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

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

Hello Peter, Gabriele,

On 7/1/2026 2:04 AM, Peter Zijlstra wrote:
>> -static inline bool set_nr_and_not_polling(struct thread_info *ti, int tif)
>> +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);
>> +
>>  	set_ti_thread_flag(ti, tif);
>>  	return true;
>>  }
> 
> This !POLLING thing also needs tracing, no?

Indeed! I'm stupid.

Here is the full diff with set_nr_if_polling() covered too on top of
tip:sched/core + Sechang's v3 at
https://lore.kernel.org/lkml/20260630084750.2792851-1-rhkrqnwk98@gmail.com/
for completeness.

Let me know if you prefer folding it in or if this needs to be a
separate patch.

  (lightly testet)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 7cbd541f656f..bd2f7fb87dc9 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -1049,9 +1049,16 @@ static inline void hrtick_schedule_exit(struct rq *rq) { }
  * this avoids any races wrt polling state changes and thereby avoids
  * spurious IPIs.
  */
-static inline bool set_nr_and_not_polling(struct thread_info *ti, int tif)
+static inline bool set_nr_and_not_polling(struct rq *rq, int tif)
 {
-	return !(fetch_or(&ti->flags, 1 << tif) & _TIF_POLLING_NRFLAG);
+	struct task_struct *curr = rq->curr;
+	struct thread_info *ti = task_thread_info(curr);
+	unsigned long old_flags = fetch_or(&ti->flags, 1 << tif);
+
+	if (trace_sched_set_need_resched_tp_enabled() && !(old_flags & (1 << tif)))
+		trace_call__sched_set_need_resched_tp(curr, cpu_of(rq), tif);
+
+	return !(old_flags & _TIF_POLLING_NRFLAG);
 }
 
 /*
@@ -1072,13 +1079,20 @@ static bool set_nr_if_polling(struct task_struct *p)
 			return true;
 	} while (!try_cmpxchg(&ti->flags, &val, val | _TIF_NEED_RESCHED));
 
+	trace_sched_set_need_resched_tp(p, task_cpu(p), TIF_NEED_RESCHED);
 	return true;
 }
 
 #else
-static inline bool set_nr_and_not_polling(struct thread_info *ti, int tif)
+static inline bool set_nr_and_not_polling(struct rq *rq, int tif)
 {
-	set_ti_thread_flag(ti, 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;
 }
 
@@ -1186,7 +1200,6 @@ static void __resched_curr(struct rq *rq, int tif)
 {
 	struct task_struct *curr = rq->curr;
 	struct thread_info *cti = task_thread_info(curr);
-	bool need_ipi;
 	int cpu;
 
 	lockdep_assert_rq_held(rq);
@@ -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;
 	}
 
-	need_ipi = set_nr_and_not_polling(cti, tif);
-	trace_sched_set_need_resched_tp(curr, cpu, tif);
-	if (need_ipi) {
+	if (set_nr_and_not_polling(rq, tif)) {
 		if (tif == TIF_NEED_RESCHED)
 			smp_send_reschedule(cpu);
 	} else {
@@ -1353,7 +1366,7 @@ static void wake_up_idle_cpu(int cpu)
 	 * and testing of the above solutions didn't appear to report
 	 * much benefits.
 	 */
-	if (set_nr_and_not_polling(task_thread_info(rq->idle), TIF_NEED_RESCHED))
+	if (set_nr_and_not_polling(rq, TIF_NEED_RESCHED))
 		smp_send_reschedule(cpu);
 	else
 		trace_sched_wake_idle_without_ipi(cpu);
---

Now, on a separate note to Gabriele, this trips up RV's sched:nrp monitor
when running sched-messaging because of the following race:

  CPU0                                          CPU1
  ====                                          ====

  <Any thread running>
  ...                                          /* Wakes up a task on CPU0 */
  <IRQ>                                        rq_lock(rq0)
                                               __resched_curr(rq0)
                                                 set = test_and_set_ti_thread_flag(curr, 0, TIF_NEED_RESCHED);
                                                 <... interrupted by something before tracepoint hits>
                                                 !!! rq0->curr has NEED_RESCHED set but trace_sched_set_need_resched_tp() is not called !!!

  </IRQ>
  raw_irqentry_exit_cond_resched()
    if (!preempt_count() /* Masks PREEMPT_NEED_RESCHED */) {
      if (need_resched() /* Only check TIF flags */) {
        preempt_schedule_irq()
          schedule(SM_PREEMPT)
            trace_sched_entry_tp(true /* SM_PREEMPT */)

              !!! SPLAT: rv: monitor nrp does not allow event schedule_entry_preempt on state any_thread_running !!!


Easy way to solve this is by moving the trace_sched_entry_tp() within
the rq_lock critical section. SM_PREEMPT has to happen on a
!POLLING CPU and __resched_curr() is always under the rq_lock() so
the __schedule() on a remote CPU cannot race with it before
trace_sched_set_need_resched_tp() is executed.

  (lightly tested again)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index bd2f7fb87dc93..961c325feca84 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -7086,9 +7086,6 @@ static void __sched notrace __schedule(int sched_mode)
 	struct rq *rq;
 	int cpu;
 
-	/* Trace preemptions consistently with task switches */
-	trace_sched_entry_tp(sched_mode == SM_PREEMPT);
-
 	cpu = smp_processor_id();
 	rq = cpu_rq(cpu);
 	prev = rq->curr;
@@ -7121,6 +7118,9 @@ static void __sched notrace __schedule(int sched_mode)
 	rq_lock(rq, &rf);
 	smp_mb__after_spinlock();
 
+	/* Trace preemptions consistently with task switches */
+	trace_sched_entry_tp(sched_mode == SM_PREEMPT);
+
 	hrtick_schedule_enter(rq);
 
 	/* Promote REQ to ACT */
---

I'm not sure if this is a problem with current mainline but I did trip
it in my testing and the above seems to solve it - your mileage may vary
:-)

--  
Thanks and Regards,
Prateek