metrale_telemetry/
launch_trace.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Launch trace: record the kernel launches, async memsets and
4//! device-to-device copies a region enqueues, so two executions of it can be
5//! diffed op by op ([`end_and_diff`]). The verify paths trace their captured
6//! region this way to find a host value that differs between two steps. Off
7//! until [`begin`] is called.
8//!
9//! Owner: telemetry.
10//! Invariants: [`record`] stores an op only between [`begin`] and the next
11//! [`end_and_diff`].
12
13use std::collections::HashMap;
14use std::sync::Mutex;
15use std::sync::atomic::{AtomicBool, Ordering};
16
17/// 2026-09-26: One enqueued op. `kind` separates kernels from memsets and
18/// copies, so an op that appears or vanishes shows as an op mismatch.
19#[derive(Clone, PartialEq, Eq)]
20pub struct Entry {
21    pub kind: &'static str,
22    pub func: u64,
23    pub grid: [u32; 3],
24    pub block: [u32; 3],
25    pub smem: u32,
26    /// 2026-09-26: Args as u64 words: a buffer is its raw address, a byte
27    /// argument its first 8 bytes, little-endian.
28    pub args: Vec<u64>,
29}
30
31static ON: AtomicBool = AtomicBool::new(false);
32static TRACE: Mutex<Vec<Entry>> = Mutex::new(Vec::new());
33static PREV: Mutex<Option<Vec<Entry>>> = Mutex::new(None);
34static NAMES: Mutex<Option<HashMap<u64, String>>> = Mutex::new(None);
35
36#[inline(always)]
37pub fn on() -> bool {
38    ON.load(Ordering::Relaxed)
39}
40
41/// 2026-09-26: Remember a kernel handle's name. The CUDA backend's
42/// `GpuBackend::kernel` calls it on every successful lookup.
43pub fn name_kernel(handle: u64, module: &str, func: &str) {
44    let mut g = NAMES.lock().unwrap();
45    g.get_or_insert_with(HashMap::new)
46        .insert(handle, format!("{module}::{func}"));
47}
48
49/// 2026-09-26: The `module::func` a kernel handle was looked up as, if it was named.
50pub fn kernel_name(handle: u64) -> Option<String> {
51    NAMES
52        .lock()
53        .unwrap()
54        .as_ref()
55        .and_then(|m| m.get(&handle).cloned())
56}
57
58fn name_of(handle: u64) -> String {
59    kernel_name(handle).unwrap_or_else(|| format!("fn@{handle:#x}"))
60}
61
62pub fn begin() {
63    TRACE.lock().unwrap().clear();
64    ON.store(true, Ordering::Relaxed);
65}
66
67#[inline(always)]
68pub fn record(e: Entry) {
69    if on() {
70        TRACE.lock().unwrap().push(e);
71    }
72}
73
74/// 2026-09-26: Stop recording and diff this trace against the previous one.
75/// Returns `None` on the first call, else a report of how many ops differ,
76/// detailing at most `max_report` of them.
77pub fn end_and_diff(max_report: usize) -> Option<String> {
78    ON.store(false, Ordering::Relaxed);
79    let cur = std::mem::take(&mut *TRACE.lock().unwrap());
80    let prev = PREV.lock().unwrap().replace(cur.clone())?;
81
82    let mut out = String::new();
83    let mut n = 0usize;
84    if prev.len() != cur.len() {
85        out.push_str(&format!(
86            "OP COUNT differs: prev {} vs cur {}\n",
87            prev.len(),
88            cur.len()
89        ));
90    }
91    for (i, (p, c)) in prev.iter().zip(cur.iter()).enumerate() {
92        if p == c {
93            continue;
94        }
95        n += 1;
96        if n > max_report {
97            continue;
98        }
99        let mut what = Vec::new();
100        if p.kind != c.kind || p.func != c.func {
101            what.push(format!("op {} -> {}", name_of(p.func), name_of(c.func)));
102        }
103        if p.grid != c.grid {
104            what.push(format!("grid {:?} -> {:?}", p.grid, c.grid));
105        }
106        if p.block != c.block {
107            what.push(format!("block {:?} -> {:?}", p.block, c.block));
108        }
109        if p.smem != c.smem {
110            what.push(format!("smem {} -> {}", p.smem, c.smem));
111        }
112        for (a, (pv, cv)) in p.args.iter().zip(c.args.iter()).enumerate() {
113            if pv != cv {
114                what.push(format!(
115                    "arg{a} {pv:#x} -> {cv:#x} (Δ {})",
116                    *cv as i64 - *pv as i64
117                ));
118            }
119        }
120        if p.args.len() != c.args.len() {
121            what.push(format!("argc {} -> {}", p.args.len(), c.args.len()));
122        }
123        out.push_str(&format!(
124            "#{i} {} [{}] {}\n",
125            name_of(c.func),
126            c.kind,
127            what.join("; ")
128        ));
129    }
130    Some(format!(
131        "{n} differing op(s) of {}\n{out}",
132        cur.len().min(prev.len())
133    ))
134}
metrale_telemetry/launch_trace.rs

metrale_telemetry/
launch_trace.rs
