metrale_gpu_runtime/
timing.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: L1 GPU spans: CUDA event pairs around each serve lane on every
4//! step, and around each eager kernel launch on the steps the telemetry hub arms
5//! (one in `kernel_span_every`), at telemetry level `Kernel` only.
6//!
7//! The bookkeeping is `metrale_telemetry::kernel::SpanRing`; this module is the
8//! CUDA half: the event pool (`cuda_events`) and the process-wide recorder that
9//! `MetraleCudaBackend::new` installs.
10//!
11//! Lane spans are recorded on the registry's stream (`MetraleRegistry::raw_stream`),
12//! which cudarc creates non-blocking, so a lane pair times the work queued on that
13//! stream only.
14//!
15//! Owner: gpu-runtime (telemetry).
16//! Invariants:
17//! - No span synchronises: a completed pair is harvested with a non-blocking
18//!   `cuEventQuery`, and a pair still in flight is left for a later harvest.
19//! - Below level `Kernel` no entry point reaches CUDA: each first reads the hub's
20//!   level or armed flag (a relaxed atomic load), and `kernel_end` acts only on a
21//!   span `kernel_begin` opened.
22//! - No kernel span is opened on a stream that is being captured.
23
24#[cfg(feature = "cuda")]
25mod cuda_events;
26
27#[cfg(feature = "cuda")]
28pub use imp::*;
29
30/// 2026-09-25: Event pairs in the pool. The ring hands out slots in turn; a span
31/// whose slot is still in flight is dropped and counted (`spans_dropped`).
32pub const SPAN_SLOTS: usize = 4096;
33
34#[cfg(feature = "cuda")]
35mod imp {
36    use std::sync::atomic::{AtomicUsize, Ordering};
37    use std::sync::{Arc, RwLock};
38
39    use metrale_telemetry::kernel::{LANES, SpanKey, SpanRing, SpanToken};
40    use metrale_telemetry::{Level, global};
41
42    use super::SPAN_SLOTS;
43    use super::cuda_events::CudaEvents;
44
45    /// 2026-09-25: The recorder `install_for_stream` installs; a later install replaces it.
46    pub struct GpuSpans {
47        ring: SpanRing<CudaEvents>,
48        stream: u64,
49        open_lanes: [AtomicUsize; LANES.len()],
50    }
51
52    const NONE: usize = usize::MAX;
53
54    static INSTALLED: RwLock<Option<Arc<GpuSpans>>> = RwLock::new(None);
55
56    fn installed() -> Option<Arc<GpuSpans>> {
57        INSTALLED.read().ok()?.clone()
58    }
59
60    /// 2026-09-25: Build and install the recorder for a backend whose stream is
61    /// `stream`, when telemetry is at level `Kernel`. Called from
62    /// `MetraleCudaBackend::new` after the registry loads. A pool that cannot be
63    /// created leaves L1 off and logs why.
64    pub fn install_for_stream(stream: u64) {
65        if global().level() != Level::Kernel {
66            return;
67        }
68        match CudaEvents::create(SPAN_SLOTS) {
69            Ok(events) => {
70                let spans = GpuSpans {
71                    ring: SpanRing::new(events, SPAN_SLOTS),
72                    stream,
73                    open_lanes: [const { AtomicUsize::new(NONE) }; LANES.len()],
74                };
75                if let Ok(mut slot) = INSTALLED.write() {
76                    *slot = Some(Arc::new(spans));
77                }
78            }
79            Err(e) => tracing::warn!("telemetry: GPU spans disabled — {e}"),
80        }
81    }
82
83    /// 2026-09-25: Open lane `lane`'s span; an index outside `LANES` is ignored.
84    pub fn lane_begin(lane: usize) {
85        if !global().lane_spans() {
86            return;
87        }
88        let Some(s) = installed() else { return };
89        let Some(slot) = s.open_lanes.get(lane) else {
90            return;
91        };
92        if let Some(tok) = s
93            .ring
94            .begin(SpanKey::Lane(lane), s.stream, &global().kernel)
95        {
96            slot.store(tok.slot(), Ordering::Relaxed);
97        }
98    }
99
100    /// 2026-09-25: Close lane `lane`'s span, if one is open.
101    pub fn lane_end(lane: usize) {
102        if !global().lane_spans() {
103            return;
104        }
105        let Some(s) = installed() else { return };
106        let Some(slot) = s.open_lanes.get(lane) else {
107            return;
108        };
109        let i = slot.swap(NONE, Ordering::Relaxed);
110        if i != NONE {
111            s.ring
112                .end(SpanToken::from_slot(i), s.stream, &global().kernel);
113        }
114    }
115
116    /// 2026-09-25: Land every completed span. The scheduler calls it at each step begin.
117    pub fn harvest() {
118        if global().level() != Level::Kernel {
119            return;
120        }
121        if let Some(s) = installed() {
122            s.ring.harvest(&global().kernel);
123        }
124    }
125
126    /// 2026-09-25: A kernel span, when this step is armed for kernel spans and
127    /// `stream` is not being captured (an event recorded into a capture would be a
128    /// graph node, not a timestamp).
129    pub fn kernel_begin(
130        gpu: &dyn crate::gpu::GpuBackend,
131        func: u64,
132        stream: u64,
133    ) -> Option<KernelSpan> {
134        if !global().kernel_spans_armed() || gpu.stream_is_capturing(stream) {
135            return None;
136        }
137        let s = installed()?;
138        let tok = s
139            .ring
140            .begin(SpanKey::Kernel(func), stream, &global().kernel)?;
141        Some(KernelSpan { spans: s, tok })
142    }
143
144    /// 2026-09-25: Close a kernel span opened by [`kernel_begin`].
145    pub fn kernel_end(span: Option<KernelSpan>, stream: u64) {
146        if let Some(k) = span {
147            k.spans.ring.end(k.tok, stream, &global().kernel);
148        }
149    }
150
151    /// 2026-09-25: An open kernel span.
152    pub struct KernelSpan {
153        spans: Arc<GpuSpans>,
154        tok: SpanToken,
155    }
156}
157
158/// 2026-09-25: Non-CUDA builds carry no GPU spans. The entry points the scheduler
159/// calls stay, so callers need no cfg of their own.
160#[cfg(not(feature = "cuda"))]
161mod imp_absent {
162    pub fn lane_begin(_lane: usize) {}
163    pub fn lane_end(_lane: usize) {}
164    pub fn harvest() {}
165}
166#[cfg(not(feature = "cuda"))]
167pub use imp_absent::*;
metrale_gpu_runtime/timing.rs

metrale_gpu_runtime/
timing.rs
