void cpu_do_switch_mm(phys_addr_t pgd_phys, struct mm_struct *mm)
{
	#ifndef CONFIG_IEE
	unsigned long ttbr1 = read_sysreg(ttbr1_el1);
	#endif
	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);

	#ifdef CONFIG_IEE
	ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, asid+1);
    iee_rwx_gate_entry(IEE_CONTEXT_SWITCH, ttbr0);

	#else
	/* Set ASID in TTBR0 since TCR.A1 is set 0*/

	#ifdef CONFIG_KOI
	ttbr0 |= FIELD_PREP(TTBR_ASID_MASK, asid+1);
    ttbr1 &= ~TTBR_ASID_MASK;
	ttbr1 |= FIELD_PREP(TTBR_ASID_MASK, asid);
	#else
	ttbr1 &= ~TTBR_ASID_MASK;
	ttbr1 |= FIELD_PREP(TTBR_ASID_MASK, asid);
	#endif
	write_sysreg(ttbr1, ttbr1_el1);
	isb();
	write_sysreg(ttbr0, ttbr0_el1);
	isb();
    #endif

	post_ttbr_update_workaround();
}
