scx_p2dq: Add backwards sticky scheduling

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

@@ -7,6 +7,9 @@
 
 #ifndef __KERNEL__
 typedef unsigned long long u64;
+#ifndef __kptr
+#define __kptr
+#endif
 #endif
 
 enum chaos_consts {
@@ -30,6 +33,7 @@ enum chaos_trait_kind {
 	CHAOS_TRAIT_RANDOM_DELAYS,
 	CHAOS_TRAIT_CPU_FREQ,
 	CHAOS_TRAIT_DEGRADATION,
+	CHAOS_TRAIT_BACKWARDS_STICKY,
 	CHAOS_TRAIT_MAX,
 };
 
@@ -40,6 +44,8 @@ struct chaos_task_ctx {
 	enum chaos_trait_kind	next_trait;
 	u64			enq_flags;
 	u64			p2dq_vtime;
+
+	struct bpf_cpumask __kptr	*bs_cpumask;
 };
 
 #endif /* __CHAOS_INTF_H */

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

@@ -50,6 +50,8 @@ const volatile u32 cpu_freq_max = SCX_CPUPERF_ONE;
 const volatile u32 degradation_freq_frac32 = 1;
 const volatile u64 degradation_frac7 = 0;
 
+const volatile u32 backwards_sticky_freq_frac32 = 1;
+
 #define MIN(x, y) ((x) < (y) ? (x) : (y))
 #define MAX(x, y) ((x) > (y) ? (x) : (y))
 
@@ -77,11 +79,21 @@ struct chaos_task_ctx *lookup_create_chaos_task_ctx(struct task_struct *p)
 	return bpf_task_storage_get(&chaos_task_ctxs, p, NULL, BPF_LOCAL_STORAGE_GET_F_CREATE);
 }
 
+static __always_inline bool keep_backwards_sticky(struct chaos_task_ctx *taskc)
+{
+	return taskc->next_trait == CHAOS_TRAIT_BACKWARDS_STICKY &&
+		taskc->bs_cpumask &&
+		bpf_cpumask_weight(cast_mask(taskc->bs_cpumask)) != nr_cpus;
+}
+
 static __always_inline enum chaos_trait_kind choose_chaos(struct chaos_task_ctx *taskc)
 {
 	if (taskc->match & CHAOS_MATCH_EXCLUDED)
 		return CHAOS_TRAIT_NONE;
 
+	if (keep_backwards_sticky(taskc))
+		return -1;
+
 	u32 roll = bpf_get_prandom_u32();
 
 	#pragma unroll
@@ -113,6 +125,24 @@ static __always_inline u64 get_cpu_delay_dsq(int cpu_idx)
 	return CHAOS_DSQ_BASE | cpu_idx;
 }
 
+static __always_inline s32 backwards_sticky_cpu(struct chaos_task_ctx *taskc)
+{
+	s32 cpu;
+	if (!taskc->bs_cpumask || bpf_cpumask_weight(cast_mask(taskc->bs_cpumask)) == nr_cpus)
+		return -ENOENT;
+
+	if (!taskc->bs_cpumask)
+		return -ENOENT;
+
+	cpu = bpf_cpumask_first_zero(cast_mask(taskc->bs_cpumask));
+	if (!taskc->bs_cpumask || cpu == nr_cpus)
+		return -ENOENT;
+	if (taskc->bs_cpumask)
+		bpf_cpumask_set_cpu(cpu, taskc->bs_cpumask);
+
+	return cpu;
+}
+
 static __always_inline s32 calculate_chaos_match(struct task_struct *p)
 {
 	struct chaos_task_ctx *taskc;
@@ -234,7 +264,7 @@ static __always_inline s32 calculate_chaos_match(struct task_struct *p)
 	return ret;
 }
 
-__weak s32 enqueue_random_delay(struct task_struct *p __arg_trusted, u64 enq_flags,
+static __always_inline s32 enqueue_random_delay(struct task_struct *p __arg_trusted, u64 enq_flags,
 				struct chaos_task_ctx *taskc __arg_nonnull)
 {
 	u64 rand64 = ((u64)bpf_get_prandom_u32() << 32) | bpf_get_prandom_u32();
@@ -249,7 +279,7 @@ __weak s32 enqueue_random_delay(struct task_struct *p __arg_trusted, u64 enq_fla
 	return true;
 }
 
-__weak s32 enqueue_chaotic(struct task_struct *p __arg_trusted, u64 enq_flags,
+static __always_inline s32 enqueue_chaotic(struct task_struct *p __arg_trusted, u64 enq_flags,
 			   struct chaos_task_ctx *taskc __arg_nonnull)
 {
 	bool out;
@@ -259,6 +289,9 @@ __weak s32 enqueue_chaotic(struct task_struct *p __arg_trusted, u64 enq_flags,
 		out = enqueue_random_delay(p, enq_flags, taskc);
 		break;
 
+	case CHAOS_TRAIT_BACKWARDS_STICKY:
+		out = false;
+		return out;
 	case CHAOS_TRAIT_NONE:
 	case CHAOS_TRAIT_CPU_FREQ:
 	case CHAOS_TRAIT_DEGRADATION:
@@ -474,6 +507,27 @@ void BPF_STRUCT_OPS(chaos_enqueue, struct task_struct *p __arg_trusted, u64 enq_
 		}
 	}
 
+	// Backwards sticky means the task will cycle through all CPUs before
+	// becoming sticky again. Do a direct dispatch to ensure the task lands
+	// on the correct CPU and trashes the caches.
+	if (taskc->next_trait == CHAOS_TRAIT_BACKWARDS_STICKY) {
+		s32 cpu = backwards_sticky_cpu(taskc);
+		if (cpu >= 0) {
+			u64 slice_ns;
+			if (promise.kind == P2DQ_ENQUEUE_PROMISE_VTIME)
+				slice_ns = promise.vtime.slice_ns;
+			if (promise.kind == P2DQ_ENQUEUE_PROMISE_FIFO)
+				slice_ns = promise.fifo.slice_ns;
+
+			scx_bpf_dsq_insert(p,
+					   SCX_DSQ_LOCAL_ON | cpu,
+					   slice_ns,
+					   enq_flags);
+			promise.kind = P2DQ_ENQUEUE_PROMISE_COMPLETE;
+			return;
+		}
+	}
+
 	complete_p2dq_enqueue(&promise, p);
 }
 
@@ -532,6 +586,30 @@ s32 BPF_STRUCT_OPS(chaos_select_cpu, struct task_struct *p, s32 prev_cpu, u64 wa
 s32 BPF_STRUCT_OPS_SLEEPABLE(chaos_init_task, struct task_struct *p,
 			     struct scx_init_task_args *args)
 {
+	struct bpf_cpumask *bs_cpumask;
+	struct chaos_task_ctx *taskc;
+
+	if (!(bs_cpumask = bpf_cpumask_create())) {
+		scx_bpf_error("chaos_task_ctx bs_cpumask allocation failure");
+		return -ENOMEM;
+	}
+
+	if (!(taskc = lookup_create_chaos_task_ctx(p))) {
+		bpf_cpumask_release(bs_cpumask);
+		scx_bpf_error("failed to lookup task context in init_task");
+		return -ENOMEM;
+	}
+
+	if ((bs_cpumask = bpf_kptr_xchg(&taskc->bs_cpumask, bs_cpumask))) {
+		bpf_cpumask_release(bs_cpumask);
+		scx_bpf_error("chaos_task_ctx bs_cpumask allocation failure");
+		return -EINVAL;
+	}
+	bpf_rcu_read_lock();
+	if (taskc->bs_cpumask)
+		bpf_cpumask_clear(taskc->bs_cpumask);
+	bpf_rcu_read_unlock();
+
 	s32 ret = p2dq_init_task_impl(p, args);
 	if (ret)
 		return ret;

scheds/rust/scx_chaos/src/lib.rs

@@ -95,6 +95,9 @@ pub enum Trait {
         frequency: f64,
         degradation_frac7: u64,
     },
+    BackwardsSticky {
+        frequency: f64,
+    },
 }
 
 impl Trait {
@@ -103,6 +106,7 @@ impl Trait {
             Self::RandomDelays { .. } => bpf_intf::chaos_trait_kind_CHAOS_TRAIT_RANDOM_DELAYS,
             Self::CpuFreq { .. } => bpf_intf::chaos_trait_kind_CHAOS_TRAIT_CPU_FREQ,
             Self::PerfDegradation { .. } => bpf_intf::chaos_trait_kind_CHAOS_TRAIT_DEGRADATION,
+            Self::BackwardsSticky { .. } => bpf_intf::chaos_trait_kind_CHAOS_TRAIT_BACKWARDS_STICKY,
         }
     }
 
@@ -111,6 +115,7 @@ impl Trait {
             Self::RandomDelays { frequency, .. } => *frequency,
             Self::CpuFreq { frequency, .. } => *frequency,
             Self::PerfDegradation { frequency, .. } => *frequency,
+            Self::BackwardsSticky { frequency, .. } => *frequency,
         }
     }
 }
@@ -409,6 +414,10 @@ impl Builder<'_> {
                         (frequency * 2_f64.powf(32_f64)) as u32;
                     open_skel.maps.rodata_data.degradation_frac7 = *degradation_frac7;
                 }
+                Trait::BackwardsSticky { frequency } => {
+                    open_skel.maps.rodata_data.backwards_sticky_freq_frac32 =
+                        (frequency * 2_f64.powf(32_f64)) as u32;
+                }
             }
         }
 
@@ -499,6 +508,15 @@ pub struct PerfDegradationArgs {
     pub degradation_frac7: u64,
 }
 
+/// Backwards Sticky scheduling is where a task is scheduled on all CPUs iteratively to break
+/// caches and then returned to sticky scheduling.
+#[derive(Debug, Parser)]
+pub struct BackwardsStickyArgs {
+    /// Chance of applying backward sticky to a process.
+    #[clap(long)]
+    pub bs_frequency: Option<f64>,
+}
+
 /// scx_chaos: A general purpose sched_ext scheduler designed to amplify race conditions
 ///
 /// WARNING: This scheduler is a very early alpha, and hasn't been production tested yet. The CLI
@@ -540,6 +558,9 @@ pub struct Args {
     #[command(flatten, next_help_heading = "Perf Degradation")]
     pub perf_degradation: PerfDegradationArgs,
 
+    #[command(flatten, next_help_heading = "Backward Sticky")]
+    pub backward_sticky: BackwardsStickyArgs,
+
     #[command(flatten, next_help_heading = "CPU Frequency")]
     pub cpu_freq: CpuFreqArgs,
 

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

@@ -60,6 +60,7 @@ const volatile u32 interactive_ratio = 10;
 const volatile u32 min_nr_queued_pick2 = 10;
 
 const volatile bool autoslice = true;
+const volatile bool backwards_sticky = false;
 const volatile bool dispatch_pick2_disable = false;
 const volatile bool eager_load_balance = true;
 const volatile bool interactive_sticky = false;
@@ -434,6 +435,19 @@ static s32 pick_idle_cpu(struct task_struct *p, task_ctx *taskc,
 		goto found_cpu;
 	}
 
+	if (backwards_sticky &&
+	    ((taskc->bs_mask & nr_cpus) != nr_cpus)) {
+		s32 i;
+		bpf_for(i, 0, nr_cpus)
+			if (taskc->bs_mask == 0 || (taskc->bs_mask & i) != i) {
+				if (i == nr_cpus - 1)
+					taskc->bs_mask |= nr_cpus;
+				cpu = i;
+				*is_idle = true;
+				goto found_cpu;
+			}
+	}
+
 	// First check if last CPU is idle
 	if (taskc->all_cpus &&
 	    bpf_cpumask_test_cpu(prev_cpu, (smt_enabled && !interactive) ?
@@ -770,6 +784,8 @@ static __always_inline int p2dq_running_impl(struct task_struct *p)
 	if (taskc->node_id != cpuc->node_id) {
 		stat_inc(P2DQ_STAT_NODE_MIGRATION);
 	}
+	if (backwards_sticky)
+		taskc->bs_mask |= task_cpu;
 
 	taskc->llc_id = llcx->id;
 	taskc->node_id = llcx->node_id;
@@ -1123,6 +1139,7 @@ static __always_inline s32 p2dq_init_task_impl(struct task_struct *p,
 		return -EINVAL;
 	}
 
+	taskc->bs_mask = 0;
 	taskc->llc_id = cpuc->llc_id;
 	taskc->node_id = cpuc->node_id;
 	taskc->dsq_index = init_dsq_index;

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

@@ -58,6 +58,7 @@ struct task_p2dq {
 	u64 			last_run_at;
 	u64			llc_runs; /* how many runs on the current LLC */
 	int			last_dsq_index;
+	s32			bs_mask;
 
 	/* The task is a workqueue worker thread */
 	bool			is_kworker;

scheds/rust/scx_p2dq/src/lib.rs

@@ -45,6 +45,10 @@ pub struct SchedulerOpts {
     #[clap(short = 'e', long, help="DEPRECATED", action = clap::ArgAction::SetTrue)]
     pub eager_load_balance: bool,
 
+    /// Enables backward sticky scheduling.
+    #[clap(long= "bs", action = clap::ArgAction::SetTrue)]
+    pub backwards_sticky: bool,
+
     /// Enables CPU frequency control.
     #[clap(short = 'f', long, action = clap::ArgAction::SetTrue)]
     pub freq_control: bool,
@@ -207,6 +211,7 @@ macro_rules! init_open_skel {
             $skel.maps.rodata_data.lb_slack_factor = opts.lb_slack_factor;
 
             $skel.maps.rodata_data.autoslice = opts.autoslice;
+            $skel.maps.rodata_data.backwards_sticky = opts.backwards_sticky;
             $skel.maps.rodata_data.debug = verbose as u32;
             $skel.maps.rodata_data.dispatch_pick2_disable = opts.dispatch_pick2_disable;
             $skel.maps.rodata_data.dispatch_lb_busy = opts.dispatch_lb_busy;