#include <linux/overflow.h>
 #include <linux/pgtable.h>
 #include <linux/uaccess.h>
+#include <linux/hugetlb.h>
 #include <asm/tlbflush.h>
 #include <asm/shmparam.h>
 
 /*** Page table manipulation functions ***/
 static int vmap_pte_range(pmd_t *pmd, unsigned long addr, unsigned long end,
                        phys_addr_t phys_addr, pgprot_t prot,
-                       pgtbl_mod_mask *mask)
+                       unsigned int max_page_shift, pgtbl_mod_mask *mask)
 {
        pte_t *pte;
        u64 pfn;
+       unsigned long size = PAGE_SIZE;
 
        pfn = phys_addr >> PAGE_SHIFT;
        pte = pte_alloc_kernel_track(pmd, addr, mask);
                return -ENOMEM;
        do {
                BUG_ON(!pte_none(*pte));
+
+#ifdef CONFIG_HUGETLB_PAGE
+               size = arch_vmap_pte_range_map_size(addr, end, pfn, max_page_shift);
+               if (size != PAGE_SIZE) {
+                       pte_t entry = pfn_pte(pfn, prot);
+
+                       entry = pte_mkhuge(entry);
+                       entry = arch_make_huge_pte(entry, ilog2(size), 0);
+                       set_huge_pte_at(&init_mm, addr, pte, entry);
+                       pfn += PFN_DOWN(size);
+                       continue;
+               }
+#endif
                set_pte_at(&init_mm, addr, pte, pfn_pte(pfn, prot));
                pfn++;
-       } while (pte++, addr += PAGE_SIZE, addr != end);
+       } while (pte += PFN_DOWN(size), addr += size, addr != end);
        *mask |= PGTBL_PTE_MODIFIED;
        return 0;
 }
                        continue;
                }
 
-               if (vmap_pte_range(pmd, addr, next, phys_addr, prot, mask))
+               if (vmap_pte_range(pmd, addr, next, phys_addr, prot, max_page_shift, mask))
                        return -ENOMEM;
        } while (pmd++, phys_addr += (next - addr), addr = next, addr != end);
        return 0;