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    } else {
73        e.refused += 1;
74    }
75}
76
77/// One line per shape, widest first — the shapes that dominate a render
78/// are the ones worth reading.
79pub fn report() -> String {
80    if SEEN.load(Ordering::Relaxed) == 0 {
81        return "no q4tp matmat calls were eligible for the device arm".into();
82    }
83    let t = table().lock().unwrap();
84    let mut rows: Vec<_> = t.iter().collect();
85    rows.sort_by_key(|(_, r)| std::cmp::Reverse(r.cpu_ns));
86    let mut s = String::from(
87        "\n  q4tp matmat, both arms per call (CMF_MM_AB=1)\n\
88         \x20   b     rows    cols  calls    gpu ms    cpu ms   ratio  worst\n",
89    );
90    let (mut tg, mut tc) = (0u128, 0u128);
91    for ((b, r, c), e) in rows {
92        let took = e.calls - e.refused;
93        let g = e.gpu_ns as f64 / 1e6;
94        let cp = e.cpu_ns as f64 / 1e6;
95        tg += e.gpu_ns;
96        tc += e.cpu_ns;
97        let ratio = if took > 0 && g > 0.0 {
98            format!("{:.2}x", cp / g)
99        } else {
100            "refused".into()
101        };
102        s.push_str(&format!(
103            "  {b:>5} {r:>8} {c:>7} {:>6} {g:>9.1} {cp:>9.1} {ratio:>7} {:>6.4}\n",
104            e.calls, e.worst
105        ));
106    }
107    s.push_str(&format!(
108        "  total gpu {:.1} ms, cpu {:.1} ms — the device arm is {:.2}x the host's\n",
109        tg as f64 / 1e6,
110        tc as f64 / 1e6,
111        tc as f64 / (tg as f64).max(1.0),
112    ));
113    s
114}
115
116/// The second question, once the ratio is known: of the device arm's
117/// time, how much is the GEMM and how much is moving the result back.
118/// `CMF_MM_SPLIT=1` submits the kernel alone, waits, then times the
119/// readback separately. The answer decides whether the next work is a
120/// faster kernel or a resident chain that stops shipping activations.
121pub fn split_on() -> bool {
122    static ON: OnceLock<bool> = OnceLock::new();
123    *ON.get_or_init(|| std::env::var("CMF_MM_SPLIT").as_deref() == Ok("1"))
124}
125
126#[derive(Default, Clone)]
127struct Split {
128    calls: u64,
129    kernel_ns: u128,
130    rb_ns: u128,
131}
132
133fn splits() -> &'static Mutex<BTreeMap<(usize, usize, usize), Split>> {
134    static T: OnceLock<Mutex<BTreeMap<(usize, usize, usize), Split>>> = OnceLock::new();
135    T.get_or_init(|| Mutex::new(BTreeMap::new()))
136}
137
138pub fn split_note(b: usize, rows: usize, cols: usize, kernel: Duration, rb: Duration) {
139    let mut t = splits().lock().unwrap();
140    let e = t.entry((b, rows, cols)).or_default();
141    e.calls += 1;
142    e.kernel_ns += kernel.as_nanos();
143    e.rb_ns += rb.as_nanos();
144}
145
146pub fn split_report() -> String {
147    let t = splits().lock().unwrap();
148    if t.is_empty() {
149        return "no split-timed calls".into();
150    }
151    let mut s = String::from(
152        "\n  device arm, kernel vs readback (CMF_MM_SPLIT=1)\n\
153         \x20   b     rows    cols  calls  kernel ms  readbk ms  rb MB/call\n",
154    );
155    for ((b, r, c), e) in t.iter() {
156        s.push_str(&format!(
157            "  {b:>5} {r:>8} {c:>7} {:>6} {:>10.1} {:>10.1} {:>10.1}\n",
158            e.calls,
159            e.kernel_ns as f64 / 1e6,
160            e.rb_ns as f64 / 1e6,
161            (b * r * 4) as f64 / 1e6,
162        ));
163    }
164    s
165}
cortiq_engine/mm_ab.rs

cortiq_engine/
mm_ab.rs
