[scx_lavd] modify lavd to use new peek operation

Summary:

The per-cpu queues mode in LAVD checks multiple queues for potential steal targets.
These checks can be made more efficient with the forthcoming O(1) lockless peek operation.
This patch updates LAVD to use peek in three places:

- try_to_steal: picking cpu with lowest vtime
- on consume_task: checking the cpu DSQ
- on consume_task: checking the domain DSQ

Note that these usages are necessarily racy because no lock is held. That is, we
could peek a low vtime, but the task may have moved or completed by the time we
attempt to steal it.

Test Plan:
 - ran schbench under a VM [1 cpdom, 50cpus]

Reviewers: daidavid

Author: rrnewton

scheds/include/scx/common.bpf.h

@@ -75,6 +75,7 @@ u32 scx_bpf_reenqueue_local(void) __ksym;
 void scx_bpf_kick_cpu(s32 cpu, u64 flags) __ksym;
 s32 scx_bpf_dsq_nr_queued(u64 dsq_id) __ksym;
 void scx_bpf_destroy_dsq(u64 dsq_id) __ksym;
+struct task_struct *scx_bpf_dsq_peek(u64 dsq_id) __ksym __weak;
 int bpf_iter_scx_dsq_new(struct bpf_iter_scx_dsq *it, u64 dsq_id, u64 flags) __ksym __weak;
 struct task_struct *bpf_iter_scx_dsq_next(struct bpf_iter_scx_dsq *it) __ksym __weak;
 void bpf_iter_scx_dsq_destroy(struct bpf_iter_scx_dsq *it) __ksym __weak;

scheds/include/scx/compat.bpf.h

@@ -230,6 +230,13 @@ static inline bool __COMPAT_is_enq_cpu_selected(u64 enq_flags)
 	 scx_bpf_pick_any_cpu_node(cpus_allowed, node, flags) :			\
 	 scx_bpf_pick_any_cpu(cpus_allowed, flags))
 
+#define __COMPAT_scx_bpf_dsq_peek(dsq_id)                                      \
+  (bpf_ksym_exists(scx_bpf_dsq_peek) ? scx_bpf_dsq_peek(dsq_id) : ({           \
+    struct task_struct *p = NULL;                                              \
+    bpf_for_each(scx_dsq, p, dsq_id, 0) { break; }                             \
+    p;                                                                         \
+  }))
+
 /*
  * Define sched_ext_ops. This may be expanded to define multiple variants for
  * backward compatibility. See compat.h::SCX_OPS_LOAD/ATTACH().

scheds/rust/scx_lavd/src/bpf/balance.bpf.c

@@ -100,7 +100,7 @@ int plan_x_cpdom_migration(void)
 		return 0;
 	}
 	if ((stealee_threshold <= max_sc_load || overflow_running) &&
-	    (stealer_threshold < min_sc_load)) {
+		(stealer_threshold < min_sc_load)) {
 		/*
 		 * If there is a overloaded domain, always try to steal.
 		 *  <~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~>
@@ -123,7 +123,7 @@ int plan_x_cpdom_migration(void)
 		 * Under-loaded active domains become a stealer.
 		 */
 		if (cpdomc->nr_active_cpus &&
-		    cpdomc->sc_load <= stealer_threshold) {
+			cpdomc->sc_load <= stealer_threshold) {
 			WRITE_ONCE(cpdomc->is_stealer, true);
 			WRITE_ONCE(cpdomc->is_stealee, false);
 			continue;
@@ -133,7 +133,7 @@ int plan_x_cpdom_migration(void)
 		 * Over-loaded or non-active domains become a stealee.
 		 */
 		if (!cpdomc->nr_active_cpus ||
-		    cpdomc->sc_load >= stealee_threshold) {
+			cpdomc->sc_load >= stealee_threshold) {
 			WRITE_ONCE(cpdomc->is_stealer, false);
 			WRITE_ONCE(cpdomc->is_stealee, true);
 			nr_stealee++;
@@ -173,6 +173,48 @@ static bool consume_dsq(struct cpdom_ctx *cpdomc, u64 dsq_id)
 	return ret;
 }
 
+/*
+ * Attempts to peek the vtime of the task at the head of the DSQ, or returns U64_MAX if the DSQ is empty.
+ */
+static int peek_vtime(u64 dsq_id)
+{
+	struct task_struct *task;
+	task = __COMPAT_scx_bpf_dsq_peek(dsq_id);
+	return task ? task->scx.dsq_vtime : U64_MAX;
+}
+
+/*
+ * Racy operation that returns the cpu that which appears to have the lowest vtime at its head.
+ */
+static int pick_cpu_with_lowest_vtime(struct cpdom_ctx *cpdomc)
+{
+	u64 lowest_vtime = U64_MAX;
+	int pick_cpu = -ENOENT;
+	int cpu, i, j;
+
+	if (!per_cpu_dsq)
+		return -ENOENT;
+
+	bpf_for(i, 0, LAVD_CPU_ID_MAX/64) {
+		u64 cpumask = cpdomc->__cpumask[i];
+		bpf_for(j, 0, 64) {
+			if (cpumask & 0x1LLU << j) {
+				u64 vtime;
+				cpu = (i * 64) + j;
+				if (cpu >= __nr_cpu_ids)
+					break;
+				vtime = peek_vtime(cpu_to_dsq(cpu));
+				if (vtime < lowest_vtime) {
+					lowest_vtime = vtime;
+					pick_cpu = cpu;
+				}
+			}
+		}
+	}
+
+	return pick_cpu;
+}
+
 /*
  * For simplicity, try to just steal from the CPU with
  * the highest number of queued_tasks in this domain.
@@ -256,7 +298,7 @@ static bool try_to_steal_task(struct cpdom_ctx *cpdomc)
 			if (!READ_ONCE(cpdomc_pick->is_stealee) || !cpdomc_pick->is_valid)
 				continue;
 
-			pick_cpu = pick_most_loaded_cpu(cpdomc_pick);
+			pick_cpu = pick_cpu_with_lowest_vtime(cpdomc_pick);
 			if (pick_cpu >= 0)
 				dsq_id = cpu_to_dsq(pick_cpu);
 			else
@@ -333,7 +375,7 @@ static bool force_to_steal_task(struct cpdom_ctx *cpdomc)
 			if (!cpdomc_pick->is_valid)
 				continue;
 
-			pick_cpu = pick_most_loaded_cpu(cpdomc_pick);
+			pick_cpu = pick_cpu_with_lowest_vtime(cpdomc_pick);
 			if (pick_cpu >= 0)
 				dsq_id = cpu_to_dsq(pick_cpu);
 			else
@@ -350,7 +392,6 @@ static bool force_to_steal_task(struct cpdom_ctx *cpdomc)
 static bool consume_task(u64 cpu_dsq_id, u64 cpdom_dsq_id)
 {
 	struct cpdom_ctx *cpdomc;
-	struct task_struct *p;
 	u64 vtime = U64_MAX, dsq_id = cpu_dsq_id;
 
 	cpdomc = MEMBER_VPTR(cpdom_ctxs, [dsq_to_cpdom(cpdom_dsq_id)]);
@@ -364,20 +405,13 @@ static bool consume_task(u64 cpu_dsq_id, u64 cpdom_dsq_id)
 	 * a task from any of stealee domains probabilistically.
 	 */
 	if (nr_cpdoms > 1 && READ_ONCE(cpdomc->is_stealer) &&
-	    try_to_steal_task(cpdomc))
+		try_to_steal_task(cpdomc))
 		goto x_domain_migration_out;
 
 	if (per_cpu_dsq) {
-		bpf_for_each(scx_dsq, p, cpu_dsq_id, 0) {
-			vtime = p->scx.dsq_vtime;
-			break;
-		}
-
-		bpf_for_each(scx_dsq, p, cpdom_dsq_id, 0) {
-			if (p->scx.dsq_vtime < vtime)
-				dsq_id = cpdom_dsq_id;
-			break;
-		}
+		vtime = peek_vtime(cpu_dsq_id);
+		if (peek_vtime(cpdom_dsq_id) < vtime)
+			dsq_id = cpdom_dsq_id;
 	} else {
 		dsq_id = cpdom_dsq_id;
 	}