rlx_cpu/
dispatch.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Dispatch table — calibration-aware kernel selection (plan #2).
6//!
7//! Two layers:
8//!   - **Defaults** from `kernel_config` (compile-time, per-arch).
9//!   - **Overrides** filled at runtime by autotune / calibration
10//!     when a measured number disagrees with the default.
11//!
12//! Borrowed from MAX's `dispatch_table_a100_gpu.mojo` /
13//! `dispatch_table_amd.mojo` pattern: kernel-variant selection is a
14//! data lookup, not scattered match arms in dispatch sites.
15//!
16//! Today the table is consulted by the cost model and (when an
17//! override is set) used to override a `kernel_config` default.
18//! Future work wires fusion patterns through the same table so
19//! schedule decisions are uniformly data-driven.
20
21use crate::kernel_config::{CpuArch, KernelConfig, OpClass, kernel_config_for};
22use std::collections::HashMap;
23use std::sync::{Mutex, OnceLock};
24
25/// One-line override for a default `KernelConfig` field. Add cases
26/// as autotune learns more.
27#[derive(Debug, Clone, Copy)]
28pub enum Override {
29    NeonSeqThreshold(usize),
30    ParGrain(usize),
31    ParThreshold(usize),
32    FuseAttnThreshold(usize),
33}
34
35#[derive(Debug, Default)]
36struct Table {
37    overrides: Mutex<HashMap<(CpuArch, OpClass), Vec<Override>>>,
38}
39
40fn table() -> &'static Table {
41    static T: OnceLock<Table> = OnceLock::new();
42    T.get_or_init(Table::default)
43}
44
45/// Set an override for `(arch, op)`. Idempotent — re-set replaces.
46pub fn set_override(arch: CpuArch, op: OpClass, ov: Override) {
47    let t = table();
48    let mut m = t.overrides.lock().expect("dispatch table poisoned");
49    let entry = m.entry((arch, op)).or_default();
50    // Replace any existing override of the same field tag.
51    entry.retain(|existing| std::mem::discriminant(existing) != std::mem::discriminant(&ov));
52    entry.push(ov);
53}
54
55/// Reset all overrides (test hook).
56#[doc(hidden)]
57pub fn clear_overrides_for_tests() {
58    let t = table();
59    let mut m = t.overrides.lock().expect("dispatch table poisoned");
60    m.clear();
61}
62
63/// Resolve a `KernelConfig` for `(arch, op)`, applying any
64/// overrides on top of the const-time defaults.
65pub fn resolve(arch: CpuArch, op: OpClass) -> KernelConfig {
66    let mut cfg = kernel_config_for(arch, op);
67    let t = table();
68    let m = t.overrides.lock().expect("dispatch table poisoned");
69    if let Some(list) = m.get(&(arch, op)) {
70        for ov in list {
71            match ov {
72                Override::NeonSeqThreshold(v) => cfg.neon_seq_threshold = *v,
73                Override::ParGrain(v) => cfg.par_grain = *v,
74                Override::ParThreshold(v) => cfg.par_threshold = *v,
75                Override::FuseAttnThreshold(v) => cfg.fuse_attn_threshold = *v,
76            }
77        }
78    }
79    cfg
80}
81
82/// Convenience: resolve for the running target.
83pub fn resolve_current(op: OpClass) -> KernelConfig {
84    resolve(CpuArch::current(), op)
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use std::sync::Mutex;
91
92    /// Serialize tests — `resolve` reads a process-global override table.
93    static DISPATCH_TEST_LOCK: Mutex<()> = Mutex::new(());
94
95    fn with_clean_table(f: impl FnOnce()) {
96        let _guard = DISPATCH_TEST_LOCK
97            .lock()
98            .expect("dispatch test lock poisoned");
99        clear_overrides_for_tests();
100        f();
101    }
102
103    #[test]
104    fn defaults_pass_through() {
105        with_clean_table(|| {
106            let arch = CpuArch::AppleSilicon;
107            let op = OpClass::Matmul;
108            let resolved = resolve(arch, op);
109            let default = kernel_config_for(arch, op);
110            assert_eq!(resolved.neon_seq_threshold, default.neon_seq_threshold);
111            assert_eq!(resolved.par_threshold, default.par_threshold);
112        });
113    }
114
115    #[test]
116    fn override_replaces_field() {
117        with_clean_table(|| {
118            let arch = CpuArch::AppleSilicon;
119            let op = OpClass::Matmul;
120            set_override(arch, op, Override::NeonSeqThreshold(7));
121            let r = resolve(arch, op);
122            assert_eq!(r.neon_seq_threshold, 7);
123            // Other fields untouched.
124            let d = kernel_config_for(arch, op);
125            assert_eq!(r.par_threshold, d.par_threshold);
126        });
127    }
128
129    #[test]
130    fn override_for_one_field_replaces_just_that() {
131        with_clean_table(|| {
132            let arch = CpuArch::X86_64;
133            let op = OpClass::Attention;
134            set_override(arch, op, Override::NeonSeqThreshold(5));
135            set_override(arch, op, Override::NeonSeqThreshold(9)); // replaces
136            let r = resolve(arch, op);
137            assert_eq!(r.neon_seq_threshold, 9);
138        });
139    }
140}
rlx_cpu/dispatch.rs

rlx_cpu/
dispatch.rs
