cortiq_engine/
mm_ab.rs

1//! Per-shape A/B of the two q4tp matmat arms, measured inside one
2//! process (`CMF_MM_AB=1`).
3//!
4//! It exists because wall-clock A/B on a shared machine lies. Three
5//! separate runs of the same change on the same stand read 2.4x, 1.0x
6//! and 1.0x, and one 8-second render's denoise drifted 44.9 s -> 52.5 s
7//! across six back-to-back repeats — a 25% band that swallows any real
8//! effect smaller than itself. Interleaving whole processes does not
9//! help: the drift is slower than a process.
10//!
11//! So both arms run back to back on the same activations inside the
12//! same call, and what is reported is their RATIO per shape. Whatever
13//! the machine is doing to one, it is doing to the other. The maximum
14//! disagreement between the two outputs comes along for free, which is
15//! the check that says the faster arm is also the right one.
16
17use std::collections::BTreeMap;
18use std::sync::atomic::{AtomicU64, Ordering};
19use std::sync::{Mutex, OnceLock};
20use std::time::Duration;
21
22#[derive(Default, Clone)]
23struct Row {
24    calls: u64,
25    gpu_ns: u128,
26    cpu_ns: u128,
27    refused: u64,
28    worst: f32,
29}
30
31fn table() -> &'static Mutex<BTreeMap<(usize, usize, usize), Row>> {
32    static T: OnceLock<Mutex<BTreeMap<(usize, usize, usize), Row>>> = OnceLock::new();
33    T.get_or_init(|| Mutex::new(BTreeMap::new()))
34}
35
36pub fn on() -> bool {
37    static ON: OnceLock<bool> = OnceLock::new();
38    *ON.get_or_init(|| std::env::var("CMF_MM_AB").as_deref() == Ok("1"))
39}
40
41/// Calls seen, so a report can say "nothing measured" rather than
42/// print an empty table that reads like "no difference".
43static SEEN: AtomicU64 = AtomicU64::new(0);
44
45#[allow(clippy::too_many_arguments)]
46pub fn record(
47    b: usize,
48    rows: usize,
49    cols: usize,
50    gpu_took_it: bool,
51    gpu: Duration,
52    cpu: Duration,
53    g: &[f32],
54    c: &[f32],
55) {
56    SEEN.fetch_add(1, Ordering::Relaxed);
57    let mut t = table().lock().unwrap();
58    let e = t.entry((b, rows, cols)).or_default();
59    e.calls += 1;
60    e.cpu_ns += cpu.as_nanos();
61    if gpu_took_it {
62        e.gpu_ns += gpu.as_nanos();
63        // Relative to the CPU arm's own scale, so a quiet row and a loud
64        // one are comparable.
65        let scale = c.iter().fold(0f32, |m, &v| m.max(v.abs())).max(1e-6);
66        let d = g
67            .iter()
68            .zip(c)
69            .fold(0f32, |m, (&a, &b)| m.max((a - b).abs()))
70            / scale;
71        e.worst = e.worst.max(d);
72        if d > 1e-2 {
73            eprintln!("mm-ab: b={b} {rows}x{cols}: GPU vs CPU max rel diff {d:.3e} (call {})", e.calls);
74        }
75    } else {
76        e.refused += 1;
77    }
78}
79
80/// One line per shape, widest first — the shapes that dominate a render
81/// are the ones worth reading.
82pub fn report() -> String {
83    if SEEN.load(Ordering::Relaxed) == 0 {
84        return "no q4tp matmat calls were eligible for the device arm".into();
85    }
86    let t = table().lock().unwrap();
87    let mut rows: Vec<_> = t.iter().collect();
88    rows.sort_by_key(|(_, r)| std::cmp::Reverse(r.cpu_ns));
89    let mut s = String::from(
90        "\n  q4tp matmat, both arms per call (CMF_MM_AB=1)\n\
91         \x20   b     rows    cols  calls    gpu ms    cpu ms   ratio  worst\n",
92    );
93    let (mut tg, mut tc) = (0u128, 0u128);
94    for ((b, r, c), e) in rows {
95        let took = e.calls - e.refused;
96        let g = e.gpu_ns as f64 / 1e6;
97        let cp = e.cpu_ns as f64 / 1e6;
98        tg += e.gpu_ns;
99        tc += e.cpu_ns;
100        let ratio = if took > 0 && g > 0.0 {
101            format!("{:.2}x", cp / g)
102        } else {
103            "refused".into()
104        };
105        s.push_str(&format!(
106            "  {b:>5} {r:>8} {c:>7} {:>6} {g:>9.1} {cp:>9.1} {ratio:>7} {:>6.4}\n",
107            e.calls, e.worst
108        ));
109    }
110    s.push_str(&format!(
111        "  total gpu {:.1} ms, cpu {:.1} ms — the device arm is {:.2}x the host's\n",
112        tg as f64 / 1e6,
113        tc as f64 / 1e6,
114        tc as f64 / (tg as f64).max(1.0),
115    ));
116    s
117}
118
119/// The second question, once the ratio is known: of the device arm's
120/// time, how much is the GEMM and how much is moving the result back.
121/// `CMF_MM_SPLIT=1` submits the kernel alone, waits, then times the
122/// readback separately. The answer decides whether the next work is a
123/// faster kernel or a resident chain that stops shipping activations.
124pub fn split_on() -> bool {
125    static ON: OnceLock<bool> = OnceLock::new();
126    *ON.get_or_init(|| std::env::var("CMF_MM_SPLIT").as_deref() == Ok("1"))
127}
128
129#[derive(Default, Clone)]
130struct Split {
131    calls: u64,
132    kernel_ns: u128,
133    rb_ns: u128,
134}
135
136fn splits() -> &'static Mutex<BTreeMap<(usize, usize, usize), Split>> {
137    static T: OnceLock<Mutex<BTreeMap<(usize, usize, usize), Split>>> = OnceLock::new();
138    T.get_or_init(|| Mutex::new(BTreeMap::new()))
139}
140
141pub fn split_note(b: usize, rows: usize, cols: usize, kernel: Duration, rb: Duration) {
142    let mut t = splits().lock().unwrap();
143    let e = t.entry((b, rows, cols)).or_default();
144    e.calls += 1;
145    e.kernel_ns += kernel.as_nanos();
146    e.rb_ns += rb.as_nanos();
147}
148
149pub fn split_report() -> String {
150    let t = splits().lock().unwrap();
151    if t.is_empty() {
152        return "no split-timed calls".into();
153    }
154    let mut s = String::from(
155        "\n  device arm, kernel vs readback (CMF_MM_SPLIT=1)\n\
156         \x20   b     rows    cols  calls  kernel ms  readbk ms  rb MB/call\n",
157    );
158    for ((b, r, c), e) in t.iter() {
159        s.push_str(&format!(
160            "  {b:>5} {r:>8} {c:>7} {:>6} {:>10.1} {:>10.1} {:>10.1}\n",
161            e.calls,
162            e.kernel_ns as f64 / 1e6,
163            e.rb_ns as f64 / 1e6,
164            (b * r * 4) as f64 / 1e6,
165        ));
166    }
167    s
168}
cortiq_engine/mm_ab.rs

cortiq_engine/
mm_ab.rs
