/* SPDX-License-Identifier: GPL-2.0 */
/*
 * Copyright (c) 2023, 2024 Valve Corporation.
 * Author: Changwoo Min <changwoo@igalia.com>
 */
#ifndef __LAVD_H
#define __LAVD_H

#include <scx/common.bpf.h>
#include <bpf_arena_common.bpf.h>
#include <lib/ravg.h>
#include <lib/sdt_task.h>
#include <lib/atq.h>
#include <lib/cgroup.h>

/*
 * common macros
 */
#define U64_MAX		((u64)~0ULL)
#define S64_MAX		((s64)(U64_MAX >> 1))
#define U32_MAX		((u32)~0U)
#define S32_MAX		((s32)(U32_MAX >> 1))

#define MAX_RT_PRIO	100

#define LAVD_SHIFT			10
#define LAVD_SCALE			(1L << LAVD_SHIFT)
#define p2s(percent)			(((percent) << LAVD_SHIFT) / 100)
#define s2p(scale)			(((scale) * 100) >> LAVD_SHIFT)

#define cpdom_to_dsq(cpdom_id)		((cpdom_id) | LAVD_DSQ_TYPE_CPDOM << LAVD_DSQ_TYPE_SHFT)
#define cpdom_to_turb_dsq(cpdom_id)	((cpdom_id) | LAVD_DSQ_TYPE_CPDOM_TURB << LAVD_DSQ_TYPE_SHFT)
#define dsq_to_cpdom(dsq_id)		((dsq_id) & LAVD_DSQ_ID_MASK)
#define dsq_to_cpu(dsq_id)		((dsq_id) & LAVD_DSQ_ID_MASK)
#define dsq_type(dsq_id)		(((dsq_id) & LAVD_DSQ_TYPE_MASK) >> LAVD_DSQ_TYPE_SHFT)

/*
 *  DSQ (dispatch queue) IDs are 64bit of the format:
 *  Lower 63 bits are reserved by users
 *
 *   Bits: [63] [62 .. 14] [13 .. 12] [11 .. 0]
 *         [ B] [    R   ] [    T   ] [   ID  ]
 *
 *    B: Sched_ext built-in ID bit, see include/linux/sched/ext.h
 *    R: Reserved
 *    T: Type of LAVD DSQ
 *   ID: DSQ ID
 */
enum {
	LAVD_DSQ_TYPE_SHFT		= 12,
	LAVD_DSQ_TYPE_MASK		= 0x3 << LAVD_DSQ_TYPE_SHFT,
	LAVD_DSQ_ID_SHFT		= 0,
	LAVD_DSQ_ID_MASK		= 0xfff << LAVD_DSQ_ID_SHFT,
	LAVD_DSQ_NR_TYPES		= 3,
	LAVD_DSQ_TYPE_CPDOM		= 1,
	LAVD_DSQ_TYPE_CPDOM_TURB		= 2,
	LAVD_DSQ_TYPE_CPU		= 0,
};

/*
 * common constants
 */
enum consts_internal {
	CLOCK_BOOTTIME			= 7,
	CACHELINE_SIZE			= 64,

	NSEC_PER_USEC			= 1000ULL,
	NSEC_PER_MSEC			= (1000ULL * NSEC_PER_USEC),

	LAVD_TIME_ONE_SEC		= (1000ULL * NSEC_PER_MSEC),
	LAVD_MAX_RETRY			= 3,

	LAVD_TARGETED_LATENCY_NS	= (10ULL * NSEC_PER_MSEC),
	LAVD_SLICE_MIN_NS_DFL		= (500ULL * NSEC_PER_USEC), /* min time slice */
	LAVD_SLICE_MAX_NS_DFL		= (5ULL * NSEC_PER_MSEC), /* max time slice */
	LAVD_SLICE_BOOST_BONUS		= LAVD_SLICE_MIN_NS_DFL,
	LAVD_SLICE_BOOST_MAX		= (500ULL * NSEC_PER_MSEC),
	LAVD_SLICE_BOOST_UTIL_WALL	= p2s(95), /* < 95%: cpu utilization threshold for slice boost */
	LAVD_ACC_RUNTIME_MAX		= LAVD_SLICE_MAX_NS_DFL,
	LAVD_TASK_LAG_MAX		= (500ULL * NSEC_PER_MSEC),
	LAVD_DL_COMPETE_WINDOW		= ((300ULL * NSEC_PER_MSEC) >> 16), /* assuming task's latency
									       criticality is around 1000. */

	LAVD_LC_FREQ_MAX                = 100000, /* shortest interval: 10usec */
	LAVD_LC_RUNTIME_MAX		= LAVD_TIME_ONE_SEC,
	LAVD_LC_WEIGHT_BOOST_REGULAR	= 128, /* 2^7 */
	LAVD_LC_WEIGHT_BOOST_MEDIUM	= (2 * LAVD_LC_WEIGHT_BOOST_REGULAR),
	LAVD_LC_WEIGHT_BOOST_HIGH	= (2 * LAVD_LC_WEIGHT_BOOST_MEDIUM),
	LAVD_LC_WEIGHT_BOOST_HIGHEST	= (2 * LAVD_LC_WEIGHT_BOOST_HIGH),
	LAVD_LC_GREEDY_SHIFT		= 1, /* 50% */
	LAVD_LC_WAKE_INTERVAL_MIN	= LAVD_SLICE_MIN_NS_DFL,
	LAVD_LC_INH_RECEIVER_SHIFT	= 2, /* 25.0% of receiver's latency criticality */
	LAVD_LC_INH_GIVER_SHIFT		= 3, /* 12.5 of giver's latency criticality */
	LAVD_LC_LATENCY_SENSITIVE_THRESH = LAVD_SCALE - (LAVD_SCALE >> 3), /* top 12.5% most latency-critical tasks */
	LAVD_VULN_THRESH_STEP_SIZE	= 64, /* granularity for lat and util in threshold space */
	LAVD_VULN_THRESH_UTIL_STEPS	= LAVD_SCALE / LAVD_VULN_THRESH_STEP_SIZE, /* util sub-steps per lat level (16) */
	LAVD_VULN_THRESH_MAX		= (LAVD_SCALE / LAVD_VULN_THRESH_STEP_SIZE) * (LAVD_SCALE / LAVD_VULN_THRESH_STEP_SIZE), /* 16 lat × 16 util = 256 */
	LAVD_VULN_THRESH_INIT		= 32, /* initial threshold */

	LAVD_RAVG_HALFLIFE_NS		= (128ULL * NSEC_PER_MSEC),

	LAVD_SYS_STAT_INTERVAL_NS	= (10ULL * NSEC_PER_MSEC),
	LAVD_SYS_STAT_DECAY_TIMES	= ((2ULL * LAVD_TIME_ONE_SEC) / LAVD_SYS_STAT_INTERVAL_NS),

	LAVD_CPU_UTIL_MAX_FOR_CPUPERF	= p2s(85), /* 85.0% */
	LAVD_CPU_UTIL_THR_FOR_MAX_FREQ	= p2s(80), /* cpu utilization threshold to update max freq */

	LAVD_CC_REQ_CAPACITY_HEADROOM	= p2s(25), /* 25%: inflate required capacity by 25% to handle sudden spikes */
	LAVD_CC_PER_CPU_UTIL		= p2s(50), /* 50%: maximum per-CPU utilization */
	LAVD_CC_UTIL_SPIKE		= p2s(90), /* When the CPU utilization is almost full (90%),
						      it is likely that the actual utilization is even
						      higher than that. */
	LAVD_CC_CPU_PIN_INTERVAL	= (250ULL * NSEC_PER_MSEC),
	LAVD_CC_CPU_PIN_INTERVAL_DIV	= (LAVD_CC_CPU_PIN_INTERVAL / LAVD_SYS_STAT_INTERVAL_NS),

	LAVD_AP_HIGH_UTIL_DFL_SMT_RT	= p2s(25),
	LAVD_AP_HIGH_UTIL_DFL_NO_SMT_RT	= p2s(50), /* 50%: balanced mode when 10% < cpu util <= 50%,
							  performance mode when cpu util > 50% */

	LAVD_CPDOM_MIG_SHIFT_UL		= 2, /* when under-loaded:  1/2**2 = [-25.0%, +25.0%] */
	LAVD_CPDOM_MIG_SHIFT		= 3, /* when mildly loaded: 1/2**3 = [-12.5%, +12.5%] */
	LAVD_CPDOM_MIG_SHIFT_OL		= 4, /* when over-loaded:   1/2**4 = [-6.25%, +6.25%] */
	LAVD_CPDOM_MIG_PROB_FT		= (LAVD_SYS_STAT_INTERVAL_NS / LAVD_SLICE_MAX_NS_DFL), /* roughly twice per interval */

	LAVD_FUTEX_OP_INVALID		= -1,
};

/*
 * Per-CPU warmth: a proxy for the microarchitectural state a task builds up on
 * a CPU (L1/L2 cache, TLB entries). Heat accrues with on-CPU residence and
 * decays once the task leaves; a warmer previous CPU earns a longer wait
 * before the task migrates. Full heat is LAVD_SCALE.
 */
enum consts_cpu_warmth {
	LAVD_CPU_ID_NONE		= 0xffff, /* no warm CPU, for introspection */
	LAVD_CPU_WARM_SAT_NS		= (500ULL * NSEC_PER_USEC), /* residence to full heat */
	LAVD_CPU_WARM_LIFETIME_NS	= (5ULL * NSEC_PER_MSEC), /* linear decay to 0 once away */
};

enum consts_flags {
	LAVD_FLAG_FUTEX_BOOST		= (0x1 << 0), /* futex acquired or not */
	LAVD_FLAG_NEED_LOCK_BOOST	= (0x1 << 1), /* need to boost lock for deadline calculation */
	LAVD_FLAG_IS_GREEDY		= (0x1 << 2), /* task's overscheduling ratio compared to its nice priority */
	LAVD_FLAG_IS_AFFINITIZED	= (0x1 << 3), /* is this task pinned to a subset of all CPUs? */
	LAVD_FLAG_IS_WAKEUP		= (0x1 << 4), /* is this a wake up? */
	LAVD_FLAG_IS_SYNC_WAKEUP	= (0x1 << 5), /* is this a sync wake up? */
	LAVD_FLAG_ON_BIG		= (0x1 << 6), /* can a task run on a big core? */
	LAVD_FLAG_ON_LITTLE		= (0x1 << 7), /* can a task run on a little core? */
	LAVD_FLAG_SLICE_BOOST		= (0x1 << 8), /* task's time slice is boosted. */
	LAVD_FLAG_IDLE_CPU_PICKED	= (0x1 << 9), /* an idle CPU is picked at ops.select_cpu() */
	LAVD_FLAG_KSOFTIRQD		= (0x1 << 10), /* ksoftirqd/%u thread */
	LAVD_FLAG_WOKEN_BY_RT_DL	= (0x1 << 11), /* woken by a RT/DL task */
	LAVD_FLAG_WOKEN_BY_HARDIRQ	= (0x1 << 12), /* woken by a hardware interrupt */
	LAVD_FLAG_WOKEN_BY_SOFTIRQ	= (0x1 << 13), /* woken by a softirq */
	LAVD_FLAG_MIGRATION_AGGRESSIVE  = (0x1 << 14), /* immediate task migration is necessary. */
	LAVD_FLAG_DOMAIN_PINNED		= (0x1 << 15), /* task's cpumask is confined to a single compute domain */
	LAVD_FLAG_IS_EFFECTIVELY_PINNED	= (0x1 << 16), /* effective cpumask weight is 1 (permanent or migrate-disable) */
	LAVD_FLAG_WARM_CPU		= (0x1 << 17), /* wait on the previous CPU: enqueue on its per-CPU DSQ */
};

#define LAVD_MASK_MIGRATION		(LAVD_FLAG_MIGRATION_AGGRESSIVE)

/*
 * Suffix convention for time-related variables
 * --------------------------------------------
 *  - _wall: wall clock time
 *  - _invr: CPU capacity and frequency-invariant time
 *  - _wwgt: weighted wall clock time scaled by task's weight
 *  - _iwgt: weighted invariant time scaled by task's weight
 */

/*
 * Task context
 */
struct task_ctx {
	/* --- cacheline 0 boundary (0 bytes) --- */
	/*
	 * Do NOT change the position of atq. It should be at the beginning
	 * of the task_ctx.
	 */
	struct scx_task_cgroup_bw atq __attribute__((aligned(CACHELINE_SIZE)));

	/* --- cacheline 1 boundary (64 bytes): running/stopping hot --- */
	/*
	 * Accessed together in account_task_runtime() and lavd_stopping().
	 * Grouped here to avoid cross-cacheline traffic on every task switch.
	 */
	volatile u64	flags;		/* LAVD_FLAG_* */
	u64	slice_wall;		/* time slice (wall clock time) */
	u64	last_measured_wall_clk;	/* last time when running time was measured (wall clock) */
	u64	last_measured_pelt_clk;	/* last time when running time was measured (pelt clock) */
	u64	last_measured_task_clk;	/* last time when running time was measured (task clock) */
	u64	last_measured_exec;	/* p->se.sum_exec_runtime at last measurement */
	/*
	 * Accumulated runtime from runnable to quiescent state.
	 * Used to calculate avg_runtime_wall/invr and latency criticality.
	 */
	u64	acc_runtime_wall;
	u64	acc_runtime_invr;
	/*
	 * Total invariant CPU time consumed for this task scaled by task's weight.
	 * Used to calculate avg_svc_time_iwgt.
	 */
	u64	svc_time_iwgt;

	/* --- cacheline 2 boundary (128 bytes): enqueue/wakeup hot --- */
	/*
	 * Accessed together in calc_lat_cri(), calc_when_to_run(), and
	 * ops.enqueue()/ops.select_cpu(). Grouped here to avoid cross-cacheline
	 * traffic on the critical scheduling path.
	 */
	u16	lat_cri;		/* final context-aware latency criticality */
	u16	normalized_lat_cri;	/* lat_cri normalized to [0, 1024] scale */
	u16	lat_cri_waker;		/* waker's latency criticality */
	u16	lat_cri_wakee;		/* wakee's latency criticality */
	u16	perf_cri;		/* performance criticality of a task */
	volatile u32	cpdom_id;		/* chosen compute domain id at ops.enqueue() */
	volatile u32	suggested_cpu_id;	/* suggested CPU ID at ops.enqueue() and ops.select_cpu() */
	volatile s32	pinned_cpu_id;		/* pinned CPU id. -ENOENT if not pinned or not runnable. */
	u32	__pad0;
	u64	last_running_clk;	/* last time when scheduled in */
	u64	run_freq;		/* scheduling frequency in a second */
	u64	wait_freq;		/* waiting frequency in a second */
	u64	wake_freq;		/* waking-up frequency in a second */
	/*
	 * Average invariant runtime per schedule.
	 * Used to calculate latency criticality.
	 */
	u64	avg_runtime_invr;

	/* --- cacheline 3 boundary (192 bytes): lower-frequency / monitoring --- */
	/*
	 * Average wall-clock runtime per schedule.
	 * Used for reporting.
	 */
	u64	avg_runtime_wall;
	u64	last_runnable_clk;	/* last time when a task became runnable */
	u64	last_quiescent_clk;	/* last time when a task became asleep */
	u64	cgrp_id;		/* cgroup id of this task */
	u64	resched_interval_wall;	/* reschedule interval in ns: [last running, this running] */
	u64	last_slice_used_wall;	/* time(ns) used in last scheduled interval: [last running, last stopping] */
	u32	cpu_id;			/* where a task is running now */
	u32	prev_cpu_id;		/* where a task ran last time */
	u8	queued_in_cpdom_id;	/* cpdom this task's load is counted in; LAVD_CPDOM_MAX_NR = not queued */
	s16	queued_on_cpu_id;	/* primary CPU id this task's load is counted on; -1 = not queued */
	u32	queued_load_snapshot;	/* task_load_metric() value snapshotted at enqueue time for the per-cpdom counter */
	u32	queued_load_snapshot_cpu; /* task_load_metric() value snapshotted at enqueue time for the per-CPU counter */
	pid_t	pid;			/* pid for this task */
	pid_t	waker_pid;		/* last waker's PID */

	/* --- cacheline 5 boundary (320 bytes): ravg/util read-mostly group --- */
	u32	util_est __attribute__((aligned(CACHELINE_SIZE)));
					/* Estimated task util using ravg duty cycle */
	struct ravg_data avg_util_ravg;	/* Running average of task utilization using ravg */
	char	waker_comm[TASK_COMM_LEN + 1]; /* last waker's comm */

	/* --- per-CPU warmth (cache/TLB state) --- */
	u64	last_stopping_clk;	/* when cpu_heat was last integrated (task stopped) */
	u16	cpu_heat;		/* residence-integrated heat for cpu_id [0, LAVD_SCALE] */
} __attribute__((aligned(CACHELINE_SIZE)));

/*
 * Compute domain context
 * - system > numa node > llc domain > compute domain per core type (P or E)
 */
struct cpdom_ctx {
	/* --- cacheline 0 boundary (0 bytes): read-only --- */
	u64	id;				    /* id of this compute domain */
	u64	alt_id;				    /* id of the closest compute domain of alternative type */
	u8	numa_id;			    /* numa domain id */
	u8	llc_id;				    /* llc domain id */
	u8	is_big;				    /* is it a big core or little core? */
	u8	is_valid;			    /* is this a valid compute domain? */
	u8	nr_neighbors[LAVD_CPDOM_MAX_DIST];  /* number of neighbors per distance */
	u64	__cpumask[LAVD_CPU_ID_MAX/64];	    /* cpumasks belongs to this compute domain */
	u8	neighbor_ids[LAVD_CPDOM_MAX_DIST * LAVD_CPDOM_MAX_NR]; /* neighbor IDs per distance in circular distance order */

	/* --- cacheline 8 boundary (512 bytes): read-write, read-mostly --- */
	u8	is_stealer __attribute__((aligned(CACHELINE_SIZE))); /* this domain should steal tasks from others */
	u8	is_stealee;			    /* stealer domain should steal tasks from this domain */
	u16	nr_active_cpus;			    /* the number of active CPUs in this compute domain */
	u16	nr_acpus_temp;			    /* temp for nr_active_cpus */
	u64	qload_invr;			    /* queued load: sum of task_load_metric() for all queued tasks, tracked atomically */
	u64	load_invr;			    /* domain load for balancing: avg_util_invr_sum + qload_invr */
	u32	nr_queued_task;			    /* the number of queued tasks in this domain */
	u32	cur_util_wall_sum;		    /* the sum of CPU utilization in the current interval */
	u32	avg_util_wall_sum;		    /* the sum of average CPU utilization */
	u32	cur_util_invr_sum;		    /* the sum of invariant CPU utilization in the current interval */
	u32	avg_util_invr_sum;		    /* the sum of average invariant CPU utilization */
	u32	cur_steal_util_wall_sum;	    /* the sum of steal utilization in the current interval */
	u32	avg_steal_util_wall_sum;	    /* the sum of average steal utilization */
	u32	cur_steal_util_invr_sum;	    /* the sum of invariant steal utilization in the current interval */
	u32	avg_steal_util_invr_sum;	    /* the sum of average invariant steal utilization */
	u32	cur_dom_pinned_util_wall_sum;	    /* the sum of domain-pinned task utilization in the current interval */
	u32	avg_dom_pinned_util_wall_sum;	    /* the sum of average domain-pinned task utilization */
	u32	cur_dom_pinned_util_invr_sum;	    /* the sum of invariant domain-pinned task utilization in the current interval */
	u32	avg_dom_pinned_util_invr_sum;	    /* the sum of average invariant domain-pinned task utilization */
	u32	cap_sum_active_cpus;		    /* the sum of capacities of active CPUs in this domain */
	u32	cap_sum_temp;			    /* temp for cap_sum_active_cpus */
	u32	dsq_consume_lat;		    /* latency to consume from dsq, shows how contended the dsq is */

	/* per-cpdom preemption vulnerability threshold tracking */
	u32	vuln_thresh;			    /* unified lat/util threshold step [0, LAVD_VULN_THRESH_MAX] */
	u32	util_sum_steady;		    /* sum of util_est for steady CPUs in this cpdom */
	u32	util_sum_turb;		    /* sum of util_est for turbulent CPUs in this cpdom */
	u32	cap_sum_steady;		    /* sum of capacity for steady CPUs in this cpdom */
	u32	cap_sum_turb;		    /* sum of capacity for turbulent CPUs in this cpdom */
	u16	nr_steady_cpus;		    /* count of steady CPUs in this cpdom */
	u16	nr_turb_cpus;		    /* count of turbulent CPUs in this cpdom */

	s64	stealee_budget_invr;		    /* egress budget: how much load can leave this domain per round */
	s64	stealer_budget_invr;		    /* ingress budget: how much additional load this stealer can accept */
} __attribute__((aligned(CACHELINE_SIZE)));

#define get_neighbor_id(cpdomc, d, i) ((cpdomc)->neighbor_ids[((d) * LAVD_CPDOM_MAX_NR) + (i)])

/*
 * Atomically subtract @amount from the stealee's egress budget. Concurrent
 * stealers on other CPUs may call this in parallel, so use __sync_fetch_and_sub
 * to avoid race conditions. The signed s64 field lets the counter go slightly
 * negative on underflow; once it happens, the <= 0 check clears the stealee
 * flag so further stealers skip this domain for the rest of the round. Budgets
 * are recomputed from scratch each LB round. Hopefully, transient negativity is
 * harmless.
 */
static __always_inline void decrement_stealee_budget(struct cpdom_ctx *cpdomc,
						     u64 amount)
{
	__sync_fetch_and_sub(&cpdomc->stealee_budget_invr, amount);

	if (READ_ONCE(cpdomc->stealee_budget_invr) <= 0)
		WRITE_ONCE(cpdomc->is_stealee, false);
}

/*
 * Atomically subtract @amount from the stealer's ingress budget.
 * Same rationale as decrement_stealee_budget().
 */
static __always_inline void decrement_stealer_budget(struct cpdom_ctx *cpdomc,
						     u64 amount)
{
	__sync_fetch_and_sub(&cpdomc->stealer_budget_invr, amount);

	if (READ_ONCE(cpdomc->stealer_budget_invr) <= 0)
		WRITE_ONCE(cpdomc->is_stealer, false);
}

extern struct cpdom_ctx		cpdom_ctxs[LAVD_CPDOM_MAX_NR];
extern struct bpf_cpumask	cpdom_cpumask[LAVD_CPDOM_MAX_NR];
extern int			nr_cpdoms;

typedef struct task_ctx __arena task_ctx;

struct cpu_ctx *get_cpu_ctx(void);
struct cpu_ctx *get_cpu_ctx_id(s32 cpu_id);
struct cpu_ctx *get_cpu_ctx_task(const struct task_struct *p);

/*
 * CPU context
 */
struct cpu_ctx {
	/* --- cacheline 0 boundary (0 bytes): IPI-read + topology --- */
	/*
	 * Fields read by remote CPUs during IPI-based preemption decisions
	 * (flags, est_stopping_clk, running_clk, lat_cri), plus topology and
	 * classification fields set at init time. Keeping them in a dedicated
	 * cacheline prevents the write-heavy accumulators in CL1 from causing
	 * false-sharing with remote readers.
	 */
	volatile u64	flags;		/* cached copy of task's flags */
	volatile u64	est_stopping_clk; /* estimated stopping time */
	volatile u64	running_clk;	/* when a task starts running */
	volatile u16	lat_cri;	/* latency criticality */
	volatile u16	effective_capacity;/* the capacity that CPU can do right now */
	u16		cpu_id;		/* cpu id */
	u16		max_capacity;	/* the maximum capacity that CPU can do */
	volatile u64	sum_lat_cri;	/* sum of latency criticality */
	volatile u32	max_lat_cri;	/* maximum latency criticality */
	volatile u32	nr_pinned_tasks; /* the number of pinned tasks waiting for running on this CPU */
	u8		cpdom_id;	/* compute domain id */
	u8		big_core;	/* is it a big core? */
	u8		turbo_core;	/* is it a turbo core? */
	u8		llc_id;		/* llc domain id */
	u8		cpdom_alt_id;	/* compute domain id of alternative type */
	u8		is_online;	/* is this CPU online? */
	u8		__pad0[2];
	u32		cpuperf_cur;	/* CPU's current performance target */
	volatile s32	futex_op;	/* futex op in futex V1 */

	/* --- cacheline 1 boundary (64 bytes): write accumulators --- */
	/*
	 * Updated on every lavd_stopping() call. Isolated in their own
	 * cacheline so that the frequent local writes do not invalidate
	 * the IPI-read fields in CL0 for remote CPUs.
	 */
	/*
	 * Total wall-clock time this CPU has spent running scx tasks so far.
	 * Used to calculate non_scx_time.
	 */
	volatile u64	tot_task_time_wall;
	/*
	 * Total scx tasks' invariant service time on a CPU scaled by tasks' weights.
	 * Used to calculate avg_svc_time_iwgt.
	 */
	volatile u64	tot_task_time_iwgt;
	/*
	 * Total invariant CPU time consumed by SCX tasks in the current
	 * interval (capacity and frequency scaled). Used to calculate
	 * util_invr.
	 */
	volatile u64	tot_task_time_invr;
	/*
	 * Total wall-clock and invariant time spent running
	 * LAVD_FLAG_DOMAIN_PINNED tasks (tasks whose cpumask is confined
	 * to a single compute domain). Used to calculate
	 * dom_pinned_util_wall/invr.
	 */
	volatile u64	tot_dom_pinned_task_time_wall;
	volatile u64	tot_dom_pinned_task_time_invr;
	volatile u64	sum_perf_cri;	/* sum of performance criticality */
	volatile u32	min_perf_cri;	/* minimum performance criticality */
	volatile u32	max_perf_cri;	/* maximum performance criticality */
	volatile u32	nr_sched;	/* number of schedules */
	volatile u32	nr_preempt;

	/*
	 * Per-CPU task_ctx lookup cache. Local-only writes/reads, never
	 * accessed remotely. Used by get_task_ctx_curcpu() / get_task_ctx()
	 * to skip bpf_task_storage_get() when consecutive ops callbacks
	 * reference the same task on the same CPU.
	 *
	 * Keyed by (task_struct *, pid). Either alone is unsafe:
	 *  - task_struct * alone: SLUB can recycle a freed task_struct
	 *    address; a stale cache entry on an idle CPU would then alias
	 *    the new task and return the freed taskc.
	 *  - pid alone: a non-leader thread's pid changes when it calls
	 *    execve() -- de_thread() / exchange_tids() swaps the calling
	 *    thread's pid with the leader's, and the leader is then
	 *    released; entries keyed on the old pid would become stale
	 *    hits for the surviving thread.
	 *
	 * Together: ABA defeated by pid mismatch, execve swap defeated
	 * by address mismatch. cached_pid == 0 means invalid.
	 */
	u64		cached_task;		/* (struct task_struct *) as u64 */
	u64		cached_taskc_raw;	/* (task_ctx __arena *) as u64 */
	u32		cached_pid;

	/* --- cacheline 2 boundary (128 bytes): per-interval results --- */
	/*
	 * Updated once per sys_stat collection interval. Read by userspace
	 * monitoring and sys_stat aggregation.
	 */
	volatile u32	nr_x_migration;
	volatile u32	nr_perf_cri;
	volatile u32	nr_lat_cri;
	volatile u32	avg_util_wall;	/* average of the CPU utilization (based on wall clock time) */
	volatile u32	cur_util_wall;	/* CPU utilization of the current interval (based on wall clock time) */
	volatile u32	avg_util_invr;	/* average of the scaled CPU utilization, which is capacity and frequency invariant. */
	volatile u32	cur_util_invr;	/* the scaled CPU utilization of the current interval, which is capacity and frequency invariant. */
	volatile u32	lat_headroom;	/* latency headroom available to this CPU (inversely related to irq/steal time) */
	/*
	 * Steal utilization: steal_time as a fraction of duration_wall,
	 * in LAVD_SHIFT fixed-point. cur_* is the current interval value;
	 * avg_* is the asymmetric EWMA (fast-rise, slow-decay).
	 * Used for load balancing: a CPU with high steal utilization has
	 * less capacity remaining for SCX tasks.
	 */
	u32		cur_steal_util_wall;
	u32		avg_steal_util_wall;
	u32		cur_steal_util_invr;
	u32		avg_steal_util_invr;

	/*
	 * Domain-pinned task utilization: the fraction of duration_wall
	 * spent running LAVD_FLAG_DOMAIN_PINNED tasks, in LAVD_SHIFT
	 * fixed-point. cur_* is the current interval value; avg_* is the
	 * asymmetric EWMA (fast-rise, slow-decay).
	 */
	u32		cur_dom_pinned_util_wall;
	u32		avg_dom_pinned_util_wall;
	u32		cur_dom_pinned_util_invr;
	u32		avg_dom_pinned_util_invr;

	/* --- cacheline 3 boundary (192 bytes): sys_stat raw inputs --- */
	/*
	 * Updated at collect_sys_stat() and used as raw inputs for deriving
	 * utilization and other per-interval metrics.
	 */
	/*
	 * Steal time for the current interval: time the CPU was not running
	 * SCX tasks and not idle (= IRQ + hypervisor steal + RT/DL).
	 */
	u64		steal_time_wall;	/* wall clock */
	u64		steal_time_invr;	/* capacity + frequency invariant */
	volatile u64	idle_total_wall;/* total idle time so far (wall clock time) */
	volatile u64	idle_start_clk;	/* when the CPU becomes idle */
	/*
	 * Exponential weighted moving average of the observed performance
	 * factor (delta_pelt / task_wall) in LAVD_SHIFT fixed-point format.
	 * Updated each collect_sys_stat() interval when task_wall > 0. Used
	 * as a fallback for conv_wall_to_invr_obs() when the CPU has no active
	 * time in the current interval (e.g., mostly idle with only IRQ
	 * traffic), so that irq_steal_invr is estimated from recent history
	 * rather than defaulting to zero or assuming max frequency.
	 * Initialized to LAVD_SCALE at init_per_cpu_ctx() and
	 * cpu_ctx_init_online().
	 */
	u32		avg_perf_factor;
	volatile u32	max_freq_observed; /* maximum CPU frequency observed within an interval scaled to 1024 */
	volatile u32	max_freq;	/* maximum CPU frequency averaged across multiple intervals */
	u32		__pad2;
	/*
	 * Snapshot of scx_clock_task() taken at the end of the last
	 * collect_sys_stat() interval. scx_clock_task() advances during tasks
	 * (SCX, RT/DL, idle) but is frozen during IRQ and hypervisor steal.
	 * With NO_HZ_IDLE, rq->clock_task is only updated at scheduling
	 * events; it is stale for a currently-idle remote CPU. When a CPU
	 * wakes from idle, rq->clock_task catches up to include the elapsed
	 * idle duration. See collect_sys_stat() for how task_wall and
	 * irq_steal_wall are correctly derived from delta_task.
	 * Initialized at init_per_cpu_ctx() and cpu_ctx_init_online(),
	 * updated each collect_sys_stat().
	 */
	u64		prev_task_clk;
	/*
	 * Snapshot of scx_clock_pelt() taken at the end of the last
	 * collect_sys_stat() interval. scx_clock_pelt() advances only during
	 * active execution (SCX + RT/DL), normalized by CPU capacity and
	 * frequency. It is frozen during IRQ, steal, and idle.
	 * The delta over an interval satisfies:
	 *   delta_pelt = scx_task_time_invr + rt_dl_time_invr
	 * so subtracting tot_task_time_invr gives rt_dl_time_invr exactly, and
	 * the ratio delta_pelt / task_wall is the observed performance factor
	 * used by conv_wall_to_invr_obs().
	 * Initialized at init_per_cpu_ctx() and cpu_ctx_init_online(),
	 * updated each collect_sys_stat().
	 */
	u64		prev_pelt_clk;

	/* --- cacheline 4 boundary (256 bytes): cpumask scratch --- */
	struct bpf_cpumask __kptr *a_mask;	/* scratch: task & active */
	struct bpf_cpumask __kptr *o_mask;	/* scratch: task & overflow */
	struct bpf_cpumask __kptr *temp_mask;	/* scratch: general-purpose */
	struct bpf_cpumask __kptr *i_mask;	/* scratch: task & idle */
	struct bpf_cpumask __kptr *ia_mask;	/* scratch: idle & active */
	struct bpf_cpumask __kptr *io_mask;	/* scratch: idle & overflow */
	struct bpf_cpumask __kptr *iat_mask;	/* scratch: idle & active & turbo */

	struct ravg_data avg_irq_steal_ravg;	/* Running average of IRQ steal utilization using ravg */
	struct ravg_data avg_util_ravg;	/* Running average of CPU utilization using ravg */
	volatile u32	util_est;	/* Estimated CPU utilization from ravg tracking */

	/* --- cacheline 7 boundary (448 bytes): cross-CPU write-heavy accumulator --- */
	/*
	 * qload_invr is incremented from any CPU on enqueue (remote write
	 * via bpf_map_lookup_percpu_elem + __sync_fetch_and_add) and
	 * decremented likewise on dispatch/dequeue/exit.
	 */
	u64	qload_invr __attribute__((aligned(CACHELINE_SIZE)));
} __attribute__((aligned(CACHELINE_SIZE)));

extern const volatile u64	nr_llcs;	/* number of LLC domains */
const extern volatile u32	nr_cpu_ids;
extern volatile u64		nr_cpus_onln;	/* current number of online CPUs */

extern const volatile u16	cpu_capacity[LAVD_CPU_ID_MAX];
extern const volatile u8	cpu_big[LAVD_CPU_ID_MAX];
extern const volatile u8	cpu_turbo[LAVD_CPU_ID_MAX];

/* Logging helpers. */

extern const volatile bool	no_wake_sync;
extern const volatile bool	no_slice_boost;
extern const volatile u8	verbose;

#define debugln(fmt, ...)						\
({									\
	if (verbose > 0)						\
		bpf_printk("[%s:%d] " fmt, __func__, __LINE__,		\
					##__VA_ARGS__);			\
})

#define traceln(fmt, ...)						\
({									\
	if (verbose > 1)						\
		bpf_printk("[%s:%d] " fmt, __func__, __LINE__,		\
					##__VA_ARGS__);			\
})

/* Arithmetic helpers. */

#ifndef min
#define min(X, Y) (((X) < (Y)) ? (X) : (Y))
#endif

#ifndef max
#define max(X, Y) (((X) < (Y)) ? (Y) : (X))
#endif

#ifndef clamp
#define clamp(val, lo, hi) min(max(val, lo), hi)
#endif

u64 calc_avg(u64 old_val, u64 new_val);
u64 calc_asym_avg(u64 old_val, u64 new_val);

/* Bitmask helpers. */
static __always_inline int cpumask_next_set_bit(u64 *cpumask)
{
	/*
	 * Check the cpumask is not empty. ctzll(x) is only well-defined
	 * for nonzero x; that's why we check for zero earlier to avoid
	 * undefined behavior.
	 */
	if (!*cpumask)
		return -ENOENT;

	/* Find the next set bit. */
	int bit = ctzll(*cpumask);

	/*
	 * This is equivalent to finding and clearing the least significant set
	 * bit.  The statement works because subtracting one from a nonzero bit
	 * flips all bits from the lowest set bit (inclusive) to the rightmost
	 * position; Then, The logic here ANDing it with the original value
	 * clears the lowest set bit.
	 */
	*cpumask &= *cpumask - 1;
	return bit;
}

/* System statistics module .*/
extern struct sys_stat		sys_stat;

s32 init_sys_stat(u64 now);
int update_sys_stat(void);

extern volatile u64		performance_mode_ns;
extern volatile u64		balanced_mode_ns;
extern volatile u64		powersave_mode_ns;

/* Helpers from util.bpf.c for querying CPU/task state. */
extern const volatile bool	per_cpu_dsq;
extern const volatile u64	pinned_slice_ns;

extern volatile bool		reinit_cpumask_for_performance;
extern volatile bool		no_preemption;
extern volatile bool		no_core_compaction;
extern volatile bool		no_freq_scaling;

bool test_cpu_flag(struct cpu_ctx *cpuc, u64 flag);
void set_cpu_flag(struct cpu_ctx *cpuc, u64 flag);
void reset_cpu_flag(struct cpu_ctx *cpuc, u64 flag);

bool is_lock_holder(task_ctx *taskc);
bool is_lock_holder_running(struct cpu_ctx *cpuc);
bool have_scheduled(task_ctx *taskc);
bool have_pending_tasks(struct cpu_ctx *cpuc);
bool can_boost_slice(void);
bool is_lat_cri(task_ctx *taskc);
u16 get_nice_prio(struct task_struct *p);
u32 cpu_to_dsq(u32 cpu);

void set_task_flag(task_ctx *taskc, u64 flag);
void reset_task_flag(task_ctx *taskc, u64 flag);
bool test_task_flag(task_ctx *taskc, u64 flag);
bool test_task_flag_mask(task_ctx __arg_arena *taskc, u64 flag);

extern const volatile u64	warm_cpu_ns;	/* warm-CPU wait budget (ns) */

/* Per-CPU warmth clock (util.bpf.c). */
u64 task_cpu_warmth(task_ctx __arg_arena *taskc, u32 cpu_id, u64 now);
void task_update_cpu_warmth(task_ctx __arg_arena *taskc, struct cpu_ctx *cpuc,
			    u64 slice_used, u64 now);

static __always_inline bool use_per_cpu_dsq(void)
{
	return per_cpu_dsq || pinned_slice_ns || warm_cpu_ns;
}

static __always_inline  bool is_per_cpu_dsq_migratable(void)
{
	/*
	 * When per_cpu-dsq is on, all tasks go to the per-CPU DSQ.
	 * So a task on a per-CPU DSQ can be migrated to another CPU.
	 * However, when pinned_slice_ns is on but per_cpu-dsq is not,
	 * only pinned tasks go to the per-CPU DSQ.
	 * Hence, tasks in a per-CPU DSQ are not migratable.
	 */
	return per_cpu_dsq;
}

static __always_inline bool use_cpdom_dsq(void)
{
	return !per_cpu_dsq;
}

bool queued_on_cpu(struct cpu_ctx *cpuc);
u64 get_target_dsq_id(struct task_struct *p, struct cpu_ctx *cpuc, task_ctx *taskc);
u16 normalize_lat_cri(u16 lat_cri);

/*
 * Compute a task's preemption vulnerability — how likely it is to be
 * routed to the turbulent DSQ. Lat is the major axis (weighted by
 * LAVD_VULN_THRESH_UTIL_STEPS) and util is the minor axis, so the
 * threshold naturally cascades: util sub-steps carry into lat levels
 * when vuln_thresh is incremented/decremented.
 */
static __always_inline
u32 preemption_vulnerability(u16 normalized_lat_cri, u32 util_est)
{
	u32 lat_step = normalized_lat_cri / LAVD_VULN_THRESH_STEP_SIZE;
	u32 util_step = util_est / LAVD_VULN_THRESH_STEP_SIZE;
	return lat_step * LAVD_VULN_THRESH_UTIL_STEPS + util_step;
}

/*
 * Does the predicted wait for @cpu to free up fit the task's warmth-extended
 * budget? The base budget is warm_cpu_ns; warm cache and TLB state on @cpu
 * stretch it up to 2x, so a task still warm there waits rather than migrate to
 * a cold CPU and refill. The wait is the time until the running task stops plus
 * the service time of tasks already queued ahead on that CPU's DSQ.
 */
static __always_inline
bool warm_cpu_wait_ok(task_ctx *taskc, s32 cpu, u64 now)
{
	struct cpu_ctx *cpuc = get_cpu_ctx_id(cpu);
	u64 heat, budget, est, wait;

	if (!cpuc)
		return false;

	heat = task_cpu_warmth(taskc, cpu, now);
	budget = (warm_cpu_ns * (LAVD_SCALE + heat)) >> LAVD_SHIFT;
	est = READ_ONCE(cpuc->est_stopping_clk);
	wait = time_delta(est, now);

	/*
	 * Add the wait for tasks already queued ahead on @cpu. This is rough: it
	 * assumes @p is served last and that every queued task runs the
	 * system-average slice. A per-core qload_invr would sharpen the latter;
	 * revisit once per-core queued load is tracked.
	 */
	wait += (u64)scx_bpf_dsq_nr_queued(cpu_to_dsq(cpu)) * sys_stat.slice_wall;

	return wait <= budget;
}

/*
 * Return the task's load contribution for queued load tracking.
 * Uses RAVG util_est: [0, 1024], 128ms half-life, tracks duty cycle.
 * Updated in lavd_runnable/lavd_quiescent — active under sched_ext.
 *
 * Note: kernel PELT p->se.avg.util_avg is NOT updated when sched_ext
 * is enabled (tasks bypass CFS), so it cannot be used here.
 */
static __always_inline u32 task_load_metric(task_ctx *taskc)
{
	return taskc->util_est;
}

extern struct bpf_cpumask __kptr *turbo_cpumask; /* CPU mask for turbo CPUs */
extern struct bpf_cpumask __kptr *big_cpumask; /* CPU mask for big CPUs */
extern struct bpf_cpumask __kptr *active_cpumask; /* CPU mask for active CPUs */
extern struct bpf_cpumask __kptr *ovrflw_cpumask; /* CPU mask for overflow CPUs */
extern struct bpf_cpumask __kptr *steady_cpumask; /* CPU mask for non-turbulent (steady) CPUs */

/* DSQ helpers. */

struct dsq_entry {
	u64 dsq_id;
	u64 vtime;
	bool eligible;
};

u64 peek_dsq_vtime(u64 dsq_id);
void sort_dsqs(struct dsq_entry *a, struct dsq_entry *b, struct dsq_entry *c);

/* Load balancer helpers. */

int plan_x_cpdom_migration(void);

/* Preemption management helpers. */
void shrink_slice_at_tick(struct task_struct *p, struct cpu_ctx *cpuc, u64 now);

/* Futex lock-related helpers. */

void reset_lock_futex_boost(task_ctx *taskc, struct cpu_ctx *cpuc);

/* Scheduler introspection-related helpers. */

u64 get_est_stopping_clk(task_ctx *taskc, u64 now);
void try_proc_introspec_cmd(struct task_struct *p, task_ctx *taskc);
void reset_cpu_preemption_info(struct cpu_ctx *cpuc);
int shrink_boosted_slice_remote(struct cpu_ctx *cpuc, u64 now);
void shrink_boosted_slice_at_tick(struct task_struct *p,
					 struct cpu_ctx *cpuc, u64 now);
void preempt_at_tick(struct task_struct *p, struct cpu_ctx *cpuc);
void try_find_and_kick_victim_cpu(struct task_struct *p,
					 task_ctx *taskc,
					 s32 preferred_cpu,
					 u64 dsq_id);

extern volatile bool is_monitored;


/* Idle CPU pick helpers */

struct pick_ctx {
	/*
	 * Input arguments for pick_idle_cpu().
	 */
	const struct task_struct *p;
	task_ctx *taskc;
	u64 wake_flags;
	s32 prev_cpu;
	/*
	 * Additional input arguments for find_sticky_cpu_and_cpdom().
	 */
	s32 sync_waker_cpu;
	/*
	 * Additional output arguments for init_active_ovrflw_masks().
	 */
	struct bpf_cpumask *active; /* global active mask */
	struct bpf_cpumask *ovrflw; /* global overflow mask */
	/*
	 * Additional output arguments for init_ao_masks().
	 * Additional input arguments for find_sticky_cpu_and_cpdom().
	 */
	struct cpu_ctx *cpuc_cur;
	struct bpf_cpumask *a_mask; /* task's active mask */
	struct bpf_cpumask *o_mask; /* task's overflow mask */
	/*
	 * Additional input arguments for init_idle_i_mask().
	 */
	const struct cpumask *i_mask;
	/*
	 * Additional input arguments for init_idle_ato_masks().
	 * Additional input arguments for pick_idle_cpu_at_cpdom().
	 */
	struct bpf_cpumask *ia_mask;
	struct bpf_cpumask *iat_mask;
	struct bpf_cpumask *io_mask;
	struct bpf_cpumask *temp_mask;
	/*
	 * Flags.
	 */
	bool a_empty:1;
	bool o_empty:1;
	bool is_task_big:1;
	bool i_empty:1;
	bool ia_empty:1;
	bool iat_empty:1;
	bool io_empty:1;
};


s32 find_cpu_in(const struct cpumask *src_mask, struct cpu_ctx *cpuc_cur);
s32  pick_idle_cpu(struct pick_ctx *ctx, bool *is_idle);

bool consume_task(u64 cpu_dsq_id, u64 cpdom_dsq_id);

extern u64 cur_logical_clk;
u64 calc_when_to_run(struct task_struct *p, task_ctx *taskc);

#endif /* __LAVD_H */
