diff --git a/arch/arm64/include/asm/kernel-pgtable.h b/arch/arm64/include/asm/kernel-pgtable.h
index 85d26143faa5..796bf3d8c181 100644
--- a/arch/arm64/include/asm/kernel-pgtable.h
+++ b/arch/arm64/include/asm/kernel-pgtable.h
@@ -91,6 +91,7 @@
 #else
 #define INIT_IDMAP_DIR_SIZE	(INIT_IDMAP_DIR_PAGES * PAGE_SIZE)
 #endif
+//
 #define INIT_IDMAP_DIR_PAGES	EARLY_PAGES(KIMAGE_VADDR, _end + MAX_FDT_SIZE + SWAPPER_BLOCK_SIZE, 1)
 
 /* Initial memory map size */
diff --git a/arch/arm64/include/asm/mmu.h b/arch/arm64/include/asm/mmu.h
index 2fcf51231d6e..b817b694d1ba 100644
--- a/arch/arm64/include/asm/mmu.h
+++ b/arch/arm64/include/asm/mmu.h
@@ -12,6 +12,10 @@
 #define USER_ASID_FLAG	(UL(1) << USER_ASID_BIT)
 #define TTBR_ASID_MASK	(UL(0xffff) << 48)
 
+#define NO_BLOCK_MAPPINGS	BIT(0)
+#define NO_CONT_MAPPINGS	BIT(1)
+#define NO_EXEC_MAPPINGS	BIT(2)	/* assumes FEAT_HPDS is not used */
+
 #ifndef __ASSEMBLY__
 
 #include <linux/refcount.h>
diff --git a/arch/arm64/kernel/head.S b/arch/arm64/kernel/head.S
index 7b236994f0e1..e2fa6b95f809 100644
--- a/arch/arm64/kernel/head.S
+++ b/arch/arm64/kernel/head.S
@@ -27,6 +27,7 @@
 #include <asm/kernel-pgtable.h>
 #include <asm/kvm_arm.h>
 #include <asm/memory.h>
+#include <asm/mmu.h>
 #include <asm/pgtable-hwdef.h>
 #include <asm/page.h>
 #include <asm/scs.h>
@@ -332,79 +333,69 @@ SYM_FUNC_START_LOCAL(remap_region)
 SYM_FUNC_END(remap_region)
 
 SYM_FUNC_START_LOCAL(create_idmap)
-	mov	x28, lr
-	/*
-	 * The ID map carries a 1:1 mapping of the physical address range
-	 * covered by the loaded image, which could be anywhere in DRAM. This
-	 * means that the required size of the VA (== PA) space is decided at
-	 * boot time, and could be more than the configured size of the VA
-	 * space for ordinary kernel and user space mappings.
-	 *
-	 * There are three cases to consider here:
-	 * - 39 <= VA_BITS < 48, and the ID map needs up to 48 VA bits to cover
-	 *   the placement of the image. In this case, we configure one extra
-	 *   level of translation on the fly for the ID map only. (This case
-	 *   also covers 42-bit VA/52-bit PA on 64k pages).
-	 *
-	 * - VA_BITS == 48, and the ID map needs more than 48 VA bits. This can
-	 *   only happen when using 64k pages, in which case we need to extend
-	 *   the root level table rather than add a level. Note that we can
-	 *   treat this case as 'always extended' as long as we take care not
-	 *   to program an unsupported T0SZ value into the TCR register.
-	 *
-	 * - Combinations that would require two additional levels of
-	 *   translation are not supported, e.g., VA_BITS==36 on 16k pages, or
-	 *   VA_BITS==39/4k pages with 5-level paging, where the input address
-	 *   requires more than 47 or 48 bits, respectively.
-	 */
-#if (VA_BITS < 48)
-#define IDMAP_PGD_ORDER	(VA_BITS - PGDIR_SHIFT)
-#define EXTRA_SHIFT	(PGDIR_SHIFT + PAGE_SHIFT - 3)
+	adr_l	x0, init_stack
+	add	sp, x0, #THREAD_SIZE
+	sub	sp, sp, #16
+	stp	lr, x0, [sp, #0]		// x0 is useless, just to keep stack 16-bytes align
 
-	/*
-	 * If VA_BITS < 48, we have to configure an additional table level.
-	 * First, we have to verify our assumption that the current value of
-	 * VA_BITS was chosen such that all translation levels are fully
-	 * utilised, and that lowering T0SZ will always result in an additional
-	 * translation level to be configured.
-	 */
-#if VA_BITS != EXTRA_SHIFT
-#error "Mismatch between VA_BITS and page size/number of translation levels"
-#endif
-#else
-#define IDMAP_PGD_ORDER	(PHYS_MASK_SHIFT - PGDIR_SHIFT)
-#define EXTRA_SHIFT
-	/*
-	 * If VA_BITS == 48, we don't have to configure an additional
-	 * translation level, but the top-level table has more entries.
-	 */
-#endif
 	adrp	x0, init_idmap_pg_dir
-	adrp	x3, _text
-	adrp	x6, _end + MAX_FDT_SIZE + SWAPPER_BLOCK_SIZE
-	mov_q	x7, SWAPPER_RX_MMUFLAGS
-
-	map_memory x0, x1, x3, x6, x7, x3, IDMAP_PGD_ORDER, x10, x11, x12, x13, x14, EXTRA_SHIFT
-
-	/* Remap the kernel page tables r/w in the ID map */
-	adrp	x1, _text
-	adrp	x2, init_pg_dir
-	adrp	x3, init_pg_end
-	bic	x4, x2, #SWAPPER_BLOCK_SIZE - 1
-	mov_q	x5, SWAPPER_RW_MMUFLAGS
-	mov	x6, #SWAPPER_BLOCK_SHIFT
-	bl	remap_region
-
-	/* Remap the FDT after the kernel image */
-	adrp	x1, _text
-	adrp	x22, _end + SWAPPER_BLOCK_SIZE
-	bic	x2, x22, #SWAPPER_BLOCK_SIZE - 1
+	adrp	x1, init_idmap_pg_end
+	sub     x1, x1, x0
+	bl      headpool_init
+	mov     x0, #0
+	bl      headpool_pgtable_alloc          // return x0, containing init_idmap_pg_dir
+	mov	x27, x0				// bake in case of flush
+
+	adr_l	x1, _text			// phys
+	mov	x2, x1				// virt for idmap
+	adr_l	x3, _etext - 1
+	sub	x3, x3, x1			// size
+	ldr     x4, =SWAPPER_RX_MMUFLAGS
+	adr_l   x5, headpool_pgtable_alloc
+	mov     x6, #0
+	bl      mmu_head_create_pgd_mapping
+
+	mov	x0, x27			// pgd
+	adr_l	x1, init_pg_dir		// phys
+	mov	x2, x1			// virt for idmap
+	adr_l	x3, init_pg_end
+	sub	x3, x3, x1
+	ldr     x4, =SWAPPER_RW_MMUFLAGS
+	adr_l   x5, headpool_pgtable_alloc
+	mov     x6, #0
+	bl      mmu_head_create_pgd_mapping
+
+	mov	x0, x27			// pgd
+	adr_l	x1, init_stack		// kernel mapping need write-permission to use this stack
+	mov	x2, x1			// virt for idmap
+	ldr	x3, =THREAD_SIZE
+	ldr     x4, =SWAPPER_RW_MMUFLAGS
+	adr_l   x5, headpool_pgtable_alloc
+	mov     x6, #0
+	bl      mmu_head_create_pgd_mapping
+
+	mov	x0, x27			// pgd
+	adr_l	x1, __initdata_begin	// kernel mapping need write-permission to it
+	mov	x2, x1			// virt for idmap
+	adr_l	x3, __initdata_end
+	sub	x3, x3, x1
+	ldr     x4, =SWAPPER_RW_MMUFLAGS
+	adr_l   x5, headpool_pgtable_alloc
+	mov     x6, #0
+	bl      mmu_head_create_pgd_mapping
+
+
+	mov	x0, x27							// pgd
+	mov	x1, x21							// FDT phys
+	adr_l	x2, _end + SWAPPER_BLOCK_SIZE				// virt
+	mov	x3, #(MAX_FDT_SIZE + SWAPPER_BLOCK_SIZE)		// size
+	ldr     x4, =SWAPPER_RW_MMUFLAGS
+	adr_l   x5, headpool_pgtable_alloc
+	mov     x6, #0
+	bl      mmu_head_create_pgd_mapping
+
+	adr_l	x22, _end + SWAPPER_BLOCK_SIZE
 	bfi	x22, x21, #0, #SWAPPER_BLOCK_SHIFT		// remapped FDT address
-	add	x3, x2, #MAX_FDT_SIZE + SWAPPER_BLOCK_SIZE
-	bic	x4, x21, #SWAPPER_BLOCK_SIZE - 1
-	mov_q	x5, SWAPPER_RW_MMUFLAGS
-	mov	x6, #SWAPPER_BLOCK_SHIFT
-	bl	remap_region
 
 	/*
 	 * Since the page tables have been populated with non-cacheable
@@ -417,22 +408,42 @@ SYM_FUNC_START_LOCAL(create_idmap)
 	adrp	x0, init_idmap_pg_dir
 	adrp	x1, init_idmap_pg_end
 	bl	dcache_inval_poc
-0:	ret	x28
+	ldp     lr, x0, [sp], #16
+0:	ret
 SYM_FUNC_END(create_idmap)
 
 SYM_FUNC_START_LOCAL(create_kernel_mapping)
+	sub	sp, sp, #80
+	stp	x0, x1, [sp, #0]
+	stp     x2, x3, [sp, #16]
+	stp     x4, x5, [sp, #32]
+	stp	x6, x7, [sp, #48]
+	stp	lr, xzr, [sp, #64]
+
 	adrp	x0, init_pg_dir
-	mov_q	x5, KIMAGE_VADDR		// compile time __va(_text)
+	adrp    x1, init_pg_end
+	sub     x1, x1, x0
+	bl      headpool_init
+	mov     x0, #0
+	bl      headpool_pgtable_alloc          // return x0, containing init_pg_dir
+
+	adrp    x1, _text                       // runtime __pa(_text)
+	mov_q   x2, KIMAGE_VADDR                // compile time __va(_text)
 #ifdef CONFIG_RELOCATABLE
-	add	x5, x5, x23			// add KASLR displacement
+	add     x2, x2, x23                     // add KASLR displacement
 #endif
-	adrp	x6, _end			// runtime __pa(_end)
-	adrp	x3, _text			// runtime __pa(_text)
-	sub	x6, x6, x3			// _end - _text
-	add	x6, x6, x5			// runtime __va(_end)
-	mov_q	x7, SWAPPER_RW_MMUFLAGS
-
-	map_memory x0, x1, x5, x6, x7, x3, (VA_BITS - PGDIR_SHIFT), x10, x11, x12, x13, x14
+	adrp    x3, _end                        // runtime __pa(_end)
+	sub     x3, x3, x1                      // _end - _text
+	ldr     x4, =SWAPPER_RW_MMUFLAGS
+	adr_l   x5, headpool_pgtable_alloc
+	mov     x6, #0
+	bl      mmu_head_create_pgd_mapping
+
+	ldp	lr, xzr, [sp, #64]
+	ldp     x6, x7, [sp, #48]
+	ldp     x4, x5, [sp, #32]
+	ldp     x2, x3, [sp, #16]
+	ldp     x0, x1, [sp], #80
 
 	dsb	ishst				// sync with page table walker
 	ret
diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index 80e49faaf066..e9748c7017dd 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -41,10 +41,6 @@
 #include <asm/pgalloc.h>
 #include <asm/kfence.h>
 
-#define NO_BLOCK_MAPPINGS	BIT(0)
-#define NO_CONT_MAPPINGS	BIT(1)
-#define NO_EXEC_MAPPINGS	BIT(2)	/* assumes FEAT_HPDS is not used */
-
 int idmap_t0sz __ro_after_init;
 
 #if VA_BITS > 48
