ldr     \rd, [\rn, #VMA_VM_MM]
        .endm
 
-/*
- * mmid - get context id from mm pointer (mm->context.id)
- */
-       .macro  mmid, rd, rn
-       ldr     \rd, [\rn, #MM_CONTEXT_ID]
-       .endm
 /*
  * read_ctr - read CTR_EL0. If the system has mismatched register fields,
  * provide the system wide safe value from arm64_ftr_reg_ctrel0.sys_val
 
        isb();
 }
 
+void cpu_do_switch_mm(phys_addr_t pgd_phys, struct mm_struct *mm);
+
 static inline void cpu_switch_mm(pgd_t *pgd, struct mm_struct *mm)
 {
        BUG_ON(pgd == swapper_pg_dir);
 
 
 #include <asm/page.h>
 
-struct mm_struct;
 struct cpu_suspend_ctx;
 
 extern void cpu_do_idle(void);
-extern void cpu_do_switch_mm(unsigned long pgd_phys, struct mm_struct *mm);
 extern void cpu_do_suspend(struct cpu_suspend_ctx *ptr);
 extern u64 cpu_do_resume(phys_addr_t ptr, u64 idmap_ttbr);
 
 
  * Copyright (C) 2012 ARM Ltd.
  */
 
+#include <linux/bitfield.h>
 #include <linux/bitops.h>
 #include <linux/sched.h>
 #include <linux/slab.h>
 /* Errata workaround post TTBRx_EL1 update. */
 asmlinkage void post_ttbr_update_workaround(void)
 {
+       if (!IS_ENABLED(CONFIG_CAVIUM_ERRATUM_27456))
+               return;
+
        asm(ALTERNATIVE("nop; nop; nop",
                        "ic iallu; dsb nsh; isb",
-                       ARM64_WORKAROUND_CAVIUM_27456,
-                       CONFIG_CAVIUM_ERRATUM_27456));
+                       ARM64_WORKAROUND_CAVIUM_27456));
+}
+
+void cpu_do_switch_mm(phys_addr_t pgd_phys, struct mm_struct *mm)
+{
+       unsigned long ttbr1 = read_sysreg(ttbr1_el1);
+       unsigned long asid = ASID(mm);
+       unsigned long ttbr0 = phys_to_ttbr(pgd_phys);
+
+       /* Skip CNP for the reserved ASID */
+       if (system_supports_cnp() && asid)
+               ttbr0 |= TTBR_CNP_BIT;
+
+       /* SW PAN needs a copy of the ASID in TTBR0 for entry */
+       if (IS_ENABLED(CONFIG_ARM64_SW_TTBR0_PAN))
+               ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, asid);
+
+       /* Set ASID in TTBR1 since TCR.A1 is set */
+       ttbr1 &= ~TTBR_ASID_MASK;
+       ttbr1 |= FIELD_PREP(TTBR_ASID_MASK, asid);
+
+       write_sysreg(ttbr1, ttbr1_el1);
+       isb();
+       write_sysreg(ttbr0, ttbr0_el1);
+       isb();
+       post_ttbr_update_workaround();
 }
 
 static int asids_init(void)
 
        .popsection
 #endif
 
-/*
- *     cpu_do_switch_mm(pgd_phys, tsk)
- *
- *     Set the translation table base pointer to be pgd_phys.
- *
- *     - pgd_phys - physical address of new TTB
- */
-SYM_FUNC_START(cpu_do_switch_mm)
-       mrs     x2, ttbr1_el1
-       mmid    x1, x1                          // get mm->context.id
-       phys_to_ttbr x3, x0
-
-alternative_if ARM64_HAS_CNP
-       cbz     x1, 1f                          // skip CNP for reserved ASID
-       orr     x3, x3, #TTBR_CNP_BIT
-1:
-alternative_else_nop_endif
-#ifdef CONFIG_ARM64_SW_TTBR0_PAN
-       bfi     x3, x1, #48, #16                // set the ASID field in TTBR0
-#endif
-       bfi     x2, x1, #48, #16                // set the ASID
-       msr     ttbr1_el1, x2                   // in TTBR1 (since TCR.A1 is set)
-       isb
-       msr     ttbr0_el1, x3                   // now update TTBR0
-       isb
-       b       post_ttbr_update_workaround     // Back to C code...
-SYM_FUNC_END(cpu_do_switch_mm)
-
        .pushsection ".idmap.text", "awx"
 
 .macro __idmap_cpu_set_reserved_ttbr1, tmp1, tmp2