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 22 static bool is_cpu_idle(s32 cpu, int node) 23 { 24 const struct cpumask *idle_cpumask; 25 bool idle; 26 27 idle_cpumask = __COMPAT_scx_bpf_get_idle_cpumask_node(node); 28 idle = bpf_cpumask_test_cpu(cpu, idle_cpumask); 29 scx_bpf_put_cpumask(idle_cpumask); 30 31 return idle; 32 } 33 34 s32 BPF_STRUCT_OPS(numa_select_cpu, 35 struct task_struct *p, s32 prev_cpu, u64 wake_flags) 36 { 37 s32 task_cpu = scx_bpf_task_cpu(p); 38 int node = __COMPAT_scx_bpf_cpu_node(task_cpu); 39 s32 cpu; 40 41 /* 42 * We could just use __COMPAT_scx_bpf_pick_any_cpu_node() here, 43 * since it already tries to pick an idle CPU within the node 44 * first, but let's use both functions for better testing coverage. 45 */ 46 cpu = __COMPAT_scx_bpf_pick_idle_cpu_node(p->cpus_ptr, node, 47 __COMPAT_SCX_PICK_IDLE_IN_NODE); 48 if (cpu < 0) 49 cpu = __COMPAT_scx_bpf_pick_any_cpu_node(p->cpus_ptr, node, 50 __COMPAT_SCX_PICK_IDLE_IN_NODE); 51 52 /* 53 * @task_cpu may be outside of p->cpus_ptr if @p's affinity 54 * changed while it was sleeping. This means it's possible for 55 * p->cpus_ptr to not include any CPUs from @node. 56 * If we failed to find a cpu in @node, check if @task_cpu 57 * is outside of p->cpus_ptr and just return @prev_cpu if it is. 58 */ 59 if (cpu < 0 && !bpf_cpumask_test_cpu(task_cpu, p->cpus_ptr)) 60 return prev_cpu; 61 62 if (is_cpu_idle(cpu, node)) 63 scx_bpf_error("CPU %d should be marked as busy", cpu); 64 65 if (__COMPAT_scx_bpf_cpu_node(cpu) != node) 66 scx_bpf_error("CPU %d should be in node %d", cpu, node); 67 68 return cpu; 69 } 70 71 void BPF_STRUCT_OPS(numa_enqueue, struct task_struct *p, u64 enq_flags) 72 { 73 int node = __COMPAT_scx_bpf_cpu_node(scx_bpf_task_cpu(p)); 74 75 scx_bpf_dsq_insert(p, node, SCX_SLICE_DFL, enq_flags); 76 } 77 78 void BPF_STRUCT_OPS(numa_dispatch, s32 cpu, struct task_struct *prev) 79 { 80 int node = __COMPAT_scx_bpf_cpu_node(cpu); 81 82 scx_bpf_dsq_move_to_local(node, 0); 83 } 84 85 s32 BPF_STRUCT_OPS_SLEEPABLE(numa_init) 86 { 87 int node, err; 88 89 bpf_for(node, 0, __COMPAT_scx_bpf_nr_node_ids()) { 90 err = scx_bpf_create_dsq(node, node); 91 if (err) 92 return err; 93 } 94 95 return 0; 96 } 97 98 void BPF_STRUCT_OPS(numa_exit, struct scx_exit_info *ei) 99 { 100 UEI_RECORD(uei, ei); 101 } 102 103 SEC(".struct_ops.link") 104 struct sched_ext_ops numa_ops = { 105 .select_cpu = (void *)numa_select_cpu, 106 .enqueue = (void *)numa_enqueue, 107 .dispatch = (void *)numa_dispatch, 108 .init = (void *)numa_init, 109 .exit = (void *)numa_exit, 110 .name = "numa", 111 }; 112