Re: [RFC PATCH 01/13] arm64: preempt: Simplify and optimize __preempt_count_dec_and_test()

[Jinjie Ruan]

在 2026/7/28 20:38, Mark Rutland 写道:
> In arm64's __preempt_count_dec_and_test(), the final conditional load of
> 'ti->preempt_count' is always executed in the common case. The
> conditional load leads to unfortunate code generation for
> preempt_enable[_notrace](), and it would be better to unconditionally
> load 'ti->preempt_count', as described below.
> 
> On arm64, struct thread_info contains the following union:
> 
> | union {
> |         u64             preempt_count;
> |         struct {
> |                 u32     count;
> |                 u32     need_resched;
> |         } preempt;
> | };
> 
> Note: 'need_resched' is encoded so that '0' means a reschedule is
> needed, and '1' means a reschedule is NOT needed.

Indeed, that's true

> 
> The core logic of __preempt_count_dec_and_test() is:
> 
> | static inline bool __preempt_count_dec_and_test(void)
> | {
> |         struct thread_info *ti = current_thread_info();
> |         u64 pc = READ_ONCE(ti->preempt_count);
> |
> |         WRITE_ONCE(ti->preempt.count, --pc);
> |
> |         return !pc || !READ_ONCE(ti->preempt_count);
> | }
> 
> The '!pc' condition can only be true when both:
> 
> * The initial value of 'need_resched' was 0, meaning that a reschedule
>   is needed. This should be rare.

At the moment of executing --pc, the probability of it coinciding with
the process needing scheduling is indeed very small.

> 
> * The initial value of 'count' was exactly 1. This cannot be true for a
>   nested preempt_disable() ... preempt_enable() sequence.
> 
> Hence in common cases, '!pc' will be false, and it's necessary to
> execute the final READ_ONCE(ti->preempt_count).
> 
> This results in a conditional branch in the common case, as can be seen
> when __preempt_count_dec_and_test() is outlined:
> 
> | <outline___preempt_count_dec_and_test>:
> |        mrs     x2, sp_el0
> |        ldr     x1, [x2, #8]
> |        mov     w0, #0x1
> |        sub     x1, x1, #0x1
> |        str     w1, [x2, #8]
> |        cbz     x1, 1f
> |        ldr     x0, [x2, #8]
> |        cmp     x0, #0x0
> |        cset    w0, eq  // eq = none
> | 1:     ret
> 
> It would be better to avoid the special case for 'pc == 0', and to
> always load 'ti->preempt_count' after decrementing 'ti->preempt.count'.
> For the common cases this will remove a conditional branch. For the rare
> cases where preemption is needed initially, this only adds a single
> load, whose cost should be dominated by other factors.
> 
> Remove the special case for 'pc == 0', and always load the combined
> 'ti->preempt_count' after decrementing 'ti->preempt.count'.
> 
> The removal of the conditional branch helps with code generation, as the
> compiler can more easily move a dependent slow path out-of-line, as
> demonstrated with the following compiled with GCC 15.2.0:
> 
> | void outline_preempt_enable_notrace(void)
> | {
> |        preempt_enable_notrace();
> | }
> 
> Before this patch:
> 
> | <outline_preempt_enable_notrace>:
> |        mrs     x1, sp_el0
> |        ldr     x0, [x1, #8]
> |        sub     x0, x0, #0x1
> |        str     w0, [x1, #8]
> |        cbz     x0, 1f
> |        ldr     x0, [x1, #8]
> |        cbnz    x0, 2f
> | 1:     paciasp
> |        stp     x29, x30, [sp, #-16]!
> |        mov     x29, sp
> |        bl      preempt_schedule_notrace
> |        ldp     x29, x30, [sp], #16
> |        autiasp
> |        ret
> | 2:     ret
> 
> After this patch:
> 
> | <outline_preempt_enable_notrace>:
> |        mrs     x0, sp_el0
> |        ldr     w1, [x0, #8]
> |        sub     w1, w1, #0x1
> |        str     w1, [x0, #8]
> |        ldr     x0, [x0, #8]
> |        cbz     x0, 1f
> |        ret
> | 1:     paciasp
> |        stp     x29, x30, [sp, #-16]!
> |        mov     x29, sp
> |        bl      preempt_schedule_notrace
> |        ldp     x29, x30, [sp], #16
> |        autiasp
> |        ret
> 
> Signed-off-by: Mark Rutland <mark.rutland@xxxxxxx>
> Cc: Ada Couprie Diaz <ada.coupriediaz@xxxxxxx>
> Cc: Ard Biesheuvel <ardb@xxxxxxxxxx>
> Cc: Catalin Marinas <catalin.marinas@xxxxxxx>
> Cc: Jinjie Ruan <ruanjinjie@xxxxxxxxxx>
> Cc: Marc Zyngier <maz@xxxxxxxxxx>
> Cc: Peter Zijlstra <peterz@xxxxxxxxxxxxx>
> Cc: Vladimir Murzin <vladimir.murzin@xxxxxxx>
> Cc: Will Deacon <will@xxxxxxxxxx>
> Cc: Yang Shi <yang@xxxxxxxxxxxxxxxxxxxxxx>
> ---
>  arch/arm64/include/asm/preempt.h | 29 +++++++++++------------------
>  1 file changed, 11 insertions(+), 18 deletions(-)
> 
> diff --git a/arch/arm64/include/asm/preempt.h b/arch/arm64/include/asm/preempt.h
> index 932ea4b620428..ca2ad1a8db095 100644
> --- a/arch/arm64/include/asm/preempt.h
> +++ b/arch/arm64/include/asm/preempt.h
> @@ -55,30 +55,23 @@ static inline void __preempt_count_sub(int val)
>  	WRITE_ONCE(current_thread_info()->preempt.count, pc);
>  }
>  
> -static inline bool __preempt_count_dec_and_test(void)
> -{
> -	struct thread_info *ti = current_thread_info();
> -	u64 pc = READ_ONCE(ti->preempt_count);
> -
> -	/* Update only the count field, leaving need_resched unchanged */
> -	WRITE_ONCE(ti->preempt.count, --pc);
> -
> -	/*
> -	 * If we wrote back all zeroes, then we're preemptible and in
> -	 * need of a reschedule. Otherwise, we need to reload the
> -	 * preempt_count in case the need_resched flag was cleared by an
> -	 * interrupt occurring between the non-atomic READ_ONCE/WRITE_ONCE
> -	 * pair.
> -	 */
> -	return !pc || !READ_ONCE(ti->preempt_count);
> -}
> -
>  static inline bool should_resched(int preempt_offset)
>  {
>  	u64 pc = READ_ONCE(current_thread_info()->preempt_count);
>  	return pc == preempt_offset;
>  }
>  
> +static inline bool __preempt_count_dec_and_test(void)
> +{
> +	/*
> +	 * We must load the combined 'prempt_count' after decrementing
> +	 * 'preempt.count' as an interrupt could modify 'need_resched' before
> +	 * __preempt_count_sub() writes back to 'preempt.count'.
> +	 */
> +	__preempt_count_sub(1);
> +	return should_resched(0);

In general, the functions are consistent.

Reviewed-by: Jinjie Ruan <ruanjinjie@xxxxxxxxxx>

> +}
> +
>  #ifdef CONFIG_PREEMPTION
>  
>  void preempt_schedule(void);