On Mon, Jul 10, 2023 at 03:13:30PM -1000, Tejun Heo wrote:
> Implement a new scheduler class sched_ext (SCX), which allows scheduling
> policies to be implemented as BPF programs to achieve the following:
>
> 1. Ease of experimentation and exploration: Enabling rapid iteration of new
> scheduling policies.
>
> 2. Customization: Building application-specific schedulers which implement
> policies that are not applicable to general-purpose schedulers.
>
> 3. Rapid scheduler deployments: Non-disruptive swap outs of scheduling
> policies in production environments.
>
> sched_ext leverages BPF’s struct_ops feature to define a structure which
> exports function callbacks and flags to BPF programs that wish to implement
> scheduling policies. The struct_ops structure exported by sched_ext is
> struct sched_ext_ops, and is conceptually similar to struct sched_class. The
> role of sched_ext is to map the complex sched_class callbacks to the more
> simple and ergonomic struct sched_ext_ops callbacks.
>
> For more detailed discussion on the motivations and overview, please refer
> to the cover letter.
>
> Later patches will also add several example schedulers and documentation.
>
> This patch implements the minimum core framework to enable implementation of
> BPF schedulers. Subsequent patches will gradually add functionalities
> including safety guarantee mechanisms, nohz and cgroup support.
>
> include/linux/sched/ext.h defines struct sched_ext_ops. With the comment on
> top, each operation should be self-explanatory. The followings are worth
> noting:
>
> * Both "sched_ext" and its shorthand "scx" are used. If the identifier
> already has "sched" in it, "ext" is used; otherwise, "scx".
>
> * In sched_ext_ops, only .name is mandatory. Every operation is optional and
> if omitted a simple but functional default behavior is provided.
>
> * A new policy constant SCHED_EXT is added and a task can select sched_ext
> by invoking sched_setscheduler(2) with the new policy constant. However,
> if the BPF scheduler is not loaded, SCHED_EXT is the same as SCHED_NORMAL
> and the task is scheduled by CFS. When the BPF scheduler is loaded, all
> tasks which have the SCHED_EXT policy are switched to sched_ext.
>
> * To bridge the workflow imbalance between the scheduler core and
> sched_ext_ops callbacks, sched_ext uses simple FIFOs called dispatch
> queues (dsq's). By default, there is one global dsq (SCX_DSQ_GLOBAL), and
> one local per-CPU dsq (SCX_DSQ_LOCAL). SCX_DSQ_GLOBAL is provided for
> convenience and need not be used by a scheduler that doesn't require it.
> SCX_DSQ_LOCAL is the per-CPU FIFO that sched_ext pulls from when putting
> the next task on the CPU. The BPF scheduler can manage an arbitrary number
> of dsq's using scx_bpf_create_dsq() and scx_bpf_destroy_dsq().
>
> * sched_ext guarantees system integrity no matter what the BPF scheduler
> does. To enable this, each task's ownership is tracked through
> p->scx.ops_state and all tasks are put on scx_tasks list. The disable path
> can always recover and revert all tasks back to CFS. See p->scx.ops_state
> and scx_tasks.
>
> * A task is not tied to its rq while enqueued. This decouples CPU selection
> from queueing and allows sharing a scheduling queue across an arbitrary
> subset of CPUs. This adds some complexities as a task may need to be
> bounced between rq's right before it starts executing. See
> dispatch_to_local_dsq() and move_task_to_local_dsq().
>
> * One complication that arises from the above weak association between task
> and rq is that synchronizing with dequeue() gets complicated as dequeue()
> may happen anytime while the task is enqueued and the dispatch path might
> need to release the rq lock to transfer the task. Solving this requires a
> bit of complexity. See the logic around p->scx.sticky_cpu and
> p->scx.ops_qseq.
>
> * Both enable and disable paths are a bit complicated. The enable path
> switches all tasks without blocking to avoid issues which can arise from
> partially switched states (e.g. the switching task itself being starved).
> The disable path can't trust the BPF scheduler at all, so it also has to
> guarantee forward progress without blocking. See scx_ops_enable() and
> scx_ops_disable_workfn().
>
> * When sched_ext is disabled, static_branches are used to shut down the
> entry points from hot paths.
>
> v4: * SCHED_CHANGE_BLOCK replaced with the previous
> sched_deq_and_put_task()/sched_enq_and_set_tsak() pair. This is
> because upstream is adaopting a different generic cleanup mechanism.
> Once that lands, the code will be adapted accordingly.
>
> * task_on_scx() used to test whether a task should be switched into SCX,
> which is confusing. Renamed to task_should_scx(). task_on_scx() now
> tests whether a task is currently on SCX.
>
> * scx_has_idle_cpus is barely used anymore and replaced with direct
> check on the idle cpumask.
>
> * SCX_PICK_IDLE_CORE added and scx_pick_idle_cpu() improved to prefer
> fully idle cores.
>
> * ops.enable() now sees up-to-date p->scx.weight value.
>
> * ttwu_queue path is disabled for tasks on SCX to avoid confusing BPF
> schedulers expecting ->select_cpu() call.
>
> * Use cpu_smt_mask() instead of topology_sibling_cpumask() like the rest
> of the scheduler.
>
> v3: * ops.set_weight() added to allow BPF schedulers to track weight changes
> without polling p->scx.weight.
>
> * move_task_to_local_dsq() was losing SCX-specific enq_flags when
> enqueueing the task on the target dsq because it goes through
> activate_task() which loses the upper 32bit of the flags. Carry the
> flags through rq->scx.extra_enq_flags.
>
> * scx_bpf_dispatch(), scx_bpf_pick_idle_cpu(), scx_bpf_task_running()
> and scx_bpf_task_cpu() now use the new KF_RCU instead of
> KF_TRUSTED_ARGS to make it easier for BPF schedulers to call them.
>
> * The kfunc helper access control mechanism implemented through
> sched_ext_entity.kf_mask is improved. Now SCX_CALL_OP*() is always
> used when invoking scx_ops operations.
>
> v2: * balance_scx_on_up() is dropped. Instead, on UP, balance_scx() is
> called from put_prev_taks_scx() and pick_next_task_scx() as necessary.
> To determine whether balance_scx() should be called from
> put_prev_task_scx(), SCX_TASK_DEQD_FOR_SLEEP flag is added. See the
> comment in put_prev_task_scx() for details.
>
> * sched_deq_and_put_task() / sched_enq_and_set_task() sequences replaced
> with SCHED_CHANGE_BLOCK().
>
> * Unused all_dsqs list removed. This was a left-over from previous
> iterations.
>
> * p->scx.kf_mask is added to track and enforce which kfunc helpers are
> allowed. Also, init/exit sequences are updated to make some kfuncs
> always safe to call regardless of the current BPF scheduler state.
> Combined, this should make all the kfuncs safe.
>
> * BPF now supports sleepable struct_ops operations. Hacky workaround
> removed and operations and kfunc helpers are tagged appropriately.
>
> * BPF now supports bitmask / cpumask helpers. scx_bpf_get_idle_cpumask()
> and friends are added so that BPF schedulers can use the idle masks
> with the generic helpers. This replaces the hacky kfunc helpers added
> by a separate patch in V1.
>
> * CONFIG_SCHED_CLASS_EXT can no longer be enabled if SCHED_CORE is
> enabled. This restriction will be removed by a later patch which adds
> core-sched support.
>
> * Add MAINTAINERS entries and other misc changes.
>
> Signed-off-by: Tejun Heo <tj@xxxxxxxxxx>
> Co-authored-by: David Vernet <dvernet@xxxxxxxx>
> Acked-by: Josh Don <joshdon@xxxxxxxxxx>
> Acked-by: Hao Luo <haoluo@xxxxxxxxxx>
> Acked-by: Barret Rhoden <brho@xxxxxxxxxx>
> Cc: Andrea Righi <andrea.righi@xxxxxxxxxxxxx>
> ---
> MAINTAINERS | 3 +
> include/asm-generic/vmlinux.lds.h | 1 +
> include/linux/sched.h | 5 +
> include/linux/sched/ext.h | 400 +++-
> include/uapi/linux/sched.h | 1 +
> init/init_task.c | 10 +
> kernel/Kconfig.preempt | 22 +-
> kernel/bpf/bpf_struct_ops_types.h | 4 +
> kernel/sched/build_policy.c | 4 +
> kernel/sched/core.c | 70 +
> kernel/sched/debug.c | 6 +
> kernel/sched/ext.c | 3140 +++++++++++++++++++++++++++++
> kernel/sched/ext.h | 118 +-
> kernel/sched/sched.h | 16 +
> 14 files changed, 3796 insertions(+), 4 deletions(-)
> create mode 100644 kernel/sched/ext.c
>
> diff --git a/MAINTAINERS b/MAINTAINERS
> index c904dba1733b..5c301e22ff74 100644
> --- a/MAINTAINERS
> +++ b/MAINTAINERS
> @@ -18767,6 +18767,8 @@ R: Ben Segall <bsegall@xxxxxxxxxx> (CONFIG_CFS_BANDWIDTH)
> R: Mel Gorman <mgorman@xxxxxxx> (CONFIG_NUMA_BALANCING)
> R: Daniel Bristot de Oliveira <bristot@xxxxxxxxxx> (SCHED_DEADLINE)
> R: Valentin Schneider <vschneid@xxxxxxxxxx> (TOPOLOGY)
> +R: Tejun Heo <tj@xxxxxxxxxx> (SCHED_EXT)
> +R: David Vernet <void@xxxxxxxxxxxxx> (SCHED_EXT)
> L: linux-kernel@xxxxxxxxxxxxxxx
> S: Maintained
> T: git git://git.kernel.org/pub/scm/linux/kernel/git/tip/tip.git sched/core
> @@ -18775,6 +18777,7 @@ F: include/linux/sched.h
> F: include/linux/wait.h
> F: include/uapi/linux/sched.h
> F: kernel/sched/
> +F: tools/sched_ext/
>
> SCSI RDMA PROTOCOL (SRP) INITIATOR
> M: Bart Van Assche <bvanassche@xxxxxxx>
> diff --git a/include/asm-generic/vmlinux.lds.h b/include/asm-generic/vmlinux.lds.h
> index d1f57e4868ed..cd5a718ba49f 100644
> --- a/include/asm-generic/vmlinux.lds.h
> +++ b/include/asm-generic/vmlinux.lds.h
> @@ -131,6 +131,7 @@
> *(__dl_sched_class) \
> *(__rt_sched_class) \
> *(__fair_sched_class) \
> + *(__ext_sched_class) \
> *(__idle_sched_class) \
> __sched_class_lowest = .;
>
> diff --git a/include/linux/sched.h b/include/linux/sched.h
> index eed5d65b8d1f..00d4ce3af52a 100644
> --- a/include/linux/sched.h
> +++ b/include/linux/sched.h
> @@ -72,6 +72,8 @@ struct task_delay_info;
> struct task_group;
> struct user_event_mm;
>
> +#include <linux/sched/ext.h>
> +
> /*
> * Task state bitmask. NOTE! These bits are also
> * encoded in fs/proc/array.c: get_task_state().
> @@ -790,6 +792,9 @@ struct task_struct {
> struct sched_entity se;
> struct sched_rt_entity rt;
> struct sched_dl_entity dl;
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + struct sched_ext_entity scx;
> +#endif
> const struct sched_class *sched_class;
>
> #ifdef CONFIG_SCHED_CORE
> diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
> index a05dfcf533b0..92011a63cc15 100644
> --- a/include/linux/sched/ext.h
> +++ b/include/linux/sched/ext.h
> @@ -1,9 +1,407 @@
> /* SPDX-License-Identifier: GPL-2.0 */
> +/*
> + * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
> + * Copyright (c) 2022 Tejun Heo <tj@xxxxxxxxxx>
> + * Copyright (c) 2022 David Vernet <dvernet@xxxxxxxx>
> + */
> #ifndef _LINUX_SCHED_EXT_H
> #define _LINUX_SCHED_EXT_H
>
> #ifdef CONFIG_SCHED_CLASS_EXT
> -#error "NOT IMPLEMENTED YET"
> +
> +#include <linux/rhashtable.h>
> +#include <linux/llist.h>
> +
> +enum scx_consts {
> + SCX_OPS_NAME_LEN = 128,
> + SCX_EXIT_REASON_LEN = 128,
> + SCX_EXIT_BT_LEN = 64,
> + SCX_EXIT_MSG_LEN = 1024,
> +
> + SCX_SLICE_DFL = 20 * NSEC_PER_MSEC,
> +};
> +
> +/*
> + * DSQ (dispatch queue) IDs are 64bit of the format:
> + *
> + * Bits: [63] [62 .. 0]
> + * [ B] [ ID ]
> + *
> + * B: 1 for IDs for built-in DSQs, 0 for ops-created user DSQs
> + * ID: 63 bit ID
> + *
> + * Built-in IDs:
> + *
> + * Bits: [63] [62] [61..32] [31 .. 0]
> + * [ 1] [ L] [ R ] [ V ]
> + *
> + * 1: 1 for built-in DSQs.
> + * L: 1 for LOCAL_ON DSQ IDs, 0 for others
> + * V: For LOCAL_ON DSQ IDs, a CPU number. For others, a pre-defined value.
> + */
> +enum scx_dsq_id_flags {
> + SCX_DSQ_FLAG_BUILTIN = 1LLU << 63,
> + SCX_DSQ_FLAG_LOCAL_ON = 1LLU << 62,
> +
> + SCX_DSQ_INVALID = SCX_DSQ_FLAG_BUILTIN | 0,
> + SCX_DSQ_GLOBAL = SCX_DSQ_FLAG_BUILTIN | 1,
> + SCX_DSQ_LOCAL = SCX_DSQ_FLAG_BUILTIN | 2,
> + SCX_DSQ_LOCAL_ON = SCX_DSQ_FLAG_BUILTIN | SCX_DSQ_FLAG_LOCAL_ON,
> + SCX_DSQ_LOCAL_CPU_MASK = 0xffffffffLLU,
> +};
> +
> +enum scx_exit_type {
> + SCX_EXIT_NONE,
> + SCX_EXIT_DONE,
> +
> + SCX_EXIT_UNREG = 64, /* BPF unregistration */
> +
> + SCX_EXIT_ERROR = 1024, /* runtime error, error msg contains details */
> + SCX_EXIT_ERROR_BPF, /* ERROR but triggered through scx_bpf_error() */
> +};
> +
> +/*
> + * scx_exit_info is passed to ops.exit() to describe why the BPF scheduler is
> + * being disabled.
> + */
> +struct scx_exit_info {
> + /* %SCX_EXIT_* - broad category of the exit reason */
> + enum scx_exit_type type;
> + /* textual representation of the above */
> + char reason[SCX_EXIT_REASON_LEN];
> + /* number of entries in the backtrace */
> + u32 bt_len;
> + /* backtrace if exiting due to an error */
> + unsigned long bt[SCX_EXIT_BT_LEN];
> + /* extra message */
> + char msg[SCX_EXIT_MSG_LEN];
> +};
> +
> +/* sched_ext_ops.flags */
> +enum scx_ops_flags {
> + /*
> + * Keep built-in idle tracking even if ops.update_idle() is implemented.
> + */
> + SCX_OPS_KEEP_BUILTIN_IDLE = 1LLU << 0,
> +
> + /*
> + * By default, if there are no other task to run on the CPU, ext core
> + * keeps running the current task even after its slice expires. If this
> + * flag is specified, such tasks are passed to ops.enqueue() with
> + * %SCX_ENQ_LAST. See the comment above %SCX_ENQ_LAST for more info.
> + */
> + SCX_OPS_ENQ_LAST = 1LLU << 1,
> +
> + /*
> + * An exiting task may schedule after PF_EXITING is set. In such cases,
> + * bpf_task_from_pid() may not be able to find the task and if the BPF
> + * scheduler depends on pid lookup for dispatching, the task will be
> + * lost leading to various issues including RCU grace period stalls.
> + *
> + * To mask this problem, by default, unhashed tasks are automatically
> + * dispatched to the local DSQ on enqueue. If the BPF scheduler doesn't
> + * depend on pid lookups and wants to handle these tasks directly, the
> + * following flag can be used.
> + */
> + SCX_OPS_ENQ_EXITING = 1LLU << 2,
> +
> + SCX_OPS_ALL_FLAGS = SCX_OPS_KEEP_BUILTIN_IDLE |
> + SCX_OPS_ENQ_LAST |
> + SCX_OPS_ENQ_EXITING,
> +};
> +
> +/* argument container for ops.enable() and friends */
> +struct scx_enable_args {
> + /* empty for now */
> +};
> +
> +/**
> + * struct sched_ext_ops - Operation table for BPF scheduler implementation
> + *
> + * Userland can implement an arbitrary scheduling policy by implementing and
> + * loading operations in this table.
> + */
> +struct sched_ext_ops {
> + /**
> + * select_cpu - Pick the target CPU for a task which is being woken up
> + * @p: task being woken up
> + * @prev_cpu: the cpu @p was on before sleeping
> + * @wake_flags: SCX_WAKE_*
> + *
> + * Decision made here isn't final. @p may be moved to any CPU while it
> + * is getting dispatched for execution later. However, as @p is not on
> + * the rq at this point, getting the eventual execution CPU right here
> + * saves a small bit of overhead down the line.
> + *
> + * If an idle CPU is returned, the CPU is kicked and will try to
> + * dispatch. While an explicit custom mechanism can be added,
> + * select_cpu() serves as the default way to wake up idle CPUs.
> + */
> + s32 (*select_cpu)(struct task_struct *p, s32 prev_cpu, u64 wake_flags);
> +
> + /**
> + * enqueue - Enqueue a task on the BPF scheduler
> + * @p: task being enqueued
> + * @enq_flags: %SCX_ENQ_*
> + *
> + * @p is ready to run. Dispatch directly by calling scx_bpf_dispatch()
> + * or enqueue on the BPF scheduler. If not directly dispatched, the bpf
> + * scheduler owns @p and if it fails to dispatch @p, the task will
> + * stall.
> + */
> + void (*enqueue)(struct task_struct *p, u64 enq_flags);
> +
> + /**
> + * dequeue - Remove a task from the BPF scheduler
> + * @p: task being dequeued
> + * @deq_flags: %SCX_DEQ_*
> + *
> + * Remove @p from the BPF scheduler. This is usually called to isolate
> + * the task while updating its scheduling properties (e.g. priority).
> + *
> + * The ext core keeps track of whether the BPF side owns a given task or
> + * not and can gracefully ignore spurious dispatches from BPF side,
> + * which makes it safe to not implement this method. However, depending
> + * on the scheduling logic, this can lead to confusing behaviors - e.g.
> + * scheduling position not being updated across a priority change.
> + */
> + void (*dequeue)(struct task_struct *p, u64 deq_flags);
> +
> + /**
> + * dispatch - Dispatch tasks from the BPF scheduler and/or consume DSQs
> + * @cpu: CPU to dispatch tasks for
> + * @prev: previous task being switched out
> + *
> + * Called when a CPU's local dsq is empty. The operation should dispatch
> + * one or more tasks from the BPF scheduler into the DSQs using
> + * scx_bpf_dispatch() and/or consume user DSQs into the local DSQ using
> + * scx_bpf_consume().
> + *
> + * The maximum number of times scx_bpf_dispatch() can be called without
> + * an intervening scx_bpf_consume() is specified by
> + * ops.dispatch_max_batch. See the comments on top of the two functions
> + * for more details.
> + *
> + * When not %NULL, @prev is an SCX task with its slice depleted. If
> + * @prev is still runnable as indicated by set %SCX_TASK_QUEUED in
> + * @prev->scx.flags, it is not enqueued yet and will be enqueued after
> + * ops.dispatch() returns. To keep executing @prev, return without
> + * dispatching or consuming any tasks. Also see %SCX_OPS_ENQ_LAST.
> + */
> + void (*dispatch)(s32 cpu, struct task_struct *prev);
> +
> + /**
> + * yield - Yield CPU
> + * @from: yielding task
> + * @to: optional yield target task
> + *
> + * If @to is NULL, @from is yielding the CPU to other runnable tasks.
> + * The BPF scheduler should ensure that other available tasks are
> + * dispatched before the yielding task. Return value is ignored in this
> + * case.
> + *
> + * If @to is not-NULL, @from wants to yield the CPU to @to. If the bpf
> + * scheduler can implement the request, return %true; otherwise, %false.
> + */
> + bool (*yield)(struct task_struct *from, struct task_struct *to);
> +
> + /**
> + * set_weight - Set task weight
> + * @p: task to set weight for
> + * @weight: new eight [1..10000]
> + *
> + * Update @p's weight to @weight.
> + */
> + void (*set_weight)(struct task_struct *p, u32 weight);
> +
> + /**
> + * set_cpumask - Set CPU affinity
> + * @p: task to set CPU affinity for
> + * @cpumask: cpumask of cpus that @p can run on
> + *
> + * Update @p's CPU affinity to @cpumask.
> + */
> + void (*set_cpumask)(struct task_struct *p, struct cpumask *cpumask);
> +
> + /**
> + * update_idle - Update the idle state of a CPU
> + * @cpu: CPU to udpate the idle state for
> + * @idle: whether entering or exiting the idle state
> + *
> + * This operation is called when @rq's CPU goes or leaves the idle
> + * state. By default, implementing this operation disables the built-in
> + * idle CPU tracking and the following helpers become unavailable:
> + *
> + * - scx_bpf_select_cpu_dfl()
> + * - scx_bpf_test_and_clear_cpu_idle()
> + * - scx_bpf_pick_idle_cpu()
> + *
> + * The user also must implement ops.select_cpu() as the default
> + * implementation relies on scx_bpf_select_cpu_dfl().
> + *
> + * Specify the %SCX_OPS_KEEP_BUILTIN_IDLE flag to keep the built-in idle
> + * tracking.
> + */
> + void (*update_idle)(s32 cpu, bool idle);
> +
> + /**
> + * prep_enable - Prepare to enable BPF scheduling for a task
> + * @p: task to prepare BPF scheduling for
> + * @args: enable arguments, see the struct definition
> + *
> + * Either we're loading a BPF scheduler or a new task is being forked.
> + * Prepare BPF scheduling for @p. This operation may block and can be
> + * used for allocations.
> + *
> + * Return 0 for success, -errno for failure. An error return while
> + * loading will abort loading of the BPF scheduler. During a fork, will
> + * abort the specific fork.
> + */
> + s32 (*prep_enable)(struct task_struct *p, struct scx_enable_args *args);
> +
> + /**
> + * enable - Enable BPF scheduling for a task
> + * @p: task to enable BPF scheduling for
> + * @args: enable arguments, see the struct definition
> + *
> + * Enable @p for BPF scheduling. @p will start running soon.
> + */
> + void (*enable)(struct task_struct *p, struct scx_enable_args *args);
> +
> + /**
> + * cancel_enable - Cancel prep_enable()
> + * @p: task being canceled
> + * @args: enable arguments, see the struct definition
> + *
> + * @p was prep_enable()'d but failed before reaching enable(). Undo the
> + * preparation.
> + */
> + void (*cancel_enable)(struct task_struct *p,
> + struct scx_enable_args *args);
> +
> + /**
> + * disable - Disable BPF scheduling for a task
> + * @p: task to disable BPF scheduling for
> + *
> + * @p is exiting, leaving SCX or the BPF scheduler is being unloaded.
> + * Disable BPF scheduling for @p.
> + */
> + void (*disable)(struct task_struct *p);
> +
> + /*
> + * All online ops must come before ops.init().
> + */
> +
> + /**
> + * init - Initialize the BPF scheduler
> + */
> + s32 (*init)(void);
> +
> + /**
> + * exit - Clean up after the BPF scheduler
> + * @info: Exit info
> + */
> + void (*exit)(struct scx_exit_info *info);
> +
> + /**
> + * dispatch_max_batch - Max nr of tasks that dispatch() can dispatch
> + */
> + u32 dispatch_max_batch;
> +
> + /**
> + * flags - %SCX_OPS_* flags
> + */
> + u64 flags;
> +
> + /**
> + * name - BPF scheduler's name
> + *
> + * Must be a non-zero valid BPF object name including only isalnum(),
> + * '_' and '.' chars. Shows up in kernel.sched_ext_ops sysctl while the
> + * BPF scheduler is enabled.
> + */
> + char name[SCX_OPS_NAME_LEN];
> +};
> +
> +/*
> + * Dispatch queue (dsq) is a simple FIFO which is used to buffer between the
> + * scheduler core and the BPF scheduler. See the documentation for more details.
> + */
> +struct scx_dispatch_q {
> + raw_spinlock_t lock;
> + struct list_head fifo; /* processed in dispatching order */
> + u32 nr;
> + u64 id;
> + struct rhash_head hash_node;
> + struct llist_node free_node;
> + struct rcu_head rcu;
> +};
> +
> +/* scx_entity.flags */
> +enum scx_ent_flags {
> + SCX_TASK_QUEUED = 1 << 0, /* on ext runqueue */
> + SCX_TASK_BAL_KEEP = 1 << 1, /* balance decided to keep current */
> + SCX_TASK_ENQ_LOCAL = 1 << 2, /* used by scx_select_cpu_dfl() to set SCX_ENQ_LOCAL */
> +
> + SCX_TASK_OPS_PREPPED = 1 << 8, /* prepared for BPF scheduler enable */
> + SCX_TASK_OPS_ENABLED = 1 << 9, /* task has BPF scheduler enabled */
> +
> + SCX_TASK_DEQD_FOR_SLEEP = 1 << 17, /* last dequeue was for SLEEP */
> +
> + SCX_TASK_CURSOR = 1 << 31, /* iteration cursor, not a task */
> +};
> +
> +/*
> + * Mask bits for scx_entity.kf_mask. Not all kfuncs can be called from
> + * everywhere and the following bits track which kfunc sets are currently
> + * allowed for %current. This simple per-task tracking works because SCX ops
> + * nest in a limited way. BPF will likely implement a way to allow and disallow
> + * kfuncs depending on the calling context which will replace this manual
> + * mechanism. See scx_kf_allow().
> + */
> +enum scx_kf_mask {
> + SCX_KF_UNLOCKED = 0, /* not sleepable, not rq locked */
> + /* all non-sleepables may be nested inside INIT and SLEEPABLE */
> + SCX_KF_INIT = 1 << 0, /* running ops.init() */
> + SCX_KF_SLEEPABLE = 1 << 1, /* other sleepable init operations */
> + /* ops.dequeue (in REST) may be nested inside DISPATCH */
> + SCX_KF_DISPATCH = 1 << 3, /* ops.dispatch() */
> + SCX_KF_ENQUEUE = 1 << 4, /* ops.enqueue() */
> + SCX_KF_REST = 1 << 5, /* other rq-locked operations */
> +
> + __SCX_KF_RQ_LOCKED = SCX_KF_DISPATCH | SCX_KF_ENQUEUE | SCX_KF_REST,
> +};
> +
> +/*
> + * The following is embedded in task_struct and contains all fields necessary
> + * for a task to be scheduled by SCX.
> + */
> +struct sched_ext_entity {
> + struct scx_dispatch_q *dsq;
> + struct list_head dsq_node;
> + u32 flags; /* protected by rq lock */
> + u32 weight;
> + s32 sticky_cpu;
> + s32 holding_cpu;
> + u32 kf_mask; /* see scx_kf_mask above */
> + atomic64_t ops_state;
> +
> + /* BPF scheduler modifiable fields */
> +
> + /*
> + * Runtime budget in nsecs. This is usually set through
> + * scx_bpf_dispatch() but can also be modified directly by the BPF
> + * scheduler. Automatically decreased by SCX as the task executes. On
> + * depletion, a scheduling event is triggered.
> + */
> + u64 slice;
> +
> + /* cold fields */
> + struct list_head tasks_node;
> +};
> +
> +void sched_ext_free(struct task_struct *p);
> +
> #else /* !CONFIG_SCHED_CLASS_EXT */
>
> static inline void sched_ext_free(struct task_struct *p) {}
> diff --git a/include/uapi/linux/sched.h b/include/uapi/linux/sched.h
> index 3bac0a8ceab2..359a14cc76a4 100644
> --- a/include/uapi/linux/sched.h
> +++ b/include/uapi/linux/sched.h
> @@ -118,6 +118,7 @@ struct clone_args {
> /* SCHED_ISO: reserved but not implemented yet */
> #define SCHED_IDLE 5
> #define SCHED_DEADLINE 6
> +#define SCHED_EXT 7
>
> /* Can be ORed in to make sure the process is reverted back to SCHED_NORMAL on fork */
> #define SCHED_RESET_ON_FORK 0x40000000
> diff --git a/init/init_task.c b/init/init_task.c
> index ff6c4b9bfe6b..bdbc663107bf 100644
> --- a/init/init_task.c
> +++ b/init/init_task.c
> @@ -6,6 +6,7 @@
> #include <linux/sched/sysctl.h>
> #include <linux/sched/rt.h>
> #include <linux/sched/task.h>
> +#include <linux/sched/ext.h>
> #include <linux/init.h>
> #include <linux/fs.h>
> #include <linux/mm.h>
> @@ -101,6 +102,15 @@ struct task_struct init_task
> #endif
> #ifdef CONFIG_CGROUP_SCHED
> .sched_task_group = &root_task_group,
> +#endif
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + .scx = {
> + .dsq_node = LIST_HEAD_INIT(init_task.scx.dsq_node),
> + .sticky_cpu = -1,
> + .holding_cpu = -1,
> + .ops_state = ATOMIC_INIT(0),
> + .slice = SCX_SLICE_DFL,
> + },
> #endif
> .ptraced = LIST_HEAD_INIT(init_task.ptraced),
> .ptrace_entry = LIST_HEAD_INIT(init_task.ptrace_entry),
> diff --git a/kernel/Kconfig.preempt b/kernel/Kconfig.preempt
> index c2f1fd95a821..0afcda19bc50 100644
> --- a/kernel/Kconfig.preempt
> +++ b/kernel/Kconfig.preempt
> @@ -133,4 +133,24 @@ config SCHED_CORE
> which is the likely usage by Linux distributions, there should
> be no measurable impact on performance.
>
> -
> +config SCHED_CLASS_EXT
> + bool "Extensible Scheduling Class"
> + depends on BPF_SYSCALL && BPF_JIT && !SCHED_CORE
> + help
> + This option enables a new scheduler class sched_ext (SCX), which
> + allows scheduling policies to be implemented as BPF programs to
> + achieve the following:
> +
> + - Ease of experimentation and exploration: Enabling rapid
> + iteration of new scheduling policies.
> + - Customization: Building application-specific schedulers which
> + implement policies that are not applicable to general-purpose
> + schedulers.
> + - Rapid scheduler deployments: Non-disruptive swap outs of
> + scheduling policies in production environments.
> +
> + sched_ext leverages BPF’s struct_ops feature to define a structure
> + which exports function callbacks and flags to BPF programs that
> + wish to implement scheduling policies. The struct_ops structure
> + exported by sched_ext is struct sched_ext_ops, and is conceptually
> + similar to struct sched_class.
> diff --git a/kernel/bpf/bpf_struct_ops_types.h b/kernel/bpf/bpf_struct_ops_types.h
> index 5678a9ddf817..3618769d853d 100644
> --- a/kernel/bpf/bpf_struct_ops_types.h
> +++ b/kernel/bpf/bpf_struct_ops_types.h
> @@ -9,4 +9,8 @@ BPF_STRUCT_OPS_TYPE(bpf_dummy_ops)
> #include <net/tcp.h>
> BPF_STRUCT_OPS_TYPE(tcp_congestion_ops)
> #endif
> +#ifdef CONFIG_SCHED_CLASS_EXT
> +#include <linux/sched/ext.h>
> +BPF_STRUCT_OPS_TYPE(sched_ext_ops)
> +#endif
> #endif
> diff --git a/kernel/sched/build_policy.c b/kernel/sched/build_policy.c
> index d9dc9ab3773f..4c658b21f603 100644
> --- a/kernel/sched/build_policy.c
> +++ b/kernel/sched/build_policy.c
> @@ -28,6 +28,7 @@
> #include <linux/suspend.h>
> #include <linux/tsacct_kern.h>
> #include <linux/vtime.h>
> +#include <linux/percpu-rwsem.h>
>
> #include <uapi/linux/sched/types.h>
>
> @@ -52,3 +53,6 @@
> #include "cputime.c"
> #include "deadline.c"
>
> +#ifdef CONFIG_SCHED_CLASS_EXT
> +# include "ext.c"
> +#endif
> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index c8a2c99248b7..a182a9adec3e 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -3925,6 +3925,15 @@ bool cpus_share_cache(int this_cpu, int that_cpu)
>
> static inline bool ttwu_queue_cond(struct task_struct *p, int cpu)
> {
> + /*
> + * The BPF scheduler may depend on select_task_rq() being invoked during
> + * wakeups. In addition, @p may end up executing on a different CPU
> + * regardless of what happens in the wakeup path making the ttwu_queue
> + * optimization less meaningful. Skip if on SCX.
> + */
> + if (task_on_scx(p))
> + return false;
> +
> /*
> * Do not complicate things with the async wake_list while the CPU is
> * in hotplug state.
> @@ -4496,6 +4505,18 @@ static void __sched_fork(unsigned long clone_flags, struct task_struct *p)
> p->rt.on_rq = 0;
> p->rt.on_list = 0;
>
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + p->scx.dsq = NULL;
> + INIT_LIST_HEAD(&p->scx.dsq_node);
> + p->scx.flags = 0;
> + p->scx.weight = 0;
> + p->scx.sticky_cpu = -1;
> + p->scx.holding_cpu = -1;
> + p->scx.kf_mask = 0;
> + atomic64_set(&p->scx.ops_state, 0);
> + p->scx.slice = SCX_SLICE_DFL;
> +#endif
> +
> #ifdef CONFIG_PREEMPT_NOTIFIERS
> INIT_HLIST_HEAD(&p->preempt_notifiers);
> #endif
> @@ -4744,6 +4765,10 @@ int sched_fork(unsigned long clone_flags, struct task_struct *p)
> goto out_cancel;
> } else if (rt_prio(p->prio)) {
> p->sched_class = &rt_sched_class;
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + } else if (task_should_scx(p)) {
> + p->sched_class = &ext_sched_class;
> +#endif
> } else {
> p->sched_class = &fair_sched_class;
> }
> @@ -7032,6 +7057,10 @@ void __setscheduler_prio(struct task_struct *p, int prio)
> p->sched_class = &dl_sched_class;
> else if (rt_prio(prio))
> p->sched_class = &rt_sched_class;
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + else if (task_should_scx(p))
> + p->sched_class = &ext_sched_class;
> +#endif
> else
> p->sched_class = &fair_sched_class;
>
> @@ -9036,6 +9065,7 @@ SYSCALL_DEFINE1(sched_get_priority_max, int, policy)
> case SCHED_NORMAL:
> case SCHED_BATCH:
> case SCHED_IDLE:
> + case SCHED_EXT:
> ret = 0;
> break;
> }
> @@ -9063,6 +9093,7 @@ SYSCALL_DEFINE1(sched_get_priority_min, int, policy)
> case SCHED_NORMAL:
> case SCHED_BATCH:
> case SCHED_IDLE:
> + case SCHED_EXT:
> ret = 0;
> }
> return ret;
> @@ -9918,6 +9949,10 @@ void __init sched_init(void)
> BUG_ON(!sched_class_above(&dl_sched_class, &rt_sched_class));
> BUG_ON(!sched_class_above(&rt_sched_class, &fair_sched_class));
> BUG_ON(!sched_class_above(&fair_sched_class, &idle_sched_class));
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + BUG_ON(!sched_class_above(&fair_sched_class, &ext_sched_class));
> + BUG_ON(!sched_class_above(&ext_sched_class, &idle_sched_class));
> +#endif
>
> wait_bit_init();
>
> @@ -12046,3 +12081,38 @@ void sched_mm_cid_fork(struct task_struct *t)
> t->mm_cid_active = 1;
> }
> #endif
> +
> +#ifdef CONFIG_SCHED_CLASS_EXT
> +void sched_deq_and_put_task(struct task_struct *p, int queue_flags,
> + struct sched_enq_and_set_ctx *ctx)
> +{
> + struct rq *rq = task_rq(p);
> +
> + lockdep_assert_rq_held(rq);
> +
> + *ctx = (struct sched_enq_and_set_ctx){
> + .p = p,
> + .queue_flags = queue_flags,
> + .queued = task_on_rq_queued(p),
> + .running = task_current(rq, p),
> + };
> +
> + update_rq_clock(rq);
> + if (ctx->queued)
> + dequeue_task(rq, p, queue_flags | DEQUEUE_NOCLOCK);
> + if (ctx->running)
> + put_prev_task(rq, p);
> +}
> +
> +void sched_enq_and_set_task(struct sched_enq_and_set_ctx *ctx)
> +{
> + struct rq *rq = task_rq(ctx->p);
> +
> + lockdep_assert_rq_held(rq);
> +
> + if (ctx->queued)
> + enqueue_task(rq, ctx->p, ctx->queue_flags | ENQUEUE_NOCLOCK);
> + if (ctx->running)
> + set_next_task(rq, ctx->p);
> +}
> +#endif /* CONFIG_SCHED_CLASS_EXT */
> diff --git a/kernel/sched/debug.c b/kernel/sched/debug.c
> index 0b2340a79b65..79fac9c92a22 100644
> --- a/kernel/sched/debug.c
> +++ b/kernel/sched/debug.c
> @@ -377,6 +377,9 @@ static __init int sched_init_debug(void)
>
> debugfs_create_file("debug", 0444, debugfs_sched, NULL, &sched_debug_fops);
>
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + debugfs_create_file("ext", 0444, debugfs_sched, NULL, &sched_ext_fops);
> +#endif
> return 0;
> }
> late_initcall(sched_init_debug);
> @@ -1093,6 +1096,9 @@ void proc_sched_show_task(struct task_struct *p, struct pid_namespace *ns,
> P(dl.runtime);
> P(dl.deadline);
> }
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + __PS("ext.enabled", task_on_scx(p));
> +#endif
> #undef PN_SCHEDSTAT
> #undef P_SCHEDSTAT
>
> diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
> new file mode 100644
> index 000000000000..51d77459d208
> --- /dev/null
> +++ b/kernel/sched/ext.c
> @@ -0,0 +1,3140 @@
> +/* SPDX-License-Identifier: GPL-2.0 */
> +/*
> + * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
> + * Copyright (c) 2022 Tejun Heo <tj@xxxxxxxxxx>
> + * Copyright (c) 2022 David Vernet <dvernet@xxxxxxxx>
> + */
> +#define SCX_OP_IDX(op) (offsetof(struct sched_ext_ops, op) / sizeof(void (*)(void)))
> +
> +enum scx_internal_consts {
> + SCX_NR_ONLINE_OPS = SCX_OP_IDX(init),
> + SCX_DSP_DFL_MAX_BATCH = 32,
> +};
> +
> +enum scx_ops_enable_state {
> + SCX_OPS_PREPPING,
> + SCX_OPS_ENABLING,
> + SCX_OPS_ENABLED,
> + SCX_OPS_DISABLING,
> + SCX_OPS_DISABLED,
> +};
> +
> +/*
> + * sched_ext_entity->ops_state
> + *
> + * Used to track the task ownership between the SCX core and the BPF scheduler.
> + * State transitions look as follows:
> + *
> + * NONE -> QUEUEING -> QUEUED -> DISPATCHING
> + * ^ | |
> + * | v v
> + * \-------------------------------/
> + *
> + * QUEUEING and DISPATCHING states can be waited upon. See wait_ops_state() call
> + * sites for explanations on the conditions being waited upon and why they are
> + * safe. Transitions out of them into NONE or QUEUED must store_release and the
> + * waiters should load_acquire.
> + *
> + * Tracking scx_ops_state enables sched_ext core to reliably determine whether
> + * any given task can be dispatched by the BPF scheduler at all times and thus
> + * relaxes the requirements on the BPF scheduler. This allows the BPF scheduler
> + * to try to dispatch any task anytime regardless of its state as the SCX core
> + * can safely reject invalid dispatches.
> + */
> +enum scx_ops_state {
> + SCX_OPSS_NONE, /* owned by the SCX core */
> + SCX_OPSS_QUEUEING, /* in transit to the BPF scheduler */
> + SCX_OPSS_QUEUED, /* owned by the BPF scheduler */
> + SCX_OPSS_DISPATCHING, /* in transit back to the SCX core */
> +
> + /*
> + * QSEQ brands each QUEUED instance so that, when dispatch races
> + * dequeue/requeue, the dispatcher can tell whether it still has a claim
> + * on the task being dispatched.
> + */
> + SCX_OPSS_QSEQ_SHIFT = 2,
> + SCX_OPSS_STATE_MASK = (1LLU << SCX_OPSS_QSEQ_SHIFT) - 1,
> + SCX_OPSS_QSEQ_MASK = ~SCX_OPSS_STATE_MASK,
> +};
> +
> +/*
> + * During exit, a task may schedule after losing its PIDs. When disabling the
> + * BPF scheduler, we need to be able to iterate tasks in every state to
> + * guarantee system safety. Maintain a dedicated task list which contains every
> + * task between its fork and eventual free.
> + */
> +static DEFINE_SPINLOCK(scx_tasks_lock);
> +static LIST_HEAD(scx_tasks);
> +
> +/* ops enable/disable */
> +static struct kthread_worker *scx_ops_helper;
> +static DEFINE_MUTEX(scx_ops_enable_mutex);
> +DEFINE_STATIC_KEY_FALSE(__scx_ops_enabled);
> +DEFINE_STATIC_PERCPU_RWSEM(scx_fork_rwsem);
> +static atomic_t scx_ops_enable_state_var = ATOMIC_INIT(SCX_OPS_DISABLED);
> +static struct sched_ext_ops scx_ops;
> +static bool scx_warned_zero_slice;
> +
> +static DEFINE_STATIC_KEY_FALSE(scx_ops_enq_last);
> +static DEFINE_STATIC_KEY_FALSE(scx_ops_enq_exiting);
> +static DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_enabled);
> +
> +struct static_key_false scx_has_op[SCX_NR_ONLINE_OPS] =
> + { [0 ... SCX_NR_ONLINE_OPS-1] = STATIC_KEY_FALSE_INIT };
> +
> +static atomic_t scx_exit_type = ATOMIC_INIT(SCX_EXIT_DONE);
> +static struct scx_exit_info scx_exit_info;
> +
> +static atomic64_t scx_nr_rejected = ATOMIC64_INIT(0);
> +
> +/* idle tracking */
> +#ifdef CONFIG_SMP
> +#ifdef CONFIG_CPUMASK_OFFSTACK
> +#define CL_ALIGNED_IF_ONSTACK
> +#else
> +#define CL_ALIGNED_IF_ONSTACK __cacheline_aligned_in_smp
> +#endif
> +
> +static struct {
> + cpumask_var_t cpu;
> + cpumask_var_t smt;
> +} idle_masks CL_ALIGNED_IF_ONSTACK;
> +
> +#endif /* CONFIG_SMP */
> +
> +/*
> + * Direct dispatch marker.
> + *
> + * Non-NULL values are used for direct dispatch from enqueue path. A valid
> + * pointer points to the task currently being enqueued. An ERR_PTR value is used
> + * to indicate that direct dispatch has already happened.
> + */
> +static DEFINE_PER_CPU(struct task_struct *, direct_dispatch_task);
> +
> +/* dispatch queues */
> +static struct scx_dispatch_q __cacheline_aligned_in_smp scx_dsq_global;
> +
> +static const struct rhashtable_params dsq_hash_params = {
> + .key_len = 8,
> + .key_offset = offsetof(struct scx_dispatch_q, id),
> + .head_offset = offsetof(struct scx_dispatch_q, hash_node),
> +};
> +
> +static struct rhashtable dsq_hash;
> +static LLIST_HEAD(dsqs_to_free);
> +
> +/* dispatch buf */
> +struct scx_dsp_buf_ent {
> + struct task_struct *task;
> + u64 qseq;
> + u64 dsq_id;
> + u64 enq_flags;
> +};
> +
> +static u32 scx_dsp_max_batch;
> +static struct scx_dsp_buf_ent __percpu *scx_dsp_buf;
> +
> +struct scx_dsp_ctx {
> + struct rq *rq;
> + struct rq_flags *rf;
> + u32 buf_cursor;
> + u32 nr_tasks;
> +};
> +
> +static DEFINE_PER_CPU(struct scx_dsp_ctx, scx_dsp_ctx);
> +
> +void scx_bpf_dispatch(struct task_struct *p, u64 dsq_id, u64 slice,
> + u64 enq_flags);
> +__printf(2, 3) static void scx_ops_error_type(enum scx_exit_type type,
> + const char *fmt, ...);
> +#define scx_ops_error(fmt, args...) \
> + scx_ops_error_type(SCX_EXIT_ERROR, fmt, ##args)
> +
> +struct scx_task_iter {
> + struct sched_ext_entity cursor;
> + struct task_struct *locked;
> + struct rq *rq;
> + struct rq_flags rf;
> +};
> +
> +#define SCX_HAS_OP(op) static_branch_likely(&scx_has_op[SCX_OP_IDX(op)])
> +
> +/* if the highest set bit is N, return a mask with bits [N+1, 31] set */
> +static u32 higher_bits(u32 flags)
> +{
> + return ~((1 << fls(flags)) - 1);
> +}
> +
> +/* return the mask with only the highest bit set */
> +static u32 highest_bit(u32 flags)
> +{
> + int bit = fls(flags);
> + return bit ? 1 << (bit - 1) : 0;
> +}
> +
> +/*
> + * scx_kf_mask enforcement. Some kfuncs can only be called from specific SCX
> + * ops. When invoking SCX ops, SCX_CALL_OP[_RET]() should be used to indicate
> + * the allowed kfuncs and those kfuncs should use scx_kf_allowed() to check
> + * whether it's running from an allowed context.
> + *
> + * @mask is constant, always inline to cull the mask calculations.
> + */
> +static __always_inline void scx_kf_allow(u32 mask)
> +{
> + /* nesting is allowed only in increasing scx_kf_mask order */
> + WARN_ONCE((mask | higher_bits(mask)) & current->scx.kf_mask,
> + "invalid nesting current->scx.kf_mask=0x%x mask=0x%x\n",
> + current->scx.kf_mask, mask);
> + current->scx.kf_mask |= mask;
> +}
> +
> +static void scx_kf_disallow(u32 mask)
> +{
> + current->scx.kf_mask &= ~mask;
> +}
> +
> +#define SCX_CALL_OP(mask, op, args...) \
> +do { \
> + if (mask) { \
> + scx_kf_allow(mask); \
> + scx_ops.op(args); \
> + scx_kf_disallow(mask); \
> + } else { \
> + scx_ops.op(args); \
> + } \
> +} while (0)
> +
> +#define SCX_CALL_OP_RET(mask, op, args...) \
> +({ \
> + __typeof__(scx_ops.op(args)) __ret; \
> + if (mask) { \
> + scx_kf_allow(mask); \
> + __ret = scx_ops.op(args); \
> + scx_kf_disallow(mask); \
> + } else { \
> + __ret = scx_ops.op(args); \
> + } \
> + __ret; \
> +})
> +
> +/* @mask is constant, always inline to cull unnecessary branches */
> +static __always_inline bool scx_kf_allowed(u32 mask)
> +{
> + if (unlikely(!(current->scx.kf_mask & mask))) {
> + scx_ops_error("kfunc with mask 0x%x called from an operation only allowing 0x%x",
> + mask, current->scx.kf_mask);
> + return false;
> + }
> +
> + if (unlikely((mask & (SCX_KF_INIT | SCX_KF_SLEEPABLE)) &&
> + in_interrupt())) {
> + scx_ops_error("sleepable kfunc called from non-sleepable context");
> + return false;
> + }
> +
> + /*
> + * Enforce nesting boundaries. e.g. A kfunc which can be called from
> + * DISPATCH must not be called if we're running DEQUEUE which is nested
> + * inside ops.dispatch(). We don't need to check the SCX_KF_SLEEPABLE
> + * boundary thanks to the above in_interrupt() check.
> + */
> + if (unlikely(highest_bit(mask) == SCX_KF_DISPATCH &&
> + (current->scx.kf_mask & higher_bits(SCX_KF_DISPATCH)))) {
> + scx_ops_error("dispatch kfunc called from a nested operation");
> + return false;
> + }
> +
> + return true;
> +}
> +
> +/**
> + * scx_task_iter_init - Initialize a task iterator
> + * @iter: iterator to init
> + *
> + * Initialize @iter. Must be called with scx_tasks_lock held. Once initialized,
> + * @iter must eventually be exited with scx_task_iter_exit().
> + *
> + * scx_tasks_lock may be released between this and the first next() call or
> + * between any two next() calls. If scx_tasks_lock is released between two
> + * next() calls, the caller is responsible for ensuring that the task being
> + * iterated remains accessible either through RCU read lock or obtaining a
> + * reference count.
> + *
> + * All tasks which existed when the iteration started are guaranteed to be
> + * visited as long as they still exist.
> + */
> +static void scx_task_iter_init(struct scx_task_iter *iter)
> +{
> + lockdep_assert_held(&scx_tasks_lock);
> +
> + iter->cursor = (struct sched_ext_entity){ .flags = SCX_TASK_CURSOR };
> + list_add(&iter->cursor.tasks_node, &scx_tasks);
> + iter->locked = NULL;
> +}
> +
> +/**
> + * scx_task_iter_exit - Exit a task iterator
> + * @iter: iterator to exit
> + *
> + * Exit a previously initialized @iter. Must be called with scx_tasks_lock held.
> + * If the iterator holds a task's rq lock, that rq lock is released. See
> + * scx_task_iter_init() for details.
> + */
> +static void scx_task_iter_exit(struct scx_task_iter *iter)
> +{
> + struct list_head *cursor = &iter->cursor.tasks_node;
> +
> + lockdep_assert_held(&scx_tasks_lock);
> +
> + if (iter->locked) {
> + task_rq_unlock(iter->rq, iter->locked, &iter->rf);
> + iter->locked = NULL;
> + }
> +
> + if (list_empty(cursor))
> + return;
> +
> + list_del_init(cursor);
> +}
> +
> +/**
> + * scx_task_iter_next - Next task
> + * @iter: iterator to walk
> + *
> + * Visit the next task. See scx_task_iter_init() for details.
> + */
> +static struct task_struct *scx_task_iter_next(struct scx_task_iter *iter)
> +{
> + struct list_head *cursor = &iter->cursor.tasks_node;
> + struct sched_ext_entity *pos;
> +
> + lockdep_assert_held(&scx_tasks_lock);
> +
> + list_for_each_entry(pos, cursor, tasks_node) {
> + if (&pos->tasks_node == &scx_tasks)
> + return NULL;
> + if (!(pos->flags & SCX_TASK_CURSOR)) {
> + list_move(cursor, &pos->tasks_node);
> + return container_of(pos, struct task_struct, scx);
> + }
> + }
> +
> + /* can't happen, should always terminate at scx_tasks above */
> + BUG();
> +}
> +
> +/**
> + * scx_task_iter_next_filtered - Next non-idle task
> + * @iter: iterator to walk
> + *
> + * Visit the next non-idle task. See scx_task_iter_init() for details.
> + */
> +static struct task_struct *
> +scx_task_iter_next_filtered(struct scx_task_iter *iter)
> +{
> + struct task_struct *p;
> +
> + while ((p = scx_task_iter_next(iter))) {
> + if (!is_idle_task(p))
> + return p;
> + }
> + return NULL;
> +}
> +
> +/**
> + * scx_task_iter_next_filtered_locked - Next non-idle task with its rq locked
> + * @iter: iterator to walk
> + *
> + * Visit the next non-idle task with its rq lock held. See scx_task_iter_init()
> + * for details.
> + */
> +static struct task_struct *
> +scx_task_iter_next_filtered_locked(struct scx_task_iter *iter)
> +{
> + struct task_struct *p;
> +
> + if (iter->locked) {
> + task_rq_unlock(iter->rq, iter->locked, &iter->rf);
> + iter->locked = NULL;
> + }
> +
> + p = scx_task_iter_next_filtered(iter);
> + if (!p)
> + return NULL;
> +
> + iter->rq = task_rq_lock(p, &iter->rf);
> + iter->locked = p;
> + return p;
> +}
> +
> +static enum scx_ops_enable_state scx_ops_enable_state(void)
> +{
> + return atomic_read(&scx_ops_enable_state_var);
> +}
> +
> +static enum scx_ops_enable_state
> +scx_ops_set_enable_state(enum scx_ops_enable_state to)
> +{
> + return atomic_xchg(&scx_ops_enable_state_var, to);
> +}
> +
> +static bool scx_ops_tryset_enable_state(enum scx_ops_enable_state to,
> + enum scx_ops_enable_state from)
> +{
> + int from_v = from;
> +
> + return atomic_try_cmpxchg(&scx_ops_enable_state_var, &from_v, to);
> +}
> +
> +static bool scx_ops_disabling(void)
> +{
> + return unlikely(scx_ops_enable_state() == SCX_OPS_DISABLING);
> +}
> +
> +/**
> + * wait_ops_state - Busy-wait the specified ops state to end
> + * @p: target task
> + * @opss: state to wait the end of
> + *
> + * Busy-wait for @p to transition out of @opss. This can only be used when the
> + * state part of @opss is %SCX_QUEUEING or %SCX_DISPATCHING. This function also
> + * has load_acquire semantics to ensure that the caller can see the updates made
> + * in the enqueueing and dispatching paths.
> + */
> +static void wait_ops_state(struct task_struct *p, u64 opss)
> +{
> + do {
> + cpu_relax();
> + } while (atomic64_read_acquire(&p->scx.ops_state) == opss);
> +}
> +
> +/**
> + * ops_cpu_valid - Verify a cpu number
> + * @cpu: cpu number which came from a BPF ops
> + *
> + * @cpu is a cpu number which came from the BPF scheduler and can be any value.
> + * Verify that it is in range and one of the possible cpus.
> + */
> +static bool ops_cpu_valid(s32 cpu)
> +{
> + return likely(cpu >= 0 && cpu < nr_cpu_ids && cpu_possible(cpu));
> +}
> +
> +/**
> + * ops_sanitize_err - Sanitize a -errno value
> + * @ops_name: operation to blame on failure
> + * @err: -errno value to sanitize
> + *
> + * Verify @err is a valid -errno. If not, trigger scx_ops_error() and return
> + * -%EPROTO. This is necessary because returning a rogue -errno up the chain can
> + * cause misbehaviors. For an example, a large negative return from
> + * ops.prep_enable() triggers an oops when passed up the call chain because the
> + * value fails IS_ERR() test after being encoded with ERR_PTR() and then is
> + * handled as a pointer.
> + */
> +static int ops_sanitize_err(const char *ops_name, s32 err)
> +{
> + if (err < 0 && err >= -MAX_ERRNO)
> + return err;
> +
> + scx_ops_error("ops.%s() returned an invalid errno %d", ops_name, err);
> + return -EPROTO;
> +}
> +
> +static void update_curr_scx(struct rq *rq)
> +{
> + struct task_struct *curr = rq->curr;
> + u64 now = rq_clock_task(rq);
> + u64 delta_exec;
> +
> + if (time_before_eq64(now, curr->se.exec_start))
> + return;
> +
> + delta_exec = now - curr->se.exec_start;
> + curr->se.exec_start = now;
> + curr->se.sum_exec_runtime += delta_exec;
> + account_group_exec_runtime(curr, delta_exec);
> + cgroup_account_cputime(curr, delta_exec);
> +
> + curr->scx.slice -= min(curr->scx.slice, delta_exec);
> +}
> +
> +static void dispatch_enqueue(struct scx_dispatch_q *dsq, struct task_struct *p,
> + u64 enq_flags)
> +{
> + bool is_local = dsq->id == SCX_DSQ_LOCAL;
> +
> + WARN_ON_ONCE(p->scx.dsq || !list_empty(&p->scx.dsq_node));
> +
> + if (!is_local) {
> + raw_spin_lock(&dsq->lock);
> + if (unlikely(dsq->id == SCX_DSQ_INVALID)) {
> + scx_ops_error("attempting to dispatch to a destroyed dsq");
> + /* fall back to the global dsq */
> + raw_spin_unlock(&dsq->lock);
> + dsq = &scx_dsq_global;
> + raw_spin_lock(&dsq->lock);
> + }
> + }
> +
> + if (enq_flags & SCX_ENQ_HEAD)
> + list_add(&p->scx.dsq_node, &dsq->fifo);
> + else
> + list_add_tail(&p->scx.dsq_node, &dsq->fifo);
> + dsq->nr++;
> + p->scx.dsq = dsq;
> +
> + /*
> + * We're transitioning out of QUEUEING or DISPATCHING. store_release to
> + * match waiters' load_acquire.
> + */
> + if (enq_flags & SCX_ENQ_CLEAR_OPSS)
> + atomic64_set_release(&p->scx.ops_state, SCX_OPSS_NONE);
> +
> + if (is_local) {
> + struct rq *rq = container_of(dsq, struct rq, scx.local_dsq);
> +
> + if (sched_class_above(&ext_sched_class, rq->curr->sched_class))
> + resched_curr(rq);
> + } else {
> + raw_spin_unlock(&dsq->lock);
> + }
> +}
> +
> +static void dispatch_dequeue(struct scx_rq *scx_rq, struct task_struct *p)
> +{
> + struct scx_dispatch_q *dsq = p->scx.dsq;
> + bool is_local = dsq == &scx_rq->local_dsq;
> +
> + if (!dsq) {
> + WARN_ON_ONCE(!list_empty(&p->scx.dsq_node));
> + /*
> + * When dispatching directly from the BPF scheduler to a local
> + * DSQ, the task isn't associated with any DSQ but
> + * @p->scx.holding_cpu may be set under the protection of
> + * %SCX_OPSS_DISPATCHING.
> + */
> + if (p->scx.holding_cpu >= 0)
> + p->scx.holding_cpu = -1;
> + return;
> + }
> +
> + if (!is_local)
> + raw_spin_lock(&dsq->lock);
> +
> + /*
> + * Now that we hold @dsq->lock, @p->holding_cpu and @p->scx.dsq_node
> + * can't change underneath us.
> + */
> + if (p->scx.holding_cpu < 0) {
> + /* @p must still be on @dsq, dequeue */
> + WARN_ON_ONCE(list_empty(&p->scx.dsq_node));
> + list_del_init(&p->scx.dsq_node);
> + dsq->nr--;
> + } else {
> + /*
> + * We're racing against dispatch_to_local_dsq() which already
> + * removed @p from @dsq and set @p->scx.holding_cpu. Clear the
> + * holding_cpu which tells dispatch_to_local_dsq() that it lost
> + * the race.
> + */
> + WARN_ON_ONCE(!list_empty(&p->scx.dsq_node));
> + p->scx.holding_cpu = -1;
> + }
> + p->scx.dsq = NULL;
> +
> + if (!is_local)
> + raw_spin_unlock(&dsq->lock);
> +}
> +
> +static struct scx_dispatch_q *find_non_local_dsq(u64 dsq_id)
> +{
> + lockdep_assert(rcu_read_lock_any_held());
> +
> + if (dsq_id == SCX_DSQ_GLOBAL)
> + return &scx_dsq_global;
> + else
> + return rhashtable_lookup_fast(&dsq_hash, &dsq_id,
> + dsq_hash_params);
> +}
> +
> +static struct scx_dispatch_q *find_dsq_for_dispatch(struct rq *rq, u64 dsq_id,
> + struct task_struct *p)
> +{
> + struct scx_dispatch_q *dsq;
> +
> + if (dsq_id == SCX_DSQ_LOCAL)
> + return &rq->scx.local_dsq;
> +
> + dsq = find_non_local_dsq(dsq_id);
> + if (unlikely(!dsq)) {
> + scx_ops_error("non-existent DSQ 0x%llx for %s[%d]",
> + dsq_id, p->comm, p->pid);
> + return &scx_dsq_global;
> + }
> +
> + return dsq;
> +}
> +
> +static void direct_dispatch(struct task_struct *ddsp_task, struct task_struct *p,
> + u64 dsq_id, u64 enq_flags)
> +{
> + struct scx_dispatch_q *dsq;
> +
> + /* @p must match the task which is being enqueued */
> + if (unlikely(p != ddsp_task)) {
> + if (IS_ERR(ddsp_task))
> + scx_ops_error("%s[%d] already direct-dispatched",
> + p->comm, p->pid);
> + else
> + scx_ops_error("enqueueing %s[%d] but trying to direct-dispatch %s[%d]",
> + ddsp_task->comm, ddsp_task->pid,
> + p->comm, p->pid);
> + return;
> + }
> +
> + /*
> + * %SCX_DSQ_LOCAL_ON is not supported during direct dispatch because
> + * dispatching to the local DSQ of a different CPU requires unlocking
> + * the current rq which isn't allowed in the enqueue path. Use
> + * ops.select_cpu() to be on the target CPU and then %SCX_DSQ_LOCAL.
> + */
> + if (unlikely((dsq_id & SCX_DSQ_LOCAL_ON) == SCX_DSQ_LOCAL_ON)) {
> + scx_ops_error("SCX_DSQ_LOCAL_ON can't be used for direct-dispatch");
> + return;
> + }
> +
> + dsq = find_dsq_for_dispatch(task_rq(p), dsq_id, p);
> + dispatch_enqueue(dsq, p, enq_flags | SCX_ENQ_CLEAR_OPSS);
> +
> + /*
> + * Mark that dispatch already happened by spoiling direct_dispatch_task
> + * with a non-NULL value which can never match a valid task pointer.
> + */
> + __this_cpu_write(direct_dispatch_task, ERR_PTR(-ESRCH));
> +}
> +
> +static bool test_rq_online(struct rq *rq)
> +{
> +#ifdef CONFIG_SMP
> + return rq->online;
> +#else
> + return true;
> +#endif
> +}
> +
> +static void do_enqueue_task(struct rq *rq, struct task_struct *p, u64 enq_flags,
> + int sticky_cpu)
> +{
> + struct task_struct **ddsp_taskp;
> + u64 qseq;
> +
> + WARN_ON_ONCE(!(p->scx.flags & SCX_TASK_QUEUED));
> +
> + if (p->scx.flags & SCX_TASK_ENQ_LOCAL) {
> + enq_flags |= SCX_ENQ_LOCAL;
> + p->scx.flags &= ~SCX_TASK_ENQ_LOCAL;
> + }
> +
> + /* rq migration */
> + if (sticky_cpu == cpu_of(rq))
> + goto local_norefill;
> +
> + /*
> + * If !rq->online, we already told the BPF scheduler that the CPU is
> + * offline. We're just trying to on/offline the CPU. Don't bother the
> + * BPF scheduler.
> + */
> + if (unlikely(!test_rq_online(rq)))
> + goto local;
> +
> + /* see %SCX_OPS_ENQ_EXITING */
> + if (!static_branch_unlikely(&scx_ops_enq_exiting) &&
> + unlikely(p->flags & PF_EXITING))
> + goto local;
> +
> + /* see %SCX_OPS_ENQ_LAST */
> + if (!static_branch_unlikely(&scx_ops_enq_last) &&
> + (enq_flags & SCX_ENQ_LAST))
> + goto local;
> +
> + if (!SCX_HAS_OP(enqueue)) {
> + if (enq_flags & SCX_ENQ_LOCAL)
> + goto local;
> + else
> + goto global;
> + }
> +
> + /* DSQ bypass didn't trigger, enqueue on the BPF scheduler */
> + qseq = rq->scx.ops_qseq++ << SCX_OPSS_QSEQ_SHIFT;
> +
> + WARN_ON_ONCE(atomic64_read(&p->scx.ops_state) != SCX_OPSS_NONE);
> + atomic64_set(&p->scx.ops_state, SCX_OPSS_QUEUEING | qseq);
> +
> + ddsp_taskp = this_cpu_ptr(&direct_dispatch_task);
> + WARN_ON_ONCE(*ddsp_taskp);
> + *ddsp_taskp = p;
> +
> + SCX_CALL_OP(SCX_KF_ENQUEUE, enqueue, p, enq_flags);
> +
> + /*
> + * If not directly dispatched, QUEUEING isn't clear yet and dispatch or
> + * dequeue may be waiting. The store_release matches their load_acquire.
> + */
> + if (*ddsp_taskp == p)
> + atomic64_set_release(&p->scx.ops_state, SCX_OPSS_QUEUED | qseq);
> + *ddsp_taskp = NULL;
> + return;
> +
> +local:
> + p->scx.slice = SCX_SLICE_DFL;
> +local_norefill:
> + dispatch_enqueue(&rq->scx.local_dsq, p, enq_flags);
> + return;
> +
> +global:
> + p->scx.slice = SCX_SLICE_DFL;
> + dispatch_enqueue(&scx_dsq_global, p, enq_flags);
> +}
> +
> +static void enqueue_task_scx(struct rq *rq, struct task_struct *p, int enq_flags)
> +{
> + int sticky_cpu = p->scx.sticky_cpu;
> +
> + enq_flags |= rq->scx.extra_enq_flags;
> +
> + if (sticky_cpu >= 0)
> + p->scx.sticky_cpu = -1;
> +
> + /*
> + * Restoring a running task will be immediately followed by
> + * set_next_task_scx() which expects the task to not be on the BPF
> + * scheduler as tasks can only start running through local DSQs. Force
> + * direct-dispatch into the local DSQ by setting the sticky_cpu.
> + */
> + if (unlikely(enq_flags & ENQUEUE_RESTORE) && task_current(rq, p))
> + sticky_cpu = cpu_of(rq);
> +
> + if (p->scx.flags & SCX_TASK_QUEUED)
> + return;
> +
> + p->scx.flags |= SCX_TASK_QUEUED;
> + rq->scx.nr_running++;
> + add_nr_running(rq, 1);
> +
> + do_enqueue_task(rq, p, enq_flags, sticky_cpu);
> +}
> +
> +static void ops_dequeue(struct task_struct *p, u64 deq_flags)
> +{
> + u64 opss;
> +
> + /* acquire ensures that we see the preceding updates on QUEUED */
> + opss = atomic64_read_acquire(&p->scx.ops_state);
> +
> + switch (opss & SCX_OPSS_STATE_MASK) {
> + case SCX_OPSS_NONE:
> + break;
> + case SCX_OPSS_QUEUEING:
> + /*
> + * QUEUEING is started and finished while holding @p's rq lock.
> + * As we're holding the rq lock now, we shouldn't see QUEUEING.
> + */
> + BUG();
> + case SCX_OPSS_QUEUED:
> + if (SCX_HAS_OP(dequeue))
> + SCX_CALL_OP(SCX_KF_REST, dequeue, p, deq_flags);
> +
> + if (atomic64_try_cmpxchg(&p->scx.ops_state, &opss,
> + SCX_OPSS_NONE))
> + break;
> + fallthrough;
> + case SCX_OPSS_DISPATCHING:
> + /*
> + * If @p is being dispatched from the BPF scheduler to a DSQ,
> + * wait for the transfer to complete so that @p doesn't get
> + * added to its DSQ after dequeueing is complete.
> + *
> + * As we're waiting on DISPATCHING with the rq locked, the
> + * dispatching side shouldn't try to lock the rq while
> + * DISPATCHING is set. See dispatch_to_local_dsq().
> + *
> + * DISPATCHING shouldn't have qseq set and control can reach
> + * here with NONE @opss from the above QUEUED case block.
> + * Explicitly wait on %SCX_OPSS_DISPATCHING instead of @opss.
> + */
> + wait_ops_state(p, SCX_OPSS_DISPATCHING);
> + BUG_ON(atomic64_read(&p->scx.ops_state) != SCX_OPSS_NONE);
> + break;
> + }
> +}
> +
> +static void dequeue_task_scx(struct rq *rq, struct task_struct *p, int deq_flags)
> +{
> + struct scx_rq *scx_rq = &rq->scx;
> +
> + if (!(p->scx.flags & SCX_TASK_QUEUED))
> + return;
> +
> + ops_dequeue(p, deq_flags);
> +
> + if (deq_flags & SCX_DEQ_SLEEP)
> + p->scx.flags |= SCX_TASK_DEQD_FOR_SLEEP;
> + else
> + p->scx.flags &= ~SCX_TASK_DEQD_FOR_SLEEP;
> +
> + p->scx.flags &= ~SCX_TASK_QUEUED;
> + scx_rq->nr_running--;
> + sub_nr_running(rq, 1);
> +
> + dispatch_dequeue(scx_rq, p);
> +}
> +
> +static void yield_task_scx(struct rq *rq)
> +{
> + struct task_struct *p = rq->curr;
> +
> + if (SCX_HAS_OP(yield))
> + SCX_CALL_OP_RET(SCX_KF_REST, yield, p, NULL);
> + else
> + p->scx.slice = 0;
> +}
> +
> +static bool yield_to_task_scx(struct rq *rq, struct task_struct *to)
> +{
> + struct task_struct *from = rq->curr;
> +
> + if (SCX_HAS_OP(yield))
> + return SCX_CALL_OP_RET(SCX_KF_REST, yield, from, to);
> + else
> + return false;
> +}
> +
> +#ifdef CONFIG_SMP
> +/**
> + * move_task_to_local_dsq - Move a task from a different rq to a local DSQ
> + * @rq: rq to move the task into, currently locked
> + * @p: task to move
> + * @enq_flags: %SCX_ENQ_*
> + *
> + * Move @p which is currently on a different rq to @rq's local DSQ. The caller
> + * must:
> + *
> + * 1. Start with exclusive access to @p either through its DSQ lock or
> + * %SCX_OPSS_DISPATCHING flag.
> + *
> + * 2. Set @p->scx.holding_cpu to raw_smp_processor_id().
> + *
> + * 3. Remember task_rq(@p). Release the exclusive access so that we don't
> + * deadlock with dequeue.
> + *
> + * 4. Lock @rq and the task_rq from #3.
> + *
> + * 5. Call this function.
> + *
> + * Returns %true if @p was successfully moved. %false after racing dequeue and
> + * losing.
> + */
> +static bool move_task_to_local_dsq(struct rq *rq, struct task_struct *p,
> + u64 enq_flags)
> +{
> + struct rq *task_rq;
> +
> + lockdep_assert_rq_held(rq);
> +
> + /*
> + * If dequeue got to @p while we were trying to lock both rq's, it'd
> + * have cleared @p->scx.holding_cpu to -1. While other cpus may have
> + * updated it to different values afterwards, as this operation can't be
> + * preempted or recurse, @p->scx.holding_cpu can never become
> + * raw_smp_processor_id() again before we're done. Thus, we can tell
> + * whether we lost to dequeue by testing whether @p->scx.holding_cpu is
> + * still raw_smp_processor_id().
> + *
> + * See dispatch_dequeue() for the counterpart.
> + */
> + if (unlikely(p->scx.holding_cpu != raw_smp_processor_id()))
> + return false;
> +
> + /* @p->rq couldn't have changed if we're still the holding cpu */
> + task_rq = task_rq(p);
> + lockdep_assert_rq_held(task_rq);
> +
> + WARN_ON_ONCE(!cpumask_test_cpu(cpu_of(rq), p->cpus_ptr));
> + deactivate_task(task_rq, p, 0);
> + set_task_cpu(p, cpu_of(rq));
> + p->scx.sticky_cpu = cpu_of(rq);
> +
> + /*
> + * We want to pass scx-specific enq_flags but activate_task() will
> + * truncate the upper 32 bit. As we own @rq, we can pass them through
> + * @rq->scx.extra_enq_flags instead.
> + */
> + WARN_ON_ONCE(rq->scx.extra_enq_flags);
> + rq->scx.extra_enq_flags = enq_flags;
> + activate_task(rq, p, 0);
> + rq->scx.extra_enq_flags = 0;
> +
> + return true;
> +}
> +
> +/**
> + * dispatch_to_local_dsq_lock - Ensure source and desitnation rq's are locked
> + * @rq: current rq which is locked
> + * @rf: rq_flags to use when unlocking @rq
> + * @src_rq: rq to move task from
> + * @dst_rq: rq to move task to
> + *
> + * We're holding @rq lock and trying to dispatch a task from @src_rq to
> + * @dst_rq's local DSQ and thus need to lock both @src_rq and @dst_rq. Whether
> + * @rq stays locked isn't important as long as the state is restored after
> + * dispatch_to_local_dsq_unlock().
> + */
> +static void dispatch_to_local_dsq_lock(struct rq *rq, struct rq_flags *rf,
> + struct rq *src_rq, struct rq *dst_rq)
> +{
> + rq_unpin_lock(rq, rf);
> +
> + if (src_rq == dst_rq) {
> + raw_spin_rq_unlock(rq);
> + raw_spin_rq_lock(dst_rq);
> + } else if (rq == src_rq) {
> + double_lock_balance(rq, dst_rq);
> + rq_repin_lock(rq, rf);
> + } else if (rq == dst_rq) {
> + double_lock_balance(rq, src_rq);
> + rq_repin_lock(rq, rf);
> + } else {
> + raw_spin_rq_unlock(rq);
> + double_rq_lock(src_rq, dst_rq);
> + }
> +}
> +
> +/**
> + * dispatch_to_local_dsq_unlock - Undo dispatch_to_local_dsq_lock()
> + * @rq: current rq which is locked
> + * @rf: rq_flags to use when unlocking @rq
> + * @src_rq: rq to move task from
> + * @dst_rq: rq to move task to
> + *
> + * Unlock @src_rq and @dst_rq and ensure that @rq is locked on return.
> + */
> +static void dispatch_to_local_dsq_unlock(struct rq *rq, struct rq_flags *rf,
> + struct rq *src_rq, struct rq *dst_rq)
> +{
> + if (src_rq == dst_rq) {
> + raw_spin_rq_unlock(dst_rq);
> + raw_spin_rq_lock(rq);
> + rq_repin_lock(rq, rf);
> + } else if (rq == src_rq) {
> + double_unlock_balance(rq, dst_rq);
> + } else if (rq == dst_rq) {
> + double_unlock_balance(rq, src_rq);
> + } else {
> + double_rq_unlock(src_rq, dst_rq);
> + raw_spin_rq_lock(rq);
> + rq_repin_lock(rq, rf);
> + }
> +}
> +#endif /* CONFIG_SMP */
> +
> +
> +static bool consume_dispatch_q(struct rq *rq, struct rq_flags *rf,
> + struct scx_dispatch_q *dsq)
> +{
> + struct scx_rq *scx_rq = &rq->scx;
> + struct task_struct *p;
> + struct rq *task_rq;
> + bool moved = false;
> +retry:
> + if (list_empty(&dsq->fifo))
> + return false;
> +
> + raw_spin_lock(&dsq->lock);
> + list_for_each_entry(p, &dsq->fifo, scx.dsq_node) {
> + task_rq = task_rq(p);
> + if (rq == task_rq)
> + goto this_rq;
> + if (likely(test_rq_online(rq)) && !is_migration_disabled(p) &&
> + cpumask_test_cpu(cpu_of(rq), p->cpus_ptr))
> + goto remote_rq;
> + }
> + raw_spin_unlock(&dsq->lock);
> + return false;
> +
> +this_rq:
> + /* @dsq is locked and @p is on this rq */
> + WARN_ON_ONCE(p->scx.holding_cpu >= 0);
> + list_move_tail(&p->scx.dsq_node, &scx_rq->local_dsq.fifo);
> + dsq->nr--;
> + scx_rq->local_dsq.nr++;
> + p->scx.dsq = &scx_rq->local_dsq;
> + raw_spin_unlock(&dsq->lock);
> + return true;
> +
> +remote_rq:
> +#ifdef CONFIG_SMP
> + /*
> + * @dsq is locked and @p is on a remote rq. @p is currently protected by
> + * @dsq->lock. We want to pull @p to @rq but may deadlock if we grab
> + * @task_rq while holding @dsq and @rq locks. As dequeue can't drop the
> + * rq lock or fail, do a little dancing from our side. See
> + * move_task_to_local_dsq().
> + */
> + WARN_ON_ONCE(p->scx.holding_cpu >= 0);
> + list_del_init(&p->scx.dsq_node);
> + dsq->nr--;
> + p->scx.holding_cpu = raw_smp_processor_id();
> + raw_spin_unlock(&dsq->lock);
> +
> + rq_unpin_lock(rq, rf);
> + double_lock_balance(rq, task_rq);
> + rq_repin_lock(rq, rf);
> +
> + moved = move_task_to_local_dsq(rq, p, 0);
> +
> + double_unlock_balance(rq, task_rq);
> +#endif /* CONFIG_SMP */
> + if (likely(moved))
> + return true;
> + goto retry;
> +}
> +
> +enum dispatch_to_local_dsq_ret {
> + DTL_DISPATCHED, /* successfully dispatched */
> + DTL_LOST, /* lost race to dequeue */
> + DTL_NOT_LOCAL, /* destination is not a local DSQ */
> + DTL_INVALID, /* invalid local dsq_id */
> +};
> +
> +/**
> + * dispatch_to_local_dsq - Dispatch a task to a local dsq
> + * @rq: current rq which is locked
> + * @rf: rq_flags to use when unlocking @rq
> + * @dsq_id: destination dsq ID
> + * @p: task to dispatch
> + * @enq_flags: %SCX_ENQ_*
> + *
> + * We're holding @rq lock and want to dispatch @p to the local DSQ identified by
> + * @dsq_id. This function performs all the synchronization dancing needed
> + * because local DSQs are protected with rq locks.
> + *
> + * The caller must have exclusive ownership of @p (e.g. through
> + * %SCX_OPSS_DISPATCHING).
> + */
> +static enum dispatch_to_local_dsq_ret
> +dispatch_to_local_dsq(struct rq *rq, struct rq_flags *rf, u64 dsq_id,
> + struct task_struct *p, u64 enq_flags)
> +{
> + struct rq *src_rq = task_rq(p);
> + struct rq *dst_rq;
> +
> + /*
> + * We're synchronized against dequeue through DISPATCHING. As @p can't
> + * be dequeued, its task_rq and cpus_allowed are stable too.
> + */
> + if (dsq_id == SCX_DSQ_LOCAL) {
> + dst_rq = rq;
> + } else if ((dsq_id & SCX_DSQ_LOCAL_ON) == SCX_DSQ_LOCAL_ON) {
> + s32 cpu = dsq_id & SCX_DSQ_LOCAL_CPU_MASK;
> +
> + if (!ops_cpu_valid(cpu)) {
> + scx_ops_error("invalid cpu %d in SCX_DSQ_LOCAL_ON verdict for %s[%d]",
> + cpu, p->comm, p->pid);
> + return DTL_INVALID;
> + }
> + dst_rq = cpu_rq(cpu);
> + } else {
> + return DTL_NOT_LOCAL;
> + }
> +
> + /* if dispatching to @rq that @p is already on, no lock dancing needed */
> + if (rq == src_rq && rq == dst_rq) {
> + dispatch_enqueue(&dst_rq->scx.local_dsq, p,
> + enq_flags | SCX_ENQ_CLEAR_OPSS);
> + return DTL_DISPATCHED;
> + }
> +
> +#ifdef CONFIG_SMP
> + if (cpumask_test_cpu(cpu_of(dst_rq), p->cpus_ptr)) {
> + struct rq *locked_dst_rq = dst_rq;
> + bool dsp;
> +
> + /*
> + * @p is on a possibly remote @src_rq which we need to lock to
> + * move the task. If dequeue is in progress, it'd be locking
> + * @src_rq and waiting on DISPATCHING, so we can't grab @src_rq
> + * lock while holding DISPATCHING.
> + *
> + * As DISPATCHING guarantees that @p is wholly ours, we can
> + * pretend that we're moving from a DSQ and use the same
> + * mechanism - mark the task under transfer with holding_cpu,
> + * release DISPATCHING and then follow the same protocol.
> + */
> + p->scx.holding_cpu = raw_smp_processor_id();
> +
> + /* store_release ensures that dequeue sees the above */
> + atomic64_set_release(&p->scx.ops_state, SCX_OPSS_NONE);
> +
> + dispatch_to_local_dsq_lock(rq, rf, src_rq, locked_dst_rq);
> +
> + /*
> + * We don't require the BPF scheduler to avoid dispatching to
> + * offline CPUs mostly for convenience but also because CPUs can
> + * go offline between scx_bpf_dispatch() calls and here. If @p
> + * is destined to an offline CPU, queue it on its current CPU
> + * instead, which should always be safe. As this is an allowed
> + * behavior, don't trigger an ops error.
> + */
> + if (unlikely(!test_rq_online(dst_rq)))
> + dst_rq = src_rq;
> +
> + if (src_rq == dst_rq) {
> + /*
> + * As @p is staying on the same rq, there's no need to
> + * go through the full deactivate/activate cycle.
> + * Optimize by abbreviating the operations in
> + * move_task_to_local_dsq().
> + */
> + dsp = p->scx.holding_cpu == raw_smp_processor_id();
> + if (likely(dsp)) {
> + p->scx.holding_cpu = -1;
> + dispatch_enqueue(&dst_rq->scx.local_dsq, p,
> + enq_flags);
> + }
> + } else {
> + dsp = move_task_to_local_dsq(dst_rq, p, enq_flags);
> + }
> +
> + /* if the destination CPU is idle, wake it up */
> + if (dsp && p->sched_class > dst_rq->curr->sched_class)
> + resched_curr(dst_rq);
> +
> + dispatch_to_local_dsq_unlock(rq, rf, src_rq, locked_dst_rq);
> +
> + return dsp ? DTL_DISPATCHED : DTL_LOST;
> + }
> +#endif /* CONFIG_SMP */
> +
> + scx_ops_error("SCX_DSQ_LOCAL[_ON] verdict target cpu %d not allowed for %s[%d]",
> + cpu_of(dst_rq), p->comm, p->pid);
> + return DTL_INVALID;
> +}
> +
> +/**
> + * finish_dispatch - Asynchronously finish dispatching a task
> + * @rq: current rq which is locked
> + * @rf: rq_flags to use when unlocking @rq
> + * @p: task to finish dispatching
> + * @qseq_at_dispatch: qseq when @p started getting dispatched
> + * @dsq_id: destination DSQ ID
> + * @enq_flags: %SCX_ENQ_*
> + *
> + * Dispatching to local DSQs may need to wait for queueing to complete or
> + * require rq lock dancing. As we don't wanna do either while inside
> + * ops.dispatch() to avoid locking order inversion, we split dispatching into
> + * two parts. scx_bpf_dispatch() which is called by ops.dispatch() records the
> + * task and its qseq. Once ops.dispatch() returns, this function is called to
> + * finish up.
> + *
> + * There is no guarantee that @p is still valid for dispatching or even that it
> + * was valid in the first place. Make sure that the task is still owned by the
> + * BPF scheduler and claim the ownership before dispatching.
> + */
> +static void finish_dispatch(struct rq *rq, struct rq_flags *rf,
> + struct task_struct *p, u64 qseq_at_dispatch,
> + u64 dsq_id, u64 enq_flags)
> +{
> + struct scx_dispatch_q *dsq;
> + u64 opss;
> +
> +retry:
> + /*
> + * No need for _acquire here. @p is accessed only after a successful
> + * try_cmpxchg to DISPATCHING.
> + */
> + opss = atomic64_read(&p->scx.ops_state);
> +
> + switch (opss & SCX_OPSS_STATE_MASK) {
> + case SCX_OPSS_DISPATCHING:
> + case SCX_OPSS_NONE:
> + /* someone else already got to it */
> + return;
> + case SCX_OPSS_QUEUED:
> + /*
> + * If qseq doesn't match, @p has gone through at least one
> + * dispatch/dequeue and re-enqueue cycle between
> + * scx_bpf_dispatch() and here and we have no claim on it.
> + */
> + if ((opss & SCX_OPSS_QSEQ_MASK) != qseq_at_dispatch)
> + return;
> +
> + /*
> + * While we know @p is accessible, we don't yet have a claim on
> + * it - the BPF scheduler is allowed to dispatch tasks
> + * spuriously and there can be a racing dequeue attempt. Let's
> + * claim @p by atomically transitioning it from QUEUED to
> + * DISPATCHING.
> + */
> + if (likely(atomic64_try_cmpxchg(&p->scx.ops_state, &opss,
> + SCX_OPSS_DISPATCHING)))
> + break;
> + goto retry;
> + case SCX_OPSS_QUEUEING:
> + /*
> + * do_enqueue_task() is in the process of transferring the task
> + * to the BPF scheduler while holding @p's rq lock. As we aren't
> + * holding any kernel or BPF resource that the enqueue path may
> + * depend upon, it's safe to wait.
> + */
> + wait_ops_state(p, opss);
> + goto retry;
> + }
> +
> + BUG_ON(!(p->scx.flags & SCX_TASK_QUEUED));
> +
> + switch (dispatch_to_local_dsq(rq, rf, dsq_id, p, enq_flags)) {
> + case DTL_DISPATCHED:
> + break;
> + case DTL_LOST:
> + break;
> + case DTL_INVALID:
> + dsq_id = SCX_DSQ_GLOBAL;
> + fallthrough;
> + case DTL_NOT_LOCAL:
> + dsq = find_dsq_for_dispatch(cpu_rq(raw_smp_processor_id()),
> + dsq_id, p);
> + dispatch_enqueue(dsq, p, enq_flags | SCX_ENQ_CLEAR_OPSS);
> + break;
> + }
> +}
> +
> +static void flush_dispatch_buf(struct rq *rq, struct rq_flags *rf)
> +{
> + struct scx_dsp_ctx *dspc = this_cpu_ptr(&scx_dsp_ctx);
> + u32 u;
> +
> + for (u = 0; u < dspc->buf_cursor; u++) {
> + struct scx_dsp_buf_ent *ent = &this_cpu_ptr(scx_dsp_buf)[u];
> +
> + finish_dispatch(rq, rf, ent->task, ent->qseq, ent->dsq_id,
> + ent->enq_flags);
> + }
> +
> + dspc->nr_tasks += dspc->buf_cursor;
> + dspc->buf_cursor = 0;
> +}
> +
> +static int balance_scx(struct rq *rq, struct task_struct *prev,
> + struct rq_flags *rf)
> +{
> + struct scx_rq *scx_rq = &rq->scx;
> + struct scx_dsp_ctx *dspc = this_cpu_ptr(&scx_dsp_ctx);
> + bool prev_on_scx = prev->sched_class == &ext_sched_class;
> +
> + lockdep_assert_rq_held(rq);
> +
> + if (prev_on_scx) {
> + WARN_ON_ONCE(prev->scx.flags & SCX_TASK_BAL_KEEP);
> + update_curr_scx(rq);
> +
> + /*
> + * If @prev is runnable & has slice left, it has priority and
> + * fetching more just increases latency for the fetched tasks.
> + * Tell put_prev_task_scx() to put @prev on local_dsq.
> + *
> + * See scx_ops_disable_workfn() for the explanation on the
> + * disabling() test.
> + */
> + if ((prev->scx.flags & SCX_TASK_QUEUED) &&
> + prev->scx.slice && !scx_ops_disabling()) {
> + prev->scx.flags |= SCX_TASK_BAL_KEEP;
> + return 1;
> + }
> + }
> +
> + /* if there already are tasks to run, nothing to do */
> + if (scx_rq->local_dsq.nr)
> + return 1;
> +
> + if (consume_dispatch_q(rq, rf, &scx_dsq_global))
> + return 1;
> +
> + if (!SCX_HAS_OP(dispatch))
> + return 0;
> +
> + dspc->rq = rq;
> + dspc->rf = rf;
> +
> + /*
> + * The dispatch loop. Because flush_dispatch_buf() may drop the rq lock,
> + * the local DSQ might still end up empty after a successful
> + * ops.dispatch(). If the local DSQ is empty even after ops.dispatch()
> + * produced some tasks, retry. The BPF scheduler may depend on this
> + * looping behavior to simplify its implementation.
> + */
> + do {
> + dspc->nr_tasks = 0;
> +
> + SCX_CALL_OP(SCX_KF_DISPATCH, dispatch, cpu_of(rq),
> + prev_on_scx ? prev : NULL);
> +
> + flush_dispatch_buf(rq, rf);
> +
> + if (scx_rq->local_dsq.nr)
> + return 1;
> + if (consume_dispatch_q(rq, rf, &scx_dsq_global))
> + return 1;
> + } while (dspc->nr_tasks);
> +
> + return 0;
> +}
> +
> +static void set_next_task_scx(struct rq *rq, struct task_struct *p, bool first)
> +{
> + if (p->scx.flags & SCX_TASK_QUEUED) {
> + WARN_ON_ONCE(atomic64_read(&p->scx.ops_state) != SCX_OPSS_NONE);
> + dispatch_dequeue(&rq->scx, p);
> + }
> +
> + p->se.exec_start = rq_clock_task(rq);
> +}
> +
> +static void put_prev_task_scx(struct rq *rq, struct task_struct *p)
> +{
> +#ifndef CONFIG_SMP
> + /*
> + * UP workaround.
> + *
> + * Because SCX may transfer tasks across CPUs during dispatch, dispatch
> + * is performed from its balance operation which isn't called in UP.
> + * Let's work around by calling it from the operations which come right
> + * after.
> + *
> + * 1. If the prev task is on SCX, pick_next_task() calls
> + * .put_prev_task() right after. As .put_prev_task() is also called
> + * from other places, we need to distinguish the calls which can be
> + * done by looking at the previous task's state - if still queued or
> + * dequeued with %SCX_DEQ_SLEEP, the caller must be pick_next_task().
> + * This case is handled here.
> + *
> + * 2. If the prev task is not on SCX, the first following call into SCX
> + * will be .pick_next_task(), which is covered by calling
> + * balance_scx() from pick_next_task_scx().
> + *
> + * Note that we can't merge the first case into the second as
> + * balance_scx() must be called before the previous SCX task goes
> + * through put_prev_task_scx().
> + *
> + * As UP doesn't transfer tasks around, balance_scx() doesn't need @rf.
> + * Pass in %NULL.
> + */
> + if (p->scx.flags & (SCX_TASK_QUEUED | SCX_TASK_DEQD_FOR_SLEEP))
> + balance_scx(rq, p, NULL);
> +#endif
> +
> + update_curr_scx(rq);
> +
> + /*
> + * If we're being called from put_prev_task_balance(), balance_scx() may
> + * have decided that @p should keep running.
> + */
> + if (p->scx.flags & SCX_TASK_BAL_KEEP) {
> + p->scx.flags &= ~SCX_TASK_BAL_KEEP;
> + dispatch_enqueue(&rq->scx.local_dsq, p, SCX_ENQ_HEAD);
> + return;
> + }
> +
> + if (p->scx.flags & SCX_TASK_QUEUED) {
> + /*
> + * If @p has slice left and balance_scx() didn't tag it for
> + * keeping, @p is getting preempted by a higher priority
> + * scheduler class. Leave it at the head of the local DSQ.
> + */
> + if (p->scx.slice && !scx_ops_disabling()) {
> + dispatch_enqueue(&rq->scx.local_dsq, p, SCX_ENQ_HEAD);
> + return;
> + }
> +
> + /*
> + * If we're in the pick_next_task path, balance_scx() should
> + * have already populated the local DSQ if there are any other
> + * available tasks. If empty, tell ops.enqueue() that @p is the
> + * only one available for this cpu. ops.enqueue() should put it
> + * on the local DSQ so that the subsequent pick_next_task_scx()
> + * can find the task unless it wants to trigger a separate
> + * follow-up scheduling event.
> + */
> + if (list_empty(&rq->scx.local_dsq.fifo))
> + do_enqueue_task(rq, p, SCX_ENQ_LAST | SCX_ENQ_LOCAL, -1);
> + else
> + do_enqueue_task(rq, p, 0, -1);
> + }
> +}
> +
> +static struct task_struct *first_local_task(struct rq *rq)
> +{
> + return list_first_entry_or_null(&rq->scx.local_dsq.fifo,
> + struct task_struct, scx.dsq_node);
> +}
> +
> +static struct task_struct *pick_next_task_scx(struct rq *rq)
> +{
> + struct task_struct *p;
> +
> +#ifndef CONFIG_SMP
> + /* UP workaround - see the comment at the head of put_prev_task_scx() */
> + if (unlikely(rq->curr->sched_class != &ext_sched_class))
> + balance_scx(rq, rq->curr, NULL);
> +#endif
> +
> + p = first_local_task(rq);
> + if (!p)
> + return NULL;
> +
> + if (unlikely(!p->scx.slice)) {
> + if (!scx_ops_disabling() && !scx_warned_zero_slice) {
> + printk_deferred(KERN_WARNING "sched_ext: %s[%d] has zero slice in pick_next_task_scx()\n",
> + p->comm, p->pid);
> + scx_warned_zero_slice = true;
> + }
> + p->scx.slice = SCX_SLICE_DFL;
> + }
> +
> + set_next_task_scx(rq, p, true);
> +
> + return p;
> +}
> +
> +#ifdef CONFIG_SMP
> +
> +static bool test_and_clear_cpu_idle(int cpu)
> +{
> +#ifdef CONFIG_SCHED_SMT
> + /*
> + * SMT mask should be cleared whether we can claim @cpu or not. The SMT
> + * cluster is not wholly idle either way. This also prevents
> + * scx_pick_idle_cpu() from getting caught in an infinite loop.
> + */
> + if (sched_smt_active()) {
> + const struct cpumask *smt = cpu_smt_mask(cpu);
> +
> + /*
> + * If offline, @cpu is not its own sibling and
> + * scx_pick_idle_cpu() can get caught in an infinite loop as
> + * @cpu is never cleared from idle_masks.smt. Ensure that @cpu
> + * is eventually cleared.
> + */
> + if (cpumask_intersects(smt, idle_masks.smt))
> + cpumask_andnot(idle_masks.smt, idle_masks.smt, smt);
> + else if (cpumask_test_cpu(cpu, idle_masks.smt))
> + __cpumask_clear_cpu(cpu, idle_masks.smt);
> + }
> +#endif
> + return cpumask_test_and_clear_cpu(cpu, idle_masks.cpu);
> +}
> +
> +static s32 scx_pick_idle_cpu(const struct cpumask *cpus_allowed, u64 flags)
> +{
> + int cpu;
> +
> +retry:
> + if (sched_smt_active()) {
> + cpu = cpumask_any_and_distribute(idle_masks.smt, cpus_allowed);
> + if (cpu < nr_cpu_ids)
> + goto found;
> +
> + if (flags & SCX_PICK_IDLE_CORE)
> + return -EBUSY;
> + }
> +
> + cpu = cpumask_any_and_distribute(idle_masks.cpu, cpus_allowed);
> + if (cpu >= nr_cpu_ids)
> + return -EBUSY;
> +
> +found:
> + if (test_and_clear_cpu_idle(cpu))
> + return cpu;
> + else
> + goto retry;
> +}
> +
> +static s32 scx_select_cpu_dfl(struct task_struct *p, s32 prev_cpu, u64 wake_flags)
> +{
> + s32 cpu;
> +
> + if (!static_branch_likely(&scx_builtin_idle_enabled)) {
> + scx_ops_error("built-in idle tracking is disabled");
> + return prev_cpu;
> + }
> +
> + /*
> + * If WAKE_SYNC and the machine isn't fully saturated, wake up @p to the
> + * local DSQ of the waker.
> + */
> + if ((wake_flags & SCX_WAKE_SYNC) && p->nr_cpus_allowed > 1 &&
> + !cpumask_empty(idle_masks.cpu) && !(current->flags & PF_EXITING)) {
> + cpu = smp_processor_id();
> + if (cpumask_test_cpu(cpu, p->cpus_ptr)) {
> + p->scx.flags |= SCX_TASK_ENQ_LOCAL;
> + return cpu;
> + }
> + }
> +
> + if (p->nr_cpus_allowed == 1)
> + return prev_cpu;
> +
> + /*
> + * If CPU has SMT, any wholly idle CPU is likely a better pick than
> + * partially idle @prev_cpu.
> + */
> + if (sched_smt_active()) {
> + if (cpumask_test_cpu(prev_cpu, idle_masks.smt) &&
> + test_and_clear_cpu_idle(prev_cpu)) {
> + p->scx.flags |= SCX_TASK_ENQ_LOCAL;
> + return prev_cpu;
> + }
> +
> + cpu = scx_pick_idle_cpu(p->cpus_ptr, SCX_PICK_IDLE_CORE);
> + if (cpu >= 0) {
> + p->scx.flags |= SCX_TASK_ENQ_LOCAL;
> + return cpu;
> + }
> + }
> +
> + if (test_and_clear_cpu_idle(prev_cpu)) {
> + p->scx.flags |= SCX_TASK_ENQ_LOCAL;
> + return prev_cpu;
> + }
> +
> + cpu = scx_pick_idle_cpu(p->cpus_ptr, 0);
> + if (cpu >= 0) {
> + p->scx.flags |= SCX_TASK_ENQ_LOCAL;
> + return cpu;
> + }
> +
> + return prev_cpu;
> +}
> +
> +static int select_task_rq_scx(struct task_struct *p, int prev_cpu, int wake_flags)
> +{
> + if (SCX_HAS_OP(select_cpu)) {
> + s32 cpu;
> +
> + cpu = SCX_CALL_OP_RET(SCX_KF_REST, select_cpu, p, prev_cpu,
> + wake_flags);
> + if (ops_cpu_valid(cpu)) {
> + return cpu;
> + } else {
> + scx_ops_error("select_cpu returned invalid cpu %d", cpu);
> + return prev_cpu;
> + }
> + } else {
> + return scx_select_cpu_dfl(p, prev_cpu, wake_flags);
> + }
> +}
> +
> +static void set_cpus_allowed_scx(struct task_struct *p,
> + struct affinity_context *ac)
> +{
> + set_cpus_allowed_common(p, ac);
> +
> + /*
> + * The effective cpumask is stored in @p->cpus_ptr which may temporarily
> + * differ from the configured one in @p->cpus_mask. Always tell the bpf
> + * scheduler the effective one.
> + *
> + * Fine-grained memory write control is enforced by BPF making the const
> + * designation pointless. Cast it away when calling the operation.
> + */
> + if (SCX_HAS_OP(set_cpumask))
> + SCX_CALL_OP(SCX_KF_REST, set_cpumask, p,
> + (struct cpumask *)p->cpus_ptr);
> +}
> +
> +static void reset_idle_masks(void)
> +{
> + /* consider all cpus idle, should converge to the actual state quickly */
> + cpumask_setall(idle_masks.cpu);
> + cpumask_setall(idle_masks.smt);
> +}
> +
> +void __scx_update_idle(struct rq *rq, bool idle)
> +{
> + int cpu = cpu_of(rq);
> +
> + if (SCX_HAS_OP(update_idle)) {
> + SCX_CALL_OP(SCX_KF_REST, update_idle, cpu_of(rq), idle);
> + if (!static_branch_unlikely(&scx_builtin_idle_enabled))
> + return;
> + }
> +
> + if (idle)
> + cpumask_set_cpu(cpu, idle_masks.cpu);
> + else
> + cpumask_clear_cpu(cpu, idle_masks.cpu);
> +
> +#ifdef CONFIG_SCHED_SMT
> + if (sched_smt_active()) {
> + const struct cpumask *smt = cpu_smt_mask(cpu);
> +
> + if (idle) {
> + /*
> + * idle_masks.smt handling is racy but that's fine as
> + * it's only for optimization and self-correcting.
> + */
> + for_each_cpu(cpu, smt) {
> + if (!cpumask_test_cpu(cpu, idle_masks.cpu))
> + return;
> + }
> + cpumask_or(idle_masks.smt, idle_masks.smt, smt);
> + } else {
> + cpumask_andnot(idle_masks.smt, idle_masks.smt, smt);
> + }
> + }
> +#endif
> +}
> +
> +#else /* !CONFIG_SMP */
> +
> +static bool test_and_clear_cpu_idle(int cpu) { return false; }
> +static s32 scx_pick_idle_cpu(const struct cpumask *cpus_allowed, u64 flags) { return -EBUSY; }
> +static void reset_idle_masks(void) {}
> +
> +#endif /* CONFIG_SMP */
> +
> +static void task_tick_scx(struct rq *rq, struct task_struct *curr, int queued)
> +{
> + update_curr_scx(rq);
> +
> + /*
> + * While disabling, always resched as we can't trust the slice
> + * management.
> + */
> + if (scx_ops_disabling())
> + curr->scx.slice = 0;
> +
> + if (!curr->scx.slice)
> + resched_curr(rq);
> +}
> +
> +static int scx_ops_prepare_task(struct task_struct *p, struct task_group *tg)
> +{
> + int ret;
> +
> + WARN_ON_ONCE(p->scx.flags & SCX_TASK_OPS_PREPPED);
> +
> + if (SCX_HAS_OP(prep_enable)) {
> + struct scx_enable_args args = { };
> +
> + ret = SCX_CALL_OP_RET(SCX_KF_SLEEPABLE, prep_enable, p, &args);
> + if (unlikely(ret)) {
> + ret = ops_sanitize_err("prep_enable", ret);
> + return ret;
> + }
> + }
> +
> + p->scx.flags |= SCX_TASK_OPS_PREPPED;
> + return 0;
> +}
> +
> +static void scx_ops_enable_task(struct task_struct *p)
> +{
> + lockdep_assert_rq_held(task_rq(p));
> + WARN_ON_ONCE(!(p->scx.flags & SCX_TASK_OPS_PREPPED));
> +
> + if (SCX_HAS_OP(enable)) {
> + struct scx_enable_args args = { };
> + SCX_CALL_OP(SCX_KF_REST, enable, p, &args);
> + }
> + p->scx.flags &= ~SCX_TASK_OPS_PREPPED;
> + p->scx.flags |= SCX_TASK_OPS_ENABLED;
> +}
> +
> +static void scx_ops_disable_task(struct task_struct *p)
> +{
> + lockdep_assert_rq_held(task_rq(p));
> +
> + if (p->scx.flags & SCX_TASK_OPS_PREPPED) {
> + if (SCX_HAS_OP(cancel_enable)) {
> + struct scx_enable_args args = { };
> + SCX_CALL_OP(SCX_KF_REST, cancel_enable, p, &args);
> + }
> + p->scx.flags &= ~SCX_TASK_OPS_PREPPED;
> + } else if (p->scx.flags & SCX_TASK_OPS_ENABLED) {
> + if (SCX_HAS_OP(disable))
> + SCX_CALL_OP(SCX_KF_REST, disable, p);
> + p->scx.flags &= ~SCX_TASK_OPS_ENABLED;
> + }
> +}
> +
> +static void set_task_scx_weight(struct task_struct *p)
> +{
> + u32 weight = sched_prio_to_weight[p->static_prio - MAX_RT_PRIO];
> +
> + p->scx.weight = sched_weight_to_cgroup(weight);
> +}
> +
> +/**
> + * refresh_scx_weight - Refresh a task's ext weight
> + * @p: task to refresh ext weight for
> + *
> + * @p->scx.weight carries the task's static priority in cgroup weight scale to
> + * enable easy access from the BPF scheduler. To keep it synchronized with the
> + * current task priority, this function should be called when a new task is
> + * created, priority is changed for a task on sched_ext, and a task is switched
> + * to sched_ext from other classes.
> + */
> +static void refresh_scx_weight(struct task_struct *p)
> +{
> + lockdep_assert_rq_held(task_rq(p));
> + set_task_scx_weight(p);
> + if (SCX_HAS_OP(set_weight))
> + SCX_CALL_OP(SCX_KF_REST, set_weight, p, p->scx.weight);
> +}
> +
> +void scx_pre_fork(struct task_struct *p)
> +{
> + /*
> + * BPF scheduler enable/disable paths want to be able to iterate and
> + * update all tasks which can become complex when racing forks. As
> + * enable/disable are very cold paths, let's use a percpu_rwsem to
> + * exclude forks.
> + */
> + percpu_down_read(&scx_fork_rwsem);
> +}
> +
> +int scx_fork(struct task_struct *p)
> +{
> + percpu_rwsem_assert_held(&scx_fork_rwsem);
> +
> + if (scx_enabled())
> + return scx_ops_prepare_task(p, task_group(p));
> + else
> + return 0;
> +}
> +
> +void scx_post_fork(struct task_struct *p)
> +{
> + if (scx_enabled()) {
> + struct rq_flags rf;
> + struct rq *rq;
> +
> + rq = task_rq_lock(p, &rf);
> + /*
> + * Set the weight manually before calling ops.enable() so that
> + * the scheduler doesn't see a stale value if they inspect the
> + * task struct. We'll invoke ops.set_weight() afterwards, as it
> + * would be odd to receive a callback on the task before we
> + * tell the scheduler that it's been fully enabled.
> + */
> + set_task_scx_weight(p);
> + scx_ops_enable_task(p);
> + refresh_scx_weight(p);
> + task_rq_unlock(rq, p, &rf);
> + }
> +
> + spin_lock_irq(&scx_tasks_lock);
> + list_add_tail(&p->scx.tasks_node, &scx_tasks);
> + spin_unlock_irq(&scx_tasks_lock);
> +
> + percpu_up_read(&scx_fork_rwsem);
> +}
> +
> +void scx_cancel_fork(struct task_struct *p)
> +{
> + if (scx_enabled())
> + scx_ops_disable_task(p);
> + percpu_up_read(&scx_fork_rwsem);
> +}
> +
> +void sched_ext_free(struct task_struct *p)
> +{
> + unsigned long flags;
> +
> + spin_lock_irqsave(&scx_tasks_lock, flags);
> + list_del_init(&p->scx.tasks_node);
> + spin_unlock_irqrestore(&scx_tasks_lock, flags);
> +
> + /*
> + * @p is off scx_tasks and wholly ours. scx_ops_enable()'s PREPPED ->
> + * ENABLED transitions can't race us. Disable ops for @p.
> + */
> + if (p->scx.flags & (SCX_TASK_OPS_PREPPED | SCX_TASK_OPS_ENABLED)) {
> + struct rq_flags rf;
> + struct rq *rq;
> +
> + rq = task_rq_lock(p, &rf);
> + scx_ops_disable_task(p);
> + task_rq_unlock(rq, p, &rf);
> + }
> +}
> +
> +static void reweight_task_scx(struct rq *rq, struct task_struct *p, int newprio)
> +{
> + refresh_scx_weight(p);
> +}
> +
> +static void prio_changed_scx(struct rq *rq, struct task_struct *p, int oldprio)
> +{
> +}
> +
> +static void switching_to_scx(struct rq *rq, struct task_struct *p)
> +{
> + refresh_scx_weight(p);
> +
> + /*
> + * set_cpus_allowed_scx() is not called while @p is associated with a
> + * different scheduler class. Keep the BPF scheduler up-to-date.
> + */
> + if (SCX_HAS_OP(set_cpumask))
> + SCX_CALL_OP(SCX_KF_REST, set_cpumask, p,
> + (struct cpumask *)p->cpus_ptr);
> +}
> +
> +static void check_preempt_curr_scx(struct rq *rq, struct task_struct *p,int wake_flags) {}
> +static void switched_to_scx(struct rq *rq, struct task_struct *p) {}
> +
> +/*
> + * Omitted operations:
> + *
> + * - check_preempt_curr: NOOP as it isn't useful in the wakeup path because the
> + * task isn't tied to the CPU at that point.
> + *
> + * - migrate_task_rq: Unncessary as task to cpu mapping is transient.
> + *
> + * - task_fork/dead: We need fork/dead notifications for all tasks regardless of
> + * their current sched_class. Call them directly from sched core instead.
> + *
> + * - task_woken, switched_from: Unnecessary.
> + */
> +DEFINE_SCHED_CLASS(ext) = {
> + .enqueue_task = enqueue_task_scx,
> + .dequeue_task = dequeue_task_scx,
> + .yield_task = yield_task_scx,
> + .yield_to_task = yield_to_task_scx,
> +
> + .check_preempt_curr = check_preempt_curr_scx,
> +
> + .pick_next_task = pick_next_task_scx,
> +
> + .put_prev_task = put_prev_task_scx,
> + .set_next_task = set_next_task_scx,
> +
> +#ifdef CONFIG_SMP
> + .balance = balance_scx,
> + .select_task_rq = select_task_rq_scx,
> + .set_cpus_allowed = set_cpus_allowed_scx,
> +#endif
> +
> + .task_tick = task_tick_scx,
> +
> + .switching_to = switching_to_scx,
> + .switched_to = switched_to_scx,
> + .reweight_task = reweight_task_scx,
> + .prio_changed = prio_changed_scx,
> +
> + .update_curr = update_curr_scx,
> +
> +#ifdef CONFIG_UCLAMP_TASK
> + .uclamp_enabled = 0,
> +#endif
> +};
> +
> +static void init_dsq(struct scx_dispatch_q *dsq, u64 dsq_id)
> +{
> + memset(dsq, 0, sizeof(*dsq));
> +
> + raw_spin_lock_init(&dsq->lock);
> + INIT_LIST_HEAD(&dsq->fifo);
> + dsq->id = dsq_id;
> +}
> +
> +static struct scx_dispatch_q *create_dsq(u64 dsq_id, int node)
> +{
> + struct scx_dispatch_q *dsq;
> + int ret;
> +
> + if (dsq_id & SCX_DSQ_FLAG_BUILTIN)
> + return ERR_PTR(-EINVAL);
> +
> + dsq = kmalloc_node(sizeof(*dsq), GFP_KERNEL, node);
> + if (!dsq)
> + return ERR_PTR(-ENOMEM);
> +
> + init_dsq(dsq, dsq_id);
> +
> + ret = rhashtable_insert_fast(&dsq_hash, &dsq->hash_node,
> + dsq_hash_params);
> + if (ret) {
> + kfree(dsq);
> + return ERR_PTR(ret);
> + }
> + return dsq;
> +}
> +
> +static void free_dsq_irq_workfn(struct irq_work *irq_work)
> +{
> + struct llist_node *to_free = llist_del_all(&dsqs_to_free);
> + struct scx_dispatch_q *dsq, *tmp_dsq;
> +
> + llist_for_each_entry_safe(dsq, tmp_dsq, to_free, free_node)
> + kfree_rcu(dsq);
> +}
> +
> +static DEFINE_IRQ_WORK(free_dsq_irq_work, free_dsq_irq_workfn);
> +
> +static void destroy_dsq(u64 dsq_id)
> +{
> + struct scx_dispatch_q *dsq;
> + unsigned long flags;
> +
> + rcu_read_lock();
> +
> + dsq = rhashtable_lookup_fast(&dsq_hash, &dsq_id, dsq_hash_params);
> + if (!dsq)
> + goto out_unlock_rcu;
> +
> + raw_spin_lock_irqsave(&dsq->lock, flags);
> +
> + if (dsq->nr) {
> + scx_ops_error("attempting to destroy in-use dsq 0x%016llx (nr=%u)",
> + dsq->id, dsq->nr);
> + goto out_unlock_dsq;
> + }
> +
> + if (rhashtable_remove_fast(&dsq_hash, &dsq->hash_node, dsq_hash_params))
> + goto out_unlock_dsq;
> +
> + /*
> + * Mark dead by invalidating ->id to prevent dispatch_enqueue() from
> + * queueing more tasks. As this function can be called from anywhere,
> + * freeing is bounced through an irq work to avoid nesting RCU
> + * operations inside scheduler locks.
> + */
> + dsq->id = SCX_DSQ_INVALID;
> + llist_add(&dsq->free_node, &dsqs_to_free);
> + irq_work_queue(&free_dsq_irq_work);
> +
> +out_unlock_dsq:
> + raw_spin_unlock_irqrestore(&dsq->lock, flags);
> +out_unlock_rcu:
> + rcu_read_unlock();
> +}
> +
> +/*
> + * Used by sched_fork() and __setscheduler_prio() to pick the matching
> + * sched_class. dl/rt are already handled.
> + */
> +bool task_should_scx(struct task_struct *p)
> +{
> + if (!scx_enabled() || scx_ops_disabling())
> + return false;
> + return p->policy == SCHED_EXT;
> +}
> +
> +static void scx_ops_fallback_enqueue(struct task_struct *p, u64 enq_flags)
> +{
> + if (enq_flags & SCX_ENQ_LAST)
> + scx_bpf_dispatch(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, enq_flags);
> + else
> + scx_bpf_dispatch(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags);
> +}
> +
> +static void scx_ops_fallback_dispatch(s32 cpu, struct task_struct *prev) {}
> +
> +static void scx_ops_disable_workfn(struct kthread_work *work)
> +{
> + struct scx_exit_info *ei = &scx_exit_info;
> + struct scx_task_iter sti;
> + struct task_struct *p;
> + struct rhashtable_iter rht_iter;
> + struct scx_dispatch_q *dsq;
> + const char *reason;
> + int i, type;
> +
> + type = atomic_read(&scx_exit_type);
> + while (true) {
> + /*
> + * NONE indicates that a new scx_ops has been registered since
> + * disable was scheduled - don't kill the new ops. DONE
> + * indicates that the ops has already been disabled.
> + */
> + if (type == SCX_EXIT_NONE || type == SCX_EXIT_DONE)
> + return;
> + if (atomic_try_cmpxchg(&scx_exit_type, &type, SCX_EXIT_DONE))
> + break;
> + }
> +
> + switch (type) {
> + case SCX_EXIT_UNREG:
> + reason = "BPF scheduler unregistered";
> + break;
> + case SCX_EXIT_ERROR:
> + reason = "runtime error";
> + break;
> + case SCX_EXIT_ERROR_BPF:
> + reason = "scx_bpf_error";
> + break;
> + default:
> + reason = "<UNKNOWN>";
> + }
> +
> + ei->type = type;
> + strlcpy(ei->reason, reason, sizeof(ei->reason));
> +
> + switch (scx_ops_set_enable_state(SCX_OPS_DISABLING)) {
> + case SCX_OPS_DISABLED:
> + pr_warn("sched_ext: ops error detected without ops (%s)\n",
> + scx_exit_info.msg);
> + WARN_ON_ONCE(scx_ops_set_enable_state(SCX_OPS_DISABLED) !=
> + SCX_OPS_DISABLING);
> + return;
> + case SCX_OPS_PREPPING:
> + goto forward_progress_guaranteed;
> + case SCX_OPS_DISABLING:
> + /* shouldn't happen but handle it like ENABLING if it does */
> + WARN_ONCE(true, "sched_ext: duplicate disabling instance?");
> + fallthrough;
> + case SCX_OPS_ENABLING:
> + case SCX_OPS_ENABLED:
> + break;
> + }
> +
> + /*
> + * DISABLING is set and ops was either ENABLING or ENABLED indicating
> + * that the ops and static branches are set.
> + *
> + * We must guarantee that all runnable tasks make forward progress
> + * without trusting the BPF scheduler. We can't grab any mutexes or
> + * rwsems as they might be held by tasks that the BPF scheduler is
> + * forgetting to run, which unfortunately also excludes toggling the
> + * static branches.
> + *
> + * Let's work around by overriding a couple ops and modifying behaviors
> + * based on the DISABLING state and then cycling the tasks through
> + * dequeue/enqueue to force global FIFO scheduling.
> + *
> + * a. ops.enqueue() and .dispatch() are overridden for simple global
> + * FIFO scheduling.
> + *
> + * b. balance_scx() never sets %SCX_TASK_BAL_KEEP as the slice value
> + * can't be trusted. Whenever a tick triggers, the running task is
> + * rotated to the tail of the queue.
> + *
> + * c. pick_next_task() suppresses zero slice warning.
> + */
> + scx_ops.enqueue = scx_ops_fallback_enqueue;
> + scx_ops.dispatch = scx_ops_fallback_dispatch;
> +
> + spin_lock_irq(&scx_tasks_lock);
> + scx_task_iter_init(&sti);
> + while ((p = scx_task_iter_next_filtered_locked(&sti))) {
> + if (READ_ONCE(p->__state) != TASK_DEAD) {
> + struct sched_enq_and_set_ctx ctx;
> +
> + /* cycling deq/enq is enough, see above */
> + sched_deq_and_put_task(p, DEQUEUE_SAVE | DEQUEUE_MOVE, &ctx);
> + sched_enq_and_set_task(&ctx);
> + }
> + }
> + scx_task_iter_exit(&sti);
> + spin_unlock_irq(&scx_tasks_lock);
> +
> +forward_progress_guaranteed:
> + /*
> + * Here, every runnable task is guaranteed to make forward progress and
> + * we can safely use blocking synchronization constructs. Actually
> + * disable ops.
> + */
> + mutex_lock(&scx_ops_enable_mutex);
> +
> + /* avoid racing against fork */
> + cpus_read_lock();
> + percpu_down_write(&scx_fork_rwsem);
> +
> + spin_lock_irq(&scx_tasks_lock);
> + scx_task_iter_init(&sti);
> + while ((p = scx_task_iter_next_filtered_locked(&sti))) {
> + const struct sched_class *old_class = p->sched_class;
> + struct sched_enq_and_set_ctx ctx;
> + bool alive = READ_ONCE(p->__state) != TASK_DEAD;
> +
> + sched_deq_and_put_task(p, DEQUEUE_SAVE | DEQUEUE_MOVE, &ctx);
> +
> + p->scx.slice = min_t(u64, p->scx.slice, SCX_SLICE_DFL);
> +
> + __setscheduler_prio(p, p->prio);
> + if (alive)
> + check_class_changing(task_rq(p), p, old_class);
> +
> + sched_enq_and_set_task(&ctx);
> +
> + if (alive)
> + check_class_changed(task_rq(p), p, old_class, p->prio);
> +
> + scx_ops_disable_task(p);
> + }
> + scx_task_iter_exit(&sti);
> + spin_unlock_irq(&scx_tasks_lock);
> +
> + /* no task is on scx, turn off all the switches and flush in-progress calls */
> + static_branch_disable_cpuslocked(&__scx_ops_enabled);
> + for (i = 0; i < SCX_NR_ONLINE_OPS; i++)
> + static_branch_disable_cpuslocked(&scx_has_op[i]);
> + static_branch_disable_cpuslocked(&scx_ops_enq_last);
> + static_branch_disable_cpuslocked(&scx_ops_enq_exiting);
> + static_branch_disable_cpuslocked(&scx_builtin_idle_enabled);
> + synchronize_rcu();
> +
> + percpu_up_write(&scx_fork_rwsem);
> + cpus_read_unlock();
> +
> + if (ei->type >= SCX_EXIT_ERROR) {
> + printk(KERN_ERR "sched_ext: BPF scheduler \"%s\" errored, disabling\n", scx_ops.name);
> +
> + if (ei->msg[0] == '\0')
> + printk(KERN_ERR "sched_ext: %s\n", ei->reason);
> + else
> + printk(KERN_ERR "sched_ext: %s (%s)\n", ei->reason, ei->msg);
> +
> + stack_trace_print(ei->bt, ei->bt_len, 2);
> + }
> +
> + if (scx_ops.exit)
> + SCX_CALL_OP(SCX_KF_UNLOCKED, exit, ei);
> +
> + memset(&scx_ops, 0, sizeof(scx_ops));
> +
> + rhashtable_walk_enter(&dsq_hash, &rht_iter);
> + do {
> + rhashtable_walk_start(&rht_iter);
> +
> + while ((dsq = rhashtable_walk_next(&rht_iter)) && !IS_ERR(dsq))
> + destroy_dsq(dsq->id);
> +
> + rhashtable_walk_stop(&rht_iter);
> + } while (dsq == ERR_PTR(-EAGAIN));
> + rhashtable_walk_exit(&rht_iter);
> +
> + free_percpu(scx_dsp_buf);
> + scx_dsp_buf = NULL;
> + scx_dsp_max_batch = 0;
> +
> + mutex_unlock(&scx_ops_enable_mutex);
> +
> + WARN_ON_ONCE(scx_ops_set_enable_state(SCX_OPS_DISABLED) !=
> + SCX_OPS_DISABLING);
> +}
> +
> +static DEFINE_KTHREAD_WORK(scx_ops_disable_work, scx_ops_disable_workfn);
> +
> +static void schedule_scx_ops_disable_work(void)
> +{
> + struct kthread_worker *helper = READ_ONCE(scx_ops_helper);
> +
> + /*
> + * We may be called spuriously before the first bpf_sched_ext_reg(). If
> + * scx_ops_helper isn't set up yet, there's nothing to do.
> + */
> + if (helper)
> + kthread_queue_work(helper, &scx_ops_disable_work);
> +}
> +
> +static void scx_ops_disable(enum scx_exit_type type)
> +{
> + int none = SCX_EXIT_NONE;
> +
> + if (WARN_ON_ONCE(type == SCX_EXIT_NONE || type == SCX_EXIT_DONE))
> + type = SCX_EXIT_ERROR;
> +
> + atomic_try_cmpxchg(&scx_exit_type, &none, type);
> +
> + schedule_scx_ops_disable_work();
> +}
> +
> +static void scx_ops_error_irq_workfn(struct irq_work *irq_work)
> +{
> + schedule_scx_ops_disable_work();
> +}
> +
> +static DEFINE_IRQ_WORK(scx_ops_error_irq_work, scx_ops_error_irq_workfn);
> +
> +__printf(2, 3) static void scx_ops_error_type(enum scx_exit_type type,
> + const char *fmt, ...)
> +{
> + struct scx_exit_info *ei = &scx_exit_info;
> + int none = SCX_EXIT_NONE;
> + va_list args;
> +
> + if (!atomic_try_cmpxchg(&scx_exit_type, &none, type))
> + return;
> +
> + ei->bt_len = stack_trace_save(ei->bt, ARRAY_SIZE(ei->bt), 1);
> +
> + va_start(args, fmt);
> + vscnprintf(ei->msg, ARRAY_SIZE(ei->msg), fmt, args);
> + va_end(args);
> +
> + irq_work_queue(&scx_ops_error_irq_work);
> +}
> +
> +static struct kthread_worker *scx_create_rt_helper(const char *name)
> +{
> + struct kthread_worker *helper;
> +
> + helper = kthread_create_worker(0, name);
> + if (helper)
> + sched_set_fifo(helper->task);
> + return helper;
> +}
> +
> +static int scx_ops_enable(struct sched_ext_ops *ops)
> +{
> + struct scx_task_iter sti;
> + struct task_struct *p;
> + int i, ret;
> +
> + mutex_lock(&scx_ops_enable_mutex);
> +
> + if (!scx_ops_helper) {
> + WRITE_ONCE(scx_ops_helper,
> + scx_create_rt_helper("sched_ext_ops_helper"));
> + if (!scx_ops_helper) {
> + ret = -ENOMEM;
> + goto err_unlock;
> + }
> + }
> +
> + if (scx_ops_enable_state() != SCX_OPS_DISABLED) {
> + ret = -EBUSY;
> + goto err_unlock;
> + }
> +
> + /*
> + * Set scx_ops, transition to PREPPING and clear exit info to arm the
> + * disable path. Failure triggers full disabling from here on.
> + */
> + scx_ops = *ops;
> +
> + WARN_ON_ONCE(scx_ops_set_enable_state(SCX_OPS_PREPPING) !=
> + SCX_OPS_DISABLED);
> +
> + memset(&scx_exit_info, 0, sizeof(scx_exit_info));
> + atomic_set(&scx_exit_type, SCX_EXIT_NONE);
> + scx_warned_zero_slice = false;
> +
> + atomic64_set(&scx_nr_rejected, 0);
> +
> + /*
> + * Keep CPUs stable during enable so that the BPF scheduler can track
> + * online CPUs by watching ->on/offline_cpu() after ->init().
> + */
> + cpus_read_lock();
> +
> + if (scx_ops.init) {
> + ret = SCX_CALL_OP_RET(SCX_KF_INIT, init);
> + if (ret) {
> + ret = ops_sanitize_err("init", ret);
> + goto err_disable;
> + }
> +
> + /*
> + * Exit early if ops.init() triggered scx_bpf_error(). Not
> + * strictly necessary as we'll fail transitioning into ENABLING
> + * later but that'd be after calling ops.prep_enable() on all
> + * tasks and with -EBUSY which isn't very intuitive. Let's exit
> + * early with success so that the condition is notified through
> + * ops.exit() like other scx_bpf_error() invocations.
> + */
> + if (atomic_read(&scx_exit_type) != SCX_EXIT_NONE)
> + goto err_disable;
> + }
> +
> + WARN_ON_ONCE(scx_dsp_buf);
> + scx_dsp_max_batch = ops->dispatch_max_batch ?: SCX_DSP_DFL_MAX_BATCH;
> + scx_dsp_buf = __alloc_percpu(sizeof(scx_dsp_buf[0]) * scx_dsp_max_batch,
> + __alignof__(scx_dsp_buf[0]));
> + if (!scx_dsp_buf) {
> + ret = -ENOMEM;
> + goto err_disable;
> + }
> +
> + /*
> + * Lock out forks before opening the floodgate so that they don't wander
> + * into the operations prematurely.
> + */
> + percpu_down_write(&scx_fork_rwsem);
> +
> + for (i = 0; i < SCX_NR_ONLINE_OPS; i++)
> + if (((void (**)(void))ops)[i])
> + static_branch_enable_cpuslocked(&scx_has_op[i]);
> +
> + if (ops->flags & SCX_OPS_ENQ_LAST)
> + static_branch_enable_cpuslocked(&scx_ops_enq_last);
> +
> + if (ops->flags & SCX_OPS_ENQ_EXITING)
> + static_branch_enable_cpuslocked(&scx_ops_enq_exiting);
> +
> + if (!ops->update_idle || (ops->flags & SCX_OPS_KEEP_BUILTIN_IDLE)) {
> + reset_idle_masks();
> + static_branch_enable_cpuslocked(&scx_builtin_idle_enabled);
> + } else {
> + static_branch_disable_cpuslocked(&scx_builtin_idle_enabled);
> + }
> +
> + static_branch_enable_cpuslocked(&__scx_ops_enabled);
> +
> + /*
> + * Enable ops for every task. Fork is excluded by scx_fork_rwsem
> + * preventing new tasks from being added. No need to exclude tasks
> + * leaving as sched_ext_free() can handle both prepped and enabled
> + * tasks. Prep all tasks first and then enable them with preemption
> + * disabled.
> + */
> + spin_lock_irq(&scx_tasks_lock);
> +
> + scx_task_iter_init(&sti);
> + while ((p = scx_task_iter_next_filtered(&sti))) {
> + get_task_struct(p);
> + spin_unlock_irq(&scx_tasks_lock);
> +
> + ret = scx_ops_prepare_task(p, task_group(p));
> + if (ret) {
> + put_task_struct(p);
> + spin_lock_irq(&scx_tasks_lock);
> + scx_task_iter_exit(&sti);
> + spin_unlock_irq(&scx_tasks_lock);
> + pr_err("sched_ext: ops.prep_enable() failed (%d) for %s[%d] while loading\n",
> + ret, p->comm, p->pid);
> + goto err_disable_unlock;
> + }
> +
> + put_task_struct(p);
> + spin_lock_irq(&scx_tasks_lock);
> + }
> + scx_task_iter_exit(&sti);
> +
> + /*
> + * All tasks are prepped but are still ops-disabled. Ensure that
> + * %current can't be scheduled out and switch everyone.
> + * preempt_disable() is necessary because we can't guarantee that
> + * %current won't be starved if scheduled out while switching.
> + */
> + preempt_disable();
> +
> + /*
> + * From here on, the disable path must assume that tasks have ops
> + * enabled and need to be recovered.
> + */
> + if (!scx_ops_tryset_enable_state(SCX_OPS_ENABLING, SCX_OPS_PREPPING)) {
> + preempt_enable();
> + spin_unlock_irq(&scx_tasks_lock);
> + ret = -EBUSY;
> + goto err_disable_unlock;
> + }
> +
> + /*
> + * We're fully committed and can't fail. The PREPPED -> ENABLED
> + * transitions here are synchronized against sched_ext_free() through
> + * scx_tasks_lock.
> + */
> + scx_task_iter_init(&sti);
> + while ((p = scx_task_iter_next_filtered_locked(&sti))) {
> + if (READ_ONCE(p->__state) != TASK_DEAD) {
> + const struct sched_class *old_class = p->sched_class;
> + struct sched_enq_and_set_ctx ctx;
> +
> + sched_deq_and_put_task(p, DEQUEUE_SAVE | DEQUEUE_MOVE,
> + &ctx);
> +
> + scx_ops_enable_task(p);
> + __setscheduler_prio(p, p->prio);
> + check_class_changing(task_rq(p), p, old_class);
> +
> + sched_enq_and_set_task(&ctx);
> +
> + check_class_changed(task_rq(p), p, old_class, p->prio);
> + } else {
> + scx_ops_disable_task(p);
> + }
> + }
> + scx_task_iter_exit(&sti);
> +
> + spin_unlock_irq(&scx_tasks_lock);
> + preempt_enable();
> + percpu_up_write(&scx_fork_rwsem);
> +
> + if (!scx_ops_tryset_enable_state(SCX_OPS_ENABLED, SCX_OPS_ENABLING)) {
> + ret = -EBUSY;
> + goto err_disable;
> + }
> +
> + cpus_read_unlock();
> + mutex_unlock(&scx_ops_enable_mutex);
> +
> + return 0;
> +
> +err_unlock:
> + mutex_unlock(&scx_ops_enable_mutex);
> + return ret;
> +
> +err_disable_unlock:
> + percpu_up_write(&scx_fork_rwsem);
> +err_disable:
> + cpus_read_unlock();
> + mutex_unlock(&scx_ops_enable_mutex);
> + /* must be fully disabled before returning */
> + scx_ops_disable(SCX_EXIT_ERROR);
> + kthread_flush_work(&scx_ops_disable_work);
> + return ret;
> +}
> +
> +#ifdef CONFIG_SCHED_DEBUG
> +static const char *scx_ops_enable_state_str[] = {
> + [SCX_OPS_PREPPING] = "prepping",
> + [SCX_OPS_ENABLING] = "enabling",
> + [SCX_OPS_ENABLED] = "enabled",
> + [SCX_OPS_DISABLING] = "disabling",
> + [SCX_OPS_DISABLED] = "disabled",
> +};
> +
> +static int scx_debug_show(struct seq_file *m, void *v)
> +{
> + mutex_lock(&scx_ops_enable_mutex);
> + seq_printf(m, "%-30s: %s\n", "ops", scx_ops.name);
> + seq_printf(m, "%-30s: %ld\n", "enabled", scx_enabled());
> + seq_printf(m, "%-30s: %s\n", "enable_state",
> + scx_ops_enable_state_str[scx_ops_enable_state()]);
> + seq_printf(m, "%-30s: %llu\n", "nr_rejected",
> + atomic64_read(&scx_nr_rejected));
> + mutex_unlock(&scx_ops_enable_mutex);
> + return 0;
> +}
> +
> +static int scx_debug_open(struct inode *inode, struct file *file)
> +{
> + return single_open(file, scx_debug_show, NULL);
> +}
> +
> +const struct file_operations sched_ext_fops = {
> + .open = scx_debug_open,
> + .read = seq_read,
> + .llseek = seq_lseek,
> + .release = single_release,
> +};
> +#endif
> +
> +/********************************************************************************
> + * bpf_struct_ops plumbing.
> + */
> +#include <linux/bpf_verifier.h>
> +#include <linux/bpf.h>
> +#include <linux/btf.h>
> +
> +extern struct btf *btf_vmlinux;
> +static const struct btf_type *task_struct_type;
> +
> +static bool bpf_scx_is_valid_access(int off, int size,
> + enum bpf_access_type type,
> + const struct bpf_prog *prog,
> + struct bpf_insn_access_aux *info)
> +{
> + if (off < 0 || off >= sizeof(__u64) * MAX_BPF_FUNC_ARGS)
> + return false;
> + if (type != BPF_READ)
> + return false;
> + if (off % size != 0)
> + return false;
> +
> + return btf_ctx_access(off, size, type, prog, info);
> +}
> +
> +static int bpf_scx_btf_struct_access(struct bpf_verifier_log *log,
> + const struct bpf_reg_state *reg, int off,
> + int size)
> +{
> + const struct btf_type *t;
> +
> + t = btf_type_by_id(reg->btf, reg->btf_id);
> + if (t == task_struct_type) {
> + if (off >= offsetof(struct task_struct, scx.slice) &&
> + off + size <= offsetofend(struct task_struct, scx.slice))
> + return SCALAR_VALUE;
> + }
> +
> + return 0;
> +}
> +
> +static const struct bpf_func_proto *
> +bpf_scx_get_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
> +{
> + switch (func_id) {
> + case BPF_FUNC_task_storage_get:
> + return &bpf_task_storage_get_proto;
> + case BPF_FUNC_task_storage_delete:
> + return &bpf_task_storage_delete_proto;
> + default:
> + return bpf_base_func_proto(func_id);
> + }
> +}
> +
> +const struct bpf_verifier_ops bpf_scx_verifier_ops = {
> + .get_func_proto = bpf_scx_get_func_proto,
> + .is_valid_access = bpf_scx_is_valid_access,
> + .btf_struct_access = bpf_scx_btf_struct_access,
> +};
> +
> +static int bpf_scx_init_member(const struct btf_type *t,
> + const struct btf_member *member,
> + void *kdata, const void *udata)
> +{
> + const struct sched_ext_ops *uops = udata;
> + struct sched_ext_ops *ops = kdata;
> + u32 moff = __btf_member_bit_offset(t, member) / 8;
> + int ret;
> +
> + switch (moff) {
> + case offsetof(struct sched_ext_ops, dispatch_max_batch):
> + if (*(u32 *)(udata + moff) > INT_MAX)
> + return -E2BIG;
> + ops->dispatch_max_batch = *(u32 *)(udata + moff);
> + return 1;
> + case offsetof(struct sched_ext_ops, flags):
> + if (*(u64 *)(udata + moff) & ~SCX_OPS_ALL_FLAGS)
> + return -EINVAL;
> + ops->flags = *(u64 *)(udata + moff);
> + return 1;
> + case offsetof(struct sched_ext_ops, name):
> + ret = bpf_obj_name_cpy(ops->name, uops->name,
> + sizeof(ops->name));
> + if (ret < 0)
> + return ret;
> + if (ret == 0)
> + return -EINVAL;
> + return 1;
> + }
> +
> + return 0;
> +}
> +
> +static int bpf_scx_check_member(const struct btf_type *t,
> + const struct btf_member *member,
> + const struct bpf_prog *prog)
> +{
> + u32 moff = __btf_member_bit_offset(t, member) / 8;
> +
> + switch (moff) {
> + case offsetof(struct sched_ext_ops, prep_enable):
> + case offsetof(struct sched_ext_ops, init):
> + case offsetof(struct sched_ext_ops, exit):
> + break;
> + default:
> + if (prog->aux->sleepable)
> + return -EINVAL;
> + }
> +
> + return 0;
> +}
> +
> +static int bpf_scx_reg(void *kdata)
> +{
> + return scx_ops_enable(kdata);
> +}
> +
> +static void bpf_scx_unreg(void *kdata)
> +{
> + scx_ops_disable(SCX_EXIT_UNREG);
> + kthread_flush_work(&scx_ops_disable_work);
> +}
> +
> +static int bpf_scx_init(struct btf *btf)
> +{
> + u32 type_id;
> +
> + type_id = btf_find_by_name_kind(btf, "task_struct", BTF_KIND_STRUCT);
> + if (type_id < 0)
> + return -EINVAL;
> + task_struct_type = btf_type_by_id(btf, type_id);
> +
> + return 0;
> +}
> +
> +static int bpf_scx_update(void *kdata, void *old_kdata)
> +{
> + /*
> + * sched_ext does not support updating the actively-loaded BPF
> + * scheduler, as registering a BPF scheduler can always fail if the
> + * scheduler returns an error code for e.g. ops.init(),
> + * ops.prep_enable(), etc. Similarly, we can always race with
> + * unregistration happening elsewhere, such as with sysrq.
> + */
> + return -EOPNOTSUPP;
> +}
> +
> +static int bpf_scx_validate(void *kdata)
> +{
> + return 0;
> +}
> +
> +/* "extern" to avoid sparse warning, only used in this file */
> +extern struct bpf_struct_ops bpf_sched_ext_ops;
> +
> +struct bpf_struct_ops bpf_sched_ext_ops = {
> + .verifier_ops = &bpf_scx_verifier_ops,
> + .reg = bpf_scx_reg,
> + .unreg = bpf_scx_unreg,
> + .check_member = bpf_scx_check_member,
> + .init_member = bpf_scx_init_member,
> + .init = bpf_scx_init,
> + .update = bpf_scx_update,
> + .validate = bpf_scx_validate,
> + .name = "sched_ext_ops",
> +};
> +
> +void __init init_sched_ext_class(void)
> +{
> + int cpu;
> + u32 v;
> +
> + /*
> + * The following is to prevent the compiler from optimizing out the enum
> + * definitions so that BPF scheduler implementations can use them
> + * through the generated vmlinux.h.
> + */
> + WRITE_ONCE(v, SCX_WAKE_EXEC | SCX_ENQ_WAKEUP | SCX_DEQ_SLEEP);
> +
> + BUG_ON(rhashtable_init(&dsq_hash, &dsq_hash_params));
> + init_dsq(&scx_dsq_global, SCX_DSQ_GLOBAL);
> +#ifdef CONFIG_SMP
> + BUG_ON(!alloc_cpumask_var(&idle_masks.cpu, GFP_KERNEL));
> + BUG_ON(!alloc_cpumask_var(&idle_masks.smt, GFP_KERNEL));
> +#endif
> + for_each_possible_cpu(cpu) {
> + struct rq *rq = cpu_rq(cpu);
> +
> + init_dsq(&rq->scx.local_dsq, SCX_DSQ_LOCAL);
> + }
> +}
> +
> +
> +/********************************************************************************
> + * Helpers that can be called from the BPF scheduler.
> + */
> +#include <linux/btf_ids.h>
> +
> +/* Disables missing prototype warnings for kfuncs */
> +__diag_push();
> +__diag_ignore_all("-Wmissing-prototypes",
> + "Global functions as their definitions will be in vmlinux BTF");
> +
> +/**
> + * scx_bpf_create_dsq - Create a custom DSQ
> + * @dsq_id: DSQ to create
> + * @node: NUMA node to allocate from
> + *
> + * Create a custom DSQ identified by @dsq_id. Can be called from ops.init() and
> + * ops.prep_enable().
> + */
> +s32 scx_bpf_create_dsq(u64 dsq_id, s32 node)
> +{
> + if (!scx_kf_allowed(SCX_KF_INIT | SCX_KF_SLEEPABLE))
> + return -EINVAL;
> +
> + if (unlikely(node >= (int)nr_node_ids ||
> + (node < 0 && node != NUMA_NO_NODE)))
> + return -EINVAL;
> + return PTR_ERR_OR_ZERO(create_dsq(dsq_id, node));
> +}
> +
> +BTF_SET8_START(scx_kfunc_ids_sleepable)
> +BTF_ID_FLAGS(func, scx_bpf_create_dsq, KF_SLEEPABLE)
> +BTF_SET8_END(scx_kfunc_ids_sleepable)
> +
> +static const struct btf_kfunc_id_set scx_kfunc_set_sleepable = {
> + .owner = THIS_MODULE,
> + .set = &scx_kfunc_ids_sleepable,
> +};
> +
> +static bool scx_dispatch_preamble(struct task_struct *p, u64 enq_flags)
> +{
> + if (!scx_kf_allowed(SCX_KF_ENQUEUE | SCX_KF_DISPATCH))
> + return false;
> +
> + lockdep_assert_irqs_disabled();
> +
> + if (unlikely(!p)) {
> + scx_ops_error("called with NULL task");
> + return false;
> + }
> +
> + if (unlikely(enq_flags & __SCX_ENQ_INTERNAL_MASK)) {
> + scx_ops_error("invalid enq_flags 0x%llx", enq_flags);
> + return false;
> + }
> +
> + return true;
> +}
> +
> +static void scx_dispatch_commit(struct task_struct *p, u64 dsq_id, u64 enq_flags)
> +{
> + struct task_struct *ddsp_task;
> + int idx;
> +
> + ddsp_task = __this_cpu_read(direct_dispatch_task);
> + if (ddsp_task) {
> + direct_dispatch(ddsp_task, p, dsq_id, enq_flags);
> + return;
> + }
> +
> + idx = __this_cpu_read(scx_dsp_ctx.buf_cursor);
> + if (unlikely(idx >= scx_dsp_max_batch)) {
> + scx_ops_error("dispatch buffer overflow");
> + return;
> + }
> +
> + this_cpu_ptr(scx_dsp_buf)[idx] = (struct scx_dsp_buf_ent){
> + .task = p,
> + .qseq = atomic64_read(&p->scx.ops_state) & SCX_OPSS_QSEQ_MASK,
> + .dsq_id = dsq_id,
> + .enq_flags = enq_flags,
> + };
> + __this_cpu_inc(scx_dsp_ctx.buf_cursor);
> +}
> +
> +/**
> + * scx_bpf_dispatch - Dispatch a task into the FIFO queue of a DSQ
> + * @p: task_struct to dispatch
> + * @dsq_id: DSQ to dispatch to
> + * @slice: duration @p can run for in nsecs
> + * @enq_flags: SCX_ENQ_*
> + *
> + * Dispatch @p into the FIFO queue of the DSQ identified by @dsq_id. It is safe
> + * to call this function spuriously. Can be called from ops.enqueue() and
> + * ops.dispatch().
> + *
> + * When called from ops.enqueue(), it's for direct dispatch and @p must match
> + * the task being enqueued. Also, %SCX_DSQ_LOCAL_ON can't be used to target the
> + * local DSQ of a CPU other than the enqueueing one. Use ops.select_cpu() to be
> + * on the target CPU in the first place.
> + *
> + * When called from ops.dispatch(), there are no restrictions on @p or @dsq_id
> + * and this function can be called upto ops.dispatch_max_batch times to dispatch
> + * multiple tasks. scx_bpf_dispatch_nr_slots() returns the number of the
> + * remaining slots. scx_bpf_consume() flushes the batch and resets the counter.
> + *
> + * This function doesn't have any locking restrictions and may be called under
> + * BPF locks (in the future when BPF introduces more flexible locking).
> + *
> + * @p is allowed to run for @slice. The scheduling path is triggered on slice
> + * exhaustion. If zero, the current residual slice is maintained. If
> + * %SCX_SLICE_INF, @p never expires and the BPF scheduler must kick the CPU with
> + * scx_bpf_kick_cpu() to trigger scheduling.
> + */
> +void scx_bpf_dispatch(struct task_struct *p, u64 dsq_id, u64 slice,
> + u64 enq_flags)
> +{
> + if (!scx_dispatch_preamble(p, enq_flags))
> + return;
> +
> + if (slice)
> + p->scx.slice = slice;
> + else
> + p->scx.slice = p->scx.slice ?: 1;
> +
> + scx_dispatch_commit(p, dsq_id, enq_flags);
> +}
> +
> +BTF_SET8_START(scx_kfunc_ids_enqueue_dispatch)
> +BTF_ID_FLAGS(func, scx_bpf_dispatch, KF_RCU)
> +BTF_SET8_END(scx_kfunc_ids_enqueue_dispatch)
> +
> +static const struct btf_kfunc_id_set scx_kfunc_set_enqueue_dispatch = {
> + .owner = THIS_MODULE,
> + .set = &scx_kfunc_ids_enqueue_dispatch,
> +};
> +
> +/**
> + * scx_bpf_dispatch_nr_slots - Return the number of remaining dispatch slots
> + *
> + * Can only be called from ops.dispatch().
> + */
> +u32 scx_bpf_dispatch_nr_slots(void)
> +{
> + if (!scx_kf_allowed(SCX_KF_DISPATCH))
> + return 0;
> +
> + return scx_dsp_max_batch - __this_cpu_read(scx_dsp_ctx.buf_cursor);
> +}
> +
> +/**
> + * scx_bpf_consume - Transfer a task from a DSQ to the current CPU's local DSQ
> + * @dsq_id: DSQ to consume
> + *
> + * Consume a task from the non-local DSQ identified by @dsq_id and transfer it
> + * to the current CPU's local DSQ for execution. Can only be called from
> + * ops.dispatch().
> + *
> + * This function flushes the in-flight dispatches from scx_bpf_dispatch() before
> + * trying to consume the specified DSQ. It may also grab rq locks and thus can't
> + * be called under any BPF locks.
> + *
> + * Returns %true if a task has been consumed, %false if there isn't any task to
> + * consume.
> + */
> +bool scx_bpf_consume(u64 dsq_id)
> +{
> + struct scx_dsp_ctx *dspc = this_cpu_ptr(&scx_dsp_ctx);
> + struct scx_dispatch_q *dsq;
> +
> + if (!scx_kf_allowed(SCX_KF_DISPATCH))
> + return false;
> +
> + flush_dispatch_buf(dspc->rq, dspc->rf);
> +
> + dsq = find_non_local_dsq(dsq_id);
> + if (unlikely(!dsq)) {
> + scx_ops_error("invalid DSQ ID 0x%016llx", dsq_id);
> + return false;
> + }
> +
> + if (consume_dispatch_q(dspc->rq, dspc->rf, dsq)) {
> + /*
> + * A successfully consumed task can be dequeued before it starts
> + * running while the CPU is trying to migrate other dispatched
> + * tasks. Bump nr_tasks to tell balance_scx() to retry on empty
> + * local DSQ.
> + */
> + dspc->nr_tasks++;
> + return true;
> + } else {
> + return false;
> + }
> +}
> +
> +BTF_SET8_START(scx_kfunc_ids_dispatch)
> +BTF_ID_FLAGS(func, scx_bpf_dispatch_nr_slots)
> +BTF_ID_FLAGS(func, scx_bpf_consume)
> +BTF_SET8_END(scx_kfunc_ids_dispatch)
> +
> +static const struct btf_kfunc_id_set scx_kfunc_set_dispatch = {
> + .owner = THIS_MODULE,
> + .set = &scx_kfunc_ids_dispatch,
> +};
> +
> +/**
> + * scx_bpf_dsq_nr_queued - Return the number of queued tasks
> + * @dsq_id: id of the DSQ
> + *
> + * Return the number of tasks in the DSQ matching @dsq_id. If not found,
> + * -%ENOENT is returned. Can be called from any non-sleepable online scx_ops
> + * operations.
> + */
> +s32 scx_bpf_dsq_nr_queued(u64 dsq_id)
> +{
> + struct scx_dispatch_q *dsq;
> +
> + lockdep_assert(rcu_read_lock_any_held());
> +
> + if (dsq_id == SCX_DSQ_LOCAL) {
> + return this_rq()->scx.local_dsq.nr;
> + } else if ((dsq_id & SCX_DSQ_LOCAL_ON) == SCX_DSQ_LOCAL_ON) {
> + s32 cpu = dsq_id & SCX_DSQ_LOCAL_CPU_MASK;
> +
> + if (ops_cpu_valid(cpu))
> + return cpu_rq(cpu)->scx.local_dsq.nr;
> + } else {
> + dsq = find_non_local_dsq(dsq_id);
> + if (dsq)
> + return dsq->nr;
> + }
> + return -ENOENT;
> +}
> +
> +/**
> + * scx_bpf_test_and_clear_cpu_idle - Test and clear @cpu's idle state
> + * @cpu: cpu to test and clear idle for
> + *
> + * Returns %true if @cpu was idle and its idle state was successfully cleared.
> + * %false otherwise.
> + *
> + * Unavailable if ops.update_idle() is implemented and
> + * %SCX_OPS_KEEP_BUILTIN_IDLE is not set.
> + */
> +bool scx_bpf_test_and_clear_cpu_idle(s32 cpu)
> +{
> + if (!static_branch_likely(&scx_builtin_idle_enabled)) {
> + scx_ops_error("built-in idle tracking is disabled");
> + return false;
> + }
> +
> + if (ops_cpu_valid(cpu))
> + return test_and_clear_cpu_idle(cpu);
> + else
> + return false;
> +}
> +
> +/**
> + * scx_bpf_pick_idle_cpu - Pick and claim an idle cpu
> + * @cpus_allowed: Allowed cpumask
> + * @flags: %SCX_PICK_IDLE_CPU_* flags
> + *
> + * Pick and claim an idle cpu in @cpus_allowed. Returns the picked idle cpu
> + * number on success. -%EBUSY if no matching cpu was found.
> + *
> + * Idle CPU tracking may race against CPU scheduling state transitions. For
> + * example, this function may return -%EBUSY as CPUs are transitioning into the
> + * idle state. If the caller then assumes that there will be dispatch events on
> + * the CPUs as they were all busy, the scheduler may end up stalling with CPUs
> + * idling while there are pending tasks. Use scx_bpf_pick_any_cpu() and
> + * scx_bpf_kick_cpu() to guarantee that there will be at least one dispatch
> + * event in the near future.
> + *
> + * Unavailable if ops.update_idle() is implemented and
> + * %SCX_OPS_KEEP_BUILTIN_IDLE is not set.
> + */
> +s32 scx_bpf_pick_idle_cpu(const struct cpumask *cpus_allowed, u64 flags)
> +{
> + if (!static_branch_likely(&scx_builtin_idle_enabled)) {
> + scx_ops_error("built-in idle tracking is disabled");
> + return -EBUSY;
> + }
> +
> + return scx_pick_idle_cpu(cpus_allowed, flags);
> +}
> +
> +/**
> + * scx_bpf_pick_any_cpu - Pick and claim an idle cpu if available or pick any CPU
> + * @cpus_allowed: Allowed cpumask
> + * @flags: %SCX_PICK_IDLE_CPU_* flags
> + *
> + * Pick and claim an idle cpu in @cpus_allowed. If none is available, pick any
> + * CPU in @cpus_allowed. Guaranteed to succeed and returns the picked idle cpu
> + * number if @cpus_allowed is not empty. -%EBUSY is returned if @cpus_allowed is
> + * empty.
> + *
> + * If ops.update_idle() is implemented and %SCX_OPS_KEEP_BUILTIN_IDLE is not
> + * set, this function can't tell which CPUs are idle and will always pick any
> + * CPU.
> + */
> +s32 scx_bpf_pick_any_cpu(const struct cpumask *cpus_allowed, u64 flags)
> +{
> + s32 cpu;
> +
> + if (static_branch_likely(&scx_builtin_idle_enabled)) {
> + cpu = scx_pick_idle_cpu(cpus_allowed, flags);
> + if (cpu >= 0)
> + return cpu;
> + }
> +
> + cpu = cpumask_any_distribute(cpus_allowed);
> + if (cpu < nr_cpu_ids)
> + return cpu;
> + else
> + return -EBUSY;
> +}
> +
> +/**
> + * scx_bpf_get_idle_cpumask - Get a referenced kptr to the idle-tracking
> + * per-CPU cpumask.
> + *
> + * Returns NULL if idle tracking is not enabled, or running on a UP kernel.
> + */
> +const struct cpumask *scx_bpf_get_idle_cpumask(void)
> +{
> + if (!static_branch_likely(&scx_builtin_idle_enabled)) {
> + scx_ops_error("built-in idle tracking is disabled");
> + return cpu_none_mask;
> + }
> +
> +#ifdef CONFIG_SMP
> + return idle_masks.cpu;
> +#else
> + return cpu_none_mask;
> +#endif
> +}
> +
> +/**
> + * scx_bpf_get_idle_smtmask - Get a referenced kptr to the idle-tracking,
> + * per-physical-core cpumask. Can be used to determine if an entire physical
> + * core is free.
> + *
> + * Returns NULL if idle tracking is not enabled, or running on a UP kernel.
> + */
> +const struct cpumask *scx_bpf_get_idle_smtmask(void)
> +{
> + if (!static_branch_likely(&scx_builtin_idle_enabled)) {
> + scx_ops_error("built-in idle tracking is disabled");
> + return cpu_none_mask;
> + }
> +
> +#ifdef CONFIG_SMP
> + if (sched_smt_active())
> + return idle_masks.smt;
> + else
> + return idle_masks.cpu;
> +#else
> + return cpu_none_mask;
> +#endif
> +}
> +
> +/**
> + * scx_bpf_put_idle_cpumask - Release a previously acquired referenced kptr to
> + * either the percpu, or SMT idle-tracking cpumask.
> + */
> +void scx_bpf_put_idle_cpumask(const struct cpumask *idle_mask)
> +{
> + /*
> + * Empty function body because we aren't actually acquiring or
> + * releasing a reference to a global idle cpumask, which is read-only
> + * in the caller and is never released. The acquire / release semantics
> + * here are just used to make the cpumask is a trusted pointer in the
> + * caller.
> + */
> +}
> +
> +struct scx_bpf_error_bstr_bufs {
> + u64 data[MAX_BPRINTF_VARARGS];
> + char msg[SCX_EXIT_MSG_LEN];
> +};
> +
> +static DEFINE_PER_CPU(struct scx_bpf_error_bstr_bufs, scx_bpf_error_bstr_bufs);
> +
> +/**
> + * scx_bpf_error_bstr - Indicate fatal error
> + * @fmt: error message format string
> + * @data: format string parameters packaged using ___bpf_fill() macro
> + * @data__sz: @data len, must end in '__sz' for the verifier
> + *
> + * Indicate that the BPF scheduler encountered a fatal error and initiate ops
> + * disabling.
> + */
> +void scx_bpf_error_bstr(char *fmt, unsigned long long *data, u32 data__sz)
> +{
> + struct bpf_bprintf_data bprintf_data = { .get_bin_args = true };
> + struct scx_bpf_error_bstr_bufs *bufs;
> + unsigned long flags;
> + int ret;
> +
> + local_irq_save(flags);
> + bufs = this_cpu_ptr(&scx_bpf_error_bstr_bufs);
> +
> + if (data__sz % 8 || data__sz > MAX_BPRINTF_VARARGS * 8 ||
> + (data__sz && !data)) {
> + scx_ops_error("invalid data=%p and data__sz=%u",
> + (void *)data, data__sz);
> + goto out_restore;
> + }
> +
> + ret = copy_from_kernel_nofault(bufs->data, data, data__sz);
> + if (ret) {
> + scx_ops_error("failed to read data fields (%d)", ret);
> + goto out_restore;
> + }
> +
> + ret = bpf_bprintf_prepare(fmt, UINT_MAX, bufs->data, data__sz / 8,
> + &bprintf_data);
> + if (ret < 0) {
> + scx_ops_error("failed to format prepration (%d)", ret);
> + goto out_restore;
> + }
> +
> + ret = bstr_printf(bufs->msg, sizeof(bufs->msg), fmt,
> + bprintf_data.bin_args);
> + bpf_bprintf_cleanup(&bprintf_data);
> + if (ret < 0) {
> + scx_ops_error("scx_ops_error(\"%s\", %p, %u) failed to format",
> + fmt, data, data__sz);
> + goto out_restore;
> + }
> +
> + scx_ops_error_type(SCX_EXIT_ERROR_BPF, "%s", bufs->msg);
> +out_restore:
> + local_irq_restore(flags);
> +}
> +
> +/**
> + * scx_bpf_destroy_dsq - Destroy a custom DSQ
> + * @dsq_id: DSQ to destroy
> + *
> + * Destroy the custom DSQ identified by @dsq_id. Only DSQs created with
> + * scx_bpf_create_dsq() can be destroyed. The caller must ensure that the DSQ is
> + * empty and no further tasks are dispatched to it. Ignored if called on a DSQ
> + * which doesn't exist. Can be called from any online scx_ops operations.
> + */
> +void scx_bpf_destroy_dsq(u64 dsq_id)
> +{
> + destroy_dsq(dsq_id);
> +}
> +
> +/**
> + * scx_bpf_task_running - Is task currently running?
> + * @p: task of interest
> + */
> +bool scx_bpf_task_running(const struct task_struct *p)
> +{
> + return task_rq(p)->curr == p;
> +}
> +
> +/**
> + * scx_bpf_task_cpu - CPU a task is currently associated with
> + * @p: task of interest
> + */
> +s32 scx_bpf_task_cpu(const struct task_struct *p)
> +{
> + return task_cpu(p);
> +}
> +
> +BTF_SET8_START(scx_kfunc_ids_any)
> +BTF_ID_FLAGS(func, scx_bpf_dsq_nr_queued)
> +BTF_ID_FLAGS(func, scx_bpf_test_and_clear_cpu_idle)
> +BTF_ID_FLAGS(func, scx_bpf_pick_idle_cpu, KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_pick_any_cpu, KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_get_idle_cpumask, KF_ACQUIRE)
> +BTF_ID_FLAGS(func, scx_bpf_get_idle_smtmask, KF_ACQUIRE)
> +BTF_ID_FLAGS(func, scx_bpf_put_idle_cpumask, KF_RELEASE)
> +BTF_ID_FLAGS(func, scx_bpf_error_bstr, KF_TRUSTED_ARGS)
> +BTF_ID_FLAGS(func, scx_bpf_destroy_dsq)
> +BTF_ID_FLAGS(func, scx_bpf_task_running, KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_task_cpu, KF_RCU)
> +BTF_SET8_END(scx_kfunc_ids_any)
> +
> +static const struct btf_kfunc_id_set scx_kfunc_set_any = {
> + .owner = THIS_MODULE,
> + .set = &scx_kfunc_ids_any,
> +};
> +
> +__diag_pop();
> +
> +/*
> + * This can't be done from init_sched_ext_class() as register_btf_kfunc_id_set()
> + * needs most of the system to be up.
> + */
> +static int __init register_ext_kfuncs(void)
> +{
> + int ret;
> +
> + /*
> + * Some kfuncs are context-sensitive and can only be called from
> + * specific SCX ops. They are grouped into BTF sets accordingly.
> + * Unfortunately, BPF currently doesn't have a way of enforcing such
> + * restrictions. Eventually, the verifier should be able to enforce
> + * them. For now, register them the same and make each kfunc explicitly
> + * check using scx_kf_allowed().
> + */
> + if ((ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> + &scx_kfunc_set_sleepable)) ||
> + (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> + &scx_kfunc_set_enqueue_dispatch)) ||
> + (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> + &scx_kfunc_set_dispatch)) ||
> + (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> + &scx_kfunc_set_any))) {
> + pr_err("sched_ext: failed to register kfunc sets (%d)\n", ret);
> + return ret;
> + }
> +
> + return 0;
> +}
> +__initcall(register_ext_kfuncs);
> diff --git a/kernel/sched/ext.h b/kernel/sched/ext.h
> index 6a93c4825339..d78d151fdbf8 100644
> --- a/kernel/sched/ext.h
> +++ b/kernel/sched/ext.h
> @@ -1,11 +1,119 @@
> /* SPDX-License-Identifier: GPL-2.0 */
> +/*
> + * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
> + * Copyright (c) 2022 Tejun Heo <tj@xxxxxxxxxx>
> + * Copyright (c) 2022 David Vernet <dvernet@xxxxxxxx>
> + */
> +enum scx_wake_flags {
> + /* expose select WF_* flags as enums */
> + SCX_WAKE_EXEC = WF_EXEC,
> + SCX_WAKE_FORK = WF_FORK,
> + SCX_WAKE_TTWU = WF_TTWU,
> + SCX_WAKE_SYNC = WF_SYNC,
> +};
> +
> +enum scx_enq_flags {
> + /* expose select ENQUEUE_* flags as enums */
> + SCX_ENQ_WAKEUP = ENQUEUE_WAKEUP,
> + SCX_ENQ_HEAD = ENQUEUE_HEAD,
> +
> + /* high 32bits are SCX specific */
> +
> + /*
> + * The task being enqueued is the only task available for the cpu. By
> + * default, ext core keeps executing such tasks but when
> + * %SCX_OPS_ENQ_LAST is specified, they're ops.enqueue()'d with
> + * %SCX_ENQ_LAST and %SCX_ENQ_LOCAL flags set.
> + *
> + * If the BPF scheduler wants to continue executing the task,
> + * ops.enqueue() should dispatch the task to %SCX_DSQ_LOCAL immediately.
> + * If the task gets queued on a different dsq or the BPF side, the BPF
> + * scheduler is responsible for triggering a follow-up scheduling event.
> + * Otherwise, Execution may stall.
> + */
> + SCX_ENQ_LAST = 1LLU << 41,
> +
> + /*
> + * A hint indicating that it's advisable to enqueue the task on the
> + * local dsq of the currently selected CPU. Currently used by
> + * select_cpu_dfl() and together with %SCX_ENQ_LAST.
> + */
> + SCX_ENQ_LOCAL = 1LLU << 42,
> +
> + /* high 8 bits are internal */
> + __SCX_ENQ_INTERNAL_MASK = 0xffLLU << 56,
> +
> + SCX_ENQ_CLEAR_OPSS = 1LLU << 56,
> +};
> +
> +enum scx_deq_flags {
> + /* expose select DEQUEUE_* flags as enums */
> + SCX_DEQ_SLEEP = DEQUEUE_SLEEP,
> +};
> +
> +enum scx_pick_idle_cpu_flags {
> + SCX_PICK_IDLE_CORE = 1LLU << 0, /* pick a CPU whose SMT siblings are also idle */
> +};
>
> #ifdef CONFIG_SCHED_CLASS_EXT
> -#error "NOT IMPLEMENTED YET"
> +
> +struct sched_enq_and_set_ctx {
> + struct task_struct *p;
> + int queue_flags;
> + bool queued;
> + bool running;
> +};
> +
> +void sched_deq_and_put_task(struct task_struct *p, int queue_flags,
> + struct sched_enq_and_set_ctx *ctx);
> +void sched_enq_and_set_task(struct sched_enq_and_set_ctx *ctx);
> +
> +extern const struct sched_class ext_sched_class;
> +extern const struct bpf_verifier_ops bpf_sched_ext_verifier_ops;
> +extern const struct file_operations sched_ext_fops;
> +
> +DECLARE_STATIC_KEY_FALSE(__scx_ops_enabled);
> +#define scx_enabled() static_branch_unlikely(&__scx_ops_enabled)
> +
> +static inline bool task_on_scx(struct task_struct *p)
> +{
> + return scx_enabled() && p->sched_class == &ext_sched_class;
> +}
While building the kernel, I encountered the following warning:
{KERNEL_SRC}/kernel/sched/core.c: In function ‘__task_prio’:
{KERNEL_SRC}/kernel/sched/core.c:170:25: warning: passing argument 1 of ‘task_on_scx’ discards ‘const’ qualifier from pointer target type [-Wdiscarded-qualifiers]
170 | if (task_on_scx(p))
| ^
In file included from {KERNEL_SRC}/kernel/sched/sched.h:3593,
from {KERNEL_SRC}/kernel/sched/core.c:86:
{KERNEL_SRC}/kernel/sched/ext.h:124:52: note: expected ‘struct task_struct *’ but argument is of type ‘const struct task_struct *’
124 | static inline bool task_on_scx(struct task_struct *p)
| ~~~~~~~~~~~~~~~~~~~~^
To address this warning, I'd suggest modifying the signature of `task_on_scx` to
accept `task_struct` argument as `const`. The proposed change is as follows:
diff --git a/kernel/sched/ext.h b/kernel/sched/ext.h
index 405037a4e6ce..e9c699a87770 100644
--- a/kernel/sched/ext.h
+++ b/kernel/sched/ext.h
@@ -121,7 +121,7 @@ DECLARE_STATIC_KEY_FALSE(__scx_switched_all);
DECLARE_STATIC_KEY_FALSE(scx_ops_cpu_preempt);
-static inline bool task_on_scx(struct task_struct *p)
+static inline bool task_on_scx(const struct task_struct *p)
{
return scx_enabled() && p->sched_class == &ext_sched_class;
}
@@ -214,7 +214,7 @@ bool scx_prio_less(const struct task_struct *a, const struct task_struct *b,
#define scx_enabled() false
#define scx_switched_all() false
-static inline bool task_on_scx(struct task_struct *p) { return false; }
+static inline bool task_on_scx(const struct task_struct *p) { return false; }
static inline void scx_pre_fork(struct task_struct *p) {}
static inline int scx_fork(struct task_struct *p) { return 0; }
static inline void scx_post_fork(struct task_struct *p) {}
> +
> +bool task_should_scx(struct task_struct *p);
> +void scx_pre_fork(struct task_struct *p);
> +int scx_fork(struct task_struct *p);
> +void scx_post_fork(struct task_struct *p);
> +void scx_cancel_fork(struct task_struct *p);
> +void init_sched_ext_class(void);
> +
> +static inline const struct sched_class *next_active_class(const struct sched_class *class)
> +{
> + class++;
> + if (!scx_enabled() && class == &ext_sched_class)
> + class++;
> + return class;
> +}
> +
> +#define for_active_class_range(class, _from, _to) \
> + for (class = (_from); class != (_to); class = next_active_class(class))
> +
> +#define for_each_active_class(class) \
> + for_active_class_range(class, __sched_class_highest, __sched_class_lowest)
> +
> +/*
> + * SCX requires a balance() call before every pick_next_task() call including
> + * when waking up from idle.
> + */
> +#define for_balance_class_range(class, prev_class, end_class) \
> + for_active_class_range(class, (prev_class) > &ext_sched_class ? \
> + &ext_sched_class : (prev_class), (end_class))
> +
> #else /* CONFIG_SCHED_CLASS_EXT */
>
> #define scx_enabled() false
>
> +static inline bool task_on_scx(struct task_struct *p) { return false; }
> static inline void scx_pre_fork(struct task_struct *p) {}
> static inline int scx_fork(struct task_struct *p) { return 0; }
> static inline void scx_post_fork(struct task_struct *p) {}
> @@ -18,7 +126,13 @@ static inline void init_sched_ext_class(void) {}
> #endif /* CONFIG_SCHED_CLASS_EXT */
>
> #if defined(CONFIG_SCHED_CLASS_EXT) && defined(CONFIG_SMP)
> -#error "NOT IMPLEMENTED YET"
> +void __scx_update_idle(struct rq *rq, bool idle);
> +
> +static inline void scx_update_idle(struct rq *rq, bool idle)
> +{
> + if (scx_enabled())
> + __scx_update_idle(rq, idle);
> +}
> #else
> static inline void scx_update_idle(struct rq *rq, bool idle) {}
> #endif
> diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
> index 17bd277cf27a..666166908eb6 100644
> --- a/kernel/sched/sched.h
> +++ b/kernel/sched/sched.h
> @@ -185,6 +185,10 @@ static inline int idle_policy(int policy)
>
> static inline int normal_policy(int policy)
> {
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + if (policy == SCHED_EXT)
> + return true;
> +#endif
> return policy == SCHED_NORMAL;
> }
>
> @@ -681,6 +685,15 @@ struct cfs_rq {
> #endif /* CONFIG_FAIR_GROUP_SCHED */
> };
>
> +#ifdef CONFIG_SCHED_CLASS_EXT
> +struct scx_rq {
> + struct scx_dispatch_q local_dsq;
> + u64 ops_qseq;
> + u64 extra_enq_flags; /* see move_task_to_local_dsq() */
> + u32 nr_running;
> +};
> +#endif /* CONFIG_SCHED_CLASS_EXT */
> +
> static inline int rt_bandwidth_enabled(void)
> {
> return sysctl_sched_rt_runtime >= 0;
> @@ -1022,6 +1035,9 @@ struct rq {
> struct cfs_rq cfs;
> struct rt_rq rt;
> struct dl_rq dl;
> +#ifdef CONFIG_SCHED_CLASS_EXT
> + struct scx_rq scx;
> +#endif
>
> #ifdef CONFIG_FAIR_GROUP_SCHED
> /* list of leaf cfs_rq on this CPU: */
> --
> 2.41.0
>
