cpu_switch_mm(mm->pgd, mm);
 }
 
-static inline void cpu_install_idmap(void)
+static inline void __cpu_install_idmap(pgd_t *idmap)
 {
        cpu_set_reserved_ttbr0();
        local_flush_tlb_all();
        cpu_set_idmap_tcr_t0sz();
 
-       cpu_switch_mm(lm_alias(idmap_pg_dir), &init_mm);
+       cpu_switch_mm(lm_alias(idmap), &init_mm);
+}
+
+static inline void cpu_install_idmap(void)
+{
+       __cpu_install_idmap(idmap_pg_dir);
 }
 
 /*
  * Atomically replaces the active TTBR1_EL1 PGD with a new VA-compatible PGD,
  * avoiding the possibility of conflicting TLB entries being allocated.
  */
-static inline void __nocfi cpu_replace_ttbr1(pgd_t *pgdp)
+static inline void __nocfi cpu_replace_ttbr1(pgd_t *pgdp, pgd_t *idmap)
 {
        typedef void (ttbr_replace_func)(phys_addr_t);
        extern ttbr_replace_func idmap_cpu_replace_ttbr1;
 
        replace_phys = (void *)__pa_symbol(function_nocfi(idmap_cpu_replace_ttbr1));
 
-       cpu_install_idmap();
+       __cpu_install_idmap(idmap);
        replace_phys(ttbr1);
        cpu_uninstall_idmap();
 }
 
 
 static void __maybe_unused cpu_enable_cnp(struct arm64_cpu_capabilities const *cap)
 {
-       cpu_replace_ttbr1(lm_alias(swapper_pg_dir));
+       cpu_replace_ttbr1(lm_alias(swapper_pg_dir), idmap_pg_dir);
 }
 
 /*
 
 
        /* Restore CnP bit in TTBR1_EL1 */
        if (system_supports_cnp())
-               cpu_replace_ttbr1(lm_alias(swapper_pg_dir));
+               cpu_replace_ttbr1(lm_alias(swapper_pg_dir), idmap_pg_dir);
 
        /*
         * PSTATE was not saved over suspend/resume, re-enable any detected
 
         */
        memcpy(tmp_pg_dir, swapper_pg_dir, sizeof(tmp_pg_dir));
        dsb(ishst);
-       cpu_replace_ttbr1(lm_alias(tmp_pg_dir));
+       cpu_replace_ttbr1(lm_alias(tmp_pg_dir), idmap_pg_dir);
 
        clear_pgds(KASAN_SHADOW_START, KASAN_SHADOW_END);
 
                                PAGE_KERNEL_RO));
 
        memset(kasan_early_shadow_page, KASAN_SHADOW_INIT, PAGE_SIZE);
-       cpu_replace_ttbr1(lm_alias(swapper_pg_dir));
+       cpu_replace_ttbr1(lm_alias(swapper_pg_dir), idmap_pg_dir);
 }
 
 static void __init kasan_init_depth(void)
 
 
        pgd_clear_fixmap();
 
-       cpu_replace_ttbr1(lm_alias(swapper_pg_dir));
+       cpu_replace_ttbr1(lm_alias(swapper_pg_dir), idmap_pg_dir);
        init_mm.pgd = swapper_pg_dir;
 
        memblock_phys_free(__pa_symbol(init_pg_dir),