rlx_cpu/
cost.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//! Cost model — estimates execution time for kernel dispatch decisions.
6//!
7//! Instead of hardcoded thresholds scattered across thunk.rs, all dispatch
8//! decisions are routed through this cost model. The model considers:
9//! - Hardware: cache sizes, core count, AMX availability, NEON throughput
10//! - Workload: matrix dimensions, batch size, sequence length
11//! - Strategy: BLAS vs NEON, parallel vs sequential, fused vs individual
12//!
13//! The model is calibrated at compile time (platform-specific constants)
14//! and refined at runtime (detected hardware + optional autotune).
15
16use crate::config::RuntimeConfig;
17
18/// Estimated cost in nanoseconds for a kernel execution strategy.
19#[derive(Debug, Clone, Copy)]
20pub struct Cost(pub f64);
21
22impl Cost {
23    pub fn ns(self) -> f64 {
24        self.0
25    }
26}
27
28/// Hardware model — derived from RuntimeConfig + platform detection.
29pub struct HwModel {
30    /// NEON throughput: FLOP/s for element-wise (FMA chains)
31    pub neon_flops: f64,
32    /// BLAS throughput: FLOP/s for sgemm (AMX or optimized NEON)
33    pub blas_flops: f64,
34    /// BLAS call overhead in nanoseconds (function call + AMX sync)
35    pub blas_overhead_ns: f64,
36    /// par_for dispatch overhead in nanoseconds
37    pub par_for_overhead_ns: f64,
38    /// L1 data cache size in bytes
39    pub l1_bytes: usize,
40    /// L2 cache size in bytes
41    pub l2_bytes: usize,
42    /// Memory bandwidth (L2 → registers) in bytes/ns
43    pub mem_bw: f64,
44    /// Number of worker threads
45    pub num_threads: usize,
46}
47
48impl HwModel {
49    /// Build from runtime config and platform defaults.
50    pub fn from_config(cfg: &RuntimeConfig) -> Self {
51        // Platform-calibrated constants
52        #[cfg(all(target_arch = "aarch64", target_vendor = "apple"))]
53        let model = HwModel {
54            neon_flops: 72e9,            // ~72 GFLOP/s NEON FMA throughput (M4 Pro P-core)
55            blas_flops: 1000e9,          // ~1 TFLOP/s AMX peak (effective varies with tile fill)
56            blas_overhead_ns: 500.0,     // ~0.5µs per cblas_sgemm call
57            par_for_overhead_ns: 5000.0, // ~5µs spin-wait dispatch
58            l1_bytes: 65536,             // 64KB L1d (refined by sysctl at runtime)
59            l2_bytes: 4 * 1024 * 1024,   // 4MB L2 per core
60            mem_bw: 50.0,                // ~50 GB/s = 50 B/ns
61            num_threads: cfg.pool_workers + 1,
62        };
63
64        #[cfg(not(all(target_arch = "aarch64", target_vendor = "apple")))]
65        let model = HwModel {
66            neon_flops: 32e9,
67            blas_flops: 200e9,
68            blas_overhead_ns: 300.0,
69            par_for_overhead_ns: 3000.0,
70            l1_bytes: 32768,
71            l2_bytes: 1024 * 1024,
72            mem_bw: 30.0,
73            num_threads: cfg.pool_workers + 1,
74        };
75
76        model
77    }
78
79    // ── Dispatch decisions ──────────────────────────────────────────
80
81    /// Should we use NEON sgemm instead of BLAS for this matrix multiply?
82    /// Returns true when BLAS overhead dominates the compute.
83    pub fn prefer_neon_sgemm(&self, m: usize, k: usize, n: usize) -> bool {
84        let flops = 2.0 * m as f64 * k as f64 * n as f64;
85        let blas_time = flops / self.blas_flops + self.blas_overhead_ns * 1e-9;
86        let neon_time = flops / self.neon_flops;
87        neon_time < blas_time
88    }
89
90    /// Should we use par_for for this element-wise operation?
91    /// Returns true when parallelism benefit exceeds dispatch overhead.
92    pub fn prefer_parallel(&self, total_elements: usize, cost_per_element_ns: f64) -> bool {
93        let seq_time = total_elements as f64 * cost_per_element_ns;
94        let par_time = seq_time / self.num_threads as f64 + self.par_for_overhead_ns;
95        par_time < seq_time
96    }
97
98    /// Should we use strided BLAS for SDPA, or sequential NEON dots?
99    pub fn prefer_blas_sdpa(
100        &self,
101        batch: usize,
102        seq: usize,
103        num_heads: usize,
104        head_dim: usize,
105    ) -> bool {
106        let total_heads = batch * num_heads;
107        // Two sgemm per head (Q@K^T + scores@V)
108        let per_head_flops = 2.0 * seq as f64 * seq as f64 * head_dim as f64 * 2.0;
109        let blas_per_head = per_head_flops / self.blas_flops + 2.0 * self.blas_overhead_ns * 1e-9;
110        let neon_per_head = per_head_flops / self.neon_flops;
111
112        // With par_for, BLAS heads run in parallel
113        let blas_total = blas_per_head * total_heads as f64 / self.num_threads as f64
114            + self.par_for_overhead_ns * 1e-9;
115        let neon_total = neon_per_head * total_heads as f64; // sequential
116
117        blas_total < neon_total
118    }
119
120    /// Should we fuse the entire transformer layer into one thunk?
121    /// True when intermediates fit in L1 and per-thunk overhead dominates.
122    pub fn prefer_fused_layer(
123        &self,
124        batch: usize,
125        seq: usize,
126        hidden: usize,
127        intermediate: usize,
128    ) -> bool {
129        let m = batch * seq;
130        // Estimate intermediate buffer sizes
131        let qkv_bytes = m * 3 * hidden * 4;
132        let attn_bytes = m * hidden * 4;
133        let ffn_bytes = m * intermediate * 4;
134        let total_bytes = qkv_bytes + 2 * attn_bytes + ffn_bytes;
135        // Fuse if total intermediates fit in L2 (L1 would be ideal but tight)
136        total_bytes <= self.l2_bytes / 2
137    }
138}
139
140/// Global hardware model singleton.
141pub fn hw_model() -> &'static HwModel {
142    use std::sync::OnceLock;
143    static MODEL: OnceLock<HwModel> = OnceLock::new();
144    MODEL.get_or_init(|| HwModel::from_config(RuntimeConfig::global()))
145}
rlx_cpu/cost.rs

rlx_cpu/
cost.rs
