path: root/tools/sched_ext/scx_qmap.bpf.c

/* SPDX-License-Identifier: GPL-2.0 */
/*
 * A simple five-level FIFO queue scheduler.
 *
 * There are five FIFOs implemented as arena-backed doubly-linked lists
 * threaded through per-task context. A task gets assigned to one depending on
 * its compound weight. Each CPU round robins through the FIFOs and dispatches
 * more from FIFOs with higher indices - 1 from queue0, 2 from queue1, 4 from
 * queue2 and so on.
 *
 * This scheduler demonstrates:
 *
 * - BPF-side queueing using TIDs.
 * - BPF arena for scheduler state.
 * - Core-sched support.
 *
 * This scheduler is primarily for demonstration and testing of sched_ext
 * features and unlikely to be useful for actual workloads.
 *
 * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2022 Tejun Heo <tj@kernel.org>
 * Copyright (c) 2022 David Vernet <dvernet@meta.com>
 */
#include <scx/common.bpf.h>

#include "scx_qmap.h"

enum consts {
	ONE_SEC_IN_NS		= 1000000000,
	ONE_MSEC_IN_NS		= 1000000,
	LOWPRI_INTV_NS		= 10 * ONE_MSEC_IN_NS,
	SHARED_DSQ		= 0,
	HIGHPRI_DSQ		= 1,
	LOWPRI_DSQ		= 2,
	HIGHPRI_WEIGHT		= 8668,		/* this is what -20 maps to */
};

char _license[] SEC("license") = "GPL";

const volatile u64 slice_ns;
const volatile u32 stall_user_nth;
const volatile u32 stall_kernel_nth;
const volatile u32 dsp_inf_loop_after;
const volatile u32 dsp_batch;
const volatile bool highpri_boosting;
const volatile bool print_dsqs_and_events;
const volatile bool print_msgs;
const volatile u64 sub_cgroup_id;
const volatile s32 disallow_tgid;
const volatile bool suppress_dump;
const volatile bool always_enq_immed;
const volatile u32 immed_stress_nth;
const volatile u32 max_tasks;

/*
 * Optional cid-override test harness. When cid_override_mode is non-zero,
 * qmap_init() calls scx_bpf_cid_override() with the caller-supplied
 * cpu_to_cid array to exercise the kfunc's acceptance and error paths.
 *
 *   0 = disabled
 *   1 = valid reverse mapping
 *   2 = invalid: duplicate cid assignment
 *   3 = invalid: out-of-range cid
 */
const volatile u32 cid_override_mode;
/*
 * Array lives in bss (writable) because scx_bpf_cid_override()'s BPF
 * verifier signature treats its len-paired pointer as read/write - rodata
 * fails verification with "write into map forbidden". Userspace populates
 * it before SCX_OPS_LOAD, same as rodata, and nothing writes it after.
 */
s32 cid_override_cpu_to_cid[SCX_QMAP_MAX_CPUS];

UEI_DEFINE(uei);

/*
 * All scheduler state - per-cpu context, stats counters, core-sched sequence
 * numbers, sub-sched cgroup ids - lives in this single BPF arena map. Userspace
 * reaches it via skel->arena->qa.
 */
struct {
	__uint(type, BPF_MAP_TYPE_ARENA);
	__uint(map_flags, BPF_F_MMAPABLE);
	__uint(max_entries, 1 << 16);		/* upper bound in pages */
#if defined(__TARGET_ARCH_arm64) || defined(__aarch64__)
	__ulong(map_extra, 0x1ull << 32);	/* user/BPF mmap base */
#else
	__ulong(map_extra, 0x1ull << 44);
#endif
} arena SEC(".maps");

struct qmap_arena __arena_global qa;

/*
 * Global idle-cid tracking, maintained via update_idle / cpu_offline and
 * scanned by the direct-dispatch path. Allocated in qmap_init() from one
 * arena page, sized to the full cid space.
 */
struct scx_cmask __arena *qa_idle_cids;

/* Per-queue locks. Each in its own .data section as bpf_res_spin_lock requires. */
__hidden struct bpf_res_spin_lock qa_q_lock0 SEC(".data.qa_q_lock0");
__hidden struct bpf_res_spin_lock qa_q_lock1 SEC(".data.qa_q_lock1");
__hidden struct bpf_res_spin_lock qa_q_lock2 SEC(".data.qa_q_lock2");
__hidden struct bpf_res_spin_lock qa_q_lock3 SEC(".data.qa_q_lock3");
__hidden struct bpf_res_spin_lock qa_q_lock4 SEC(".data.qa_q_lock4");

static struct bpf_res_spin_lock *qa_q_lock(s32 qid)
{
	switch (qid) {
	case 0:	return &qa_q_lock0;
	case 1:	return &qa_q_lock1;
	case 2:	return &qa_q_lock2;
	case 3:	return &qa_q_lock3;
	case 4:	return &qa_q_lock4;
	default: return NULL;
	}
}

/*
 * If enabled, CPU performance target is set according to the queue index
 * according to the following table.
 */
static const u32 qidx_to_cpuperf_target[] = {
	[0] = SCX_CPUPERF_ONE * 0 / 4,
	[1] = SCX_CPUPERF_ONE * 1 / 4,
	[2] = SCX_CPUPERF_ONE * 2 / 4,
	[3] = SCX_CPUPERF_ONE * 3 / 4,
	[4] = SCX_CPUPERF_ONE * 4 / 4,
};

/*
 * Per-queue sequence numbers to implement core-sched ordering.
 *
 * Tail seq is assigned to each queued task and incremented. Head seq tracks the
 * sequence number of the latest dispatched task. The distance between the a
 * task's seq and the associated queue's head seq is called the queue distance
 * and used when comparing two tasks for ordering. See qmap_core_sched_before().
 */

/*
 * Per-task scheduling context. Allocated from the qa.task_ctxs[] slab in
 * arena. While the task is alive the entry is referenced from task_ctx_stor;
 * while it's free the entry sits on the free list singly-linked through
 * @next_free.
 *
 * When the task is queued on one of the five priority FIFOs, @q_idx is the
 * queue index and @q_next/@q_prev link it in the queue's doubly-linked list.
 * @q_idx is -1 when the task isn't on any queue.
 */
struct task_ctx {
	struct task_ctx __arena	*next_free;	/* only valid on free list */
	struct task_ctx __arena	*q_next;	/* queue link, NULL if tail */
	struct task_ctx __arena	*q_prev;	/* queue link, NULL if head */
	struct qmap_fifo __arena *fifo;		/* queue we're on, NULL if not queued */
	u64			tid;
	s32			pid;	/* for dump only */
	bool			force_local;	/* Dispatch directly to local_dsq */
	bool			highpri;
	u64			core_sched_seq;
	struct scx_cmask	cpus_allowed;	/* per-task affinity in cid space */
};

/*
 * Slab stride for task_ctx. cpus_allowed's flex array bits[] overlaps the
 * tail bytes appended per entry; struct_size() gives the actual per-entry
 * footprint.
 */
#define TASK_CTX_STRIDE							\
	struct_size_t(struct task_ctx, cpus_allowed.bits,		\
		      CMASK_NR_WORDS(SCX_QMAP_MAX_CPUS))

/* All task_ctx pointers are arena pointers. */
typedef struct task_ctx __arena task_ctx_t;

/* Holds an arena pointer to the task's slab entry. */
struct task_ctx_stor_val {
	task_ctx_t		*taskc;
};

struct {
	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
	__uint(map_flags, BPF_F_NO_PREALLOC);
	__type(key, int);
	__type(value, struct task_ctx_stor_val);
} task_ctx_stor SEC(".maps");

/* Protects the task_ctx slab free list. */
__hidden struct bpf_res_spin_lock qa_task_lock SEC(".data.qa_task_lock");

static int qmap_spin_lock(struct bpf_res_spin_lock *lock)
{
	if (bpf_res_spin_lock(lock)) {
		scx_bpf_error("res_spin_lock failed");
		return -EBUSY;
	}
	return 0;
}

/*
 * Try prev_cid, then scan taskc->cpus_allowed AND qa_idle_cids round-robin
 * from prev_cid + 1. Atomic claim retries on race; bounded by
 * IDLE_PICK_RETRIES to keep the verifier's insn budget in check.
 */
#define IDLE_PICK_RETRIES	16

static s32 pick_direct_dispatch_cid(struct task_struct *p, s32 prev_cid,
				    task_ctx_t *taskc)
{
	u32 nr_cids = scx_bpf_nr_cids();
	s32 cid;
	u32 i;

	if (!always_enq_immed && p->nr_cpus_allowed == 1)
		return prev_cid;

	if (cmask_test_and_clear(prev_cid, qa_idle_cids))
		return prev_cid;

	cid = prev_cid;
	bpf_for(i, 0, IDLE_PICK_RETRIES) {
		cid = cmask_next_and_set_wrap(&taskc->cpus_allowed,
					      qa_idle_cids, cid + 1);
		barrier_var(cid);
		if (cid >= nr_cids)
			return -1;
		if (cmask_test_and_clear(cid, qa_idle_cids))
			return cid;
	}
	return -1;
}

/*
 * Force a reference to the arena map. The verifier associates an arena with
 * a program by finding an LD_IMM64 instruction that loads the arena's BPF
 * map; programs that only use arena pointers returned from task-local
 * storage (like qmap_select_cpu) never reference @arena directly. Without
 * this, the verifier rejects addr_space_cast with "addr_space_cast insn
 * can only be used in a program that has an associated arena".
 */
#define QMAP_TOUCH_ARENA() do { asm volatile("" :: "r"(&arena)); } while (0)

static task_ctx_t *lookup_task_ctx(struct task_struct *p)
{
	struct task_ctx_stor_val *v;

	QMAP_TOUCH_ARENA();

	v = bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
	if (!v || !v->taskc)
		return NULL;
	return v->taskc;
}

/* Append @taskc to the tail of @fifo. Must not already be queued. */
static void qmap_fifo_enqueue(struct qmap_fifo __arena *fifo, task_ctx_t *taskc)
{
	struct bpf_res_spin_lock *lock = qa_q_lock(fifo->idx);

	if (!lock || qmap_spin_lock(lock))
		return;
	taskc->fifo = fifo;
	taskc->q_next = NULL;
	taskc->q_prev = fifo->tail;
	if (fifo->tail)
		fifo->tail->q_next = taskc;
	else
		fifo->head = taskc;
	fifo->tail = taskc;
	bpf_res_spin_unlock(lock);
}

/* Pop the head of @fifo. Returns NULL if empty. */
static task_ctx_t *qmap_fifo_pop(struct qmap_fifo __arena *fifo)
{
	struct bpf_res_spin_lock *lock = qa_q_lock(fifo->idx);
	task_ctx_t *taskc;

	if (!lock || qmap_spin_lock(lock))
		return NULL;
	taskc = fifo->head;
	if (taskc) {
		fifo->head = taskc->q_next;
		if (taskc->q_next)
			taskc->q_next->q_prev = NULL;
		else
			fifo->tail = NULL;
		taskc->q_next = NULL;
		taskc->q_prev = NULL;
		taskc->fifo = NULL;
	}
	bpf_res_spin_unlock(lock);
	return taskc;
}

/* Remove @taskc from its fifo. No-op if not queued. */
static void qmap_fifo_remove(task_ctx_t *taskc)
{
	struct qmap_fifo __arena *fifo = taskc->fifo;
	struct bpf_res_spin_lock *lock;

	if (!fifo)
		return;

	lock = qa_q_lock(fifo->idx);
	if (!lock || qmap_spin_lock(lock))
		return;

	/* Re-check under lock — a concurrent pop may have cleared fifo. */
	if (taskc->fifo != fifo) {
		bpf_res_spin_unlock(lock);
		return;
	}

	if (taskc->q_next)
		taskc->q_next->q_prev = taskc->q_prev;
	else
		fifo->tail = taskc->q_prev;
	if (taskc->q_prev)
		taskc->q_prev->q_next = taskc->q_next;
	else
		fifo->head = taskc->q_next;
	taskc->q_next = NULL;
	taskc->q_prev = NULL;
	taskc->fifo = NULL;
	bpf_res_spin_unlock(lock);
}

s32 BPF_STRUCT_OPS(qmap_select_cid, struct task_struct *p,
		   s32 prev_cid, u64 wake_flags)
{
	task_ctx_t *taskc;
	s32 cid;

	if (!(taskc = lookup_task_ctx(p)))
		return prev_cid;

	if (p->scx.weight < 2 && !(p->flags & PF_KTHREAD))
		return prev_cid;

	cid = pick_direct_dispatch_cid(p, prev_cid, taskc);

	if (cid >= 0) {
		taskc->force_local = true;
		return cid;
	} else {
		return prev_cid;
	}
}

static int weight_to_idx(u32 weight)
{
	/* Coarsely map the compound weight to a FIFO. */
	if (weight <= 25)
		return 0;
	else if (weight <= 50)
		return 1;
	else if (weight < 200)
		return 2;
	else if (weight < 400)
		return 3;
	else
		return 4;
}

void BPF_STRUCT_OPS(qmap_enqueue, struct task_struct *p, u64 enq_flags)
{
	static u32 user_cnt, kernel_cnt;
	task_ctx_t *taskc;
	int idx = weight_to_idx(p->scx.weight);
	s32 cid;

	if (enq_flags & SCX_ENQ_REENQ) {
		__sync_fetch_and_add(&qa.nr_reenqueued, 1);
		if (scx_bpf_task_cid(p) == 0)
			__sync_fetch_and_add(&qa.nr_reenqueued_cid0, 1);
	}

	if (p->flags & PF_KTHREAD) {
		if (stall_kernel_nth && !(++kernel_cnt % stall_kernel_nth))
			return;
	} else {
		if (stall_user_nth && !(++user_cnt % stall_user_nth))
			return;
	}

	if (qa.test_error_cnt && !--qa.test_error_cnt)
		scx_bpf_error("test triggering error");

	if (!(taskc = lookup_task_ctx(p)))
		return;

	/*
	 * All enqueued tasks must have their core_sched_seq updated for correct
	 * core-sched ordering. Also, take a look at the end of qmap_dispatch().
	 */
	taskc->core_sched_seq = qa.core_sched_tail_seqs[idx]++;

	/*
	 * IMMED stress testing: Every immed_stress_nth'th enqueue, dispatch
	 * directly to prev_cpu's local DSQ even when busy to force dsq->nr > 1
	 * and exercise the kernel IMMED reenqueue trigger paths.
	 */
	if (immed_stress_nth && !(enq_flags & SCX_ENQ_REENQ)) {
		static u32 immed_stress_cnt;

		if (!(++immed_stress_cnt % immed_stress_nth)) {
			taskc->force_local = false;
			scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | scx_bpf_task_cid(p),
					   slice_ns, enq_flags);
			return;
		}
	}

	/*
	 * If qmap_select_cid() is telling us to or this is the last runnable
	 * task on the CPU, enqueue locally.
	 */
	if (taskc->force_local) {
		taskc->force_local = false;
		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, slice_ns, enq_flags);
		return;
	}

	/* see lowpri_timerfn() */
	if (__COMPAT_has_generic_reenq() &&
	    p->scx.weight < 2 && !(p->flags & PF_KTHREAD) && !(enq_flags & SCX_ENQ_REENQ)) {
		scx_bpf_dsq_insert(p, LOWPRI_DSQ, slice_ns, enq_flags);
		return;
	}

	/* if select_cid() wasn't called, try direct dispatch */
	if (!__COMPAT_is_enq_cpu_selected(enq_flags) &&
	    (cid = pick_direct_dispatch_cid(p, scx_bpf_task_cid(p), taskc)) >= 0) {
		__sync_fetch_and_add(&qa.nr_ddsp_from_enq, 1);
		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cid, slice_ns, enq_flags);
		return;
	}

	/*
	 * If the task was re-enqueued due to the CPU being preempted by a
	 * higher priority scheduling class, just re-enqueue the task directly
	 * on the global DSQ. As we want another CPU to pick it up, find and
	 * kick an idle cid.
	 */
	if (enq_flags & SCX_ENQ_REENQ) {
		s32 cid;

		scx_bpf_dsq_insert(p, SHARED_DSQ, 0, enq_flags);
		cid = cmask_next_and_set_wrap(&taskc->cpus_allowed,
					      qa_idle_cids, 0);
		if (cid < scx_bpf_nr_cids())
			scx_bpf_kick_cid(cid, SCX_KICK_IDLE);
		return;
	}

	/* Queue on the selected FIFO. */
	qmap_fifo_enqueue(&qa.fifos[idx], taskc);

	if (highpri_boosting && p->scx.weight >= HIGHPRI_WEIGHT) {
		taskc->highpri = true;
		__sync_fetch_and_add(&qa.nr_highpri_queued, 1);
	}
	__sync_fetch_and_add(&qa.nr_enqueued, 1);
}

void BPF_STRUCT_OPS(qmap_dequeue, struct task_struct *p, u64 deq_flags)
{
	task_ctx_t *taskc;

	__sync_fetch_and_add(&qa.nr_dequeued, 1);
	if (deq_flags & SCX_DEQ_CORE_SCHED_EXEC)
		__sync_fetch_and_add(&qa.nr_core_sched_execed, 1);

	taskc = lookup_task_ctx(p);
	if (taskc && taskc->fifo) {
		if (taskc->highpri)
			__sync_fetch_and_sub(&qa.nr_highpri_queued, 1);
		qmap_fifo_remove(taskc);
	}
}

static void update_core_sched_head_seq(struct task_struct *p)
{
	int idx = weight_to_idx(p->scx.weight);
	task_ctx_t *taskc;

	if ((taskc = lookup_task_ctx(p)))
		qa.core_sched_head_seqs[idx] = taskc->core_sched_seq;
}

/*
 * To demonstrate the use of scx_bpf_dsq_move(), implement silly selective
 * priority boosting mechanism by scanning SHARED_DSQ looking for highpri tasks,
 * moving them to HIGHPRI_DSQ and then consuming them first. This makes minor
 * difference only when dsp_batch is larger than 1.
 *
 * scx_bpf_dispatch[_vtime]_from_dsq() are allowed both from ops.dispatch() and
 * non-rq-lock holding BPF programs. As demonstration, this function is called
 * from qmap_dispatch() and monitor_timerfn().
 */
static bool dispatch_highpri(bool from_timer)
{
	struct task_struct *p;
	s32 this_cid = scx_bpf_this_cid();
	u32 nr_cids = scx_bpf_nr_cids();

	/* scan SHARED_DSQ and move highpri tasks to HIGHPRI_DSQ */
	bpf_for_each(scx_dsq, p, SHARED_DSQ, 0) {
		static u64 highpri_seq;
		task_ctx_t *taskc;

		if (!(taskc = lookup_task_ctx(p)))
			return false;

		if (taskc->highpri) {
			/* exercise the set_*() and vtime interface too */
			scx_bpf_dsq_move_set_slice(BPF_FOR_EACH_ITER, slice_ns * 2);
			scx_bpf_dsq_move_set_vtime(BPF_FOR_EACH_ITER, highpri_seq++);
			scx_bpf_dsq_move_vtime(BPF_FOR_EACH_ITER, p, HIGHPRI_DSQ, 0);
		}
	}

	/*
	 * Scan HIGHPRI_DSQ and dispatch until a task that can run here is
	 * found. Prefer this_cid if the task allows it; otherwise RR-scan the
	 * task's cpus_allowed starting after this_cid.
	 */
	bpf_for_each(scx_dsq, p, HIGHPRI_DSQ, 0) {
		task_ctx_t *taskc;
		bool dispatched = false;
		s32 cid;

		if (!(taskc = lookup_task_ctx(p)))
			return false;

		if (cmask_test(this_cid, &taskc->cpus_allowed))
			cid = this_cid;
		else
			cid = cmask_next_set_wrap(&taskc->cpus_allowed,
						  this_cid + 1);
		if (cid >= nr_cids)
			continue;

		if (scx_bpf_dsq_move(BPF_FOR_EACH_ITER, p, SCX_DSQ_LOCAL_ON | cid,
				     SCX_ENQ_PREEMPT)) {
			if (cid == this_cid) {
				dispatched = true;
				__sync_fetch_and_add(&qa.nr_expedited_local, 1);
			} else {
				__sync_fetch_and_add(&qa.nr_expedited_remote, 1);
			}
			if (from_timer)
				__sync_fetch_and_add(&qa.nr_expedited_from_timer, 1);
		} else {
			__sync_fetch_and_add(&qa.nr_expedited_lost, 1);
		}

		if (dispatched)
			return true;
	}

	return false;
}

void BPF_STRUCT_OPS(qmap_dispatch, s32 cid, struct task_struct *prev)
{
	struct task_struct *p;
	struct cpu_ctx __arena *cpuc;
	task_ctx_t *taskc;
	u32 batch = dsp_batch ?: 1;
	s32 i;

	if (dispatch_highpri(false))
		return;

	if (!qa.nr_highpri_queued && scx_bpf_dsq_move_to_local(SHARED_DSQ, 0))
		return;

	if (dsp_inf_loop_after && qa.nr_dispatched > dsp_inf_loop_after) {
		/*
		 * PID 2 should be kthreadd which should mostly be idle and off
		 * the scheduler. Let's keep dispatching it to force the kernel
		 * to call this function over and over again.
		 */
		p = bpf_task_from_pid(2);
		if (p) {
			scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, slice_ns, 0);
			bpf_task_release(p);
			return;
		}
	}

	cpuc = &qa.cpu_ctxs[scx_bpf_this_cid()];

	for (i = 0; i < 5; i++) {
		/* Advance the dispatch cursor and pick the fifo. */
		if (!cpuc->dsp_cnt) {
			cpuc->dsp_idx = (cpuc->dsp_idx + 1) % 5;
			cpuc->dsp_cnt = 1 << cpuc->dsp_idx;
		}

		/* Dispatch or advance. */
		bpf_repeat(BPF_MAX_LOOPS) {
			task_ctx_t *taskc;

			taskc = qmap_fifo_pop(&qa.fifos[cpuc->dsp_idx]);
			if (!taskc)
				break;

			p = scx_bpf_tid_to_task(taskc->tid);
			if (!p)
				continue;

			if (taskc->highpri)
				__sync_fetch_and_sub(&qa.nr_highpri_queued, 1);

			update_core_sched_head_seq(p);
			__sync_fetch_and_add(&qa.nr_dispatched, 1);

			scx_bpf_dsq_insert(p, SHARED_DSQ, slice_ns, 0);

			/*
			 * scx_qmap uses a global BPF queue that any CPU's
			 * dispatch can pop from. If this CPU popped a task that
			 * can't run here, it gets stranded on SHARED_DSQ after
			 * consume_dispatch_q() skips it. Kick the task's home
			 * CPU so it drains SHARED_DSQ.
			 *
			 * There's a race between the pop and the flush of the
			 * buffered dsq_insert:
			 *
			 *  CPU 0 (dispatching)      CPU 1 (home, idle)
			 *  ~~~~~~~~~~~~~~~~~~~      ~~~~~~~~~~~~~~~~~~~
			 *  pop from BPF queue
			 *  dsq_insert(buffered)
			 *                           balance:
			 *                             SHARED_DSQ empty
			 *                             BPF queue empty
			 *                             -> goes idle
			 *  flush -> on SHARED
			 *  kick CPU 1
			 *                           wakes, drains task
			 *
			 * The kick prevents indefinite stalls but a per-CPU
			 * kthread like ksoftirqd can be briefly stranded when
			 * its home CPU enters idle with softirq pending,
			 * triggering:
			 *
			 *  "NOHZ tick-stop error: local softirq work is pending, handler #N!!!"
			 *
			 * from report_idle_softirq(). The kick lands shortly
			 * after and the home CPU drains the task. This could be
			 * avoided by e.g. dispatching pinned tasks to local or
			 * global DSQs, but the current code is left as-is to
			 * document this class of issue -- other schedulers
			 * seeing similar warnings can use this as a reference.
			 */
			if (!cmask_test(cid, &taskc->cpus_allowed))
				scx_bpf_kick_cid(scx_bpf_task_cid(p), 0);

			batch--;
			cpuc->dsp_cnt--;
			if (!batch || !scx_bpf_dispatch_nr_slots()) {
				if (dispatch_highpri(false))
					return;
				scx_bpf_dsq_move_to_local(SHARED_DSQ, 0);
				return;
			}
			if (!cpuc->dsp_cnt)
				break;
		}

		cpuc->dsp_cnt = 0;
	}

	for (i = 0; i < MAX_SUB_SCHEDS; i++) {
		if (qa.sub_sched_cgroup_ids[i] &&
		    scx_bpf_sub_dispatch(qa.sub_sched_cgroup_ids[i]))
			return;
	}

	/*
	 * No other tasks. @prev will keep running. Update its core_sched_seq as
	 * if the task were enqueued and dispatched immediately.
	 */
	if (prev) {
		taskc = lookup_task_ctx(prev);
		if (!taskc)
			return;

		taskc->core_sched_seq =
			qa.core_sched_tail_seqs[weight_to_idx(prev->scx.weight)]++;
	}
}

void BPF_STRUCT_OPS(qmap_tick, struct task_struct *p)
{
	struct cpu_ctx __arena *cpuc = &qa.cpu_ctxs[scx_bpf_this_cid()];
	int idx;

	/*
	 * Use the running avg of weights to select the target cpuperf level.
	 * This is a demonstration of the cpuperf feature rather than a
	 * practical strategy to regulate CPU frequency.
	 */
	cpuc->avg_weight = cpuc->avg_weight * 3 / 4 + p->scx.weight / 4;
	idx = weight_to_idx(cpuc->avg_weight);
	cpuc->cpuperf_target = qidx_to_cpuperf_target[idx];

	scx_bpf_cidperf_set(scx_bpf_task_cid(p), cpuc->cpuperf_target);
}

/*
 * The distance from the head of the queue scaled by the weight of the queue.
 * The lower the number, the older the task and the higher the priority.
 */
static s64 task_qdist(struct task_struct *p)
{
	int idx = weight_to_idx(p->scx.weight);
	task_ctx_t *taskc;
	s64 qdist;

	taskc = lookup_task_ctx(p);
	if (!taskc)
		return 0;

	qdist = taskc->core_sched_seq - qa.core_sched_head_seqs[idx];

	/*
	 * As queue index increments, the priority doubles. The queue w/ index 3
	 * is dispatched twice more frequently than 2. Reflect the difference by
	 * scaling qdists accordingly. Note that the shift amount needs to be
	 * flipped depending on the sign to avoid flipping priority direction.
	 */
	if (qdist >= 0)
		return qdist << (4 - idx);
	else
		return qdist << idx;
}

/*
 * This is called to determine the task ordering when core-sched is picking
 * tasks to execute on SMT siblings and should encode about the same ordering as
 * the regular scheduling path. Use the priority-scaled distances from the head
 * of the queues to compare the two tasks which should be consistent with the
 * dispatch path behavior.
 */
bool BPF_STRUCT_OPS(qmap_core_sched_before,
		    struct task_struct *a, struct task_struct *b)
{
	return task_qdist(a) > task_qdist(b);
}

/*
 * sched_switch tracepoint and cpu_release handlers are no longer needed.
 * With SCX_OPS_ALWAYS_ENQ_IMMED, wakeup_preempt_scx() reenqueues IMMED
 * tasks when a higher-priority scheduling class takes the CPU.
 */

s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init_task, struct task_struct *p,
			     struct scx_init_task_args *args)
{
	struct task_ctx_stor_val *v;
	task_ctx_t *taskc;

	if (p->tgid == disallow_tgid)
		p->scx.disallow = true;

	/* pop a slab entry off the free list */
	if (qmap_spin_lock(&qa_task_lock))
		return -EBUSY;
	taskc = qa.task_free_head;
	if (taskc)
		qa.task_free_head = taskc->next_free;
	bpf_res_spin_unlock(&qa_task_lock);
	if (!taskc) {
		scx_bpf_error("task_ctx slab exhausted (max_tasks=%u)", max_tasks);
		return -ENOMEM;
	}

	taskc->next_free = NULL;
	taskc->q_next = NULL;
	taskc->q_prev = NULL;
	taskc->fifo = NULL;
	taskc->tid = p->scx.tid;
	taskc->pid = p->pid;
	taskc->force_local = false;
	taskc->highpri = false;
	taskc->core_sched_seq = 0;
	cmask_init(&taskc->cpus_allowed, 0, scx_bpf_nr_cids());
	bpf_rcu_read_lock();
	cmask_from_cpumask(&taskc->cpus_allowed, p->cpus_ptr);
	bpf_rcu_read_unlock();

	v = bpf_task_storage_get(&task_ctx_stor, p, NULL,
				 BPF_LOCAL_STORAGE_GET_F_CREATE);
	if (!v) {
		/* push back to the free list */
		if (!qmap_spin_lock(&qa_task_lock)) {
			taskc->next_free = qa.task_free_head;
			qa.task_free_head = taskc;
			bpf_res_spin_unlock(&qa_task_lock);
		}
		return -ENOMEM;
	}
	v->taskc = taskc;
	return 0;
}

void BPF_STRUCT_OPS(qmap_exit_task, struct task_struct *p,
		    struct scx_exit_task_args *args)
{
	struct task_ctx_stor_val *v;
	task_ctx_t *taskc;

	v = bpf_task_storage_get(&task_ctx_stor, p, NULL, 0);
	if (!v || !v->taskc)
		return;
	taskc = v->taskc;
	v->taskc = NULL;

	if (qmap_spin_lock(&qa_task_lock))
		return;
	taskc->next_free = qa.task_free_head;
	qa.task_free_head = taskc;
	bpf_res_spin_unlock(&qa_task_lock);
}

void BPF_STRUCT_OPS(qmap_dump, struct scx_dump_ctx *dctx)
{
	task_ctx_t *taskc;
	s32 i;

	QMAP_TOUCH_ARENA();

	if (suppress_dump)
		return;

	/*
	 * Walk the queue lists without locking - kfunc calls (scx_bpf_dump)
	 * aren't in the verifier's kfunc_spin_allowed() list so we can't hold
	 * a lock and dump. Best-effort; racing may print stale tids but the
	 * walk is bounded by bpf_repeat() so it always terminates.
	 */
	bpf_for(i, 0, 5) {
		scx_bpf_dump("QMAP FIFO[%d]:", i);
		taskc = qa.fifos[i].head;
		bpf_repeat(4096) {
			if (!taskc)
				break;
			scx_bpf_dump(" %d:%llu", taskc->pid, taskc->tid);
			taskc = taskc->q_next;
		}
		scx_bpf_dump("\n");
	}
}

void BPF_STRUCT_OPS(qmap_dump_cid, struct scx_dump_ctx *dctx, s32 cid, bool idle)
{
	struct cpu_ctx __arena *cpuc = &qa.cpu_ctxs[cid];

	if (suppress_dump || idle)
		return;

	scx_bpf_dump("QMAP: dsp_idx=%llu dsp_cnt=%llu avg_weight=%u cpuperf_target=%u",
		     cpuc->dsp_idx, cpuc->dsp_cnt, cpuc->avg_weight,
		     cpuc->cpuperf_target);
}

void BPF_STRUCT_OPS(qmap_dump_task, struct scx_dump_ctx *dctx, struct task_struct *p)
{
	struct task_ctx_stor_val *v;
	task_ctx_t *taskc;

	QMAP_TOUCH_ARENA();

	if (suppress_dump)
		return;
	v = bpf_task_storage_get(&task_ctx_stor, p, NULL, 0);
	if (!v || !v->taskc)
		return;
	taskc = v->taskc;

	scx_bpf_dump("QMAP: force_local=%d core_sched_seq=%llu",
		     taskc->force_local, taskc->core_sched_seq);
}

s32 BPF_STRUCT_OPS(qmap_cgroup_init, struct cgroup *cgrp, struct scx_cgroup_init_args *args)
{
	if (print_msgs)
		bpf_printk("CGRP INIT %llu weight=%u period=%lu quota=%ld burst=%lu",
			   cgrp->kn->id, args->weight, args->bw_period_us,
			   args->bw_quota_us, args->bw_burst_us);
	return 0;
}

void BPF_STRUCT_OPS(qmap_cgroup_set_weight, struct cgroup *cgrp, u32 weight)
{
	if (print_msgs)
		bpf_printk("CGRP SET %llu weight=%u", cgrp->kn->id, weight);
}

void BPF_STRUCT_OPS(qmap_cgroup_set_bandwidth, struct cgroup *cgrp,
		    u64 period_us, u64 quota_us, u64 burst_us)
{
	if (print_msgs)
		bpf_printk("CGRP SET %llu period=%lu quota=%ld burst=%lu",
			   cgrp->kn->id, period_us, quota_us, burst_us);
}

void BPF_STRUCT_OPS(qmap_update_idle, s32 cid, bool idle)
{
	QMAP_TOUCH_ARENA();
	if (idle)
		cmask_set(cid, qa_idle_cids);
	else
		cmask_clear(cid, qa_idle_cids);
}

void BPF_STRUCT_OPS(qmap_set_cmask, struct task_struct *p,
		    const struct scx_cmask *cmask_in)
{
	struct scx_cmask __arena *cmask = (struct scx_cmask __arena *)(long)cmask_in;
	task_ctx_t *taskc;

	taskc = lookup_task_ctx(p);
	if (!taskc)
		return;
	cmask_copy(&taskc->cpus_allowed, cmask);
}

struct monitor_timer {
	struct bpf_timer timer;
};

struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, 1);
	__type(key, u32);
	__type(value, struct monitor_timer);
} monitor_timer SEC(".maps");

/*
 * Aggregate cidperf across the first nr_online_cids cids. Post-hotplug
 * the first-N-are-online invariant drifts, so some cap/cur values may
 * be stale. For this demo monitor that's fine; the scheduler exits on
 * the enable-time hotplug_seq mismatch and userspace restarts, which
 * rebuilds the layout.
 */
static void monitor_cpuperf(void)
{
	u32 nr_online = scx_bpf_nr_online_cids();
	u64 cap_sum = 0, cur_sum = 0, cur_min = SCX_CPUPERF_ONE, cur_max = 0;
	u64 target_sum = 0, target_min = SCX_CPUPERF_ONE, target_max = 0;
	s32 cid;

	QMAP_TOUCH_ARENA();

	bpf_for(cid, 0, nr_online) {
		struct cpu_ctx __arena *cpuc = &qa.cpu_ctxs[cid];
		u32 cap = scx_bpf_cidperf_cap(cid);
		u32 cur = scx_bpf_cidperf_cur(cid);
		u32 target;

		cur_min = cur < cur_min ? cur : cur_min;
		cur_max = cur > cur_max ? cur : cur_max;

		cur_sum += (u64)cur * cap / SCX_CPUPERF_ONE;
		cap_sum += cap;

		target = cpuc->cpuperf_target;
		target_sum += target;
		target_min = target < target_min ? target : target_min;
		target_max = target > target_max ? target : target_max;
	}

	if (!nr_online || !cap_sum)
		return;

	qa.cpuperf_min = cur_min;
	qa.cpuperf_avg = cur_sum * SCX_CPUPERF_ONE / cap_sum;
	qa.cpuperf_max = cur_max;

	qa.cpuperf_target_min = target_min;
	qa.cpuperf_target_avg = target_sum / nr_online;
	qa.cpuperf_target_max = target_max;
}

/*
 * Dump the currently queued tasks in the shared DSQ to demonstrate the usage of
 * scx_bpf_dsq_nr_queued() and DSQ iterator. Raise the dispatch batch count to
 * see meaningful dumps in the trace pipe.
 */
static void dump_shared_dsq(void)
{
	struct task_struct *p;
	s32 nr;

	if (!(nr = scx_bpf_dsq_nr_queued(SHARED_DSQ)))
		return;

	bpf_printk("Dumping %d tasks in SHARED_DSQ in reverse order", nr);

	bpf_rcu_read_lock();
	bpf_for_each(scx_dsq, p, SHARED_DSQ, SCX_DSQ_ITER_REV)
		bpf_printk("%s[%d]", p->comm, p->pid);
	bpf_rcu_read_unlock();
}

static int monitor_timerfn(void *map, int *key, struct bpf_timer *timer)
{
	bpf_rcu_read_lock();
	dispatch_highpri(true);
	bpf_rcu_read_unlock();

	monitor_cpuperf();

	if (print_dsqs_and_events) {
		struct scx_event_stats events;

		dump_shared_dsq();

		__COMPAT_scx_bpf_events(&events, sizeof(events));

		bpf_printk("%35s: %lld", "SCX_EV_SELECT_CPU_FALLBACK",
			   scx_read_event(&events, SCX_EV_SELECT_CPU_FALLBACK));
		bpf_printk("%35s: %lld", "SCX_EV_DISPATCH_LOCAL_DSQ_OFFLINE",
			   scx_read_event(&events, SCX_EV_DISPATCH_LOCAL_DSQ_OFFLINE));
		bpf_printk("%35s: %lld", "SCX_EV_DISPATCH_KEEP_LAST",
			   scx_read_event(&events, SCX_EV_DISPATCH_KEEP_LAST));
		bpf_printk("%35s: %lld", "SCX_EV_ENQ_SKIP_EXITING",
			   scx_read_event(&events, SCX_EV_ENQ_SKIP_EXITING));
		bpf_printk("%35s: %lld", "SCX_EV_REFILL_SLICE_DFL",
			   scx_read_event(&events, SCX_EV_REFILL_SLICE_DFL));
		bpf_printk("%35s: %lld", "SCX_EV_BYPASS_DURATION",
			   scx_read_event(&events, SCX_EV_BYPASS_DURATION));
		bpf_printk("%35s: %lld", "SCX_EV_BYPASS_DISPATCH",
			   scx_read_event(&events, SCX_EV_BYPASS_DISPATCH));
		bpf_printk("%35s: %lld", "SCX_EV_BYPASS_ACTIVATE",
			   scx_read_event(&events, SCX_EV_BYPASS_ACTIVATE));
	}

	bpf_timer_start(timer, ONE_SEC_IN_NS, 0);
	return 0;
}

struct lowpri_timer {
	struct bpf_timer timer;
};

struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, 1);
	__type(key, u32);
	__type(value, struct lowpri_timer);
} lowpri_timer SEC(".maps");

/*
 * Nice 19 tasks are put into the lowpri DSQ. Every 10ms, reenq is triggered and
 * the tasks are transferred to SHARED_DSQ.
 */
static int lowpri_timerfn(void *map, int *key, struct bpf_timer *timer)
{
	scx_bpf_dsq_reenq(LOWPRI_DSQ, 0);
	bpf_timer_start(timer, LOWPRI_INTV_NS, 0);
	return 0;
}

s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init)
{
	u8 __arena *slab;
	u32 nr_pages, key = 0, i;
	u32 nr_cids, nr_cpu_ids;
	struct bpf_timer *timer;
	s32 ret;

	nr_cids = scx_bpf_nr_cids();
	nr_cpu_ids = scx_bpf_nr_cpu_ids();

	if (nr_cids > SCX_QMAP_MAX_CPUS) {
		scx_bpf_error("nr_cids=%u exceeds SCX_QMAP_MAX_CPUS=%d",
			      nr_cids, SCX_QMAP_MAX_CPUS);
		return -EINVAL;
	}
	if (nr_cpu_ids > SCX_QMAP_MAX_CPUS) {
		scx_bpf_error("nr_cpu_ids=%u exceeds SCX_QMAP_MAX_CPUS=%d",
			      nr_cpu_ids, SCX_QMAP_MAX_CPUS);
		return -EINVAL;
	}

	/*
	 * cid-override test hook. Must run before anything that reads the
	 * cid space (scx_bpf_nr_cids, cmask_init, etc.). On invalid input,
	 * the kfunc calls scx_error() which aborts the scheduler.
	 */
	if (cid_override_mode) {
		scx_bpf_cid_override((const s32 *)cid_override_cpu_to_cid,
				     nr_cpu_ids * sizeof(s32));
	}

	/*
	 * Allocate the task_ctx slab in arena and thread the entire slab onto
	 * the free list. max_tasks is set by userspace before load. Each entry
	 * is TASK_CTX_STRIDE bytes - task_ctx's trailing cpus_allowed flex
	 * array extends into the stride tail.
	 */
	if (!max_tasks) {
		scx_bpf_error("max_tasks must be > 0");
		return -EINVAL;
	}

	nr_pages = (max_tasks * TASK_CTX_STRIDE + PAGE_SIZE - 1) / PAGE_SIZE;
	slab = bpf_arena_alloc_pages(&arena, NULL, nr_pages, NUMA_NO_NODE, 0);
	if (!slab) {
		scx_bpf_error("failed to allocate task_ctx slab");
		return -ENOMEM;
	}
	qa.task_ctxs = (task_ctx_t *)slab;

	bpf_for(i, 0, 5)
		qa.fifos[i].idx = i;

	bpf_for(i, 0, max_tasks) {
		task_ctx_t *cur = (task_ctx_t *)(slab + i * TASK_CTX_STRIDE);
		task_ctx_t *next = (i + 1 < max_tasks) ?
			(task_ctx_t *)(slab + (i + 1) * TASK_CTX_STRIDE) : NULL;
		cur->next_free = next;
	}
	qa.task_free_head = (task_ctx_t *)slab;

	/*
	 * Allocate and initialize the idle cmask. Starts empty - update_idle
	 * fills it as cpus enter idle.
	 */
	qa_idle_cids = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
	if (!qa_idle_cids) {
		scx_bpf_error("failed to allocate idle cmask");
		return -ENOMEM;
	}
	cmask_init(qa_idle_cids, 0, nr_cids);

	ret = scx_bpf_create_dsq(SHARED_DSQ, -1);
	if (ret) {
		scx_bpf_error("failed to create DSQ %d (%d)", SHARED_DSQ, ret);
		return ret;
	}

	ret = scx_bpf_create_dsq(HIGHPRI_DSQ, -1);
	if (ret) {
		scx_bpf_error("failed to create DSQ %d (%d)", HIGHPRI_DSQ, ret);
		return ret;
	}

	ret = scx_bpf_create_dsq(LOWPRI_DSQ, -1);
	if (ret)
		return ret;

	timer = bpf_map_lookup_elem(&monitor_timer, &key);
	if (!timer)
		return -ESRCH;
	bpf_timer_init(timer, &monitor_timer, CLOCK_MONOTONIC);
	bpf_timer_set_callback(timer, monitor_timerfn);
	ret = bpf_timer_start(timer, ONE_SEC_IN_NS, 0);
	if (ret)
		return ret;

	if (__COMPAT_has_generic_reenq()) {
		/* see lowpri_timerfn() */
		timer = bpf_map_lookup_elem(&lowpri_timer, &key);
		if (!timer)
			return -ESRCH;
		bpf_timer_init(timer, &lowpri_timer, CLOCK_MONOTONIC);
		bpf_timer_set_callback(timer, lowpri_timerfn);
		ret = bpf_timer_start(timer, LOWPRI_INTV_NS, 0);
		if (ret)
			return ret;
	}

	return 0;
}

void BPF_STRUCT_OPS(qmap_exit, struct scx_exit_info *ei)
{
	UEI_RECORD(uei, ei);
}

s32 BPF_STRUCT_OPS(qmap_sub_attach, struct scx_sub_attach_args *args)
{
	s32 i;

	for (i = 0; i < MAX_SUB_SCHEDS; i++) {
		if (!qa.sub_sched_cgroup_ids[i]) {
			qa.sub_sched_cgroup_ids[i] = args->ops->sub_cgroup_id;
			bpf_printk("attaching sub-sched[%d] on %s",
				   i, args->cgroup_path);
			return 0;
		}
	}

	return -ENOSPC;
}

void BPF_STRUCT_OPS(qmap_sub_detach, struct scx_sub_detach_args *args)
{
	s32 i;

	for (i = 0; i < MAX_SUB_SCHEDS; i++) {
		if (qa.sub_sched_cgroup_ids[i] == args->ops->sub_cgroup_id) {
			qa.sub_sched_cgroup_ids[i] = 0;
			bpf_printk("detaching sub-sched[%d] on %s",
				   i, args->cgroup_path);
			break;
		}
	}
}

SCX_OPS_CID_DEFINE(qmap_ops,
	       .flags			= SCX_OPS_ENQ_EXITING | SCX_OPS_TID_TO_TASK,
	       .select_cid		= (void *)qmap_select_cid,
	       .enqueue			= (void *)qmap_enqueue,
	       .dequeue			= (void *)qmap_dequeue,
	       .dispatch		= (void *)qmap_dispatch,
	       .tick			= (void *)qmap_tick,
	       .core_sched_before	= (void *)qmap_core_sched_before,
	       .set_cmask		= (void *)qmap_set_cmask,
	       .update_idle		= (void *)qmap_update_idle,
	       .init_task		= (void *)qmap_init_task,
	       .exit_task		= (void *)qmap_exit_task,
	       .dump			= (void *)qmap_dump,
	       .dump_cid		= (void *)qmap_dump_cid,
	       .dump_task		= (void *)qmap_dump_task,
	       .cgroup_init		= (void *)qmap_cgroup_init,
	       .cgroup_set_weight	= (void *)qmap_cgroup_set_weight,
	       .cgroup_set_bandwidth	= (void *)qmap_cgroup_set_bandwidth,
	       .sub_attach		= (void *)qmap_sub_attach,
	       .sub_detach		= (void *)qmap_sub_detach,
	       .init			= (void *)qmap_init,
	       .exit			= (void *)qmap_exit,
	       .timeout_ms		= 5000U,
	       .name			= "qmap");