[PATCH v6] sched/fair: Preserve wake-affine CPU for non-SMT reciprocal sync wakeups

[Shubhang Kaushik (Ampere)]

For WF_SYNC wakeups, wake_affine() may select or retain the waker CPU,
but select_idle_sibling() can subsequently move the wakee to another
idle CPU.

Pipe-style ping-pong workloads exercise the following handoff:

  A wakes B
  B wakes A
  A wakes B

When A wakes B for the second time, B last woke A. Check
B.last_wakee == A before record_wakee(B) records A.last_wakee == B.

Within the existing non-wide wake-affine path, on non-SMT systems, keep
this reciprocal WF_SYNC wakee on the waker CPU when wake-affine selected
or retained it and the waker has no other runnable task. This preserves
the direct handoff instead of searching for another idle CPU.

Record last_wakee history only in task context. In IRQ or softirq
context, current is the interrupted task rather than the logical waker.

This does not introduce a generic WF_SYNC placement rule. SMT systems
continue through select_idle_sibling(), which can choose an idle sibling
or idle core. On asymmetric-capacity systems, continue to require that
the wakee fits on the waker CPU.

Signed-off-by: Shubhang Kaushik (Ampere) <sh@xxxxxxxxxx>
---
Testing was performed on an 80-core non-SMT Ampere Altra system.

  perf bench sched pipe -l 1000000, 30 runs:

    tip/sched/core (4a3b51aab6e2 baseline):

    default:
      3.704 -> 2.267 usec/op mean,   about 38.8% improvement
      3.709 -> 2.264 usec/op median, about 39.0% improvement

    taskset -c 78,79:
      3.621 -> 2.255 usec/op mean,   about 37.7% improvement
      3.558 -> 2.261 usec/op median, about 36.5% improvement

Changes in v6:
  - Guard record_wakee() with in_task() so IRQ/softirq wakeups do not
    record the interrupted task as the waker.
  - Capture p->last_wakee == current before record_wakee() records p as
    current's last wakee.
  - Make the reciprocal handoff and non-SMT stacking rationale explicit
    in the code.

Link to v5: https://lore.kernel.org/r/20261008-b4-sched-sync-wakeup-v5-1-1af9fab48965@xxxxxxxxxx
---
 kernel/sched/fair.c | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 57360f5cdde4fa66f4cba2cbb53d4e322134504e..3d0d1a68ca84bebfed77a89c3c364a9562c75032 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -8513,6 +8513,10 @@ static unsigned long capacity_of(int cpu)
 
 static void record_wakee(struct task_struct *p)
 {
+	/* IRQ/softirq current is the interrupted task, not the logical waker. */
+	if (!in_task())
+		return;
+
 	/*
 	 * Only decay a single time; tasks that have less then 1 wakeup per
 	 * jiffy will not have built up many flips.
@@ -9167,6 +9171,22 @@ static inline bool asym_fits_cpu(unsigned long util,
 	return true;
 }
 
+static bool can_stack_on_waker(struct task_struct *p, int cpu)
+{
+	struct rq *rq = cpu_rq(cpu);
+
+	if ((rq->nr_running - cfs_h_nr_delayed(rq)) != 1)
+		return false;
+
+	if (sched_asym_cpucap_active()) {
+		sync_entity_load_avg(&p->se);
+		if (!task_fits_cpu(p, cpu))
+			return false;
+	}
+
+	return true;
+}
+
 /*
  * Try and locate an idle core/thread in the LLC cache domain.
  */
@@ -9920,6 +9940,7 @@ select_task_rq_fair(struct task_struct *p, int prev_cpu, int wake_flags)
 	struct sched_domain *tmp, *sd = NULL;
 	int cpu = smp_processor_id();
 	int new_cpu = prev_cpu;
+	bool reciprocal_wakeup = false;
 	int want_affine = 0;
 	/* SD_flags and WF_flags share the first nibble */
 	int sd_flag = wake_flags & 0xF;
@@ -9929,6 +9950,12 @@ select_task_rq_fair(struct task_struct *p, int prev_cpu, int wake_flags)
 	 */
 	lockdep_assert_held(&p->pi_lock);
 	if (wake_flags & WF_TTWU) {
+		/*
+		 * Check if current is p's last wakee before recording p as
+		 * current's last wakee.
+		 */
+		reciprocal_wakeup = in_task() &&
+			p->last_wakee == current;
 		record_wakee(p);
 
 		if ((wake_flags & WF_CURRENT_CPU) &&
@@ -9955,6 +9982,19 @@ select_task_rq_fair(struct task_struct *p, int prev_cpu, int wake_flags)
 			if (cpu != prev_cpu)
 				new_cpu = wake_affine(tmp, p, cpu, prev_cpu, sync);
 
+			/*
+			 * If wake-affine selected or kept the waker CPU, stack
+			 * a reciprocal WF_SYNC wakee there on non-SMT when it
+			 * has no other runnable task. On SMT, retain
+			 * select_idle_sibling() so it can consider idle
+			 * sibling/core placement.
+			 */
+			if (sync && !sched_smt_active() &&
+			    new_cpu == cpu &&
+			    reciprocal_wakeup &&
+			    can_stack_on_waker(p, cpu))
+				return cpu;
+
 			sd = NULL; /* Prefer wake_affine over balance flags */
 			break;
 		}

---
base-commit: 24a19b953e2a3d0ea5a8c4951a0ab8f32720c907
change-id: 20260721-b4-sched-sync-wakeup-04d40cbeb1da

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