1 // SPDX-License-Identifier: GPL-2.0 2 /* 3 * Copyright (c) 2023 Meta Platforms, Inc. and affiliates. 4 * Copyright (c) 2023 David Vernet <dvernet@meta.com> 5 * Copyright (c) 2023 Tejun Heo <tj@kernel.org> 6 */ 7 8 #include <scx/common.bpf.h> 9 10 char _license[] SEC("license") = "GPL"; 11 12 UEI_DEFINE(uei); 13 14 s32 BPF_STRUCT_OPS(enq_select_cpu_select_cpu, struct task_struct *p, 15 s32 prev_cpu, u64 wake_flags) 16 { 17 /* Bounce all tasks to ops.enqueue() */ 18 return prev_cpu; 19 } 20 21 void BPF_STRUCT_OPS(enq_select_cpu_enqueue, struct task_struct *p, 22 u64 enq_flags) 23 { 24 s32 cpu, prev_cpu = scx_bpf_task_cpu(p); 25 bool found = false; 26 27 cpu = scx_bpf_select_cpu_dfl(p, prev_cpu, 0, &found); 28 if (found) { 29 scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu, SCX_SLICE_DFL, enq_flags); 30 return; 31 } 32 33 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags); 34 } 35 36 void BPF_STRUCT_OPS(enq_select_cpu_exit, struct scx_exit_info *ei) 37 { 38 UEI_RECORD(uei, ei); 39 } 40 41 struct task_cpu_arg { 42 pid_t pid; 43 }; 44 45 SEC("syscall") 46 int select_cpu_from_user(struct task_cpu_arg *input) 47 { 48 struct task_struct *p; 49 bool found = false; 50 s32 cpu; 51 52 p = bpf_task_from_pid(input->pid); 53 if (!p) 54 return -EINVAL; 55 56 bpf_rcu_read_lock(); 57 cpu = scx_bpf_select_cpu_dfl(p, bpf_get_smp_processor_id(), 0, &found); 58 if (!found) 59 cpu = -EBUSY; 60 bpf_rcu_read_unlock(); 61 62 bpf_task_release(p); 63 64 return cpu; 65 } 66 67 SEC(".struct_ops.link") 68 struct sched_ext_ops enq_select_cpu_ops = { 69 .select_cpu = (void *)enq_select_cpu_select_cpu, 70 .enqueue = (void *)enq_select_cpu_enqueue, 71 .exit = (void *)enq_select_cpu_exit, 72 .name = "enq_select_cpu", 73 .timeout_ms = 1000U, 74 }; 75