Re: [PATCH v2] sched: set TIF_NEED_RESCHED before calling __trace_set_need_resched()

[Peter Zijlstra <peterz@infradead.org>]

On Wed, Jul 01, 2026 at 01:39:23PM +0530, K Prateek Nayak wrote:
> Hello Peter,
> 
> On 7/1/2026 12:24 PM, Peter Zijlstra wrote:
> >> @@ -1204,16 +1217,16 @@ static void __resched_curr(struct rq *rq, int tif)
> >>  	cpu = cpu_of(rq);
> >>  
> >>  	if (cpu == smp_processor_id()) {
> >> -		set_ti_thread_flag(cti, tif);
> >> +		int set = test_and_set_ti_thread_flag(cti, tif);
> >> +
> >> +		if (trace_sched_set_need_resched_tp_enabled() && !set)
> >> +			trace_call__sched_set_need_resched_tp(curr, cpu, tif);
> >>  		if (tif == TIF_NEED_RESCHED)
> >>  			set_preempt_need_resched();
> >> -		trace_sched_set_need_resched_tp(curr, cpu, tif);
> >>  		return;
> >>  	}
> > 
> > I can't help but notice that the local and !POLLING cases show
> > remarkable similarity. Just not sure extracting that isn't going to make
> > a mess.
> > 
> > Anyway, yes this looks about right.
> 
> If the fetch_or() based path is okay for !POLLING and local cases which
> uses a slightly (vastly?) worse instruction to set the ti->flags, we can
> instead do:
> 

Tempting, but at least arm64 would have to agree I think. That said, I
wonder why ARM64 doesn't use LDXR+WFE for idle, just like we have
MONITOR+MWAIT.

Probably too damn many idle flavours to deal with or somesuch.

> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index bd2f7fb87dc93..ea793e8a94a8f 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -1029,6 +1029,15 @@ static inline void hrtick_schedule_enter(struct rq *rq) { }
>  static inline void hrtick_schedule_exit(struct rq *rq) { }
>  #endif /* !CONFIG_SCHED_HRTICK */
>  
> +#ifndef TIF_POLLING_NRFLAG
> +/*
> + * If arch doesn't define _TIF_POLLING_NRFLAG, set it 0 to
> + * allow compilers to optimize (val & _TIF_POLLING_NRFLAG)
> + * based branches during build.
> + */
> +#define _TIF_POLLING_NRFLAG	0U
> +#endif
> +
>  /*
>   * try_cmpxchg based fetch_or() macro so it works for different integer types:
>   */
> @@ -1043,7 +1052,6 @@ static inline void hrtick_schedule_exit(struct rq *rq) { }
>  	_val;								\
>  })
>  
> -#ifdef TIF_POLLING_NRFLAG
>  /*
>   * Atomically set TIF_NEED_RESCHED and test for TIF_POLLING_NRFLAG,
>   * this avoids any races wrt polling state changes and thereby avoids
> @@ -1083,25 +1091,6 @@ static bool set_nr_if_polling(struct task_struct *p)
>  	return true;
>  }
>  
> -#else
> -static inline bool set_nr_and_not_polling(struct rq *rq, int tif)
> -{
> -	struct task_struct *curr = rq->curr;
> -	struct thread_info *ti = task_thread_info(curr);
> -	int set = test_and_set_ti_thread_flag(ti, tif);
> -
> -	if (trace_sched_set_need_resched_tp_enabled() && !set)
> -		trace_call__sched_set_need_resched_tp(curr, cpu_of(rq), tif);
> -
> -	return true;
> -}
> -
> -static inline bool set_nr_if_polling(struct task_struct *p)
> -{
> -	return false;
> -}
> -#endif
> -
>  static bool __wake_q_add(struct wake_q_head *head, struct task_struct *task)
>  {
>  	struct wake_q_node *node = &task->wake_q;
> @@ -1216,19 +1205,19 @@ static void __resched_curr(struct rq *rq, int tif)
>  
>  	cpu = cpu_of(rq);
>  
> -	if (cpu == smp_processor_id()) {
> -		int set = test_and_set_ti_thread_flag(cti, tif);
> -
> -		if (trace_sched_set_need_resched_tp_enabled() && !set)
> -			trace_call__sched_set_need_resched_tp(curr, cpu, tif);
> -		if (tif == TIF_NEED_RESCHED)
> -			set_preempt_need_resched();
> -		return;
> -	}
> -
>  	if (set_nr_and_not_polling(rq, tif)) {
> -		if (tif == TIF_NEED_RESCHED)
> -			smp_send_reschedule(cpu);
> +		if (tif != TIF_NEED_RESCHED)
> +			return;
> +		/*
> +		 * For local CPU, folding the NEED_RESCHED
> +		 * into preempt_count() is sufficient.
> +		 */
> +		if (cpu == smp_processor_id()) {
> +			set_preempt_need_resched();
> +			return;
> +		}
> +		/* Use an IPI for remote CPUs. */
> +		smp_send_reschedule(cpu);
>  	} else {
>  		trace_sched_wake_idle_without_ipi(cpu);
>  	}
> ---
> 
> set_nr_and_not_polling() will always return true for !POLLING which will
> go down the above path in __resched_curr() that deals appropriately with
> local CPU case.
> 
> set_nr_if_polling() will always return false for !POLLING from the first
> condition in the do-while loop.
> 
> Thoughts?
> 
> -- 
> Thanks and Regards,
> Prateek
>