[PATCH] sched_ext: Avoid relocking DSQ during remote consumption

[Usama Arif <usama.arif@linux.dev>]

A remote consume unlinks a task from a non-local DSQ before switching
to the task's rq. Leaving p->scx.dsq set makes the following
deactivate_task() reacquire the DSQ lock in scx_dispatch_dequeue() to
clear the transfer state.

Keep p->scx.dsq set while switching rq locks so that a racing dequeue
can find the DSQ and cancel the transfer. After acquiring the source rq
and confirming ownership, clear p->scx.dsq. The subsequent dequeue then
takes its !dsq path without reacquiring the DSQ lock.

Fleet profiles show that roughly 70% of the cycles attributed to
scx_dispatch_dequeue() are spent in queued_spin_lock_slowpath() while
taking dsq->lock.

Signed-off-by: Usama Arif <usama.arif@linux.dev>
---
 kernel/sched/ext/ext.c | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index cf8e32044ffb1..0135ff6824b03 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2612,7 +2612,7 @@ static bool task_can_run_on_remote_rq(struct scx_sched *sch,
  * points to this CPU. See scx_dispatch_dequeue() for the counterpart.
  *
  * On return, @dsq is unlocked and @src_rq is locked. Returns %true if @p is
- * still valid. %false if lost to dequeue.
+ * still valid and detached from @dsq. %false if lost to dequeue.
  */
 static bool unlink_dsq_and_switch_rq_lock(struct task_struct *p,
 					  struct scx_dispatch_q *dsq,
@@ -2632,8 +2632,19 @@ static bool unlink_dsq_and_switch_rq_lock(struct task_struct *p,
 	switch_rq_lock(locked_rq, src_rq);
 
 	/* task_rq couldn't have changed if we're still the holding cpu */
-	return likely(p->scx.holding_cpu == cpu) &&
-		!WARN_ON_ONCE(src_rq != task_rq(p));
+	if (likely(p->scx.holding_cpu == cpu) &&
+	    !WARN_ON_ONCE(src_rq != task_rq(p))) {
+		/*
+		 * Keep ->dsq set until we own @src_rq so that a racing dequeue
+		 * can find @dsq and clear holding_cpu. Once ownership is
+		 * confirmed, clear it under @src_rq so that deactivate_task()
+		 * takes the !dsq path instead of retaking @dsq->lock.
+		 */
+		p->scx.dsq = NULL;
+		return true;
+	}
+
+	return false;
 }
 
 static bool consume_remote_task(struct scx_sched *sch, struct rq *this_rq,
-- 
2.53.0-Meta