#ifndef __ASSEMBLY__
 
 #include <linux/bitfield.h>
+#include <linux/kasan-enabled.h>
 #include <linux/page-flags.h>
+#include <linux/sched.h>
 #include <linux/types.h>
 
 #include <asm/pgtable-types.h>
 
 #endif /* CONFIG_ARM64_MTE */
 
+static inline void mte_disable_tco_entry(struct task_struct *task)
+{
+       if (!system_supports_mte())
+               return;
+
+       /*
+        * Re-enable tag checking (TCO set on exception entry). This is only
+        * necessary if MTE is enabled in either the kernel or the userspace
+        * task in synchronous or asymmetric mode (SCTLR_EL1.TCF0 bit 0 is set
+        * for both). With MTE disabled in the kernel and disabled or
+        * asynchronous in userspace, tag check faults (including in uaccesses)
+        * are not reported, therefore there is no need to re-enable checking.
+        * This is beneficial on microarchitectures where re-enabling TCO is
+        * expensive.
+        */
+       if (kasan_hw_tags_enabled() ||
+           (task->thread.sctlr_user & (1UL << SCTLR_EL1_TCF0_SHIFT)))
+               asm volatile(SET_PSTATE_TCO(0));
+}
+
 #ifdef CONFIG_KASAN_HW_TAGS
 /* Whether the MTE asynchronous mode is enabled. */
 DECLARE_STATIC_KEY_FALSE(mte_async_or_asymm_mode);
 
  */
 
 #include <linux/context_tracking.h>
+#include <linux/kasan.h>
 #include <linux/linkage.h>
 #include <linux/lockdep.h>
 #include <linux/ptrace.h>
 {
        __enter_from_kernel_mode(regs);
        mte_check_tfsr_entry();
+       mte_disable_tco_entry(current);
 }
 
 /*
        CT_WARN_ON(ct_state() != CONTEXT_USER);
        user_exit_irqoff();
        trace_hardirqs_off_finish();
+       mte_disable_tco_entry(current);
 }
 
 static __always_inline void enter_from_user_mode(struct pt_regs *regs)
 
        msr_s   SYS_ICC_PMR_EL1, x20
 alternative_else_nop_endif
 
-       /* Re-enable tag checking (TCO set on exception entry) */
-#ifdef CONFIG_ARM64_MTE
-alternative_if ARM64_MTE
-       SET_PSTATE_TCO(0)
-alternative_else_nop_endif
-#endif
-
        /*
         * Registers that may be useful after this macro is invoked:
         *
 
        mte_update_sctlr_user(next);
        mte_update_gcr_excl(next);
 
+       /* TCO may not have been disabled on exception entry for the current task. */
+       mte_disable_tco_entry(next);
+
        /*
         * Check if an async tag exception occurred at EL1.
         *