* 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
+       .macro  offset_ttbr1, ttbr, tmp
 #ifdef CONFIG_ARM64_USER_VA_BITS_52
        orr     \ttbr, \ttbr, #TTBR1_BADDR_4852_OFFSET
+#endif
+
+#ifdef CONFIG_ARM64_VA_BITS_52
+       mrs_s   \tmp, SYS_ID_AA64MMFR2_EL1
+       and     \tmp, \tmp, #(0xf << ID_AA64MMFR2_LVA_SHIFT)
+       cbnz    \tmp, .Lskipoffs_\@
+       orr     \ttbr, \ttbr, #TTBR1_BADDR_4852_OFFSET
+.Lskipoffs_\@ :
 #endif
        .endm
 
  * to be nop'ed out when dealing with 52-bit kernel VAs.
  */
        .macro  restore_ttbr1, ttbr
-#ifdef CONFIG_ARM64_USER_VA_BITS_52
+#if defined(CONFIG_ARM64_USER_VA_BITS_52) || defined(CONFIG_ARM64_VA_BITS_52)
        bic     \ttbr, \ttbr, #TTBR1_BADDR_4852_OFFSET
 #endif
        .endm
 
        phys_to_ttbr x1, x1
        phys_to_ttbr x2, x2
        msr     ttbr0_el1, x2                   // load TTBR0
-       offset_ttbr1 x1
+       offset_ttbr1 x1, x3
        msr     ttbr1_el1, x1                   // load TTBR1
        isb
        msr     sctlr_el1, x0
 
  * Even switching to our copied tables will cause a changed output address at
  * each stage of the walk.
  */
-.macro break_before_make_ttbr_switch zero_page, page_table, tmp
+.macro break_before_make_ttbr_switch zero_page, page_table, tmp, tmp2
        phys_to_ttbr \tmp, \zero_page
        msr     ttbr1_el1, \tmp
        isb
        tlbi    vmalle1
        dsb     nsh
        phys_to_ttbr \tmp, \page_table
-       offset_ttbr1 \tmp
+       offset_ttbr1 \tmp, \tmp2
        msr     ttbr1_el1, \tmp
        isb
 .endm
         * We execute from ttbr0, change ttbr1 to our copied linear map tables
         * with a break-before-make via the zero page
         */
-       break_before_make_ttbr_switch   x5, x0, x6
+       break_before_make_ttbr_switch   x5, x0, x6, x8
 
        mov     x21, x1
        mov     x30, x2
        dsb     ish             /* wait for PoU cleaning to finish */
 
        /* switch to the restored kernels page tables */
-       break_before_make_ttbr_switch   x25, x21, x6
+       break_before_make_ttbr_switch   x25, x21, x6, x8
 
        ic      ialluis
        dsb     ish
 
 .macro __idmap_cpu_set_reserved_ttbr1, tmp1, tmp2
        adrp    \tmp1, empty_zero_page
        phys_to_ttbr \tmp2, \tmp1
-       offset_ttbr1 \tmp2
+       offset_ttbr1 \tmp2, \tmp1
        msr     ttbr1_el1, \tmp2
        isb
        tlbi    vmalle1
 
        __idmap_cpu_set_reserved_ttbr1 x1, x3
 
-       offset_ttbr1 x0
+       offset_ttbr1 x0, x3
        msr     ttbr1_el1, x0
        isb
 
        cbnz    w18, 1b
 
        /* All done, act like nothing happened */
-       offset_ttbr1 swapper_ttb
+       offset_ttbr1 swapper_ttb, x18
        msr     ttbr1_el1, swapper_ttb
        isb
        ret