spark_runtime/
launch_trace.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! ANOMALIES A56 diagnostic: record every GPU op a step enqueues, so two
4//! consecutive steps can be diffed.
5//!
6//! A CUDA graph BAKES the grid, block, shared-mem and argument BYTES of every
7//! launch at capture time. So a captured region is replayable if and only if
8//! two consecutive executions of it enqueue byte-identical ops. Anything that
9//! differs between step N and step N+1 is a host value the graph froze — which
10//! is exactly the class of bug where the capture pass is byte-exact but the
11//! replay is not.
12//!
13//! This turns "which host scalar leaked into a kernel argument?" from a code
14//! read into a mechanical diff. Off unless `begin()` is called.
15
16use std::collections::HashMap;
17use std::sync::Mutex;
18use std::sync::atomic::{AtomicBool, Ordering};
19
20/// One enqueued op. `kind` separates kernels from memsets/copies so an op
21/// appearing or vanishing shows up as a kind mismatch rather than arg noise.
22#[derive(Clone, PartialEq, Eq)]
23pub struct Entry {
24    pub kind: &'static str,
25    pub func: u64,
26    pub grid: [u32; 3],
27    pub block: [u32; 3],
28    pub smem: u32,
29    /// Args as u64 words: buffers are the raw address, scalars are LE-packed.
30    pub args: Vec<u64>,
31}
32
33static ON: AtomicBool = AtomicBool::new(false);
34static TRACE: Mutex<Vec<Entry>> = Mutex::new(Vec::new());
35static PREV: Mutex<Option<Vec<Entry>>> = Mutex::new(None);
36static NAMES: Mutex<Option<HashMap<u64, String>>> = Mutex::new(None);
37
38#[inline(always)]
39pub fn on() -> bool {
40    ON.load(Ordering::Relaxed)
41}
42
43/// Remember a kernel handle's name. Called from the backend's `kernel()`
44/// lookup, which runs at init only.
45pub fn name_kernel(handle: u64, module: &str, func: &str) {
46    let mut g = NAMES.lock().unwrap();
47    g.get_or_insert_with(HashMap::new)
48        .insert(handle, format!("{module}::{func}"));
49}
50
51fn name_of(handle: u64) -> String {
52    NAMES
53        .lock()
54        .unwrap()
55        .as_ref()
56        .and_then(|m| m.get(&handle).cloned())
57        .unwrap_or_else(|| format!("fn@{handle:#x}"))
58}
59
60pub fn begin() {
61    TRACE.lock().unwrap().clear();
62    ON.store(true, Ordering::Relaxed);
63}
64
65#[inline(always)]
66pub fn record(e: Entry) {
67    if on() {
68        TRACE.lock().unwrap().push(e);
69    }
70}
71
72/// Stop recording and diff this trace against the previous one. Returns
73/// `None` on the first call (nothing to compare against yet), else a report
74/// naming every op that differs.
75pub fn end_and_diff(max_report: usize) -> Option<String> {
76    ON.store(false, Ordering::Relaxed);
77    let cur = std::mem::take(&mut *TRACE.lock().unwrap());
78    let prev = PREV.lock().unwrap().replace(cur.clone())?;
79
80    let mut out = String::new();
81    let mut n = 0usize;
82    if prev.len() != cur.len() {
83        out.push_str(&format!(
84            "OP COUNT differs: prev {} vs cur {}\n",
85            prev.len(),
86            cur.len()
87        ));
88    }
89    for (i, (p, c)) in prev.iter().zip(cur.iter()).enumerate() {
90        if p == c {
91            continue;
92        }
93        n += 1;
94        if n > max_report {
95            continue;
96        }
97        let mut what = Vec::new();
98        if p.kind != c.kind || p.func != c.func {
99            what.push(format!("op {} -> {}", name_of(p.func), name_of(c.func)));
100        }
101        if p.grid != c.grid {
102            what.push(format!("grid {:?} -> {:?}", p.grid, c.grid));
103        }
104        if p.block != c.block {
105            what.push(format!("block {:?} -> {:?}", p.block, c.block));
106        }
107        if p.smem != c.smem {
108            what.push(format!("smem {} -> {}", p.smem, c.smem));
109        }
110        for (a, (pv, cv)) in p.args.iter().zip(c.args.iter()).enumerate() {
111            if pv != cv {
112                what.push(format!(
113                    "arg{a} {pv:#x} -> {cv:#x} (Δ {})",
114                    *cv as i64 - *pv as i64
115                ));
116            }
117        }
118        if p.args.len() != c.args.len() {
119            what.push(format!("argc {} -> {}", p.args.len(), c.args.len()));
120        }
121        out.push_str(&format!(
122            "#{i} {} [{}] {}\n",
123            name_of(c.func),
124            c.kind,
125            what.join("; ")
126        ));
127    }
128    Some(format!(
129        "{n} differing op(s) of {}\n{out}",
130        cur.len().min(prev.len())
131    ))
132}
spark_runtime/launch_trace.rs

spark_runtime/
launch_trace.rs
