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}
42
43const N: usize = 12;
44const NAMES: [&str; N] = [
45    "qkv",
46    "attn_core",
47    "attn_o",
48    "ffn_gate_up",
49    "ffn_down",
50    "norms",
51    "head",
52    "sampler",
53    "layers(total)",
54    "matmat(prefill)",
55    "attend(prefill)",
56    "split_act",
57];
58
59static NS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
60static CALLS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
61
62#[inline]
63pub fn on() -> bool {
64    static ON: OnceLock<bool> = OnceLock::new();
65    *ON.get_or_init(|| std::env::var("CMF_CPU_PROF").is_ok_and(|v| v != "0"))
66}
67
68/// Scope timer: records on drop. `None` when profiling is off.
69pub struct Timer(Option<(Slot, std::time::Instant)>);
70
71impl Drop for Timer {
72    #[inline]
73    fn drop(&mut self) {
74        if let Some((s, t0)) = self.0.take() {
75            NS[s as usize].fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
76            CALLS[s as usize].fetch_add(1, Ordering::Relaxed);
77        }
78    }
79}
80
81#[inline]
82pub fn time(s: Slot) -> Timer {
83    if on() {
84        Timer(Some((s, std::time::Instant::now())))
85    } else {
86        Timer(None)
87    }
88}
89
90/// Zero every slot (bench calls it right before the measured window).
91pub fn reset() {
92    for i in 0..N {
93        NS[i].store(0, Ordering::Relaxed);
94        CALLS[i].store(0, Ordering::Relaxed);
95    }
96}
97
98/// Snapshot of (name, total ms, calls) for every slot that ran.
99pub fn snapshot() -> Vec<(&'static str, f64, u64)> {
100    (0..N)
101        .filter(|&i| CALLS[i].load(Ordering::Relaxed) > 0)
102        .map(|i| {
103            (
104                NAMES[i],
105                NS[i].load(Ordering::Relaxed) as f64 / 1e6,
106                CALLS[i].load(Ordering::Relaxed),
107            )
108        })
109        .collect()
110}
111
112/// Print ms/token per slot to stderr.
113pub fn report(label: &str, tokens: usize) {
114    if !on() {
115        return;
116    }
117    let t = tokens.max(1) as f64;
118    eprintln!("cpu-prof [{label}] per token over {tokens} tokens:");
119    for (name, ms, calls) in snapshot() {
120        eprintln!(
121            "  {name:<16} {:>8.3} ms/tok  ({calls} calls, {:.1} us/call)",
122            ms / t,
123            ms * 1e3 / calls.max(1) as f64
124        );
125    }
126}
cortiq_engine/cpuprof.rs

cortiq_engine/
cpuprof.rs
