cortiq_engine/
cpuprof.rs

1//! In-process CPU op timers (`CMF_CPU_PROF=1`).
2//!
3//! Wall-clock A/B on a shared box is not evidence (see the perf notes):
4//! the same change measured 23%, 0% and 22% across whole runs. These
5//! counters accumulate nanoseconds inside the process around each CPU
6//! stage of the layer loop, so contention inflates every slot alike and
7//! the SPLIT survives it. Off by default: one cached bool per call site.
8//!
9//! `report(tokens)` prints ms per token for every slot that ran.
10
11use std::sync::OnceLock;
12use std::sync::atomic::{AtomicU64, Ordering};
13
14#[derive(Clone, Copy)]
15#[repr(usize)]
16pub enum Slot {
17    /// Q/K/V projections of one decode position.
18    Qkv,
19    /// RoPE + append + attend (the attention core, one position).
20    AttnCore,
21    /// Output projection.
22    AttnO,
23    /// Fused gate/up (+SiLU·mul).
24    FfnGateUp,
25    /// Down projection.
26    FfnDown,
27    /// RMSNorms + residual adds of the layer loop.
28    Norms,
29    /// Final norm + lm_head matvec.
30    Head,
31    /// Sampler (argmax / top-k / penalties).
32    Sampler,
33    /// Whole layer stack of one forward (decode or prefill chunk).
34    Layers,
35    /// Batched prefill projections (every matmat).
36    Matmat,
37    /// Batched prefill attention core (per-position or batched attend).
38    PrefillAttend,
39    /// Activation quantization (split_act) inside the kernels.
40    SplitAct,
41    /// One pool dispatch, entry to exit (caller thread).
42    DispTotal,
43    /// Pool dispatch bookkeeping: descriptor, tickets, unparks.
44    DispSetup,
45    /// The caller participant's share of the dispatched closure.
46    DispWork,
47    /// Barrier drain: waiting for the other participants.
48    DispWait,
49    /// Lane driver (gate+up): per-batch scale/ladder decode.
50    GuScales,
51    /// Lane driver (gate+up): the interleaved group-dot loop.
52    GuDots,
53    /// Lane driver (gate+up): outliers + silu + store.
54    GuPost,
55    /// The dense-FFN activation pass (Spark's exact-erf GELU combine).
56    FfnCombine,
57}
58
59const N: usize = 20;
60const NAMES: [&str; N] = [
61    "qkv",
62    "attn_core",
63    "attn_o",
64    "ffn_gate_up",
65    "ffn_down",
66    "norms",
67    "head",
68    "sampler",
69    "layers(total)",
70    "matmat(prefill)",
71    "attend(prefill)",
72    "split_act",
73    "disp(total)",
74    "disp(setup)",
75    "disp(work)",
76    "disp(wait)",
77    "ln.scales",
78    "ln.dots",
79    "ln.post",
80    "ffn_combine",
81];
82
83static NS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
84static CALLS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
85
86/// Total elements marked as activation outliers by `split_act` (diagnostic
87/// for the per-row outlier-correction cost). Reset by [`reset`], printed by
88/// [`report`].
89static OUTLIERS: AtomicU64 = AtomicU64::new(0);
90/// Rows routed through a lane driver (diagnostic: is the lane path live?).
91static LANE_ROWS: AtomicU64 = AtomicU64::new(0);
92/// Rows routed through the plain per-row path (diagnostic).
93static PLAIN_ROWS: AtomicU64 = AtomicU64::new(0);
94
95/// Record `n` outliers found in one activation split.
96pub fn count_outliers(n: u64) {
97    OUTLIERS.fetch_add(n, Ordering::Relaxed);
98}
99
100/// Record `n` rows about to run through a lane-interleaved driver.
101pub fn count_lane_rows(n: u64) {
102    LANE_ROWS.fetch_add(n, Ordering::Relaxed);
103}
104
105/// Record `n` rows about to run through the plain per-row path.
106pub fn count_plain_rows(n: u64) {
107    PLAIN_ROWS.fetch_add(n, Ordering::Relaxed);
108}
109
110#[inline]
111pub fn on() -> bool {
112    static ON: OnceLock<bool> = OnceLock::new();
113    *ON.get_or_init(|| std::env::var("CMF_CPU_PROF").is_ok_and(|v| v != "0"))
114}
115
116/// Scope timer: records on drop. `None` when profiling is off.
117pub struct Timer(Option<(Slot, std::time::Instant)>);
118
119impl Drop for Timer {
120    #[inline]
121    fn drop(&mut self) {
122        if let Some((s, t0)) = self.0.take() {
123            NS[s as usize].fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
124            CALLS[s as usize].fetch_add(1, Ordering::Relaxed);
125        }
126    }
127}
128
129#[inline]
130pub fn time(s: Slot) -> Timer {
131    if on() {
132        Timer(Some((s, std::time::Instant::now())))
133    } else {
134        Timer(None)
135    }
136}
137
138/// Zero every slot (bench calls it right before the measured window).
139pub fn reset() {
140    for i in 0..N {
141        NS[i].store(0, Ordering::Relaxed);
142        CALLS[i].store(0, Ordering::Relaxed);
143    }
144    OUTLIERS.store(0, Ordering::Relaxed);
145    LANE_ROWS.store(0, Ordering::Relaxed);
146    PLAIN_ROWS.store(0, Ordering::Relaxed);
147}
148
149/// Snapshot of (name, total ms, calls) for every slot that ran.
150pub fn snapshot() -> Vec<(&'static str, f64, u64)> {
151    (0..N)
152        .filter(|&i| CALLS[i].load(Ordering::Relaxed) > 0)
153        .map(|i| {
154            (
155                NAMES[i],
156                NS[i].load(Ordering::Relaxed) as f64 / 1e6,
157                CALLS[i].load(Ordering::Relaxed),
158            )
159        })
160        .collect()
161}
162
163/// Labeled stage spans (`SAM3_PROF=1`) — the generic sibling of [`time`]:
164/// a process-wide map label → (ms, calls), accumulated by RAII spans and
165/// printed (then cleared) by a `dump_on_drop` guard. The accumulator is a
166/// shared Mutex so spans taken ON POOL WORKER THREADS inside run_rows
167/// closures land in the caller's dump too. Inclusive: a child span's ms
168/// are also inside its parent's. Off by default: one cached bool per
169/// `span()` call site.
170pub mod spans {
171    use std::collections::BTreeMap;
172    use std::sync::{Mutex, OnceLock};
173    use std::time::Instant;
174
175    type Acc = Mutex<BTreeMap<&'static str, (f64, u64)>>;
176
177    #[inline]
178    fn acc() -> &'static Acc {
179        static ACC: OnceLock<Acc> = OnceLock::new();
180        ACC.get_or_init(|| Mutex::new(BTreeMap::new()))
181    }
182
183    #[inline]
184    pub fn on() -> bool {
185        static ON: OnceLock<bool> = OnceLock::new();
186        *ON.get_or_init(|| std::env::var("SAM3_PROF").is_ok_and(|v| v != "0"))
187    }
188
189    /// Records wall time from creation to drop under `label`.
190    pub struct Span(Option<(&'static str, Instant)>);
191
192    impl Drop for Span {
193        #[inline]
194        fn drop(&mut self) {
195            if let Some((label, t0)) = self.0.take() {
196                add(label, t0.elapsed());
197            }
198        }
199    }
200
201    /// Manually accumulate `d` under `label` (worker-thread friendly).
202    #[inline]
203    pub fn add(label: &'static str, d: std::time::Duration) {
204        let mut m = match acc().lock() {
205            Ok(g) => g,
206            Err(poisoned) => poisoned.into_inner(),
207        };
208        let e = m.entry(label).or_insert((0.0, 0));
209        e.0 += d.as_secs_f64() * 1e3;
210        e.1 += 1;
211    }
212
213    /// Start a span; a no-op when `SAM3_PROF` is off.
214    #[inline]
215    pub fn span(label: &'static str) -> Span {
216        if on() {
217            Span(Some((label, Instant::now())))
218        } else {
219            Span(None)
220        }
221    }
222
223    /// Print the accumulated spans under `prefix`, then clear.
224    pub fn dump(prefix: &str) {
225        if !on() {
226            return;
227        }
228        let mut m = match acc().lock() {
229            Ok(g) => g,
230            Err(poisoned) => poisoned.into_inner(),
231        };
232        if m.is_empty() {
233            return;
234        }
235        eprintln!("[sam3-prof] {prefix}:");
236        let rows: Vec<(&'static str, f64, u64)> =
237            m.iter().map(|(k, v)| (*k, v.0, v.1)).collect();
238        for (k, ms, n) in rows {
239            eprintln!("[sam3-prof]   {k:<18} {ms:>10.1} ms  ({n} calls)");
240        }
241        m.clear();
242    }
243
244    /// Prints (and clears) the accumulator when dropped — one dump per
245    /// forward even through its early-return paths.
246    pub struct Dump(String);
247
248    impl Drop for Dump {
249        fn drop(&mut self) {
250            dump(&self.0);
251        }
252    }
253
254    pub fn dump_on_drop(prefix: impl Into<String>) -> Dump {
255        Dump(prefix.into())
256    }
257}
258
259/// Print ms/token per slot to stderr.
260pub fn report(label: &str, tokens: usize) {
261    if !on() {
262        return;
263    }
264    let t = tokens.max(1) as f64;
265    eprintln!("cpu-prof [{label}] per token over {tokens} tokens:");
266    for (name, ms, calls) in snapshot() {
267        eprintln!(
268            "  {name:<16} {:>8.3} ms/tok  ({calls} calls, {:.1} us/call)",
269            ms / t,
270            ms * 1e3 / calls.max(1) as f64
271        );
272    }
273    let ol = OUTLIERS.load(Ordering::Relaxed);
274    if ol > 0 {
275        eprintln!(
276            "  outliers         {:>8.1} /split  ({} total)",
277            ol as f64 / t,
278            ol
279        );
280    }
281    let lr = LANE_ROWS.load(Ordering::Relaxed);
282    let pr = PLAIN_ROWS.load(Ordering::Relaxed);
283    if lr > 0 || pr > 0 {
284        eprintln!(
285            "  lane_rows        {:>8.1} /tok   plain_rows {:>8.1} /tok",
286            lr as f64 / t,
287            pr as f64 / t
288        );
289    }
290}
cortiq_engine/cpuprof.rs

cortiq_engine/
cpuprof.rs
