Merge pull request #2939 from sched-ext/bpfland-pref-idle-scan

scx_bpfland: Add --preferred-idle-scan option to prioritize faster CPUs

Author: arighi

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

@@ -6,6 +6,12 @@
 #include <scx/percpu.bpf.h>
 #include "intf.h"
 
+/*
+ * Maximum amount of CPUs supported by the scheduler when flat or preferred
+ * idle CPU scan is enabled.
+ */
+#define MAX_CPUS	1024
+
 /*
  * Maximum rate of task wakeups/sec (tasks with a higher rate are capped to
  * this value).
@@ -75,6 +81,21 @@ const volatile bool local_pcpu = true;
  */
 volatile s64 cpufreq_perf_lvl;
 
+/*
+ * Enable preferred cores prioritization.
+ */
+const volatile bool preferred_idle_scan;
+
+/*
+ * CPUs sorted by their capacity in descendent order.
+ */
+const volatile u64 preferred_cpus[MAX_CPUS];
+
+/*
+ * Cache CPU capacity values.
+ */
+const volatile u64 cpu_capacity[MAX_CPUS];
+
 /*
  * Scheduling statistics.
  */
@@ -252,6 +273,33 @@ static inline const struct cpumask *get_idle_cpumask(s32 cpu)
 	return __COMPAT_scx_bpf_get_idle_cpumask_node(__COMPAT_scx_bpf_cpu_node(cpu));
 }
 
+/*
+ * Return the cpumask of fully idle SMT cores within the NUMA node that
+ * contains @cpu.
+ *
+ * If NUMA support is disabled, @cpu is ignored.
+ */
+static inline const struct cpumask *get_idle_smtmask(s32 cpu)
+{
+	if (!numa_enabled)
+		return scx_bpf_get_idle_smtmask();
+
+	return __COMPAT_scx_bpf_get_idle_smtmask_node(__COMPAT_scx_bpf_cpu_node(cpu));
+}
+
+/*
+ * Return true if @cpu is valid, otherwise trigger an error and return
+ * false.
+ */
+static inline bool is_cpu_valid(s32 cpu)
+{
+	if (cpu < 0 || cpu >= MAX_CPUS) {
+		scx_bpf_error("invalid CPU id: %d", cpu);
+		return false;
+	}
+	return true;
+}
+
 /*
  * Return true if @this_cpu and @that_cpu are in the same LLC, false
  * otherwise.
@@ -261,6 +309,9 @@ static inline bool cpus_share_cache(s32 this_cpu, s32 that_cpu)
         if (this_cpu == that_cpu)
                 return true;
 
+	if (!is_cpu_valid(this_cpu) || !is_cpu_valid(that_cpu))
+		return false;
+
 	return cpu_llc_id(this_cpu) == cpu_llc_id(that_cpu);
 }
 
@@ -272,7 +323,10 @@ static inline bool is_cpu_faster(s32 this_cpu, s32 that_cpu)
         if (this_cpu == that_cpu)
                 return false;
 
-	return cpu_priority(this_cpu) > cpu_priority(that_cpu);
+	if (!is_cpu_valid(this_cpu) || !is_cpu_valid(that_cpu))
+		return false;
+
+	return cpu_capacity[this_cpu] > cpu_capacity[that_cpu];
 }
 
 /*
@@ -329,6 +383,102 @@ static inline bool is_wakeup(u64 wake_flags)
 	return wake_flags & SCX_WAKE_TTWU;
 }
 
+/*
+ * Try to pick the best idle CPU based on the @preferred_cpus ranking.
+ * Return a full-idle SMT core if @do_idle_smt is true, or any idle CPU if
+ * @do_idle_smt is false.
+ */
+static s32 pick_idle_cpu_pref_smt(struct task_struct *p, s32 prev_cpu, bool is_prev_allowed,
+				  const struct cpumask *primary, const struct cpumask *smt)
+{
+	u64 max_cpus = MIN(nr_cpu_ids, MAX_CPUS);
+	int i;
+
+	if (is_prev_allowed &&
+	    (!primary || bpf_cpumask_test_cpu(prev_cpu, primary)) &&
+	    (!smt || bpf_cpumask_test_cpu(prev_cpu, smt)) &&
+	    scx_bpf_test_and_clear_cpu_idle(prev_cpu))
+		return prev_cpu;
+
+	bpf_for(i, 0, max_cpus) {
+		s32 cpu = preferred_cpus[i];
+
+		if ((cpu == prev_cpu) || !bpf_cpumask_test_cpu(cpu, p->cpus_ptr))
+			continue;
+
+		if ((!primary || bpf_cpumask_test_cpu(cpu, primary)) &&
+		    (!smt || bpf_cpumask_test_cpu(cpu, smt)) &&
+		    scx_bpf_test_and_clear_cpu_idle(cpu))
+			return cpu;
+	}
+
+	return -EBUSY;
+}
+
+/*
+ * Return the optimal idle CPU for task @p or -EBUSY if no idle CPU is
+ * found.
+ */
+static s32 pick_idle_cpu_scan(struct task_struct *p, s32 prev_cpu)
+{
+	const struct cpumask *smt, *primary;
+	bool is_prev_allowed = bpf_cpumask_test_cpu(prev_cpu, p->cpus_ptr);
+	s32 cpu;
+
+	primary = !primary_all ? cast_mask(primary_cpumask) : NULL;
+	smt = smt_enabled ? get_idle_smtmask(prev_cpu) : NULL;
+
+	/*
+	 * If the task can't migrate, there's no point looking for other
+	 * CPUs.
+	 */
+	if (p->nr_cpus_allowed == 1 || is_migration_disabled(p)) {
+		if (scx_bpf_test_and_clear_cpu_idle(prev_cpu)) {
+			cpu = prev_cpu;
+			goto out;
+		}
+	}
+
+	if (!primary_all) {
+		if (smt_enabled) {
+			/*
+			 * Try to pick a full-idle core in the primary
+			 * domain.
+			 */
+			cpu = pick_idle_cpu_pref_smt(p, prev_cpu, is_prev_allowed, primary, smt);
+			if (cpu >= 0)
+				goto out;
+		}
+
+		/*
+		 * Try to pick any idle CPU in the primary domain.
+		 */
+		cpu = pick_idle_cpu_pref_smt(p, prev_cpu, is_prev_allowed, primary, NULL);
+		if (cpu >= 0)
+			goto out;
+	}
+
+	if (smt_enabled) {
+		/*
+		 * Try to pick any full-idle core in the system.
+		 */
+		cpu = pick_idle_cpu_pref_smt(p, prev_cpu, is_prev_allowed, NULL, smt);
+		if (cpu >= 0)
+			goto out;
+	}
+
+	/*
+	 * Try to pick any idle CPU in the system.
+	 */
+	cpu = pick_idle_cpu_pref_smt(p, prev_cpu, is_prev_allowed, NULL, NULL);
+
+out:
+	if (smt)
+		scx_bpf_put_cpumask(smt);
+
+	return cpu;
+}
+
 /*
  * Pick an optimal idle CPU for task @p (as close as possible to
  * @prev_cpu).
@@ -341,6 +491,14 @@ static s32 pick_idle_cpu(struct task_struct *p, s32 prev_cpu, s32 this_cpu,
 	const struct cpumask *primary = cast_mask(primary_cpumask);
 	s32 cpu;
 
+	/*
+	 * Use lightweight idle CPU scanning when flat or preferred idle
+	 * scan is enabled, unless the system is busy, in which case the
+	 * cpumask-based scanning is more efficient.
+	 */
+	if (preferred_idle_scan)
+		return pick_idle_cpu_scan(p, prev_cpu);
+
 	/*
 	 * Clear the wake sync bit if synchronous wakeups are disabled.
 	 */

scheds/rust/scx_bpfland/src/main.rs

@@ -11,7 +11,7 @@ pub mod bpf_intf;
 pub use bpf_intf::*;
 
 mod stats;
-use std::ffi::c_int;
+use std::ffi::{c_int, c_ulong};
 use std::fmt::Write;
 use std::mem::MaybeUninit;
 use std::sync::atomic::AtomicBool;
@@ -197,6 +197,13 @@ struct Opts {
     #[clap(short = 'm', long, default_value = "auto")]
     primary_domain: String,
 
+    /// Enable preferred idle CPU scanning.
+    ///
+    /// With this option enabled, the scheduler will prioritize assigning tasks to higher-ranked
+    /// cores before considering lower-ranked ones.
+    #[clap(short = 'P', long, action = clap::ArgAction::SetTrue)]
+    preferred_idle_scan: bool,
+
     /// Disable SMT awareness.
     #[clap(long, action = clap::ArgAction::SetTrue)]
     disable_smt: bool,
@@ -333,6 +340,21 @@ impl<'a> Scheduler<'a> {
         rodata.throttle_ns = opts.throttle_us * 1000;
         rodata.primary_all = domain.weight() == *NR_CPU_IDS;
 
+        // Generate the list of available CPUs sorted by capacity in descending order.
+        let mut cpus: Vec<_> = topo.all_cpus.values().collect();
+        cpus.sort_by_key(|cpu| std::cmp::Reverse(cpu.cpu_capacity));
+        for (i, cpu) in cpus.iter().enumerate() {
+            rodata.cpu_capacity[cpu.id] = cpu.cpu_capacity as c_ulong;
+            rodata.preferred_cpus[i] = cpu.id as u64;
+        }
+        if opts.preferred_idle_scan {
+            info!(
+                "Preferred CPUs: {:?}",
+                &rodata.preferred_cpus[0..cpus.len()]
+            );
+        }
+        rodata.preferred_idle_scan = opts.preferred_idle_scan;
+
         // Implicitly enable direct dispatch of per-CPU kthreads if CPU throttling is enabled
         // (it's never a good idea to throttle per-CPU kthreads).
         rodata.local_kthreads = opts.local_kthreads || opts.throttle_us > 0;