mrs     \rd, sp_el0
        .endm
 
+/*
+ * Offset ttbr1 to allow for 48-bit kernel VAs set with 52-bit PTRS_PER_PGD.
+ * orr is used as it can cover the immediate value (and is idempotent).
+ * In future this may be nop'ed out when dealing with 52-bit kernel VAs.
+ *     ttbr: Value of ttbr to set, modified.
+ */
+       .macro  offset_ttbr1, ttbr
+#ifdef CONFIG_ARM64_52BIT_VA
+       orr     \ttbr, \ttbr, #TTBR1_BADDR_4852_OFFSET
+#endif
+       .endm
+
+/*
+ * Perform the reverse of offset_ttbr1.
+ * bic is used as it can cover the immediate value and, in future, won't need
+ * to be nop'ed out when dealing with 52-bit kernel VAs.
+ */
+       .macro  restore_ttbr1, ttbr
+#ifdef CONFIG_ARM64_52BIT_VA
+       bic     \ttbr, \ttbr, #TTBR1_BADDR_4852_OFFSET
+#endif
+       .endm
+
 /*
  * Arrange a physical address in a TTBR register, taking care of 52-bit
  * addresses.
 
 #define PGDIR_SHIFT            ARM64_HW_PGTABLE_LEVEL_SHIFT(4 - CONFIG_PGTABLE_LEVELS)
 #define PGDIR_SIZE             (_AC(1, UL) << PGDIR_SHIFT)
 #define PGDIR_MASK             (~(PGDIR_SIZE-1))
+#ifdef CONFIG_ARM64_52BIT_VA
+#define PTRS_PER_PGD           (1 << (52 - PGDIR_SHIFT))
+#else
 #define PTRS_PER_PGD           (1 << (VA_BITS - PGDIR_SHIFT))
+#endif
 
 /*
  * Section address mask and size definitions.
 #define TTBR_BADDR_MASK_52     (((UL(1) << 46) - 1) << 2)
 #endif
 
+#ifdef CONFIG_ARM64_52BIT_VA
+/* Must be at least 64-byte aligned to prevent corruption of the TTBR */
+#define TTBR1_BADDR_4852_OFFSET        (((UL(1) << (52 - PGDIR_SHIFT)) - \
+                                (UL(1) << (48 - PGDIR_SHIFT))) * 8)
+#endif
+
 #endif
 
        phys_to_ttbr x1, x1
        phys_to_ttbr x2, x2
        msr     ttbr0_el1, x2                   // load TTBR0
+       offset_ttbr1 x1
        msr     ttbr1_el1, x1                   // load TTBR1
        isb
        msr     sctlr_el1, x0
 
        tlbi    vmalle1
        dsb     nsh
        phys_to_ttbr \tmp, \page_table
+       offset_ttbr1 \tmp
        msr     ttbr1_el1, \tmp
        isb
 .endm
 
 .macro __idmap_cpu_set_reserved_ttbr1, tmp1, tmp2
        adrp    \tmp1, empty_zero_page
        phys_to_ttbr \tmp2, \tmp1
+       offset_ttbr1 \tmp2
        msr     ttbr1_el1, \tmp2
        isb
        tlbi    vmalle1
 
        __idmap_cpu_set_reserved_ttbr1 x1, x3
 
+       offset_ttbr1 x0
        msr     ttbr1_el1, x0
        isb
 
        pte             .req    x16
 
        mrs     swapper_ttb, ttbr1_el1
+       restore_ttbr1   swapper_ttb
        adr     flag_ptr, __idmap_kpti_flag
 
        cbnz    cpu, __idmap_kpti_secondary
        cbnz    w18, 1b
 
        /* All done, act like nothing happened */
+       offset_ttbr1 swapper_ttb
        msr     ttbr1_el1, swapper_ttb
        isb
        ret