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

[K Prateek Nayak <kprateek.nayak@amd.com>]

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:

  (Lightly tested after removing "HAVE_TIF_POLLING_NRFLAG" for x86)

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