[PATCH 15/24] arm64: Refactor and remove mte_disable_tco_entry()

["Pierre-Clément Tosi" <ptosi@google.com>]

Break the function up into the fundamental steps of checking the state
and clearing PSTATE.TCO. As a result, call sites can now be optimized
to only check the conditions that matter in their context, further
clarifying the intent in the code. This results in self-documenting code
and higher maintainability, addressing the need for lengthy comments.

Cc: Peter Collingbourne <peter@pcc.me.uk>
Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
---
 arch/arm64/include/asm/mte.h     | 30 ++++++++++--------------------
 arch/arm64/include/asm/uaccess.h |  6 ++++++
 arch/arm64/kernel/entry-common.c | 13 +++++++++++--
 arch/arm64/kernel/mte.c          |  8 ++++----
 4 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/arch/arm64/include/asm/mte.h b/arch/arm64/include/asm/mte.h
index 815843079eb9..1308e37c43c6 100644
--- a/arch/arm64/include/asm/mte.h
+++ b/arch/arm64/include/asm/mte.h
@@ -18,6 +18,7 @@
 #include <linux/types.h>
 
 #include <asm/barrier.h>
+#include <asm/cpufeature.h>
 #include <asm/pgtable-types.h>
 #include <asm/sysreg.h>
 
@@ -62,6 +63,15 @@ static __always_inline bool mte_get_reset_tfsr(void)
 	return __mte_get_reset_tfsr();
 }
 
+static inline bool task_uses_sync_mte_faults(struct task_struct *task)
+{
+	if (!system_supports_mte())
+		return false;
+
+	/* TCF0[0] is set for both synchronous (0b01) and asymmetric (0b11). */
+	return task->thread.sctlr_user & (1UL << SCTLR_EL1_TCF0_SHIFT);
+}
+
 #ifdef CONFIG_ARM64_MTE
 
 /* track which pages have valid allocation tags */
@@ -254,26 +264,6 @@ static inline bool folio_try_hugetlb_mte_tagging(struct folio *folio)
 }
 #endif
 
-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)))
-		set_pstate_tco(0);
-}
-
 #ifdef CONFIG_KASAN_HW_TAGS
 void mte_check_tfsr_el1(void);
 
diff --git a/arch/arm64/include/asm/uaccess.h b/arch/arm64/include/asm/uaccess.h
index aac7676da531..e99faabd5dfc 100644
--- a/arch/arm64/include/asm/uaccess.h
+++ b/arch/arm64/include/asm/uaccess.h
@@ -123,6 +123,12 @@ static inline bool uaccess_ttbr0_enable(void)
 }
 #endif
 
+static inline bool uaccess_requires_mte(struct task_struct *task)
+{
+	/* Asynchronous faults don't get reported to userspace. */
+	return task_uses_sync_mte_faults(task);
+}
+
 static inline void __uaccess_disable_hw_pan(void)
 {
 	asm(ALTERNATIVE("nop", SET_PSTATE_PAN(0), ARM64_HAS_PAN));
diff --git a/arch/arm64/kernel/entry-common.c b/arch/arm64/kernel/entry-common.c
index fbf9bafc6962..ea385f579269 100644
--- a/arch/arm64/kernel/entry-common.c
+++ b/arch/arm64/kernel/entry-common.c
@@ -7,6 +7,7 @@
 
 #include <linux/context_tracking.h>
 #include <linux/irq-entry-common.h>
+#include <linux/kasan-enabled.h>
 #include <linux/kasan.h>
 #include <linux/linkage.h>
 #include <linux/livepatch.h>
@@ -16,6 +17,7 @@
 #include <linux/sched.h>
 #include <linux/sched/debug.h>
 #include <linux/thread_info.h>
+#include <linux/uaccess.h>
 
 #include <asm/cpufeature.h>
 #include <asm/daifflags.h>
@@ -72,7 +74,13 @@ static __always_inline void irqentry_exit_pcpu_adjust(struct pt_regs *regs)
 static __always_inline void mte_enter_from_kernel_mode(void)
 {
 	mte_check_tfsr_entry();
-	mte_disable_tco_entry(current);
+
+	/*
+	 * Taking an exception sets TCO but clearing it can be expensive on
+	 * some architectures so first check that MTE is actually required.
+	 */
+	if (kasan_hw_tags_enabled() || uaccess_requires_mte(current))
+		set_pstate_tco(0);
 }
 
 static __always_inline void mte_exit_to_kernel_mode(void)
@@ -82,7 +90,8 @@ static __always_inline void mte_exit_to_kernel_mode(void)
 
 static __always_inline void mte_enter_from_user_mode(void)
 {
-	mte_disable_tco_entry(current);
+	if (kasan_hw_tags_enabled() || uaccess_requires_mte(current))
+		set_pstate_tco(0);
 }
 
 static __always_inline void mte_exit_to_user_mode(void)
diff --git a/arch/arm64/kernel/mte.c b/arch/arm64/kernel/mte.c
index 47432c9c4de0..0d998289edcb 100644
--- a/arch/arm64/kernel/mte.c
+++ b/arch/arm64/kernel/mte.c
@@ -281,8 +281,8 @@ void mte_thread_switch(struct task_struct *next)
 	mte_update_gcr_excl(next);
 
 	/* TCO may not have been disabled on exception entry for the current task. */
-	if (!kasan_hw_tags_enabled())
-		mte_disable_tco_entry(next);
+	if (!kasan_hw_tags_enabled() && uaccess_requires_mte(next))
+		set_pstate_tco(0);
 
 	if (!system_uses_mte_async_or_asymm_mode())
 		return;
@@ -401,8 +401,8 @@ long set_mte_ctrl(struct task_struct *task, unsigned long arg)
 		mte_update_sctlr_user(task);
 		mte_update_gcr_excl(task);
 		update_sctlr_el1(task->thread.sctlr_user);
-		if (!kasan_hw_tags_enabled())
-			mte_disable_tco_entry(task);
+		if (!kasan_hw_tags_enabled() && uaccess_requires_mte(task))
+			set_pstate_tco(0);
 		preempt_enable();
 	}
 

-- 
2.56.0.455.g6c5a503b38-goog