1 // SPDX-License-Identifier: GPL-2.0
2 /*
3 * A scheduler that validates the behavior of the NUMA-aware
4 * functionalities.
5 *
6 * The scheduler creates a separate DSQ for each NUMA node, ensuring tasks
7 * are exclusively processed by CPUs within their respective nodes. Idle
8 * CPUs are selected only within the same node, so task migration can only
9 * occurs between CPUs belonging to the same node.
10 *
11 * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>
12 */
13
14 #include <scx/common.bpf.h>
15
16 char _license[] SEC("license") = "GPL";
17
18 UEI_DEFINE(uei);
19
20 const volatile unsigned int __COMPAT_SCX_PICK_IDLE_IN_NODE;
21
validate_local_idle_state(void)22 static void validate_local_idle_state(void)
23 {
24 const struct cpumask *idle_cpumask;
25 struct task_struct *curr;
26 s32 cpu = bpf_get_smp_processor_id();
27 int node = __COMPAT_scx_bpf_cpu_node(cpu);
28 bool cpu_is_idle, curr_is_idle;
29
30 bpf_rcu_read_lock();
31 curr = scx_bpf_cpu_curr(cpu);
32 curr_is_idle = curr && (curr->flags & PF_IDLE);
33 bpf_rcu_read_unlock();
34
35 idle_cpumask = __COMPAT_scx_bpf_get_idle_cpumask_node(node);
36 cpu_is_idle = bpf_cpumask_test_cpu(cpu, idle_cpumask);
37 scx_bpf_put_cpumask(idle_cpumask);
38
39 /*
40 * Unlike a remote picked CPU, the local CPU cannot go through an
41 * idle re-pick while this callback is running. If it is running a
42 * non-idle scheduling context, it must not be advertised as idle
43 * in its node's idle cpumask.
44 */
45 if (!curr_is_idle && cpu_is_idle)
46 scx_bpf_error("running CPU %d should be marked as busy", cpu);
47 }
48
BPF_STRUCT_OPS(numa_select_cpu,struct task_struct * p,s32 prev_cpu,u64 wake_flags)49 s32 BPF_STRUCT_OPS(numa_select_cpu,
50 struct task_struct *p, s32 prev_cpu, u64 wake_flags)
51 {
52 s32 task_cpu = scx_bpf_task_cpu(p);
53 int node = __COMPAT_scx_bpf_cpu_node(task_cpu);
54 s32 cpu;
55
56 validate_local_idle_state();
57
58 /*
59 * We could just use __COMPAT_scx_bpf_pick_any_cpu_node() here,
60 * since it already tries to pick an idle CPU within the node
61 * first, but let's use both functions for better testing coverage.
62 */
63 cpu = __COMPAT_scx_bpf_pick_idle_cpu_node(p->cpus_ptr, node,
64 __COMPAT_SCX_PICK_IDLE_IN_NODE);
65 if (cpu < 0)
66 cpu = __COMPAT_scx_bpf_pick_any_cpu_node(p->cpus_ptr, node,
67 __COMPAT_SCX_PICK_IDLE_IN_NODE);
68
69 /*
70 * @task_cpu may be outside of p->cpus_ptr if @p's affinity
71 * changed while it was sleeping. This means it's possible for
72 * p->cpus_ptr to not include any CPUs from @node.
73 * If we failed to find a cpu in @node, check if @task_cpu
74 * is outside of p->cpus_ptr and just return @prev_cpu if it is.
75 */
76 if (cpu < 0 && !bpf_cpumask_test_cpu(task_cpu, p->cpus_ptr))
77 return prev_cpu;
78
79 if (__COMPAT_scx_bpf_cpu_node(cpu) != node)
80 scx_bpf_error("CPU %d should be in node %d", cpu, node);
81
82 return cpu;
83 }
84
BPF_STRUCT_OPS(numa_enqueue,struct task_struct * p,u64 enq_flags)85 void BPF_STRUCT_OPS(numa_enqueue, struct task_struct *p, u64 enq_flags)
86 {
87 int node = __COMPAT_scx_bpf_cpu_node(scx_bpf_task_cpu(p));
88
89 scx_bpf_dsq_insert(p, node, SCX_SLICE_DFL, enq_flags);
90 }
91
BPF_STRUCT_OPS(numa_dispatch,s32 cpu,struct task_struct * prev)92 void BPF_STRUCT_OPS(numa_dispatch, s32 cpu, struct task_struct *prev)
93 {
94 int node = __COMPAT_scx_bpf_cpu_node(cpu);
95
96 scx_bpf_dsq_move_to_local(node, 0);
97 }
98
BPF_STRUCT_OPS_SLEEPABLE(numa_init)99 s32 BPF_STRUCT_OPS_SLEEPABLE(numa_init)
100 {
101 int node, err;
102
103 bpf_for(node, 0, __COMPAT_scx_bpf_nr_node_ids()) {
104 err = scx_bpf_create_dsq(node, node);
105 if (err)
106 return err;
107 }
108
109 return 0;
110 }
111
BPF_STRUCT_OPS(numa_exit,struct scx_exit_info * ei)112 void BPF_STRUCT_OPS(numa_exit, struct scx_exit_info *ei)
113 {
114 UEI_RECORD(uei, ei);
115 }
116
117 SEC(".struct_ops.link")
118 struct sched_ext_ops numa_ops = {
119 .select_cpu = (void *)numa_select_cpu,
120 .enqueue = (void *)numa_enqueue,
121 .dispatch = (void *)numa_dispatch,
122 .init = (void *)numa_init,
123 .exit = (void *)numa_exit,
124 .name = "numa",
125 };
126