scx_p2dq: Add NUMA aware load balancing

Add NUMA aware pick 2 load balancing. The load balancer first tries to
maintain NUMA locality before attempting to pick two across NUMA nodes.
This reduces NUMA migrations on large systems.

Signed-off-by: Daniel Hodges <hodges.daniel.scott@gmail.com>

Author: hodgesds

scheds/rust/scx_p2dq/src/bpf/intf.h

@@ -93,6 +93,9 @@ enum stat_idx {
 	P2DQ_STAT_EAS_LITTLE_SELECT,
 	P2DQ_STAT_EAS_BIG_SELECT,
 	P2DQ_STAT_EAS_FALLBACK,
+	P2DQ_STAT_PICK2_SAME_NUMA,
+	P2DQ_STAT_PICK2_CROSS_NUMA,
+	P2DQ_STAT_PICK2_NUMA_FALLBACK,
 	P2DQ_NR_STATS,
 };
 

scheds/rust/scx_p2dq/src/bpf/main.bpf.c

@@ -96,6 +96,7 @@ const volatile struct {
 
 	bool dispatch_lb_interactive;
 	bool dispatch_pick2_disable;
+	bool dispatch_pick2_numa_aware;
 	bool eager_load_balance;
 	bool max_dsq_pick2;
 	bool wakeup_llc_migrations;
@@ -110,6 +111,7 @@ const volatile struct {
 
 	.dispatch_lb_interactive = false,
 	.dispatch_pick2_disable = false,
+	.dispatch_pick2_numa_aware = true,
 	.eager_load_balance = true,
 	.max_dsq_pick2 = false,
 	.wakeup_llc_migrations = false,
@@ -206,6 +208,12 @@ u32 llcs_per_node[MAX_NUMA_NODES];
 /* Global DHQ counter for unique indexing */
 u32 global_dhq_count = 0;
 
+/* NUMA-aware pick2 load balancing data structures */
+/* Maps NUMA node ID to list of LLC IDs in that node */
+u64 node_llc_ids[MAX_NUMA_NODES][MAX_LLCS];
+/* Count of LLCs per NUMA node (for random selection) */
+u32 node_llc_counts[MAX_NUMA_NODES];
+
 u64 min_slice_ns = 500;
 
 private(A) struct bpf_cpumask __kptr *all_cpumask;
@@ -963,6 +971,127 @@ static struct llc_ctx *rand_llc_ctx(void)
 	return lookup_llc_ctx(bpf_get_prandom_u32() % topo_config.nr_llcs);
 }
 
+/*
+ * Returns a random LLC context from the specified NUMA node.
+ * Returns NULL if node_id is invalid or has no LLCs.
+ */
+static struct llc_ctx *rand_llc_ctx_from_node(u32 node_id)
+{
+	u32 llc_count, llc_index;
+	u64 llc_id;
+
+	if (node_id >= MAX_NUMA_NODES)
+		return NULL;
+
+	llc_count = node_llc_counts[node_id];
+	if (llc_count == 0 || llc_count > MAX_LLCS)
+		return NULL;
+
+	/* Random selection within this node's LLCs */
+	llc_index = bpf_get_prandom_u32() % llc_count;
+
+	/* Verifier bounds check */
+	if (llc_index >= MAX_LLCS)
+		return NULL;
+
+	llc_id = node_llc_ids[node_id][llc_index];
+	return lookup_llc_ctx(llc_id);
+}
+
+/*
+ * Pick two different LLCs from the same NUMA node.
+ * Returns 0 on success, -EINVAL if can't pick 2 different LLCs.
+ */
+static __always_inline int pick_two_same_node(u32 node_id,
+					       struct llc_ctx **left,
+					       struct llc_ctx **right)
+{
+	u32 llc_count, attempts = 0;
+	u32 left_index, right_index;
+	u64 llc_id;
+
+	if (node_id >= MAX_NUMA_NODES)
+		return -EINVAL;
+
+	llc_count = node_llc_counts[node_id];
+
+	/* Need at least 2 LLCs in this node */
+	if (llc_count < 2)
+		return -EINVAL;
+
+	*left = rand_llc_ctx_from_node(node_id);
+	if (!*left)
+		return -EINVAL;
+
+	/* Try to pick a different LLC, with bounded loop for verifier */
+	#pragma unroll
+	for (attempts = 0; attempts < 8; attempts++) {
+		*right = rand_llc_ctx_from_node(node_id);
+		if (*right && (*right)->id != (*left)->id) {
+			stat_inc(P2DQ_STAT_PICK2_SAME_NUMA);
+			return 0;
+		}
+	}
+
+	/* Fallback: deterministically pick a different LLC */
+	left_index = (*left)->index;
+	right_index = (left_index + 1) % llc_count;
+
+	if (right_index >= MAX_LLCS)
+		return -EINVAL;
+
+	llc_id = node_llc_ids[node_id][right_index];
+	*right = lookup_llc_ctx(llc_id);
+
+	if (*right && (*right)->id != (*left)->id) {
+		stat_inc(P2DQ_STAT_PICK2_SAME_NUMA);
+		return 0;
+	}
+
+	return -EINVAL;
+}
+
+/*
+ * Pick two LLCs from different NUMA nodes for cross-node balancing.
+ */
+static __always_inline int pick_two_cross_node(u32 current_node_id,
+						struct llc_ctx **left,
+						struct llc_ctx **right)
+{
+	u32 other_node, attempts = 0;
+
+	/* Pick one from current node */
+	*left = rand_llc_ctx_from_node(current_node_id);
+	if (!*left)
+		return -EINVAL;
+
+	/* Try to pick from a different node */
+	#pragma unroll
+	for (attempts = 0; attempts < 8; attempts++) {
+		other_node = bpf_get_prandom_u32() % topo_config.nr_nodes;
+		if (other_node >= MAX_NUMA_NODES)
+			continue;
+
+		if (other_node != current_node_id &&
+		    node_llc_counts[other_node] > 0) {
+			*right = rand_llc_ctx_from_node(other_node);
+			if (*right && (*right)->id != (*left)->id) {
+				stat_inc(P2DQ_STAT_PICK2_CROSS_NUMA);
+				return 0;
+			}
+		}
+	}
+
+	/* Fallback to any random LLC */
+	*right = rand_llc_ctx();
+	if (*right && (*right)->id != (*left)->id) {
+		stat_inc(P2DQ_STAT_PICK2_CROSS_NUMA);
+		return 0;
+	}
+
+	return -EINVAL;
+}
+
 static bool keep_running(struct cpu_ctx *cpuc, struct llc_ctx *llcx,
 			 struct task_struct *p)
 {
@@ -2538,18 +2667,34 @@ static __always_inline int dispatch_pick_two(s32 cpu, struct llc_ctx *cur_llcx,
 	}
 
 	/*
-	 * For pick two load balancing we randomly choose two LLCs. We then
-	 * first try to consume from the LLC with the largest load. If we are
-	 * unable to consume from the first LLC then the second LLC is consumed
-	 * from. This yields better work conservation on machines with a large
-	 * number of LLCs.
+	 * NUMA-aware pick-2: prefer same-NUMA LLCs, then try cross-NUMA,
+	 * finally fallback to random selection.
 	 */
+	if (lb_config.dispatch_pick2_numa_aware && topo_config.nr_nodes > 1) {
+		u32 current_node_id = cpu_node_ids[cpu];
+		int ret;
+
+		if (node_llc_counts[current_node_id] >= 2) {
+			ret = pick_two_same_node(current_node_id, &left, &right);
+			if (ret == 0 && left && right)
+				goto llc_selected;
+		}
+
+		ret = pick_two_cross_node(current_node_id, &left, &right);
+		if (ret == 0 && left && right)
+			goto llc_selected;
+
+		stat_inc(P2DQ_STAT_PICK2_NUMA_FALLBACK);
+	}
+
+	/* Legacy random selection */
 	left = topo_config.nr_llcs == 2 ? lookup_llc_ctx(llc_ids[0]) : rand_llc_ctx();
 	right = topo_config.nr_llcs == 2 ? lookup_llc_ctx(llc_ids[1]) : rand_llc_ctx();
 
 	if (!left || !right)
 		return -EINVAL;
 
+	/* Handle collision (same LLC picked twice) */
 	if (left->id == right->id) {
 		i = llc_get_load(cur_llcx) % topo_config.nr_llcs;
 		i &= 0x3; // verifier
@@ -2559,6 +2704,8 @@ static __always_inline int dispatch_pick_two(s32 cpu, struct llc_ctx *cur_llcx,
 			return -EINVAL;
 	}
 
+llc_selected:
+
 
 	if (llc_get_load(right) > llc_get_load(left)) {
 		first = right;
@@ -3542,9 +3689,35 @@ static s32 p2dq_init_impl()
 
 	// First we initialize LLCs because DSQs are created at the LLC level.
 	bpf_for(i, 0, topo_config.nr_llcs) {
+		u32 llc_id, node_id, node_llc_idx;
+
 		ret = init_llc(i);
 		if (ret)
 			return ret;
+
+		/* Build NUMA-to-LLC mapping for NUMA-aware pick2 */
+		if (i >= MAX_LLCS)
+			continue;
+
+		llc_id = llc_ids[i];
+		llcx = lookup_llc_ctx(llc_id);
+		if (!llcx)
+			continue;
+
+		node_id = llcx->node_id;
+		if (node_id >= MAX_NUMA_NODES) {
+			scx_bpf_error("LLC %u has invalid node_id %u", llc_id, node_id);
+			return -EINVAL;
+		}
+
+		node_llc_idx = node_llc_counts[node_id];
+		if (node_llc_idx >= MAX_LLCS) {
+			scx_bpf_error("Node %u has too many LLCs", node_id);
+			return -EINVAL;
+		}
+
+		node_llc_ids[node_id][node_llc_idx] = llc_id;
+		node_llc_counts[node_id]++;
 	}
 
 	bpf_for(i, 0, topo_config.nr_nodes) {

scheds/rust/scx_p2dq/src/lib.rs

@@ -185,6 +185,10 @@ pub struct SchedulerOpts {
     #[clap(short = 'd', long, action = clap::ArgAction::SetTrue)]
     pub dispatch_pick2_disable: bool,
 
+    /// Disables NUMA-aware pick2 load balancing.
+    #[clap(long, action = clap::ArgAction::SetTrue)]
+    pub disable_pick2_numa_aware: bool,
+
     /// Enables pick2 load balancing on the dispatch path when LLC utilization is under the
     /// specified utilization.
     #[clap(long, default_value = "75", value_parser = clap::value_parser!(u64).range(0..100))]
@@ -432,6 +436,7 @@ macro_rules! init_open_skel {
             rodata.lb_config.max_dsq_pick2 = MaybeUninit::new(opts.max_dsq_pick2);
             rodata.lb_config.eager_load_balance = MaybeUninit::new(!opts.eager_load_balance);
             rodata.lb_config.dispatch_pick2_disable = MaybeUninit::new(opts.dispatch_pick2_disable);
+            rodata.lb_config.dispatch_pick2_numa_aware = MaybeUninit::new(!opts.disable_pick2_numa_aware);
             rodata.lb_config.dispatch_lb_busy = opts.dispatch_lb_busy;
             rodata.lb_config.dispatch_lb_interactive =
                 MaybeUninit::new(opts.dispatch_lb_interactive);

scheds/rust/scx_p2dq/src/main.rs

@@ -61,6 +61,9 @@ use bpf_intf::stat_idx_P2DQ_STAT_IDLE;
 use bpf_intf::stat_idx_P2DQ_STAT_KEEP;
 use bpf_intf::stat_idx_P2DQ_STAT_LLC_MIGRATION;
 use bpf_intf::stat_idx_P2DQ_STAT_NODE_MIGRATION;
+use bpf_intf::stat_idx_P2DQ_STAT_PICK2_CROSS_NUMA;
+use bpf_intf::stat_idx_P2DQ_STAT_PICK2_NUMA_FALLBACK;
+use bpf_intf::stat_idx_P2DQ_STAT_PICK2_SAME_NUMA;
 use bpf_intf::stat_idx_P2DQ_STAT_SELECT_PICK2;
 use bpf_intf::stat_idx_P2DQ_STAT_THERMAL_AVOID;
 use bpf_intf::stat_idx_P2DQ_STAT_THERMAL_KICK;
@@ -284,6 +287,9 @@ impl<'a> Scheduler<'a> {
             eas_little_select: stats[stat_idx_P2DQ_STAT_EAS_LITTLE_SELECT as usize],
             eas_big_select: stats[stat_idx_P2DQ_STAT_EAS_BIG_SELECT as usize],
             eas_fallback: stats[stat_idx_P2DQ_STAT_EAS_FALLBACK as usize],
+            pick2_same_numa: stats[stat_idx_P2DQ_STAT_PICK2_SAME_NUMA as usize],
+            pick2_cross_numa: stats[stat_idx_P2DQ_STAT_PICK2_CROSS_NUMA as usize],
+            pick2_numa_fallback: stats[stat_idx_P2DQ_STAT_PICK2_NUMA_FALLBACK as usize],
         }
     }
 

scheds/rust/scx_p2dq/src/stats.rs

@@ -91,6 +91,12 @@ pub struct Metrics {
     pub eas_big_select: u64,
     #[stat(desc = "Number of times EAS fell back to non-preferred core type")]
     pub eas_fallback: u64,
+    #[stat(desc = "Number of times pick2 selected LLCs from same NUMA node")]
+    pub pick2_same_numa: u64,
+    #[stat(desc = "Number of times pick2 selected LLCs across NUMA nodes")]
+    pub pick2_cross_numa: u64,
+    #[stat(desc = "Number of times pick2 NUMA-aware selection fell back to random")]
+    pub pick2_numa_fallback: u64,
 }
 
 impl Metrics {
@@ -111,38 +117,47 @@ impl Metrics {
             self.enq_mig,
         )?;
 
-        // Build the stats line conditionally based on thermal tracking availability
         let mut stats_line = format!(
-            "\twake prev/llc/mig {}/{}/{}\n\tpick2 select/dispatch {}/{}\n\tmigrations llc/node: {}/{}\n\tfork balance/same {}/{}\n\texec balance/same {}/{}",
-            self.wake_prev,
-            self.wake_llc,
-            self.wake_mig,
-            self.select_pick2,
-            self.dispatch_pick2,
-            self.llc_migrations,
-            self.node_migrations,
-            self.fork_balance,
-            self.fork_same_llc,
-            self.exec_balance,
-            self.exec_same_llc,
+            "\twake prev/llc/mig {}/{}/{}",
+            self.wake_prev, self.wake_llc, self.wake_mig,
         );
 
-        // Only show thermal stats if thermal tracking is enabled
+        if crate::TOPO.all_llcs.len() > 1 {
+            stats_line.push_str(&format!(
+                "\n\tpick2 select/dispatch {}/{}\n\tmigrations llc/node: {}/{}\n\tfork balance/same {}/{}\n\texec balance/same {}/{}",
+                self.select_pick2,
+                self.dispatch_pick2,
+                self.llc_migrations,
+                self.node_migrations,
+                self.fork_balance,
+                self.fork_same_llc,
+                self.exec_balance,
+                self.exec_same_llc,
+            ));
+        }
+
         if is_thermal_tracking_enabled() {
             stats_line.push_str(&format!(
                 "\n\tthermal kick/avoid {}/{}",
                 self.thermal_kick, self.thermal_avoid,
             ));
         }
 
-        // Only show EAS stats if energy-aware scheduling is enabled
         if is_eas_enabled() {
             stats_line.push_str(&format!(
                 "\n\tEAS little/big/fallback {}/{}/{}",
                 self.eas_little_select, self.eas_big_select, self.eas_fallback,
             ));
         }
 
+        let numa_total = self.pick2_same_numa + self.pick2_cross_numa + self.pick2_numa_fallback;
+        if numa_total > 0 {
+            stats_line.push_str(&format!(
+                "\n\tpick2 same_numa/cross_numa/fallback {}/{}/{}",
+                self.pick2_same_numa, self.pick2_cross_numa, self.pick2_numa_fallback,
+            ));
+        }
+
         writeln!(w, "{}", stats_line)?;
         Ok(())
     }
@@ -176,6 +191,9 @@ impl Metrics {
             eas_little_select: self.eas_little_select - rhs.eas_little_select,
             eas_big_select: self.eas_big_select - rhs.eas_big_select,
             eas_fallback: self.eas_fallback - rhs.eas_fallback,
+            pick2_same_numa: self.pick2_same_numa - rhs.pick2_same_numa,
+            pick2_cross_numa: self.pick2_cross_numa - rhs.pick2_cross_numa,
+            pick2_numa_fallback: self.pick2_numa_fallback - rhs.pick2_numa_fallback,
         }
     }
 }