Skip to main content

hrx/execution/
trace.rs

1use std::{
2    collections::VecDeque,
3    sync::{
4        Mutex,
5        atomic::{AtomicBool, Ordering},
6    },
7    time::{Duration, Instant},
8};
9
10/// Bounded trace of host-observed synchronized region latencies, including
11/// driver and scheduling overhead. These are never GPU timestamps.
12#[derive(Clone, Debug, serde::Serialize)]
13pub struct ExecutionTrace {
14    /// Events in completion order; timestamps retain actual start order.
15    pub events: Vec<TraceEvent>,
16    /// Oldest events discarded when capacity was exceeded.
17    pub dropped: u64,
18}
19/// A synchronized native region as observed by its execution worker.
20#[derive(Clone, Debug, serde::Serialize)]
21pub struct TraceEvent {
22    /// Start relative to capture start, in microseconds.
23    pub start_us: f64,
24    /// Host-observed execution latency, in microseconds.
25    pub duration_us: f64,
26    /// Upload=0, compute=1, download=2, NPU=3, host=4.
27    pub lane: usize,
28    /// Explicit copied bytes in the region.
29    pub copied_bytes: usize,
30    /// Whether execution failed.
31    pub failed: bool,
32}
33impl ExecutionTrace {
34    /// Chrome/Perfetto trace JSON with an explicit measurement-scope label.
35    pub fn to_json(&self) -> crate::Result<String> {
36        let events = self
37            .events
38            .iter()
39            .map(|event| {
40                serde_json::json!({
41                    "name": "native region (host-observed)", "cat": "hrx.host", "ph": "X",
42                    "ts": event.start_us, "dur": event.duration_us, "pid": 0, "tid": event.lane,
43                    "args": {"copied_bytes": event.copied_bytes, "failed": event.failed},
44                })
45            })
46            .collect::<Vec<_>>();
47        serde_json::to_string(&serde_json::json!({
48            "traceEvents": events, "displayTimeUnit": "ms", "dropped_events": self.dropped,
49            "measurement": "host-observed synchronized region latency; not device timestamps",
50        }))
51        .map_err(|error| crate::Error::Message(error.to_string()))
52    }
53}
54struct Capture {
55    epoch: Instant,
56    capacity: usize,
57    events: VecDeque<TraceEvent>,
58    dropped: u64,
59}
60#[derive(Default)]
61pub(super) struct Tracer {
62    enabled: AtomicBool,
63    capture: Mutex<Option<Capture>>,
64}
65impl Tracer {
66    pub fn start(&self, capacity: usize) -> crate::Result<()> {
67        if capacity == 0 {
68            return Err(crate::Error::Message(
69                "trace capacity must be nonzero".into(),
70            ));
71        }
72        let mut capture = self.capture.lock().unwrap_or_else(|e| e.into_inner());
73        if capture.is_some() {
74            return Err(crate::Error::Busy("a trace is already running".into()));
75        }
76        *capture = Some(Capture {
77            epoch: Instant::now(),
78            capacity,
79            events: VecDeque::new(),
80            dropped: 0,
81        });
82        self.enabled.store(true, Ordering::Release);
83        Ok(())
84    }
85    pub fn finish(&self) -> Option<ExecutionTrace> {
86        let mut capture = self.capture.lock().unwrap_or_else(|e| e.into_inner());
87        self.enabled.store(false, Ordering::Release);
88        capture.take().map(|capture| ExecutionTrace {
89            events: capture.events.into(),
90            dropped: capture.dropped,
91        })
92    }
93    pub fn record(
94        &self,
95        start: Instant,
96        elapsed: Duration,
97        lane: usize,
98        copied_bytes: usize,
99        failed: bool,
100    ) {
101        if !self.enabled.load(Ordering::Acquire) {
102            return;
103        }
104        let mut capture = self.capture.lock().unwrap_or_else(|e| e.into_inner());
105        let Some(capture) = capture.as_mut() else {
106            return;
107        };
108        let Some(offset) = start.checked_duration_since(capture.epoch) else {
109            return;
110        };
111        if capture.events.len() == capture.capacity {
112            capture.events.pop_front();
113            capture.dropped += 1;
114        }
115        capture.events.push_back(TraceEvent {
116            start_us: offset.as_secs_f64() * 1e6,
117            duration_us: elapsed.as_secs_f64() * 1e6,
118            lane,
119            copied_bytes,
120            failed,
121        });
122    }
123}