diff --git a/arch/arm64/kernel/image-vars.h b/arch/arm64/kernel/image-vars.h
index 82bafa1f869c3a8b..f48b6f09d278d3bf 100644
--- a/arch/arm64/kernel/image-vars.h
+++ b/arch/arm64/kernel/image-vars.h
@@ -56,6 +56,8 @@ PROVIDE(__pi_arm64_sw_feature_override	= arm64_sw_feature_override);
 PROVIDE(__pi_arm64_use_ng_mappings	= arm64_use_ng_mappings);
 PROVIDE(__pi__ctype			= _ctype);
 
+PROVIDE(__pi_init_idmap_pg_dir		= init_idmap_pg_dir);
+PROVIDE(__pi_init_idmap_pg_end		= init_idmap_pg_end);
 PROVIDE(__pi_init_pg_dir		= init_pg_dir);
 PROVIDE(__pi_init_pg_end		= init_pg_end);
 PROVIDE(__pi_swapper_pg_dir		= swapper_pg_dir);
diff --git a/arch/arm64/kernel/pi/map_kernel.c b/arch/arm64/kernel/pi/map_kernel.c
index 3b0b3fecf2bd533b..71cc32cd2545b85a 100644
--- a/arch/arm64/kernel/pi/map_kernel.c
+++ b/arch/arm64/kernel/pi/map_kernel.c
@@ -137,6 +137,22 @@ static bool __init arm64_early_this_cpu_has_lva(void)
 						    ID_AA64MMFR2_EL1_VARange_SHIFT);
 }
 
+static bool __init arm64_early_this_cpu_has_lpa2(void)
+{
+	bool gran4k = IS_ENABLED(CONFIG_ARM64_4K_PAGES);
+	u64 mmfr0;
+	int feat;
+
+	mmfr0 = read_sysreg(id_aa64mmfr0_el1);
+	mmfr0 &= ~id_aa64mmfr0_override.mask;
+	mmfr0 |= id_aa64mmfr0_override.val;
+	feat = cpuid_feature_extract_field(mmfr0, ID_AA64MMFR0_EL1_TGRAN_SHIFT,
+					   gran4k /* signed */);
+
+	return gran4k ? feat >= ID_AA64MMFR0_EL1_TGRAN4_52_BIT
+		      : feat >= ID_AA64MMFR0_EL1_TGRAN16_52_BIT;
+}
+
 static bool __init arm64_early_this_cpu_has_pac(void)
 {
 	u64 isar1, isar2;
@@ -279,6 +295,74 @@ static void noinline __section(".idmap.text") disable_wxn(void)
 	    ::	"r"(sctlr & ~SCTLR_ELx_M), "r"(sctlr));
 }
 
+static void noinline __section(".idmap.text") set_ttbr0_for_lpa2(u64 ttbr)
+{
+	u64 sctlr = read_sysreg(sctlr_el1);
+	u64 tcr = read_sysreg(tcr_el1) | TCR_DS;
+
+	asm("	msr	sctlr_el1, %0		;"
+	    "	isb				;"
+	    "   msr     ttbr0_el1, %1		;"
+	    "   msr     tcr_el1, %2		;"
+	    "	isb				;"
+	    "	tlbi    vmalle1			;"
+	    "	dsb     nsh			;"
+	    "	isb				;"
+	    "	msr     sctlr_el1, %3		;"
+	    "	isb				;"
+	    ::	"r"(sctlr & ~SCTLR_ELx_M), "r"(ttbr), "r"(tcr), "r"(sctlr));
+}
+
+static void remap_idmap_for_lpa2(void)
+{
+	extern pgd_t init_idmap_pg_dir[], init_idmap_pg_end[];
+	pgd_t *pgdp = (void *)init_pg_dir + PAGE_SIZE;
+	pgprot_t text_prot = PAGE_KERNEL_ROX;
+	pgprot_t data_prot = PAGE_KERNEL;
+
+	/* clear the bits that change meaning once LPA2 is turned on */
+	pgprot_val(text_prot) &= ~PTE_SHARED;
+	pgprot_val(data_prot) &= ~PTE_SHARED;
+
+	/*
+	 * We have to clear bits [9:8] in all block or page descriptors in the
+	 * initial ID map, as otherwise they will be (mis)interpreted as
+	 * physical address bits once we flick the LPA2 switch (TCR.DS). Since
+	 * we cannot manipulate live descriptors in that way without creating
+	 * potential TLB conflicts, let's create another temporary ID map in a
+	 * LPA2 compatible fashion, and update the initial ID map while running
+	 * from that.
+	 */
+	map_segment(init_pg_dir, &pgdp, 0, _stext, __inittext_end, text_prot, false);
+	map_segment(init_pg_dir, &pgdp, 0, __initdata_begin, _end, data_prot, false);
+	dsb(ishst);
+	set_ttbr0_for_lpa2((u64)init_pg_dir);
+
+	/*
+	 * Recreate the initial ID map with the same granularity as before.
+	 * Don't bother with the FDT, we no longer need it after this.
+	 */
+	memset(init_idmap_pg_dir, 0,
+	       (u64)init_idmap_pg_dir - (u64)init_idmap_pg_end);
+
+	pgdp = (void *)init_idmap_pg_dir + PAGE_SIZE;
+	map_segment(init_idmap_pg_dir, &pgdp, 0,
+		    PTR_ALIGN_DOWN(&_stext[0], INIT_IDMAP_BLOCK_SIZE),
+		    PTR_ALIGN_DOWN(&__bss_start[0], INIT_IDMAP_BLOCK_SIZE),
+		    text_prot, false);
+	map_segment(init_idmap_pg_dir, &pgdp, 0,
+		    PTR_ALIGN_DOWN(&__bss_start[0], INIT_IDMAP_BLOCK_SIZE),
+		    PTR_ALIGN(&_end[0], INIT_IDMAP_BLOCK_SIZE),
+		    data_prot, false);
+	dsb(ishst);
+
+	/* switch back to the updated initial ID map */
+	set_ttbr0_for_lpa2((u64)init_idmap_pg_dir);
+
+	/* wipe the temporary ID map from memory */
+	memset(init_pg_dir, 0, (u64)init_pg_end - (u64)init_pg_dir);
+}
+
 asmlinkage void __init early_map_kernel(u64 boot_status, void *fdt)
 {
 	static char const chosen_str[] __initconst = "/chosen";
@@ -292,9 +376,6 @@ asmlinkage void __init early_map_kernel(u64 boot_status, void *fdt)
 	/* Parse the command line for CPU feature overrides */
 	init_feature_override(boot_status, fdt, chosen);
 
-	if (VA_BITS > VA_BITS_MIN && arm64_early_this_cpu_has_lva())
-		sysreg_clear_set(tcr_el1, TCR_T1SZ_MASK, TCR_T1SZ(VA_BITS));
-
 	if (IS_ENABLED(CONFIG_ARM64_WXN) &&
 	    cpuid_feature_extract_unsigned_field(arm64_sw_feature_override.val,
 						 ARM64_SW_FEATURE_OVERRIDE_NOWXN))
@@ -322,6 +403,20 @@ asmlinkage void __init early_map_kernel(u64 boot_status, void *fdt)
 			arm64_use_ng_mappings = true;
 	}
 
+	if (VA_BITS > VA_BITS_MIN) {
+		bool va52 = false;
+
+		if (IS_ENABLED(CONFIG_ARM64_64K_PAGES)) {
+			va52 = arm64_early_this_cpu_has_lva();
+		} else if (arm64_early_this_cpu_has_lpa2()) {
+			remap_idmap_for_lpa2();
+			va52 = true;
+		}
+		if (va52)
+			sysreg_clear_set(tcr_el1, TCR_T1SZ_MASK,
+					 TCR_T1SZ(VA_BITS));
+	}
+
 	va_base = KIMAGE_VADDR + kaslr_offset;
 	map_kernel(kaslr_offset, va_base - pa_base);
 }
diff --git a/arch/arm64/mm/proc.S b/arch/arm64/mm/proc.S
index c747a2ef478cabec..8197663a54f63c9d 100644
--- a/arch/arm64/mm/proc.S
+++ b/arch/arm64/mm/proc.S
@@ -443,6 +443,9 @@ SYM_FUNC_START(__cpu_setup)
 #if VA_BITS > VA_BITS_MIN
 alternative_if ARM64_HAS_LVA
 	eor		tcr, tcr, #TCR_T1SZ(VA_BITS) ^ TCR_T1SZ(VA_BITS_MIN)
+#ifndef CONFIG_ARM64_64K_PAGES
+	orr		tcr, tcr, #TCR_DS
+#endif
 alternative_else_nop_endif
 #elif VA_BITS < 48
 	idmap_get_t0sz	x9
