1 // SPDX-License-Identifier: GPL-2.0 2 /* 3 * A scheduler that validates the behavior of scx_bpf_select_cpu_and() by 4 * selecting idle CPUs strictly within a subset of allowed CPUs. 5 * 6 * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com> 7 */ 8 9 #include <scx/common.bpf.h> 10 11 char _license[] SEC("license") = "GPL"; 12 13 UEI_DEFINE(uei); 14 15 private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask; 16 17 static void 18 validate_local_idle_state(void) 19 { 20 const struct cpumask *idle; 21 struct task_struct *curr; 22 s32 cpu = bpf_get_smp_processor_id(); 23 bool cpu_is_idle, curr_is_idle; 24 25 bpf_rcu_read_lock(); 26 curr = scx_bpf_cpu_curr(cpu); 27 curr_is_idle = curr && (curr->flags & PF_IDLE); 28 bpf_rcu_read_unlock(); 29 30 idle = scx_bpf_get_idle_cpumask(); 31 cpu_is_idle = bpf_cpumask_test_cpu(cpu, idle); 32 scx_bpf_put_idle_cpumask(idle); 33 34 /* 35 * Unlike a remote selected CPU, the local CPU cannot go through an 36 * idle re-pick while this callback is running. If it is running a 37 * non-idle scheduling context, it must not be advertised as idle. 38 */ 39 if (!curr_is_idle && cpu_is_idle) 40 scx_bpf_error("running CPU %d should be marked as busy", cpu); 41 } 42 43 static void 44 validate_selected_cpu(const struct task_struct *p, s32 cpu) 45 { 46 const struct cpumask *allowed = cast_mask(allowed_cpumask); 47 48 if (!allowed) { 49 scx_bpf_error("allowed domain not initialized"); 50 return; 51 } 52 53 if (!bpf_cpumask_test_cpu(cpu, allowed)) 54 scx_bpf_error("CPU %d not in the allowed domain for %d (%s)", 55 cpu, p->pid, p->comm); 56 57 if (!bpf_cpumask_test_cpu(cpu, p->cpus_ptr)) 58 scx_bpf_error("CPU %d not in the affinity mask for %d (%s)", 59 cpu, p->pid, p->comm); 60 } 61 62 s32 BPF_STRUCT_OPS(allowed_cpus_select_cpu, 63 struct task_struct *p, s32 prev_cpu, u64 wake_flags) 64 { 65 const struct cpumask *allowed; 66 s32 cpu; 67 68 validate_local_idle_state(); 69 allowed = cast_mask(allowed_cpumask); 70 if (!allowed) { 71 scx_bpf_error("allowed domain not initialized"); 72 return -EINVAL; 73 } 74 75 /* 76 * Select an idle CPU strictly within the allowed domain. 77 */ 78 cpu = scx_bpf_select_cpu_and(p, prev_cpu, wake_flags, allowed, 0); 79 if (cpu >= 0) { 80 validate_selected_cpu(p, cpu); 81 scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0); 82 83 return cpu; 84 } 85 86 return prev_cpu; 87 } 88 89 void BPF_STRUCT_OPS(allowed_cpus_enqueue, struct task_struct *p, u64 enq_flags) 90 { 91 const struct cpumask *allowed; 92 s32 prev_cpu = scx_bpf_task_cpu(p), cpu; 93 94 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0); 95 96 validate_local_idle_state(); 97 allowed = cast_mask(allowed_cpumask); 98 if (!allowed) { 99 scx_bpf_error("allowed domain not initialized"); 100 return; 101 } 102 103 /* 104 * Use scx_bpf_select_cpu_and() to proactively kick an idle CPU 105 * within @allowed_cpumask, usable by @p. 106 */ 107 cpu = scx_bpf_select_cpu_and(p, prev_cpu, 0, allowed, 0); 108 if (cpu >= 0) { 109 validate_selected_cpu(p, cpu); 110 scx_bpf_kick_cpu(cpu, SCX_KICK_IDLE); 111 } 112 } 113 114 s32 BPF_STRUCT_OPS_SLEEPABLE(allowed_cpus_init) 115 { 116 struct bpf_cpumask *mask; 117 118 mask = bpf_cpumask_create(); 119 if (!mask) 120 return -ENOMEM; 121 122 mask = bpf_kptr_xchg(&allowed_cpumask, mask); 123 if (mask) 124 bpf_cpumask_release(mask); 125 126 bpf_rcu_read_lock(); 127 128 /* 129 * Assign the first online CPU to the allowed domain. 130 */ 131 mask = allowed_cpumask; 132 if (mask) { 133 const struct cpumask *online = scx_bpf_get_online_cpumask(); 134 135 bpf_cpumask_set_cpu(bpf_cpumask_first(online), mask); 136 scx_bpf_put_cpumask(online); 137 } 138 139 bpf_rcu_read_unlock(); 140 141 return 0; 142 } 143 144 void BPF_STRUCT_OPS(allowed_cpus_exit, struct scx_exit_info *ei) 145 { 146 UEI_RECORD(uei, ei); 147 } 148 149 struct task_cpu_arg { 150 pid_t pid; 151 }; 152 153 SEC("syscall") 154 int select_cpu_from_user(struct task_cpu_arg *input) 155 { 156 struct task_struct *p; 157 int cpu; 158 159 p = bpf_task_from_pid(input->pid); 160 if (!p) 161 return -EINVAL; 162 163 bpf_rcu_read_lock(); 164 cpu = scx_bpf_select_cpu_and(p, bpf_get_smp_processor_id(), 0, p->cpus_ptr, 0); 165 bpf_rcu_read_unlock(); 166 167 bpf_task_release(p); 168 169 return cpu; 170 } 171 172 SEC(".struct_ops.link") 173 struct sched_ext_ops allowed_cpus_ops = { 174 .select_cpu = (void *)allowed_cpus_select_cpu, 175 .enqueue = (void *)allowed_cpus_enqueue, 176 .init = (void *)allowed_cpus_init, 177 .exit = (void *)allowed_cpus_exit, 178 .name = "allowed_cpus", 179 }; 180