openvm_stark_sdk/
nvtx_tracing.rs

1use std::sync::Arc;
2
3use tracing::{
4    field::{Field, Visit},
5    Id, Subscriber,
6};
7use tracing_subscriber::{registry::LookupSpan, Layer};
8
9/// A tracing layer that maps tracing spans to NVTX ranges for Nsight Systems profiling.
10///
11/// When active, each span enter/exit maps to an NVTX range_push/range_pop, making
12/// proving phases visible in `nsys profile --trace=cuda,nvtx` timelines.
13///
14/// Only spans at or above the configured level are exported.
15pub struct NvtxLayer {
16    config: NvtxConfig,
17}
18
19pub struct NvtxConfig {
20    pub max_level: tracing::Level,
21}
22
23impl Default for NvtxConfig {
24    fn default() -> Self {
25        Self {
26            max_level: tracing::Level::INFO,
27        }
28    }
29}
30
31impl NvtxLayer {
32    pub fn new(config: NvtxConfig) -> Self {
33        Self { config }
34    }
35}
36
37/// Per-span cached NVTX label, stored in span extensions.
38/// Only inserted for spans that pass the level filter.
39struct NvtxSpanData {
40    label: Arc<str>,
41}
42
43/// Visitor that collects all span fields into a compact label suffix.
44struct FieldCollector {
45    parts: Vec<String>,
46}
47
48impl FieldCollector {
49    fn new() -> Self {
50        Self { parts: Vec::new() }
51    }
52
53    fn into_suffix(self) -> Option<String> {
54        if self.parts.is_empty() {
55            None
56        } else {
57            Some(self.parts.join(","))
58        }
59    }
60}
61
62impl Visit for FieldCollector {
63    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
64        self.parts.push(format!("{}={:?}", field.name(), value));
65    }
66
67    fn record_i64(&mut self, field: &Field, value: i64) {
68        self.parts.push(format!("{}={}", field.name(), value));
69    }
70
71    fn record_u64(&mut self, field: &Field, value: u64) {
72        self.parts.push(format!("{}={}", field.name(), value));
73    }
74
75    fn record_str(&mut self, field: &Field, value: &str) {
76        self.parts.push(format!("{}={}", field.name(), value));
77    }
78}
79
80fn build_label(meta: &tracing::Metadata<'_>, attrs: &tracing::span::Attributes<'_>) -> Arc<str> {
81    let mut collector = FieldCollector::new();
82    attrs.record(&mut collector);
83
84    let name = meta.name();
85    match collector.into_suffix() {
86        Some(suffix) => Arc::from(format!("{name}{{{suffix}}}")),
87        None => Arc::from(name),
88    }
89}
90
91impl<S> Layer<S> for NvtxLayer
92where
93    S: Subscriber + for<'a> LookupSpan<'a>,
94{
95    fn on_new_span(
96        &self,
97        attrs: &tracing::span::Attributes<'_>,
98        id: &Id,
99        ctx: tracing_subscriber::layer::Context<'_, S>,
100    ) {
101        let Some(span) = ctx.span(id) else { return };
102        let meta = span.metadata();
103
104        if meta.level() > &self.config.max_level {
105            return;
106        }
107
108        let label = build_label(meta, attrs);
109        span.extensions_mut().insert(NvtxSpanData { label });
110    }
111
112    fn on_record(
113        &self,
114        id: &Id,
115        values: &tracing::span::Record<'_>,
116        ctx: tracing_subscriber::layer::Context<'_, S>,
117    ) {
118        let Some(span) = ctx.span(id) else { return };
119        let mut exts = span.extensions_mut();
120        let Some(data) = exts.get_mut::<NvtxSpanData>() else {
121            return;
122        };
123        // Append newly recorded fields to the existing label.
124        // If the label already has fields ("name{a=1}"), strip the closing brace and
125        // append with a comma. If no fields yet ("name"), open a new brace group.
126        let mut collector = FieldCollector::new();
127        values.record(&mut collector);
128        if let Some(suffix) = collector.into_suffix() {
129            let base = data.label.trim_end_matches('}');
130            let sep = if base.len() < data.label.len() {
131                ","
132            } else {
133                "{"
134            };
135            data.label = Arc::from(format!("{base}{sep}{suffix}}}"));
136        }
137    }
138
139    fn on_enter(&self, id: &Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
140        let Some(span) = ctx.span(id) else { return };
141        let exts = span.extensions();
142        let Some(data) = exts.get::<NvtxSpanData>() else {
143            return;
144        };
145        nvtx::range_push!("{}", data.label);
146    }
147
148    fn on_exit(&self, id: &Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
149        let Some(span) = ctx.span(id) else { return };
150        let exts = span.extensions();
151        if exts.get::<NvtxSpanData>().is_some() {
152            nvtx::range_pop!();
153        }
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use tracing::info_span;
160    use tracing_subscriber::{layer::SubscriberExt, Registry};
161
162    use super::*;
163
164    #[test]
165    fn test_nvtx_layer_with_fields() {
166        let subscriber = Registry::default().with(NvtxLayer::new(NvtxConfig::default()));
167        tracing::subscriber::with_default(subscriber, || {
168            // Should produce "prove_segment{segment=42,phase=prover}"
169            let span = info_span!("prove_segment", segment = 42, phase = "prover");
170            let _guard = span.enter();
171            // No fields -> just the name
172            let inner = info_span!("stark_prove");
173            let _inner_guard = inner.enter();
174        });
175    }
176
177    #[test]
178    fn test_nvtx_layer_late_bound_fields() {
179        let subscriber = Registry::default().with(NvtxLayer::new(NvtxConfig::default()));
180        tracing::subscriber::with_default(subscriber, || {
181            // Late-bound field via span.record()
182            let span = info_span!("step", idx = tracing::field::Empty);
183            let _guard = span.enter();
184            span.record("idx", 7);
185            // Also test recording on a span that already has fields
186            let span2 = info_span!("work", phase = "init", result = tracing::field::Empty);
187            let _guard2 = span2.enter();
188            span2.record("result", 42);
189        });
190    }
191}
openvm_stark_sdk/nvtx_tracing.rs

openvm_stark_sdk/
nvtx_tracing.rs
