#include <linux/memblock.h>
 #include <linux/mmu_context.h>
+#include <linux/hugetlb.h>
 #include <asm/fixmap.h>
 #include <asm/code-patching.h>
 #include <asm/inst.h>
+#include <asm/pgalloc.h>
 
 #include <mm/mmu_decl.h>
 
        return 0;
 }
 
+static pte_t __init *early_hugepd_alloc_kernel(hugepd_t *pmdp, unsigned long va)
+{
+       if (hpd_val(*pmdp) == 0) {
+               pte_t *ptep = memblock_alloc(sizeof(pte_basic_t), SZ_4K);
+
+               if (!ptep)
+                       return NULL;
+
+               hugepd_populate_kernel((hugepd_t *)pmdp, ptep, PAGE_SHIFT_8M);
+               hugepd_populate_kernel((hugepd_t *)pmdp + 1, ptep, PAGE_SHIFT_8M);
+       }
+       return hugepte_offset(*(hugepd_t *)pmdp, va, PGDIR_SHIFT);
+}
+
+static int __ref __early_map_kernel_hugepage(unsigned long va, phys_addr_t pa,
+                                            pgprot_t prot, int psize, bool new)
+{
+       pmd_t *pmdp = pmd_ptr_k(va);
+       pte_t *ptep;
+
+       if (WARN_ON(psize != MMU_PAGE_512K && psize != MMU_PAGE_8M))
+               return -EINVAL;
+
+       if (new) {
+               if (WARN_ON(slab_is_available()))
+                       return -EINVAL;
+
+               if (psize == MMU_PAGE_512K)
+                       ptep = early_pte_alloc_kernel(pmdp, va);
+               else
+                       ptep = early_hugepd_alloc_kernel((hugepd_t *)pmdp, va);
+       } else {
+               if (psize == MMU_PAGE_512K)
+                       ptep = pte_offset_kernel(pmdp, va);
+               else
+                       ptep = hugepte_offset(*(hugepd_t *)pmdp, va, PGDIR_SHIFT);
+       }
+
+       if (WARN_ON(!ptep))
+               return -ENOMEM;
+
+       /* The PTE should never be already present */
+       if (new && WARN_ON(pte_present(*ptep) && pgprot_val(prot)))
+               return -EINVAL;
+
+       set_huge_pte_at(&init_mm, va, ptep, pte_mkhuge(pfn_pte(pa >> PAGE_SHIFT, prot)));
+
+       return 0;
+}
+
 /*
  * MMU_init_hw does the chip-specific initialization of the MMU hardware.
  */