rlx_cpu/
kernel_config.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//! Compile-time kernel-config tables (plan #14).
6//!
7//! Borrowed from MAX's `internal_utils/nvidia_configs.mojo` /
8//! `amd_configs.mojo` pattern: tile sizes, kernel-selection
9//! thresholds, etc. as compile-time data structures kernels query
10//! instead of scattered match-arms.
11//!
12//! Today the values live as `const`s here and are surfaced through
13//! [`kernel_config_for`]. The goal is one source of truth — when
14//! we want to tune for a new arch (M5 Apple Silicon, x86 Zen5,
15//! etc.) we add a row to the table, not a new match arm in 12
16//! files.
17
18/// Coarse target classification — refined as new SoCs emerge.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub enum CpuArch {
21    AppleSilicon, // M-series; AMX + NEON
22    AarchGeneric, // Other ARM (RPi, AWS Graviton)
23    X86_64,
24    Other,
25}
26
27impl CpuArch {
28    /// Pick the best label for the running target.
29    pub const fn current() -> Self {
30        #[cfg(all(target_arch = "aarch64", target_vendor = "apple"))]
31        {
32            Self::AppleSilicon
33        }
34        #[cfg(all(target_arch = "aarch64", not(target_vendor = "apple")))]
35        {
36            Self::AarchGeneric
37        }
38        #[cfg(target_arch = "x86_64")]
39        {
40            Self::X86_64
41        }
42        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
43        {
44            Self::Other
45        }
46    }
47}
48
49/// Op category that a kernel config is keyed against. Coarse —
50/// refines as we learn which thresholds actually want per-shape
51/// tuning.
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
53pub enum OpClass {
54    /// Matmul / linear-algebra hot path.
55    Matmul,
56    /// SDPA attention.
57    Attention,
58    /// Element-wise / activation / norm.
59    Elementwise,
60    /// View / shape ops (reshape, narrow, transpose).
61    Shape,
62}
63
64/// Settings the dispatch logic asks for. All in elements unless
65/// noted otherwise.
66#[derive(Debug, Clone, Copy)]
67pub struct KernelConfig {
68    /// Below this batch*seq, prefer NEON over BLAS for matmul / SDPA.
69    pub neon_seq_threshold: usize,
70    /// par_for granularity for elementwise.
71    pub par_grain: usize,
72    /// Below this total element count, run sequentially (par_for
73    /// has positive overhead even at 0 work).
74    pub par_threshold: usize,
75    /// FusedAttnBlock fires when batch*seq <= this.
76    pub fuse_attn_threshold: usize,
77}
78
79const APPLE_SILICON: KernelConfig = KernelConfig {
80    neon_seq_threshold: 32,
81    par_grain: 64,
82    par_threshold: 30_000,
83    fuse_attn_threshold: 64,
84};
85
86const AARCH_GENERIC: KernelConfig = KernelConfig {
87    neon_seq_threshold: 24,
88    par_grain: 32,
89    par_threshold: 20_000,
90    fuse_attn_threshold: 48,
91};
92
93const X86_DEFAULT: KernelConfig = KernelConfig {
94    neon_seq_threshold: 16, // AVX2 path; lower threshold reflects bigger vector unit
95    par_grain: 32,
96    par_threshold: 20_000,
97    fuse_attn_threshold: 32,
98};
99
100const FALLBACK: KernelConfig = KernelConfig {
101    neon_seq_threshold: 16,
102    par_grain: 16,
103    par_threshold: 10_000,
104    fuse_attn_threshold: 16,
105};
106
107/// Look up the canonical kernel config for `(arch, op_class)`. The
108/// table is `const`-evaluated so callers pay no lookup cost.
109pub const fn kernel_config_for(arch: CpuArch, op: OpClass) -> KernelConfig {
110    // Today the per-op variation is small — we return one row per
111    // arch and let callers read the field they care about. As more
112    // shape-specific tuning data lands, the OpClass dimension
113    // becomes load-bearing.
114    let _ = op;
115    match arch {
116        CpuArch::AppleSilicon => APPLE_SILICON,
117        CpuArch::AarchGeneric => AARCH_GENERIC,
118        CpuArch::X86_64 => X86_DEFAULT,
119        CpuArch::Other => FALLBACK,
120    }
121}
122
123/// Convenience: defaults for the running target.
124pub const fn current_config(op: OpClass) -> KernelConfig {
125    kernel_config_for(CpuArch::current(), op)
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    #[test]
133    fn current_resolves() {
134        let cfg = current_config(OpClass::Matmul);
135        // All targets at minimum produce a non-zero threshold.
136        assert!(cfg.neon_seq_threshold > 0);
137        assert!(cfg.par_threshold > 0);
138    }
139
140    #[test]
141    fn apple_silicon_picks_higher_thresholds() {
142        let m = kernel_config_for(CpuArch::AppleSilicon, OpClass::Matmul);
143        let f = kernel_config_for(CpuArch::Other, OpClass::Matmul);
144        assert!(m.neon_seq_threshold >= f.neon_seq_threshold);
145        assert!(m.par_threshold >= f.par_threshold);
146    }
147}
rlx_cpu/kernel_config.rs

rlx_cpu/
kernel_config.rs
