Re: [PATCH 05/24] arm64: Check MTE once per fault-handling section

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

On Sat, Nov 07, 2026 at 04:36:11PM +0100, Pierre-Clément Tosi wrote:
> Call system_uses_mte_async_or_asymm_mode() only once per section. This
> improves performance (fewer instructions at runtime and fewer static
> branches needing to be tracked and patched), correctness (no theoretical
> race between the execution of the section and the static branch being
> toggled, which could leave TCO in an unexpected state), and code clarity
> (no 3-level indirection with redundant checks to flip one PSTATE bit).
> 
> Signed-off-by: Pierre-Clément Tosi <ptosi@google.com>
> ---
>  arch/arm64/include/asm/mte-kasan.h      | 25 -----------
>  arch/arm64/include/asm/uaccess.h        | 80 ++++++++++++++++++++-------------
>  arch/arm64/include/asm/word-at-a-time.h | 22 +++++++--
>  3 files changed, 67 insertions(+), 60 deletions(-)
> 
> diff --git a/arch/arm64/include/asm/mte-kasan.h b/arch/arm64/include/asm/mte-kasan.h
> index 467d1d41a392..b474fca3ea56 100644
> --- a/arch/arm64/include/asm/mte-kasan.h
> +++ b/arch/arm64/include/asm/mte-kasan.h
> @@ -48,23 +48,6 @@ static inline void mte_enable_tco(void)
>  	asm volatile(ALTERNATIVE("nop", SET_PSTATE_TCO(1), ARM64_MTE));
>  }
>  
> -/*
> - * These functions disable tag checking only if in MTE async mode
> - * since the sync mode generates exceptions synchronously and the
> - * nofault or load_unaligned_zeropad can handle them.
> - */
> -static inline void __mte_disable_tco_async(void)
> -{
> -	if (system_uses_mte_async_or_asymm_mode())
> -		mte_disable_tco();
> -}
> -
> -static inline void __mte_enable_tco_async(void)
> -{
> -	if (system_uses_mte_async_or_asymm_mode())
> -		mte_enable_tco();
> -}
> -
>  /*
>   * These functions are meant to be only used from KASAN runtime through
>   * the arch_*() interface defined in asm/memory.h.
> @@ -197,14 +180,6 @@ static inline void mte_enable_tco(void)
>  {
>  }
>  
> -static inline void __mte_disable_tco_async(void)
> -{
> -}
> -
> -static inline void __mte_enable_tco_async(void)
> -{
> -}
> -
>  static inline u8 mte_get_ptr_tag(void *ptr)
>  {
>  	return 0xFF;
> diff --git a/arch/arm64/include/asm/uaccess.h b/arch/arm64/include/asm/uaccess.h
> index f0924dc0bd64..9e3957ee3642 100644
> --- a/arch/arm64/include/asm/uaccess.h
> +++ b/arch/arm64/include/asm/uaccess.h
> @@ -240,6 +240,24 @@ do {										\
>  	(x) = (__force __typeof__(*(ptr)))__gu_val;				\
>  } while (0)
>  
> +/*
> + * We must not call into the scheduler between setting and clearing TCO. As `x`
> + * and `ptr` may contain blocking functions, we must evaluate these outside of
> + * the critical section.
> + */
> +#define __raw_get_mem_tco(ldr, x, ptr, label, type)			\
> +do {									\
> +	__label__ __rgm_err;						\
> +									\
> +	asm volatile(SET_PSTATE_TCO(1));				\
> +	__raw_get_mem(ldr, x, ptr, __rgm_err, type);			\
> +	asm volatile(SET_PSTATE_TCO(0));				\
> +	break;								\
> +__rgm_err:								\
> +	asm volatile(SET_PSTATE_TCO(0));				\
> +	goto label;							\
> +} while (0)
> +
>  /*
>   * We must not call into the scheduler between uaccess_ttbr0_enable() and
>   * uaccess_ttbr0_disable(). As `x` and `ptr` could contain blocking functions,
> @@ -286,27 +304,18 @@ do {									\
>  
>  #define get_user	__get_user
>  
> -/*
> - * We must not call into the scheduler between __mte_enable_tco_async() and
> - * __mte_disable_tco_async(). As `dst` and `src` may contain blocking
> - * functions, we must evaluate these outside of the critical section.
> - */
> +/* We handle synchronous exceptions so only use PSTATE.TCO in non-sync MTE. */
>  #define __get_kernel_nofault(dst, src, type, err_label)			\
>  do {									\
>  	__typeof__(dst) __gkn_dst = (dst);				\
>  	__typeof__(src) __gkn_src = (src);				\
> -	do { 								\
> -		__label__ __gkn_label;					\
>  									\
> -		__mte_enable_tco_async();				\
> +	if (system_uses_mte_async_or_asymm_mode())			\
> +		__raw_get_mem_tco("ldr", *((type *)(__gkn_dst)),	\
> +		      (__force type *)(__gkn_src), err_label, K);	\
> +	else								\
>  		__raw_get_mem("ldr", *((type *)(__gkn_dst)),		\
> -		      (__force type *)(__gkn_src), __gkn_label, K);	\
> -		__mte_disable_tco_async();				\
> -		break;							\
> -	__gkn_label:							\
> -		__mte_disable_tco_async();				\
> -		goto err_label;						\
> -	} while (0);							\
> +		      (__force type *)(__gkn_src), err_label, K);	\
>  } while (0)
>  
>  #define __put_mem_asm(store, reg, x, addr, label, type)			\
> @@ -337,6 +346,24 @@ do {										\
>  	}									\
>  } while (0)
>  
> +/*
> + * We must not call into the scheduler between setting and clearing TCO. As `x`
> + * and `ptr` may contain blocking functions, we must evaluate these outside of
> + * the critical section.
> + */
> +#define __raw_put_mem_tco(str, x, ptr, label, type)			\
> +do {									\
> +	__label__ __rpm_err;						\
> +									\
> +	asm volatile(SET_PSTATE_TCO(1));				\
> +	__raw_put_mem(str, x, ptr, __rpm_err, type);			\
> +	asm volatile(SET_PSTATE_TCO(0));				\
> +	break;								\
> +__rpm_err:								\
> +	asm volatile(SET_PSTATE_TCO(0));				\
> +	goto label;							\
> +} while (0)
> +
>  /*
>   * We must not call into the scheduler between uaccess_ttbr0_enable() and
>   * uaccess_ttbr0_disable(). As `x` and `ptr` could contain blocking functions,
> @@ -383,28 +410,19 @@ do {									\
>  
>  #define put_user	__put_user
>  
> -/*
> - * We must not call into the scheduler between __mte_enable_tco_async() and
> - * __mte_disable_tco_async(). As `dst` and `src` may contain blocking
> - * functions, we must evaluate these outside of the critical section.
> - */
> +/* We handle synchronous exceptions so only use PSTATE.TCO in non-sync MTE. */
>  #define __put_kernel_nofault(dst, src, type, err_label)			\
>  do {									\
>  	__typeof__(dst) __pkn_dst = (dst);				\
>  	__typeof__(src) __pkn_src = (src);				\
>  									\
> -	do {								\
> -		__label__ __pkn_err;					\
> -		__mte_enable_tco_async();				\
> +	if (system_uses_mte_async_or_asymm_mode())			\
> +		__raw_put_mem_tco("str", *((type *)(__pkn_src)),	\
> +			      (__force type *)(__pkn_dst), err_label, K); \
> +	else								\
>  		__raw_put_mem("str", *((type *)(__pkn_src)),		\
> -			      (__force type *)(__pkn_dst), __pkn_err, K);	\
> -		__mte_disable_tco_async();				\
> -		break;							\
> -	__pkn_err:							\
> -		__mte_disable_tco_async();				\
> -		goto err_label;						\
> -	} while (0);							\
> -} while(0)
> +			      (__force type *)(__pkn_dst), err_label, K); \
> +} while (0)
>  
>  extern unsigned long __must_check __arch_copy_from_user(void *to, const void __user *from, unsigned long n);
>  #define raw_copy_from_user(to, from, n)					\
> diff --git a/arch/arm64/include/asm/word-at-a-time.h b/arch/arm64/include/asm/word-at-a-time.h
> index d7a3b69dc8c6..ec72126210d1 100644
> --- a/arch/arm64/include/asm/word-at-a-time.h
> +++ b/arch/arm64/include/asm/word-at-a-time.h
> @@ -10,6 +10,9 @@
>  #include <linux/bitops.h>
>  #include <linux/wordpart.h>
>  
> +#include <asm/mte-kasan.h>
> +#include <asm/sysreg.h>
> +
>  struct word_at_a_time {
>  	const unsigned long one_bits, high_bits;
>  };
> @@ -41,12 +44,10 @@ static inline unsigned long zero_bytemask(unsigned long bits)
>   * and the next page not being mapped, take the exception and
>   * return zeroes in the non-existing part.
>   */
> -static inline unsigned long load_unaligned_zeropad(const void *addr)
> +static __always_inline unsigned long __load_unaligned_zeropad(const void *addr)
>  {
>  	unsigned long ret;
>  
> -	__mte_enable_tco_async();
> -
>  	/* Load word from unaligned pointer addr */
>  	asm(
>  	"1:	ldr	%0, %2\n"
> @@ -55,7 +56,20 @@ static inline unsigned long load_unaligned_zeropad(const void *addr)
>  	: "=&r" (ret)
>  	: "r" (addr), "Q" (*(unsigned long *)addr));
>  
> -	__mte_disable_tco_async();
> +	return ret;
> +}
> +
> +static inline unsigned long load_unaligned_zeropad(const void *addr)
> +{
> +	unsigned long ret;
> +
> +	/* We handle synchronous exceptions so only use TCO in non-sync MTE. */
> +	if (!system_uses_mte_async_or_asymm_mode())
> +		return __load_unaligned_zeropad(addr);
> +
> +	asm volatile(SET_PSTATE_TCO(1));
> +	ret = __load_unaligned_zeropad(addr);
> +	asm volatile(SET_PSTATE_TCO(0));

Note that Sashiko identified that nothing prevents the compiler from re-ordering
__load_unaligned_zeropad() w.r.t. the SET_PSTATE_TCO() and the presence of the
exact same expression above is probably inciting it to do so!

I believe that the most robust approach would to use the "memory" clobber in
set_pstate_tco() so that it always acts as a full compiler barrier. I'll address
this (and reorder the patches to introduce the macro first) in the next version.

>  
>  	return ret;
>  }
> 
> -- 
> 2.56.0.455.g6c5a503b38-goog
> 

-- 
Pierre