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