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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
//! Calibration cache — measures real GPU throughput on this hardware,
//! persists results to disk, replaces hardcoded `sgemm_*_flops` defaults.
//!
//! Strategy:
//! 1. Look for cache file `~/.cache/rlx/metal-calib-<hwid>.json`
//! 2. If found and valid: use measured values
//! 3. Otherwise: run quick benchmark (~50ms total), save results, use them
//!
//! The cache is keyed by GPU registry ID, so it stays valid across runs
//! on the same machine and is invalidated automatically if hardware changes.
use rlx_ir::Tick;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use crate::blas::metal_sgemm;
use crate::device::metal_device;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Calibration {
pub gpu_name: String,
pub registry_id: u64,
/// Measured GFLOP/s for sgemm_simd_4x4 at large M (best case).
pub sgemm_simd_4x4_flops: f64,
/// Measured GFLOP/s for sgemm_simd at small-aligned M.
pub sgemm_simd_flops: f64,
/// Measured GFLOP/s for sgemm_simd_padded.
pub sgemm_padded_flops: f64,
/// Measured GFLOP/s for sgemm_tiled (scalar fp32).
pub sgemm_tiled_flops: f64,
/// Measured baseline command-buffer round-trip (ns).
pub roundtrip_overhead_ns: f64,
}
fn cache_path(registry_id: u64) -> PathBuf {
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
let dir = PathBuf::from(home).join(".cache").join("rlx");
let _ = std::fs::create_dir_all(&dir);
dir.join(format!("metal-calib-{:x}.json", registry_id))
}
impl Calibration {
pub fn load(registry_id: u64) -> Option<Self> {
let path = cache_path(registry_id);
let raw = std::fs::read_to_string(&path).ok()?;
let cal: Calibration = serde_json::from_str(&raw).ok()?;
if cal.registry_id == registry_id {
Some(cal)
} else {
None
}
}
pub fn save(&self) -> std::io::Result<()> {
let path = cache_path(self.registry_id);
let raw = serde_json::to_string_pretty(self).map_err(std::io::Error::other)?;
std::fs::write(path, raw)
}
/// Measure throughput for each kernel variant by running representative
/// matmul shapes. Total cost ~50ms; runs once per machine.
pub fn measure() -> Self {
let dev = metal_device().expect("Metal device required for calibration");
let measure = |m: usize, k: usize, n: usize| -> f64 {
// Allocate three buffers in one (m*k + k*n + m*n) * 4 bytes
let total = (m * k + k * n + m * n) * 4;
let buffer = dev.alloc_shared(total);
unsafe {
let ptr = buffer.contents() as *mut f32;
for i in 0..(m * k + k * n) {
*ptr.add(i) = ((i * 13 + 7) % 257) as f32 / 257.0;
}
}
let a_off = 0;
let b_off = m * k * 4;
let c_off = (m * k + k * n) * 4;
// Warmup (kernels JIT on first dispatch)
{
let cb = dev.queue.new_command_buffer();
let enc = cb.compute_command_encoder_with_dispatch_type(
crate::mtl::MTLDispatchType::Serial,
);
for _ in 0..2 {
metal_sgemm(enc, &buffer, a_off, b_off, c_off, m, k, n);
}
enc.end_encoding();
cb.commit();
cb.wait_until_completed();
}
// Batch many sgemm calls into ONE command buffer so compute
// dominates the single wait_until_completed (~0.8ms baseline).
// 50 iterations × ~50µs compute = ~2.5ms, dwarfing dispatch.
let n_iter = 50;
let cb = dev.queue.new_command_buffer();
let enc =
cb.compute_command_encoder_with_dispatch_type(crate::mtl::MTLDispatchType::Serial);
let t0 = Tick::now();
for _ in 0..n_iter {
metal_sgemm(enc, &buffer, a_off, b_off, c_off, m, k, n);
}
enc.end_encoding();
cb.commit();
cb.wait_until_completed();
let total_s = Tick::now().elapsed_ns(t0) as f64 / 1e9;
2.0 * (m * k * n) as f64 * (n_iter as f64) / total_s
};
// Probe shapes — sized to match production BERT FFN matmul.
// 50 iterations per probe → enough compute to dominate dispatch cost.
// Simd4x4 : 256×768×3072 (BERT FFN-up at batch=16, seq=16-ish)
// Simd : 8×512×512 (8-aligned, m<32; small variant)
// SimdPadded: 6×768×768 (batch=1 attention-out)
// Tiled : 64×128×17 (n%8 != 0 fallback)
let simd_4x4 = measure(256, 768, 3072);
let simd = measure(8, 512, 512);
let padded = measure(6, 768, 768);
let tiled = measure(64, 128, 17);
// Round-trip baseline: empty command buffer commit+wait
let roundtrip_ns = {
let n_iter = 10;
let t0 = Tick::now();
for _ in 0..n_iter {
let cb = dev.queue.new_command_buffer();
cb.commit();
cb.wait_until_completed();
}
Tick::now().elapsed_ns(t0) as f64 / n_iter as f64
};
Calibration {
gpu_name: dev.name.clone(),
registry_id: dev.registry_id,
sgemm_simd_4x4_flops: simd_4x4,
sgemm_simd_flops: simd,
sgemm_padded_flops: padded,
sgemm_tiled_flops: tiled,
roundtrip_overhead_ns: roundtrip_ns,
}
}
/// Load from cache, or measure and save. Idempotent.
pub fn load_or_measure() -> Self {
let dev = metal_device().expect("Metal device required");
if let Some(cal) = Self::load(dev.registry_id) {
return cal;
}
let verbose = rlx_ir::env::var("RLX_VERBOSE")
.and_then(|v| v.parse::<u8>().ok())
.unwrap_or(0)
>= 1;
if verbose {
eprintln!(
"[rlx-metal] no calibration cache for {}; measuring...",
dev.name
);
}
let cal = Self::measure();
if verbose {
eprintln!(
"[rlx-metal] calibrated: simd_4x4={:.0} GF/s, simd={:.0} GF/s, padded={:.0} GF/s, tiled={:.0} GF/s, rt={:.0}µs",
cal.sgemm_simd_4x4_flops / 1e9,
cal.sgemm_simd_flops / 1e9,
cal.sgemm_padded_flops / 1e9,
cal.sgemm_tiled_flops / 1e9,
cal.roundtrip_overhead_ns / 1000.0
);
}
let _ = cal.save();
cal
}
}