[PATCH v3] sched/core: Preserve wake flags across queued wakeups

[Shubhang Kaushik (Ampere)]

Queued wakeups currently save only whether the wakee migrated. When the
target CPU drains the wakelist, sched_ttwu_pending() therefore passes only
WF_MIGRATED or 0 to ttwu_do_activate().

That loses flags which are still used after CPU selection. For example,
WF_RQ_SELECTED is set after select_task_rq() chooses a runqueue, then used
by ttwu_do_activate() for ENQUEUE_RQ_SELECTED and by wakeup_preempt_fair()
for the preemption threshold:

direct wakeup:
       select_task_rq() -> WF_RQ_SELECTED
       ttwu_do_activate(WF_RQ_SELECTED)
       wakeup_preempt_fair(WF_RQ_SELECTED)

queued wakeup:
       select_task_rq() -> WF_RQ_SELECTED
       ttwu_queue_wakelist() -> save WF_MIGRATED only
       sched_ttwu_pending() -> ttwu_do_activate(WF_MIGRATED or 0)

Preserve the wake flags that still matter after queueing: WF_TTWU,
WF_SYNC, WF_MIGRATED and WF_RQ_SELECTED. Do not save WF_CURRENT_CPU,
which is only a CPU-selection hint.

Signed-off-by: Shubhang Kaushik (Ampere) <sh@xxxxxxxxxx>
---
Tested on an 80 CPU Ampere Altra system with perf bench sched messaging,
perf bench sched pipe, hackbench and SPECjBB. No material regression was
observed.

Baseline: mainline origin/master at v7.2-rc4 (248951ddc14d)
---
Changes in v3:
  - Use READ_ONCE()/WRITE_ONCE() for sched_remote_wakeup_flags.
  - Anchor the changelog on the WF_RQ_SELECTED preemption behavior.
  - Keep the cross-CPU ordering comment for the task field.

Link to v2: https://lore.kernel.org/r/20260713-b4-sched-ttwu-wake-flags-v2-1-76e23a8cc313@xxxxxxxxxx

Changes in v2:
  - Move sched_remote_wakeup_flags to a standalone u8 outside the scheduler
    bitfields.
  - Drop the unnecessary reset in sched_ttwu_pending().
  - Add WF_TTWU_QUEUE_MASK for the wake flags preserved across the wakelist.
  - Tighten the changelog around direct-vs-queued wakeup consistency.

Link to v1: https://lore.kernel.org/r/20260710-b4-sched-ttwu-wake-flags-v1-1-23182a4c5367@xxxxxxxxxx
---
 include/linux/sched.h | 26 ++++++++++----------------
 kernel/sched/core.c   | 10 ++++++++--
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 373bcc0598d10b4256a11f8c8373ece78fa5e0e5..c0de2ea02eeeb3ad9ca1a090260b00e1b35cbc39 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -983,6 +983,16 @@ struct task_struct {
 
 	/* Used for emulating ABI behavior of previous Linux versions: */
 	unsigned int			personality;
+	/*
+	 * Must not share the scheduler bitfield word because wakelist
+	 * queueing is not serialized by p->on_cpu.
+	 *
+	 * smp_load_acquire(&p->on_cpu) in ttwu_queue_wakelist() pairs
+	 * with smp_mb__after_spinlock() in schedule(), making current's
+	 * stores visible before the target CPU uses these flags from
+	 * sched_ttwu_pending().
+	 */
+	u8				sched_remote_wakeup_flags;
 
 	/* Scheduler bits, serialized by scheduler locks: */
 	unsigned			sched_reset_on_fork:1;
@@ -993,22 +1003,6 @@ struct task_struct {
 	/* Force alignment to the next boundary: */
 	unsigned			:0;
 
-	/* Unserialized, strictly 'current' */
-
-	/*
-	 * This field must not be in the scheduler word above due to wakelist
-	 * queueing no longer being serialized by p->on_cpu. However:
-	 *
-	 * p->XXX = X;			ttwu()
-	 * schedule()			  if (p->on_rq && ..) // false
-	 *   smp_mb__after_spinlock();	  if (smp_load_acquire(&p->on_cpu) && //true
-	 *   deactivate_task()		      ttwu_queue_wakelist())
-	 *     p->on_rq = 0;			p->sched_remote_wakeup = Y;
-	 *
-	 * guarantees all stores of 'current' are visible before
-	 * ->sched_remote_wakeup gets used, so it can be in this word.
-	 */
-	unsigned			sched_remote_wakeup:1;
 #ifdef CONFIG_RT_MUTEXES
 	unsigned			sched_rt_mutex:1;
 #endif
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 96226707c2f6135341aa779b8262f113e103d8ad..7ddcc6a7f1c06675ce8084cd304475d61f869eb7 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -3902,13 +3902,16 @@ void sched_ttwu_pending(void *arg)
 	update_rq_clock(rq);
 
 	llist_for_each_entry_safe(p, t, llist, wake_entry.llist) {
+		int wake_flags;
+
 		if (WARN_ON_ONCE(p->on_cpu))
 			smp_cond_load_acquire(&p->on_cpu, !VAL);
 
 		if (WARN_ON_ONCE(task_cpu(p) != cpu_of(rq)))
 			set_task_cpu(p, cpu_of(rq));
 
-		ttwu_do_activate(rq, p, p->sched_remote_wakeup ? WF_MIGRATED : 0, &rf);
+		wake_flags = READ_ONCE(p->sched_remote_wakeup_flags);
+		ttwu_do_activate(rq, p, wake_flags, &rf);
 	}
 
 	/*
@@ -3947,11 +3950,14 @@ bool call_function_single_prep_ipi(int cpu)
  * via sched_ttwu_wakeup() for activation so the wakee incurs the cost
  * of the wakeup instead of the waker.
  */
+#define WF_TTWU_QUEUE_MASK	(WF_TTWU | WF_SYNC | WF_MIGRATED | \
+				 WF_RQ_SELECTED)
 static void __ttwu_queue_wakelist(struct task_struct *p, int cpu, int wake_flags)
 {
 	struct rq *rq = cpu_rq(cpu);
 
-	p->sched_remote_wakeup = !!(wake_flags & WF_MIGRATED);
+	WRITE_ONCE(p->sched_remote_wakeup_flags,
+		   wake_flags & WF_TTWU_QUEUE_MASK);
 
 	WRITE_ONCE(rq->ttwu_pending, 1);
 #ifdef CONFIG_SMP

---
base-commit: 248951ddc14de84de3910f9b13f51491a8cd91df
change-id: 20260710-b4-sched-ttwu-wake-flags-e87ada0e65e1

Best regards,
-- 
Shubhang Kaushik (Ampere) <sh@xxxxxxxxxx>