scx_cosmos/

main.rs

// SPDX-License-Identifier: GPL-2.0
//
// Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>

// This software may be used and distributed according to the terms of the
// GNU General Public License version 2.

mod bpf_skel;
pub use bpf_skel::*;
pub mod bpf_intf;
pub use bpf_intf::*;

mod cgroup;
mod gpu;
mod stats;
use cgroup::CgroupReader;

use std::collections::{HashMap, HashSet};
use std::ffi::{c_int, c_ulong};
use std::mem::MaybeUninit;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};

use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use clap::Parser;
use crossbeam::channel::RecvTimeoutError;
use libbpf_rs::MapCore;
use libbpf_rs::MapFlags;
use libbpf_rs::OpenObject;
use libbpf_rs::ProgramInput;
use log::{debug, info, warn};
use nvml_wrapper::Nvml;
use nvml_wrapper::bitmasks::InitFlags;
use scx_stats::prelude::*;
use scx_utils::GpuIndex;
use scx_utils::NR_CPU_IDS;
use scx_utils::Powermode;
use scx_utils::Topology;
use scx_utils::UserExitInfo;
use scx_utils::build_id;
use scx_utils::compat;
use scx_utils::get_primary_cpus;
use scx_utils::libbpf_clap_opts::LibbpfOpts;
use scx_utils::perf::PerfEventSpec;
use scx_utils::perf::parse_perf_event;
use scx_utils::perf::setup_perf_events;
use scx_utils::scx_ops_attach;
use scx_utils::scx_ops_load;
use scx_utils::scx_ops_open;
use scx_utils::try_set_rlimit_infinity;
use scx_utils::uei_exited;
use scx_utils::uei_report;
use stats::Metrics;

const SCHEDULER_NAME: &str = "scx_cosmos";

#[derive(Debug, clap::Parser)]
#[command(
    name = "scx_cosmos",
    version,
    disable_version_flag = true,
    about = "Lightweight scheduler optimized for preserving task-to-CPU locality."
)]
struct Opts {
    /// Exit debug dump buffer length. 0 indicates default.
    #[clap(long, default_value = "0")]
    exit_dump_len: u32,

    /// Maximum scheduling slice duration in microseconds.
    #[clap(short = 's', long, default_value = "1000")]
    slice_us: u64,

    /// Maximum runtime (since last sleep) that can be charged to a task in microseconds.
    #[clap(short = 'l', long, default_value = "20000")]
    slice_lag_us: u64,

    /// CPU busy threshold.
    ///
    /// Specifies the user CPU utilization percentage (0-100%) at which the scheduler considers a
    /// CPU to be busy. Only user time counts: sleep-intensive workloads burn most of their CPU
    /// time in the kernel and their CPUs are deliberately not considered busy, so that tasks stay
    /// on their CPU / local dispatch queue, reducing the scheduling overhead and the balancing
    /// pressure in dispatch.
    ///
    /// When a CPU crosses this threshold, the scheduler switches it from using its per-CPU
    /// round-robin dispatch queue (which favors locality and reduced locking contention) to the
    /// global deadline-based dispatch queue (which improves load balancing).
    ///
    /// The global dispatch queue can increase task migrations and improve responsiveness for
    /// interactive tasks under heavy load. Lower values make the scheduler switch to deadline
    /// mode sooner, improving overall responsiveness at the cost of reducing single-task
    /// performance due to the additional migrations. Higher values makes task more "sticky" to
    /// their CPU, improving workloads that benefit from cache locality.
    ///
    /// A higher value is recommended for server-type workloads, while a lower value is recommended
    /// for interactive-type workloads.
    ///
    /// 0 = utilization tracking disabled: every CPU is always considered busy (pure deadline
    /// mode) with no tracking overhead.
    #[clap(short = 'c', long, default_value = "0")]
    cpu_busy_thresh: u64,

    /// Polling time (ms) to refresh the dynamic perf event thresholds (options -E / -Y in
    /// dynamic mode).
    ///
    /// Value is clamped to the range [10 .. 1000].
    ///
    /// 0 = disabled.
    #[clap(short = 'p', long, default_value = "0")]
    polling_ms: u64,

    /// Specifies a list of CPUs to prioritize.
    ///
    /// Accepts a comma-separated list of CPUs or ranges (i.e., 0-3,12-15) or the following special
    /// keywords:
    ///
    /// "turbo" = automatically detect and prioritize the CPUs with the highest max frequency,
    /// "performance" = automatically detect and prioritize the fastest CPUs,
    /// "powersave" = automatically detect and prioritize the slowest CPUs,
    /// "all" = all CPUs assigned to the primary domain.
    ///
    /// By default "all" CPUs are used.
    #[clap(short = 'm', long)]
    primary_domain: Option<String>,

    /// Maximum time (us) that tasks can wait in a shared DSQ before the primary domain is
    /// considered overloaded, allowing tasks to spill to the non-primary CPUs.
    ///
    /// Tasks are normally contained in the primary domain (option -m). When the domain's shared
    /// DSQ has been continuously backlogged for longer than this threshold, tasks are allowed to
    /// run on the non-primary CPUs until the backlog drains.
    ///
    /// 0 = tasks are always contained in the primary domain, even when it is overloaded.
    ///
    /// This option is ignored when no primary domain is defined.
    #[clap(long, default_value = "1000")]
    primary_overload_us: u64,

    /// Hardware perf event to monitor (0x0 = disabled). Accepts hex (0xN) or symbolic names
    /// (e.g. cache-misses, LLC-load-misses, page-faults, branch-misses).
    #[clap(short = 'e', long, default_value = "0x0", value_parser = parse_perf_event)]
    perf_config: PerfEventSpec,

    /// Threshold (perf events/msec) to classify a task as event heavy; exceeding it triggers migration.
    #[clap(short = 'E', default_value = "0", long)]
    perf_threshold: u64,

    /// Sticky perf event (0x0 = disabled). When a task exceeds -Y for this event, keep it on the same CPU.
    /// Accepts hex (0xN) or symbolic names (e.g. cache-misses, LLC-load-misses).
    #[clap(short = 'y', long, default_value = "0x0", value_parser = parse_perf_event)]
    perf_sticky: PerfEventSpec,

    /// Sticky perf threshold; task is kept on same CPU when its count for -y event exceeds this.
    #[clap(short = 'Y', default_value = "0", long)]
    perf_sticky_threshold: u64,

    /// Enable GPU-aware scheduling.
    #[clap(short = 'g', long, action = clap::ArgAction::SetTrue)]
    gpu: bool,

    /// Only treat a process as GPU-bound if its GPU utilization is at least this percentage (0–100).
    ///
    /// Uses NVML process utilization (SM + memory). 0 = no filter (all processes on the GPU are
    /// considered GPU-bound). Requires driver support (Maxwell or newer).
    #[clap(long, default_value = "0", value_parser = clap::value_parser!(u32).range(0..=100))]
    gpu_util_threshold: u32,

    /// Disable NUMA optimizations.
    #[clap(short = 'n', long, action = clap::ArgAction::SetTrue)]
    disable_numa: bool,

    /// Disable CPU frequency control.
    #[clap(short = 'f', long, action = clap::ArgAction::SetTrue)]
    disable_cpufreq: bool,

    /// Enable flat idle CPU scanning.
    ///
    /// This option can help reducing some overhead when trying to allocate idle CPUs and it can be
    /// quite effective with simple CPU topologies.
    #[arg(short = 'i', long, action = clap::ArgAction::SetTrue)]
    flat_idle_scan: bool,

    /// Enable preferred idle CPU scanning.
    ///
    /// With this option enabled, the scheduler will prioritize assigning tasks to higher-ranked
    /// cores before considering lower-ranked ones.
    #[clap(short = 'P', long, action = clap::ArgAction::SetTrue)]
    preferred_idle_scan: bool,

    /// Disable SMT.
    ///
    /// This option can only be used together with --flat-idle-scan or --preferred-idle-scan,
    /// otherwise it is ignored.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    disable_smt: bool,

    /// ***DEPRECATED*** SMT contention avoidance.
    #[clap(short = 'S', long, action = clap::ArgAction::SetTrue)]
    avoid_smt: bool,

    /// Disable early clearing of idle CPU state.
    ///
    /// When enabled, multiple concurrent wakeups can select the same idle CPU
    /// before it fully wakes up. This can improve performance in highly communicative
    /// workloads by aggressively stacking tasks on the same cache.
    #[clap(short = 'N', long, action = clap::ArgAction::SetTrue)]
    no_early_clear: bool,

    /// Disable direct dispatch during synchronous wakeups.
    ///
    /// Enabling this option can lead to a more uniform load distribution across available cores,
    /// potentially improving performance in certain scenarios. However, it may come at the cost of
    /// reduced efficiency for pipe-intensive workloads that benefit from tighter producer-consumer
    /// coupling.
    #[clap(short = 'w', long, action = clap::ArgAction::SetTrue)]
    no_wake_sync: bool,

    /// ***DEPRECATED*** Disable deferred wakeups.
    #[clap(short = 'd', long, action = clap::ArgAction::SetTrue)]
    no_deferred_wakeup: bool,

    /// Enable high-resolution timer preemption.
    ///
    /// By default, the scheduler preempts tasks that exceed their time slice, measuring the time
    /// slice via the tick handler. Add an option to enforce preemption based on the high-precision
    /// timer and CPU occupancy. Enable this option to improve latency-sensitive workloads.
    #[clap(long, action = clap::ArgAction::SetTrue)]
    time_preemption: bool,

    /// Enable address space affinity.
    ///
    /// This option allows to keep tasks that share the same address space (e.g., threads of the
    /// same process) on the same CPU across wakeups.
    ///
    /// This can improve locality and performance in certain cache-sensitive workloads.
    #[clap(short = 'a', long, action = clap::ArgAction::SetTrue)]
    mm_affinity: bool,

    /// Enable stats monitoring with the specified interval.
    #[clap(long)]
    stats: Option<f64>,

    /// Run in stats monitoring mode with the specified interval. Scheduler
    /// is not launched.
    #[clap(long)]
    monitor: Option<f64>,

    /// Enable verbose output, including libbpf details.
    #[clap(short = 'v', long, action = clap::ArgAction::SetTrue)]
    verbose: bool,

    /// Print scheduler version and exit.
    #[clap(short = 'V', long, action = clap::ArgAction::SetTrue)]
    version: bool,

    /// Show descriptions for statistics.
    #[clap(long)]
    help_stats: bool,

    #[clap(flatten, next_help_heading = "Libbpf Options")]
    pub libbpf: LibbpfOpts,
}

pub fn parse_cpu_list(optarg: &str) -> Result<Vec<usize>, String> {
    let mut cpus = Vec::new();
    let mut seen = HashSet::new();

    // Handle special keywords
    if let Some(mode) = match optarg {
        "powersave" => Some(Powermode::Powersave),
        "performance" => Some(Powermode::Performance),
        "turbo" => Some(Powermode::Turbo),
        "all" => Some(Powermode::Any),
        _ => None,
    } {
        return get_primary_cpus(mode).map_err(|e| e.to_string());
    }

    // Validate input characters
    if optarg
        .chars()
        .any(|c| !c.is_ascii_digit() && c != '-' && c != ',' && !c.is_whitespace())
    {
        return Err("Invalid character in CPU list".to_string());
    }

    // Replace all whitespace with tab (or just trim later)
    let cleaned = optarg.replace(' ', "\t");

    for token in cleaned.split(',') {
        let token = token.trim_matches(|c: char| c.is_whitespace());

        if token.is_empty() {
            continue;
        }

        if let Some((start_str, end_str)) = token.split_once('-') {
            let start = start_str
                .trim()
                .parse::<usize>()
                .map_err(|_| "Invalid range start")?;
            let end = end_str
                .trim()
                .parse::<usize>()
                .map_err(|_| "Invalid range end")?;

            if start > end {
                return Err(format!("Invalid CPU range: {}-{}", start, end));
            }

            for i in start..=end {
                if cpus.len() >= *NR_CPU_IDS {
                    return Err(format!("Too many CPUs specified (max {})", *NR_CPU_IDS));
                }
                if seen.insert(i) {
                    cpus.push(i);
                }
            }
        } else {
            let cpu = token
                .parse::<usize>()
                .map_err(|_| format!("Invalid CPU: {}", token))?;
            if cpus.len() >= *NR_CPU_IDS {
                return Err(format!("Too many CPUs specified (max {})", *NR_CPU_IDS));
            }
            if seen.insert(cpu) {
                cpus.push(cpu);
            }
        }
    }

    Ok(cpus)
}

/// Initial value for the dynamic threshold (in BPF units).
const DYNAMIC_THRESHOLD_INIT_VALUE: u64 = 1000;

/// Minimum value for the dynamic threshold (in BPF units).
const DYNAMIC_THRESHOLD_MIN_VALUE: u64 = 10;

/// Target event rate (per second) above which we consider migrations/sticky dispatches too high.
const DYNAMIC_THRESHOLD_RATE_HIGH: f64 = 4000.0;

/// Target event rate (per second) below which we consider migrations/sticky dispatches too low.
const DYNAMIC_THRESHOLD_RATE_LOW: f64 = 2000.0;

/// Hysteresis band: rate must move by this fraction beyond the target bounds before we act.
/// This prevents oscillation when the rate hovers near the threshold boundaries.
const DYNAMIC_THRESHOLD_HYSTERESIS: f64 = 0.1;

/// EMA smoothing factor (alpha). Higher values give more weight to recent samples.
/// 0.3 provides good balance between responsiveness and stability.
const DYNAMIC_THRESHOLD_EMA_ALPHA: f64 = 0.3;

/// Minimum scale factor when just outside the target band (slow convergence near optimal).
const DYNAMIC_THRESHOLD_SCALE_MIN: f64 = 0.0001;

/// Maximum scale factor when far from target (fast convergence when initial threshold is way off).
const DYNAMIC_THRESHOLD_SCALE_MAX: f64 = 1000.0;

/// Slope for "too high" case: scale grows with (rate/HIGH - 1) so we step much harder when rate is
/// many times over target.
const DYNAMIC_THRESHOLD_SLOPE_HIGH: f64 = 0.35;

/// Slope for "too low" case: scale grows with deficit so we step harder when rate is near zero.
const DYNAMIC_THRESHOLD_SLOPE_LOW: f64 = 0.58;

/// Minimum interval between NVML GPU PID syncs. Kept separate from CPU polling so that fast
/// polling (e.g. 100 ms) does not trigger expensive NVML calls every tick.
const GPU_SYNC_INTERVAL: Duration = Duration::from_secs(1);

/// State for EMA-based dynamic threshold adjustment with hysteresis.
///
/// This struct maintains the smoothed rate estimate and tracks whether we're
/// currently in an adjustment state (raising or lowering threshold) to implement
/// hysteresis and prevent oscillation.
#[derive(Debug, Clone)]
struct DynamicThresholdState {
    /// Current threshold value.
    threshold: u64,
    /// EMA-smoothed rate estimate.
    smoothed_rate: f64,
    /// Previous raw counter value for delta calculation.
    prev_counter: u64,
    /// Whether the EMA has been initialized with a valid sample.
    initialized: bool,
    /// Current adjustment direction: None (stable), Some(true) = raising, Some(false) = lowering.
    /// Used for hysteresis: once we start adjusting in a direction, we continue until
    /// the rate crosses back into the stable band with hysteresis margin.
    adjustment_direction: Option<bool>,
}

impl DynamicThresholdState {
    /// Create a new dynamic threshold state with the given initial threshold.
    fn new(initial_threshold: u64) -> Self {
        Self {
            threshold: initial_threshold,
            smoothed_rate: 0.0,
            prev_counter: 0,
            initialized: false,
            adjustment_direction: None,
        }
    }

    /// Update the state with a new counter sample and elapsed time.
    /// Returns the new threshold if it changed, or None if unchanged.
    fn update(
        &mut self,
        counter: u64,
        elapsed_secs: f64,
        verbose: bool,
        name: &str,
    ) -> Option<u64> {
        if elapsed_secs <= 0.0 {
            return None;
        }

        // Calculate instantaneous rate.
        let delta = counter.saturating_sub(self.prev_counter);
        self.prev_counter = counter;
        let raw_rate = delta as f64 / elapsed_secs;

        // Update EMA.
        if self.initialized {
            self.smoothed_rate = DYNAMIC_THRESHOLD_EMA_ALPHA * raw_rate
                + (1.0 - DYNAMIC_THRESHOLD_EMA_ALPHA) * self.smoothed_rate;
        } else {
            // First sample: initialize EMA directly.
            self.smoothed_rate = raw_rate;
            self.initialized = true;
        }

        // Determine if we should adjust the threshold using hysteresis.
        let rate = self.smoothed_rate;
        let old_threshold = self.threshold;

        // Calculate hysteresis-adjusted bounds based on current state.
        let (effective_high, effective_low) = match self.adjustment_direction {
            Some(true) => {
                // Currently raising threshold: need rate to drop below LOW - hysteresis to stop.
                (
                    DYNAMIC_THRESHOLD_RATE_HIGH,
                    DYNAMIC_THRESHOLD_RATE_LOW * (1.0 - DYNAMIC_THRESHOLD_HYSTERESIS),
                )
            }
            Some(false) => {
                // Currently lowering threshold: need rate to rise above HIGH + hysteresis to stop.
                (
                    DYNAMIC_THRESHOLD_RATE_HIGH * (1.0 + DYNAMIC_THRESHOLD_HYSTERESIS),
                    DYNAMIC_THRESHOLD_RATE_LOW,
                )
            }
            None => {
                // Stable state: need rate to exceed bounds + hysteresis to start adjusting.
                (
                    DYNAMIC_THRESHOLD_RATE_HIGH * (1.0 + DYNAMIC_THRESHOLD_HYSTERESIS),
                    DYNAMIC_THRESHOLD_RATE_LOW * (1.0 - DYNAMIC_THRESHOLD_HYSTERESIS),
                )
            }
        };

        // Determine new adjustment direction.
        let new_direction = if rate > effective_high {
            Some(true) // Rate too high, raise threshold.
        } else if rate < effective_low && rate >= 0.0 {
            Some(false) // Rate too low, lower threshold.
        } else {
            // Rate in stable band (considering hysteresis).
            if self.adjustment_direction.is_some() {
                // We were adjusting; check if we should stop.
                if (DYNAMIC_THRESHOLD_RATE_LOW..=DYNAMIC_THRESHOLD_RATE_HIGH).contains(&rate) {
                    None // Back in target band, stop adjusting.
                } else {
                    self.adjustment_direction // Continue current direction.
                }
            } else {
                None // Already stable.
            }
        };

        // Apply adjustment if we have a direction.
        if let Some(raising) = new_direction {
            let scale = Self::compute_scale(rate, raising);
            let factor = if raising { 1.0 + scale } else { 1.0 - scale };
            let new_threshold = ((self.threshold as f64) * factor).round() as u64;
            self.threshold = new_threshold.clamp(DYNAMIC_THRESHOLD_MIN_VALUE, u64::MAX);
        }

        self.adjustment_direction = new_direction;

        // Return new threshold only if it changed.
        if self.threshold != old_threshold {
            if verbose {
                info!(
                    "{}: {} -> {} (smoothed rate {:.1}/s, raw {:.1}/s, dir {:?})",
                    name,
                    old_threshold,
                    self.threshold,
                    self.smoothed_rate,
                    raw_rate,
                    self.adjustment_direction
                );
            }
            Some(self.threshold)
        } else {
            None
        }
    }

    /// Compute the scale factor for threshold adjustment based on how far the rate
    /// is from the target band.
    fn compute_scale(rate: f64, too_high: bool) -> f64 {
        if too_high {
            let excess = ((rate / DYNAMIC_THRESHOLD_RATE_HIGH) - 1.0).max(0.0);
            let scale =
                DYNAMIC_THRESHOLD_SCALE_MIN + DYNAMIC_THRESHOLD_SLOPE_HIGH * excess.min(4.0);
            scale.min(DYNAMIC_THRESHOLD_SCALE_MAX)
        } else {
            if rate <= 0.0 {
                return DYNAMIC_THRESHOLD_SCALE_MAX;
            }
            let deficit = (DYNAMIC_THRESHOLD_RATE_LOW - rate) / DYNAMIC_THRESHOLD_RATE_LOW;
            let t = deficit.clamp(0.0, 1.0);
            DYNAMIC_THRESHOLD_SCALE_MIN + DYNAMIC_THRESHOLD_SLOPE_LOW * t
        }
    }
}

struct Scheduler<'a> {
    skel: BpfSkel<'a>,
    opts: &'a Opts,
    struct_ops: Option<libbpf_rs::Link>,
    stats_server: StatsServer<(), Metrics>,
    /// GPU device index -> NUMA node (for NVML PID sync). Only set when --gpu and NUMA enabled.
    gpu_index_to_node: Option<HashMap<u32, u32>>,
    /// Previous (TGID, node) set so we can remove processes that stopped using the GPU.
    previous_gpu_pids: Option<HashMap<u32, u32>>,
    /// Reused NVML handle to avoid re-initializing on every sync (expensive).
    nvml: Option<Nvml>,
    /// Host cgroup v2 reader used to discover peer processes of NVML GPU processes.
    gpu_cgroup_reader: Option<CgroupReader>,
    /// Dynamic threshold state for perf event migrations (when --perf-threshold is 0/dynamic).
    perf_threshold_state: Option<DynamicThresholdState>,
    /// Dynamic threshold state for sticky perf events (when --perf-sticky-threshold is 0/dynamic).
    perf_sticky_threshold_state: Option<DynamicThresholdState>,
}

impl<'a> Scheduler<'a> {
    fn init(opts: &'a Opts, open_object: &'a mut MaybeUninit<OpenObject>) -> Result<Self> {
        try_set_rlimit_infinity();

        // Initialize CPU topology.
        let topo = Topology::new().unwrap();

        // Check host topology to determine if we need to enable SMT capabilities.
        let smt_enabled = !opts.disable_smt && topo.smt_enabled;

        // Determine the amount of non-empty NUMA nodes in the system.
        let nr_nodes = topo
            .nodes
            .values()
            .filter(|node| !node.all_cpus.is_empty())
            .count();
        info!("NUMA nodes: {}", nr_nodes);

        // Automatically disable NUMA optimizations when running on non-NUMA systems.
        let numa_enabled = !opts.disable_numa && nr_nodes > 1;
        if !numa_enabled {
            info!("Disabling NUMA optimizations");
        }

        info!(
            "{} {} {}",
            SCHEDULER_NAME,
            build_id::full_version(env!("CARGO_PKG_VERSION")),
            if smt_enabled { "SMT on" } else { "SMT off" }
        );

        // Print command line.
        info!(
            "scheduler options: {}",
            std::env::args().collect::<Vec<_>>().join(" ")
        );

        // Initialize BPF connector.
        let mut skel_builder = BpfSkelBuilder::default();
        skel_builder.obj_builder.debug(opts.verbose);
        let open_opts = opts.libbpf.clone().into_bpf_open_opts();
        let mut skel = scx_ops_open!(skel_builder, open_object, cosmos_ops, open_opts)?;

        skel.struct_ops.cosmos_ops_mut().exit_dump_len = opts.exit_dump_len;

        // Override default BPF scheduling parameters.
        let rodata = skel.maps.rodata_data.as_mut().unwrap();
        rodata.slice_ns = opts.slice_us * 1000;
        rodata.slice_lag = opts.slice_lag_us * 1000;
        rodata.cpufreq_enabled = !opts.disable_cpufreq;
        rodata.flat_idle_scan = opts.flat_idle_scan;
        rodata.smt_enabled = smt_enabled;
        rodata.numa_enabled = numa_enabled;
        rodata.nr_node_ids = topo.nodes.len() as u32;
        rodata.no_wake_sync = opts.no_wake_sync;
        rodata.no_early_clear = opts.no_early_clear;
        rodata.time_preemption = opts.time_preemption;
        rodata.mm_affinity = opts.mm_affinity;

        // Enable perf event scheduling settings.
        rodata.perf_config = opts.perf_config.event_id;
        rodata.perf_sticky = opts.perf_sticky.event_id;

        // Normalize CPU busy threshold in the range [0 .. 1024].
        rodata.busy_threshold = opts.cpu_busy_thresh * 1024 / 100;

        // Maximum shared DSQ backlog before spilling outside the primary domain.
        rodata.overload_thresh_ns = opts.primary_overload_us * 1000;

        // Generate the list of available CPUs sorted by capacity in descending order.
        let mut cpus: Vec<_> = topo.all_cpus.values().collect();
        cpus.sort_by_key(|cpu| std::cmp::Reverse(cpu.cpu_capacity));
        // Normalize CPU capacities to 1..1024 so the highest capacity is always 1024.
        let max_cap = cpus.first().map(|c| c.cpu_capacity).unwrap_or(1).max(1);
        for (i, cpu) in cpus.iter().enumerate() {
            let normalized = (cpu.cpu_capacity * 1024 / max_cap).clamp(1, 1024);
            rodata.cpu_capacity[cpu.id] = normalized as c_ulong;
            rodata.preferred_cpus[i] = cpu.id as u64;
        }
        rodata.all_cpus_same_capacity = cpus.iter().all(|cpu| cpu.cpu_capacity == max_cap);
        if opts.preferred_idle_scan {
            info!(
                "Preferred CPUs: {:?}",
                &rodata.preferred_cpus[0..cpus.len()]
            );
        }
        rodata.preferred_idle_scan = opts.preferred_idle_scan;

        // Define the primary scheduling domain.
        let primary_cpus = if let Some(ref domain) = opts.primary_domain {
            match parse_cpu_list(domain) {
                Ok(cpus) => cpus,
                Err(e) => bail!("Error parsing primary domain: {}", e),
            }
        } else {
            (0..*NR_CPU_IDS).collect()
        };
        if primary_cpus.len() < *NR_CPU_IDS {
            info!("Primary CPUs: {:?}", primary_cpus);
            rodata.primary_all = false;
        } else {
            rodata.primary_all = true;
        }

        // Enable GPU support and build GPU index -> node for NVML PID sync. Init NVML once here
        // so we reuse the handle in the run loop (re-initing every sync is very expensive).
        let (gpu_index_to_node, previous_gpu_pids, nvml) = if opts.gpu && numa_enabled {
            match Nvml::init_with_flags(InitFlags::NO_GPUS) {
                Ok(nvml) => {
                    info!("NVIDIA GPU-aware scheduling enabled (NVML PID sync)");
                    rodata.gpu_enabled = true;
                    let mut idx_to_node = HashMap::new();
                    for (id, gpu) in topo.gpus() {
                        let GpuIndex::Nvidia { nvml_id } = id;
                        idx_to_node.insert(nvml_id, gpu.node_id as u32);
                    }
                    (Some(idx_to_node), Some(HashMap::new()), Some(nvml))
                }
                Err(e) => {
                    warn!("NVML init failed, disabling GPU-aware scheduling: {}", e);
                    rodata.gpu_enabled = false;
                    (None, None, None)
                }
            }
        } else {
            rodata.gpu_enabled = false;
            (None, None, None)
        };

        let gpu_cgroup_reader = if nvml.is_some() && opts.gpu_util_threshold == 0 {
            match CgroupReader::discover() {
                Ok(reader) => {
                    info!("NVIDIA GPU workload discovery enabled (cgroup v2)");
                    Some(reader)
                }
                Err(error) => {
                    warn!(
                        "GPU workload discovery unavailable, using NVML process scope: {error:#}"
                    );
                    None
                }
            }
        } else if nvml.is_some() {
            info!(
                "NVIDIA GPU workload discovery requires --gpu-util-threshold=0; using NVML process scope"
            );
            None
        } else {
            None
        };

        // Set scheduler flags.
        skel.struct_ops.cosmos_ops_mut().flags = *compat::SCX_OPS_ENQ_EXITING
            | *compat::SCX_OPS_ENQ_LAST
            | *compat::SCX_OPS_ENQ_MIGRATION_DISABLED
            | *compat::SCX_OPS_ALLOW_QUEUED_WAKEUP
            | if numa_enabled {
                *compat::SCX_OPS_BUILTIN_IDLE_PER_NODE
            } else {
                0
            };

        info!(
            "scheduler flags: {:#x}",
            skel.struct_ops.cosmos_ops_mut().flags
        );

        // Load the BPF program for validation.
        let mut skel = scx_ops_load!(skel, cosmos_ops, uei)?;

        // Initial perf thresholds in bss. When threshold is 0 we use dynamic logic; when user
        // specifies a value > 0 we use it as a static threshold.
        let bss = skel.maps.bss_data.as_mut().unwrap();
        if opts.perf_config.event_id > 0 {
            bss.perf_threshold = if opts.perf_threshold == 0 {
                DYNAMIC_THRESHOLD_INIT_VALUE
            } else {
                opts.perf_threshold
            };
        }
        if opts.perf_sticky.event_id > 0 {
            bss.perf_sticky_threshold = if opts.perf_sticky_threshold == 0 {
                DYNAMIC_THRESHOLD_INIT_VALUE
            } else {
                opts.perf_sticky_threshold
            };
        }

        // Configure CPU->node mapping (must be done after skeleton is loaded).
        for node in topo.nodes.values() {
            for cpu in node.all_cpus.values() {
                if opts.verbose {
                    info!("CPU{} -> node{}", cpu.id, node.id);
                }
                skel.maps.cpu_node_map.update(
                    &(cpu.id as u32).to_ne_bytes(),
                    &(node.id as u32).to_ne_bytes(),
                    MapFlags::ANY,
                )?;
            }
        }

        // Setup performance events for all CPUs.
        // Counter indices must match PMU library install order: migration first (0), then sticky (1).
        // When only sticky is used, it gets index 0; when both are used, sticky gets index 1.
        let nr_cpus = *NR_CPU_IDS;
        info!("Setting up performance counters for {} CPUs...", nr_cpus);
        let mut perf_available = true;
        let sticky_counter_idx = if opts.perf_config.event_id > 0 { 1 } else { 0 };
        for cpu in 0..nr_cpus {
            if opts.perf_config.event_id > 0
                && let Err(e) =
                    setup_perf_events(&skel.maps.scx_pmu_map, cpu as i32, &opts.perf_config, 0)
                && cpu == 0
            {
                let err_str = e.to_string();
                if err_str.contains("errno 2") || err_str.contains("os error 2") {
                    warn!("Performance counters not available on this CPU architecture");
                    warn!(
                        "PMU event '{}' not supported - scheduler will run without perf monitoring",
                        opts.perf_config.display_name
                    );
                } else {
                    warn!("Failed to setup perf events: {}", e);
                }
                perf_available = false;
                break;
            }
            if opts.perf_sticky.event_id > 0
                && let Err(e) = setup_perf_events(
                    &skel.maps.scx_pmu_map,
                    cpu as i32,
                    &opts.perf_sticky,
                    sticky_counter_idx,
                )
                && cpu == 0
            {
                let err_str = e.to_string();
                if err_str.contains("errno 2") || err_str.contains("os error 2") {
                    warn!("Performance counters not available on this CPU architecture");
                    warn!(
                        "PMU event '{}' not supported - scheduler will run without perf monitoring",
                        opts.perf_sticky.display_name
                    );
                } else {
                    warn!("Failed to setup perf events: {}", e);
                }
                perf_available = false;
                break;
            }
        }
        if perf_available {
            info!("Performance counters configured successfully for all CPUs");
        }

        // Configure GPU->node mapping.
        if opts.gpu && numa_enabled {
            for (id, gpu) in topo.gpus() {
                let GpuIndex::Nvidia { nvml_id } = id;
                if opts.verbose {
                    info!("GPU{} -> node{}", nvml_id, gpu.node_id);
                }
                skel.maps.gpu_node_map.update(
                    &nvml_id.to_ne_bytes(),
                    &(gpu.node_id as u32).to_ne_bytes(),
                    MapFlags::ANY,
                )?;
            }
        }

        // Enable primary scheduling domain, if defined.
        if primary_cpus.len() < *NR_CPU_IDS {
            for cpu in primary_cpus {
                if let Err(err) = Self::enable_primary_cpu(&mut skel, cpu as i32) {
                    bail!("failed to add CPU {} to primary domain: error {}", cpu, err);
                }
            }
        }

        // Initialize SMT domains.
        if smt_enabled {
            Self::init_smt_domains(&mut skel, &topo)?;
        }

        // Attach the scheduler.
        let struct_ops = Some(scx_ops_attach!(skel, cosmos_ops)?);
        let stats_server = StatsServer::new(stats::server_data()).launch()?;

        // Initialize dynamic threshold states for perf events (only when using dynamic mode).
        let perf_threshold_state = if opts.perf_config.event_id > 0 && opts.perf_threshold == 0 {
            Some(DynamicThresholdState::new(DYNAMIC_THRESHOLD_INIT_VALUE))
        } else {
            None
        };
        let perf_sticky_threshold_state =
            if opts.perf_sticky.event_id > 0 && opts.perf_sticky_threshold == 0 {
                Some(DynamicThresholdState::new(DYNAMIC_THRESHOLD_INIT_VALUE))
            } else {
                None
            };

        Ok(Self {
            skel,
            opts,
            struct_ops,
            stats_server,
            gpu_index_to_node,
            previous_gpu_pids,
            nvml,
            gpu_cgroup_reader,
            perf_threshold_state,
            perf_sticky_threshold_state,
        })
    }

    /// Sync process TGID -> GPU (node) hints from NVML. Peer processes in the same exact,
    /// non-root cgroup inherit the hint when all observed GPUs resolve to one NUMA node.
    /// gpu_util_threshold > 0 retains process-only behavior because its NVML snapshot is
    /// intentionally filtered.
    fn sync_gpu_pids(&mut self) -> Result<()> {
        let gpu_index_to_node = match &self.gpu_index_to_node {
            Some(m) => m,
            None => return Ok(()),
        };
        let nvml = match &self.nvml {
            Some(n) => n,
            None => return Ok(()),
        };
        let threshold = self.opts.gpu_util_threshold;
        // First collect TGID -> set of nodes (GPUs) per process.
        let mut pid_to_nodes: HashMap<u32, HashSet<u32>> = HashMap::new();
        let mut snapshot_complete = true;

        // A failed NVML query is not an empty snapshot. Keep the last applied
        // hints instead of deleting processes we simply failed to observe.
        let count = nvml.device_count().context("NVML device count")?;
        for i in 0..count {
            let node = match gpu_index_to_node.get(&i) {
                Some(&n) => n,
                None => {
                    snapshot_complete = false;
                    continue;
                }
            };
            let device = match nvml.device_by_index(i) {
                Ok(device) => device,
                Err(error) => {
                    debug!("NVML device {i} lookup failed: {error:#}");
                    return Err(error).context(format!("NVML device {i} lookup failed"));
                }
            };

            if threshold > 0 {
                // Use process utilization; only add PIDs above threshold.
                match device.process_utilization_stats(None::<u64>) {
                    Ok(samples) => {
                        for sample in samples {
                            let util = sample.sm_util.max(sample.mem_util);
                            if util >= threshold {
                                pid_to_nodes.entry(sample.pid).or_default().insert(node);
                            }
                        }
                    }
                    Err(_) => {
                        // NotSupported or other: fall back to all running processes.
                        Self::add_running_gpu_processes_to_set(&device, node, &mut pid_to_nodes)
                            .with_context(|| {
                                format!("NVML device {i} process snapshot is incomplete")
                            })?;
                    }
                }
            } else {
                Self::add_running_gpu_processes_to_set(&device, node, &mut pid_to_nodes)
                    .with_context(|| format!("NVML device {i} process snapshot is incomplete"))?;
            }
        }

        let mut direct = gpu::direct_gpu_processes(&pid_to_nodes);
        direct.remove(&std::process::id());
        let mut current = if snapshot_complete {
            self.gpu_cgroup_reader
                .as_ref()
                .map(|reader| gpu::expand_gpu_processes(&pid_to_nodes, reader))
                .unwrap_or_else(|| direct.clone())
        } else {
            direct.clone()
        };
        current.remove(&std::process::id());
        let max_entries = self.skel.maps.gpu_pid_map.max_entries() as usize;

        // Workload discovery can include many peer processes. Fall back to
        // direct NVML processes rather than partially populating the map.
        if current.len() > max_entries {
            warn!(
                "GPU workload has {} processes, exceeding gpu_pid_map capacity {}; using {} direct NVML processes",
                current.len(),
                max_entries,
                direct.len()
            );
        }
        if direct.len() > max_entries {
            warn!(
                "{} direct NVML processes exceed gpu_pid_map capacity {}; clearing GPU process hints",
                direct.len(),
                max_entries
            );
        }
        current = gpu::fit_gpu_processes(current, &direct, max_entries);

        self.reconcile_gpu_pid_hints(&current)
    }

    /// Reconcile the desired GPU workload hints with the BPF map.
    fn reconcile_gpu_pid_hints(&mut self, desired: &HashMap<u32, u32>) -> Result<()> {
        let previous = self.previous_gpu_pids.as_ref().unwrap().clone();
        let map = &self.skel.maps.gpu_pid_map;

        // Track the state actually applied to BPF so a partial map operation
        // is retried and cleaned up on the next synchronization.
        let mut applied = previous.clone();
        let mut first_error = None;

        // Delete stale entries first so replacements cannot temporarily run
        // out of map capacity.
        for pid in previous.keys() {
            if desired.contains_key(pid) {
                continue;
            }
            match map.delete(&pid.to_ne_bytes()).context("gpu_pid_map delete") {
                Ok(()) => {
                    applied.remove(pid);
                }
                Err(error) => {
                    if first_error.is_none() {
                        first_error = Some(error);
                    }
                }
            }
        }
        for (pid, node) in desired {
            match map
                .update(&pid.to_ne_bytes(), &node.to_ne_bytes(), MapFlags::ANY)
                .context("gpu_pid_map update")
            {
                Ok(()) => {
                    applied.insert(*pid, *node);
                }
                Err(error) => {
                    if first_error.is_none() {
                        first_error = Some(error);
                    }
                }
            }
        }
        *self.previous_gpu_pids.as_mut().unwrap() = applied;
        if let Some(error) = first_error {
            return Err(error);
        }
        Ok(())
    }

    /// Record running compute/graphics process TGIDs and the GPU node in pid_to_nodes.
    fn add_running_gpu_processes_to_set(
        device: &nvml_wrapper::Device<'_>,
        node: u32,
        pid_to_nodes: &mut HashMap<u32, HashSet<u32>>,
    ) -> Result<()> {
        let mut errors = Vec::new();

        match device.running_compute_processes() {
            Ok(processes) => {
                for process in processes {
                    pid_to_nodes.entry(process.pid).or_default().insert(node);
                }
            }
            Err(error) => errors.push(format!("compute process query failed: {error}")),
        }
        match device.running_graphics_processes() {
            Ok(processes) => {
                for process in processes {
                    pid_to_nodes.entry(process.pid).or_default().insert(node);
                }
            }
            Err(error) => errors.push(format!("graphics process query failed: {error}")),
        }

        if errors.is_empty() {
            Ok(())
        } else {
            bail!(errors.join("; "))
        }
    }

    fn enable_primary_cpu(skel: &mut BpfSkel<'_>, cpu: i32) -> Result<(), u32> {
        let prog = &mut skel.progs.enable_primary_cpu;
        let mut args = cpu_arg {
            cpu_id: cpu as c_int,
        };
        let input = ProgramInput {
            context_in: Some(unsafe {
                std::slice::from_raw_parts_mut(
                    &mut args as *mut _ as *mut u8,
                    std::mem::size_of_val(&args),
                )
            }),
            ..Default::default()
        };
        let out = prog.test_run(input).unwrap();
        if out.return_value != 0 {
            return Err(out.return_value);
        }

        Ok(())
    }

    fn enable_sibling_cpu(
        skel: &mut BpfSkel<'_>,
        cpu: usize,
        sibling_cpu: usize,
    ) -> Result<(), u32> {
        let prog = &mut skel.progs.enable_sibling_cpu;
        let mut args = domain_arg {
            cpu_id: cpu as c_int,
            sibling_cpu_id: sibling_cpu as c_int,
        };
        let input = ProgramInput {
            context_in: Some(unsafe {
                std::slice::from_raw_parts_mut(
                    &mut args as *mut _ as *mut u8,
                    std::mem::size_of_val(&args),
                )
            }),
            ..Default::default()
        };
        let out = prog.test_run(input).unwrap();
        if out.return_value != 0 {
            return Err(out.return_value);
        }

        Ok(())
    }

    fn init_smt_domains(skel: &mut BpfSkel<'_>, topo: &Topology) -> Result<(), std::io::Error> {
        let smt_siblings = topo.sibling_cpus();

        info!("SMT sibling CPUs: {:?}", smt_siblings);
        for (cpu, sibling_cpu) in smt_siblings.iter().enumerate() {
            Self::enable_sibling_cpu(skel, cpu, *sibling_cpu as usize).unwrap();
        }

        Ok(())
    }

    fn get_metrics(&self) -> Metrics {
        let bss_data = self.skel.maps.bss_data.as_ref().unwrap();
        Metrics {
            nr_event_dispatches: bss_data.nr_event_dispatches,
            nr_ev_sticky_dispatches: bss_data.nr_ev_sticky_dispatches,
            nr_gpu_dispatches: bss_data.nr_gpu_dispatches,
            nr_overload_events: bss_data.nr_overload_events,
        }
    }

    pub fn exited(&mut self) -> bool {
        uei_exited!(&self.skel, uei)
    }

    fn run(&mut self, shutdown: Arc<AtomicBool>) -> Result<UserExitInfo> {
        let (res_ch, req_ch) = self.stats_server.channels();

        // Periodic (option -p) updates of the dynamic perf thresholds.
        let polling_time = Duration::from_millis(self.opts.polling_ms).min(Duration::from_secs(1));
        let mut last_update = Instant::now();
        let mut last_gpu_sync = Instant::now();

        while !shutdown.load(Ordering::Relaxed) && !self.exited() {
            if !polling_time.is_zero() && last_update.elapsed() >= polling_time {
                // Update dynamic perf thresholds using EMA + hysteresis.
                let elapsed_secs = last_update.elapsed().as_secs_f64();

                // Update migration threshold state if dynamic mode is enabled.
                if let Some(ref mut state) = self.perf_threshold_state {
                    let nr_event = self
                        .skel
                        .maps
                        .bss_data
                        .as_ref()
                        .unwrap()
                        .nr_event_dispatches;
                    if let Some(new_thresh) =
                        state.update(nr_event, elapsed_secs, self.opts.verbose, "perf_threshold")
                    {
                        self.skel.maps.bss_data.as_mut().unwrap().perf_threshold = new_thresh;
                    }
                }

                // Update sticky threshold state if dynamic mode is enabled.
                if let Some(ref mut state) = self.perf_sticky_threshold_state {
                    let nr_sticky = self
                        .skel
                        .maps
                        .bss_data
                        .as_ref()
                        .unwrap()
                        .nr_ev_sticky_dispatches;
                    if let Some(new_thresh) = state.update(
                        nr_sticky,
                        elapsed_secs,
                        self.opts.verbose,
                        "perf_sticky_threshold",
                    ) {
                        self.skel
                            .maps
                            .bss_data
                            .as_mut()
                            .unwrap()
                            .perf_sticky_threshold = new_thresh;
                    }
                }

                last_update = Instant::now();
            }

            // GPU PID sync is throttled to GPU_SYNC_INTERVAL.
            if self.gpu_index_to_node.is_some() && last_gpu_sync.elapsed() >= GPU_SYNC_INTERVAL {
                if let Err(e) = self.sync_gpu_pids() {
                    debug!("GPU PID sync: {}", e);
                }
                last_gpu_sync = Instant::now();
            }

            // Update statistics and check for exit condition.
            let timeout = if polling_time.is_zero() {
                Duration::from_secs(1)
            } else {
                polling_time
            };
            match req_ch.recv_timeout(timeout) {
                Ok(()) => res_ch.send(self.get_metrics())?,
                Err(RecvTimeoutError::Timeout) => {}
                Err(e) => Err(e)?,
            }
        }

        let _ = self.struct_ops.take();
        uei_report!(&self.skel, uei)
    }
}

impl Drop for Scheduler<'_> {
    fn drop(&mut self) {
        info!("Unregister {SCHEDULER_NAME} scheduler");
    }
}

fn main() -> Result<()> {
    let opts = Opts::parse();

    if opts.version {
        println!(
            "{} {}",
            SCHEDULER_NAME,
            build_id::full_version(env!("CARGO_PKG_VERSION"))
        );
        return Ok(());
    }

    if opts.help_stats {
        stats::server_data().describe_meta(&mut std::io::stdout(), None)?;
        return Ok(());
    }

    let loglevel = simplelog::LevelFilter::Info;

    let mut lcfg = simplelog::ConfigBuilder::new();
    lcfg.set_time_offset_to_local()
        .expect("Failed to set local time offset")
        .set_time_level(simplelog::LevelFilter::Error)
        .set_location_level(simplelog::LevelFilter::Off)
        .set_target_level(simplelog::LevelFilter::Off)
        .set_thread_level(simplelog::LevelFilter::Off);
    simplelog::TermLogger::init(
        loglevel,
        lcfg.build(),
        simplelog::TerminalMode::Stderr,
        simplelog::ColorChoice::Auto,
    )?;

    let shutdown = Arc::new(AtomicBool::new(false));
    let shutdown_clone = shutdown.clone();
    ctrlc::set_handler(move || {
        shutdown_clone.store(true, Ordering::Relaxed);
    })
    .context("Error setting Ctrl-C handler")?;

    if let Some(intv) = opts.monitor.or(opts.stats) {
        let shutdown_copy = shutdown.clone();
        let jh = std::thread::spawn(move || {
            match stats::monitor(Duration::from_secs_f64(intv), shutdown_copy) {
                Ok(_) => {
                    debug!("stats monitor thread finished successfully")
                }
                Err(error_object) => {
                    warn!(
                        "stats monitor thread finished because of an error {}",
                        error_object
                    )
                }
            }
        });
        if opts.monitor.is_some() {
            let _ = jh.join();
            return Ok(());
        }
    }

    let mut open_object = MaybeUninit::uninit();
    loop {
        let mut sched = Scheduler::init(&opts, &mut open_object)?;
        if !sched.run(shutdown.clone())?.should_restart() {
            break;
        }
    }

    Ok(())
}