#define _ASM_ARM64_VMALLOC_H
 
 #include <asm/page.h>
+#include <asm/pgtable.h>
 
 #ifdef CONFIG_HAVE_ARCH_HUGE_VMAP
 
 static inline bool arch_vmap_pud_supported(pgprot_t prot)
 {
        /*
-        * Only 4k granule supports level 1 block mappings.
         * SW table walks can't handle removal of intermediate entries.
         */
-       return IS_ENABLED(CONFIG_ARM64_4K_PAGES) &&
+       return pud_sect_supported() &&
               !IS_ENABLED(CONFIG_PTDUMP_DEBUGFS);
 }
 
 
 {
        int order;
 
-#ifdef CONFIG_ARM64_4K_PAGES
-       order = PUD_SHIFT - PAGE_SHIFT;
-#else
-       order = CONT_PMD_SHIFT + PMD_SHIFT - PAGE_SHIFT;
-#endif
+       if (pud_sect_supported())
+               order = PUD_SHIFT - PAGE_SHIFT;
+       else
+               order = CONT_PMD_SHIFT + PMD_SHIFT - PAGE_SHIFT;
        /*
         * HugeTLB CMA reservation is required for gigantic
         * huge pages which could not be allocated via the
        size_t pagesize = huge_page_size(h);
 
        switch (pagesize) {
-#ifdef CONFIG_ARM64_4K_PAGES
+#ifndef __PAGETABLE_PMD_FOLDED
        case PUD_SIZE:
+               return pud_sect_supported();
 #endif
        case PMD_SIZE:
        case CONT_PMD_SIZE:
        *pgsize = size;
 
        switch (size) {
-#ifdef CONFIG_ARM64_4K_PAGES
+#ifndef __PAGETABLE_PMD_FOLDED
        case PUD_SIZE:
+               if (pud_sect_supported())
+                       contig_ptes = 1;
+               break;
 #endif
        case PMD_SIZE:
                contig_ptes = 1;
 
 static int __init hugetlbpage_init(void)
 {
-#ifdef CONFIG_ARM64_4K_PAGES
-       hugetlb_add_hstate(PUD_SHIFT - PAGE_SHIFT);
-#endif
+       if (pud_sect_supported())
+               hugetlb_add_hstate(PUD_SHIFT - PAGE_SHIFT);
+
        hugetlb_add_hstate(CONT_PMD_SHIFT - PAGE_SHIFT);
        hugetlb_add_hstate(PMD_SHIFT - PAGE_SHIFT);
        hugetlb_add_hstate(CONT_PTE_SHIFT - PAGE_SHIFT);
 bool __init arch_hugetlb_valid_size(unsigned long size)
 {
        switch (size) {
-#ifdef CONFIG_ARM64_4K_PAGES
+#ifndef __PAGETABLE_PMD_FOLDED
        case PUD_SIZE:
+               return pud_sect_supported();
 #endif
        case CONT_PMD_SIZE:
        case PMD_SIZE: