diff --git a/Cargo.lock b/Cargo.lock index ce3b000..4ec8aff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -170,9 +170,9 @@ checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c" [[package]] name = "memchr" -version = "2.7.1" +version = "2.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149" +checksum = "6c8640c5d730cb13ebd907d8d04b52f55ac9a2eec55b440c8892f40d56c76c1d" [[package]] name = "minimal-lexical" @@ -190,6 +190,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + [[package]] name = "nvtx" version = "0.1.0" @@ -198,6 +208,8 @@ dependencies = [ "gettid", "libc", "nvtx-sys", + "tracing", + "tracing-subscriber", "widestring", ] @@ -216,6 +228,18 @@ version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + +[[package]] +name = "pin-project-lite" +version = "0.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8afb450f006bf6385ca15ef45d71d2288452bc3683ce2e2cacc0d18e4be60b58" + [[package]] name = "prettyplease" version = "0.2.17" @@ -292,12 +316,27 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "smallvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" + [[package]] name = "syn" version = "2.0.55" @@ -309,12 +348,85 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "thread_local" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +dependencies = [ + "cfg-if", + "once_cell", +] + +[[package]] +name = "tracing" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "nu-ansi-term", + "sharded-slab", + "smallvec", + "thread_local", + "tracing-core", + "tracing-log", +] + [[package]] name = "unicode-ident" version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + [[package]] name = "which" version = "4.4.2" diff --git a/Cargo.toml b/Cargo.toml index d4def3f..ff85a8e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,16 +25,19 @@ publish.workspace = true license.workspace = true [features] -default = ["color-name", "name-current-thread", "cuda", "cuda_runtime"] +default = ["color-name", "name-current-thread", "cuda", "cuda_runtime", "tracing"] color-name = ["dep:color-name"] name-current-thread = ["dep:gettid"] cuda = ["nvtx-sys/cuda"] cuda_runtime = ["nvtx-sys/cuda_runtime"] +tracing = ["dep:tracing", "dep:tracing-subscriber", "dep:color-name" ] [dependencies] color-name = { version = "1.1.0", optional = true } gettid = { version = "0.1.2", optional = true } nvtx-sys = { path = "crates/nvtx-sys" } +tracing = { version = "0.1.40", optional = true, features = ["attributes"] } +tracing-subscriber = { version = "0.3.18", optional = true } widestring = "1.0.2" [target.'cfg(unix)'.dependencies] @@ -59,3 +62,7 @@ required-features = ["name-current-thread"] [[example]] name = "domain_threads" required-features = ["name-current-thread"] + +[[example]] +name = "tracing" +required-features = ["tracing"] diff --git a/examples/tracing.rs b/examples/tracing.rs new file mode 100644 index 0000000..cae3e15 --- /dev/null +++ b/examples/tracing.rs @@ -0,0 +1,35 @@ +use tracing::{info, instrument}; +use tracing_subscriber::prelude::*; + +#[instrument(fields(color = "goldenrod", category = 1, payload = i))] +fn foo(i: u64) { + for j in 1..=i { + bar(j); + std::thread::sleep(std::time::Duration::from_millis(10 * i)); + } +} + +#[instrument(fields(color = "plum", category = 1, payload = j))] +fn bar(j: u64) { + for k in 1..=j { + baz(k); + std::thread::sleep(std::time::Duration::from_millis(10 * j)); + } +} + +#[instrument(fields(color = "salmon", category = 2, payload = k))] +fn baz(k: u64) { + std::thread::sleep(std::time::Duration::from_millis(10 * k)); +} + +fn main() { + let layer = nvtx::tracing::NvtxLayer::new("nvtx"); + tracing_subscriber::registry().with(layer).init(); + + info!( + message = "At the beginning of the program", + color = "blue", + category = 2 + ); + foo(10); +} diff --git a/src/color.rs b/src/color.rs index 4f9dcf9..77a0079 100644 --- a/src/color.rs +++ b/src/color.rs @@ -99,7 +99,7 @@ impl TypeValueEncodable for Color { type Value = u32; fn encode(&self) -> (Self::Type, Self::Value) { - let as_u32 = u32::from_ne_bytes([self.a, self.r, self.g, self.b]); + let as_u32 = u32::from_be_bytes([self.a, self.r, self.g, self.b]); (Self::Type::NVTX_COLOR_ARGB, as_u32) } diff --git a/src/domain/category.rs b/src/domain/category.rs index deda0af..259e0b4 100644 --- a/src/domain/category.rs +++ b/src/domain/category.rs @@ -3,6 +3,6 @@ use crate::Domain; /// Represents a category for use with event and range grouping. See [`Domain::register_category`], [`Domain::register_categories`] #[derive(Debug, Clone, Copy)] pub struct Category<'a> { - pub(super) id: u32, - pub(super) domain: &'a Domain, + pub(crate) id: u32, + pub(crate) domain: &'a Domain, } diff --git a/src/domain/event_attributes.rs b/src/domain/event_attributes.rs index 51f6646..1acffe3 100644 --- a/src/domain/event_attributes.rs +++ b/src/domain/event_attributes.rs @@ -4,10 +4,10 @@ use crate::{Color, Domain, Payload, TypeValueEncodable}; /// All attributes that are associated with marks and ranges #[derive(Debug, Clone)] pub struct EventAttributes<'a> { - pub(super) category: Option>, - pub(super) color: Option, - pub(super) payload: Option, - pub(super) message: Option>, + pub(crate) category: Option>, + pub(crate) color: Option, + pub(crate) payload: Option, + pub(crate) message: Option>, } impl<'a> EventAttributes<'a> { diff --git a/src/domain/local_range.rs b/src/domain/local_range.rs index a143438..7ad770a 100644 --- a/src/domain/local_range.rs +++ b/src/domain/local_range.rs @@ -2,7 +2,7 @@ use super::EventArgument; use crate::Domain; use std::marker::PhantomData; -/// A RAII-like object for modeling callstack Ranges within a Domain +/// A RAII-like object for modeling callstack (thread-local) Ranges within a Domain #[derive(Debug)] pub struct LocalRange<'a> { pub(super) domain: &'a Domain, diff --git a/src/domain/mod.rs b/src/domain/mod.rs index d1a4788..d193b65 100644 --- a/src/domain/mod.rs +++ b/src/domain/mod.rs @@ -224,6 +224,20 @@ impl Domain { Range::new(arg, self) } + /// Internal function for starting a range and returning a raw Range Id + pub(crate) fn range_start<'a>(&self, arg: impl Into>) -> u64 { + let arg = match arg.into() { + EventArgument::Attributes(attr) => attr, + EventArgument::Message(m) => m.into(), + }; + nvtx_sys::domain_range_start_ex(self.handle, &arg.encode()) + } + + /// Internal function for ending a range given a raw Range Id + pub(crate) fn range_end(&self, range_id: u64) { + nvtx_sys::domain_range_end(self.handle, range_id); + } + /// Name a resource /// /// ``` diff --git a/src/domain/range.rs b/src/domain/range.rs index 477e29a..54c7e0e 100644 --- a/src/domain/range.rs +++ b/src/domain/range.rs @@ -1,7 +1,7 @@ use super::EventArgument; use crate::Domain; -/// A RAII-like object for modeling start/end Ranges within a Domain +/// A RAII-like object for modeling process-wide Ranges within a Domain #[derive(Debug)] pub struct Range<'a> { pub(super) id: nvtx_sys::RangeId, @@ -10,18 +10,16 @@ pub struct Range<'a> { impl<'a> Range<'a> { pub(super) fn new(arg: impl Into>, domain: &'a Domain) -> Range<'a> { - let arg = match arg.into() { - EventArgument::Attributes(attr) => attr, - EventArgument::Message(m) => m.into(), - }; - let id = nvtx_sys::domain_range_start_ex(domain.handle, &arg.encode()); - Range { id, domain } + Range { + id: domain.range_start(arg), + domain, + } } } impl<'a> Drop for Range<'a> { fn drop(&mut self) { - nvtx_sys::domain_range_end(self.domain.handle, self.id) + self.domain.range_end(self.id); } } diff --git a/src/lib.rs b/src/lib.rs index ebe1a2e..a561a9a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -38,6 +38,9 @@ //! runtime resources such as Devices, Events, and Streams. The feature also adds `CudaRuntimeIdentifier` //! to the [`crate::domain`] module to provide an alternative naming mechanism via [`crate::Domain::name_resource`]. //! +//! * **tracing** - +//! When enabled, +//! //! ## Platform-specific types //! //! * **PThread Resource Naming** - @@ -46,12 +49,17 @@ /// color support pub mod color; + /// specialized types for use within a domain context pub mod domain; /// platform native types pub mod native_types; +#[cfg(feature = "tracing")] +/// tracing support +pub mod tracing; + mod category; mod event_argument; mod event_attributes; diff --git a/src/local_range.rs b/src/local_range.rs index 116641e..12363fb 100644 --- a/src/local_range.rs +++ b/src/local_range.rs @@ -1,7 +1,7 @@ use crate::{EventArgument, Message}; use std::marker::PhantomData; -/// A RAII-like object for modeling callstack Ranges +/// A RAII-like object for modeling callstack (thread-local) Ranges #[derive(Debug)] pub struct LocalRange { // prevent Sync + Send diff --git a/src/range.rs b/src/range.rs index 14ec341..f4a18a3 100644 --- a/src/range.rs +++ b/src/range.rs @@ -1,6 +1,6 @@ use crate::{EventArgument, Message}; -/// A RAII-like object for modeling start/end Ranges +/// A RAII-like object for modeling process-wide Ranges #[derive(Debug)] pub struct Range { id: nvtx_sys::RangeId, diff --git a/src/tracing/mod.rs b/src/tracing/mod.rs new file mode 100644 index 0000000..af8885a --- /dev/null +++ b/src/tracing/mod.rs @@ -0,0 +1,212 @@ +use crate::{ + domain::{Category, EventAttributes, Message}, + Color, Domain, Payload, Str, +}; +use std::marker::PhantomData; +use tracing::{ + field::{Field, Visit}, + span::{Attributes, Id, Record}, +}; +use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer}; + +/// The tracing layer for nvtx range and events. +/// +/// Only a subset of nvtx domain features are available. +/// +/// Unavailable functionality: +/// - registering strings for efficient reuse +/// - registering names for categories +/// +/// **Supported fields** +/// * `message` (`&str`) for Marks only -- the span name is used as the message for Spans +/// * `color` (`&str`) +/// * `payload` (one of: `f64`, `u64`, `i64`) +/// * `category` (`u32`) provides a numerical category +/// +/// `instrument` example: +/// +/// ``` +/// use tracing::instrument; +/// +/// #[instrument(fields(color = "salmon", category = 2, payload = k))] +/// fn baz (k : u64) { +/// std::thread::sleep(std::time::Duration::from_millis(10 * k)); +/// } +/// ``` +/// +/// `mark` example: +/// +/// ``` +/// use tracing::info; +/// +/// info!(message = "At the beginning of the program", color = "blue", category = 2); +/// ``` +/// +/// `span` example: +/// +/// ``` +/// use tracing::info_span; +/// +/// let span = info_span!("Running an arbitrary block!", color = "red", payload = 3.1415); +/// span.in_scope(|| { +/// // do work inside the span... +/// }); +/// ``` +pub struct NvtxLayer { + /// the domain to use + domain: crate::Domain, +} + +impl NvtxLayer { + /// Create a new layer with a given domain name + pub fn new(name: impl Into) -> NvtxLayer { + NvtxLayer { + domain: Domain::new(name), + } + } + + /// Get the layer's domain + pub fn get_domain(&self) -> &Domain { + &self.domain + } +} + +#[derive(Debug, Clone, Default)] +struct NvtxData { + message: Option, + color: Option, + category: Option, + payload: Option, +} + +struct NvtxId(u64); + +impl Layer for NvtxLayer +where + S: tracing::Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) { + let mut data = NvtxData::default(); + let mut visitor = NvtxVisitor::<'_, S>::new(&mut data); + event.record(&mut visitor); + let attr = EventAttributes { + message: data.message.as_ref().map(|s| Message::from(s.clone())), + category: data.category.map(|c| Category { + id: c, + domain: &self.domain, + }), + color: data.color, + payload: data.payload, + }; + self.domain.mark(attr); + } + + fn on_new_span<'a>(&'a self, attrs: &Attributes<'a>, id: &Id, ctx: Context<'a, S>) { + let span = ctx.span(id).unwrap(); + let mut data = NvtxData::default(); + let mut visitor = NvtxVisitor::<'_, S>::new(&mut data); + attrs.record(&mut visitor); + data.message = Some(attrs.metadata().name().to_string()); + span.extensions_mut().insert(data); + } + + fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) { + match ctx.span(id).unwrap().extensions_mut().get_mut::() { + Some(data) => { + let mut visitor = NvtxVisitor::<'_, S>::new(data); + values.record(&mut visitor); + } + None => todo!(), + } + } + + fn on_enter(&self, id: &Id, ctx: Context<'_, S>) { + let mut attr: Option = None; + if let Some(data) = ctx.span(id).unwrap().extensions().get::() { + attr = Some(EventAttributes { + message: data.message.as_ref().map(|s| Message::from(s.clone())), + category: data.category.map(|c| Category { + id: c, + domain: &self.domain, + }), + color: data.color, + payload: data.payload, + }); + } + if let Some(a) = attr { + ctx.span(id) + .unwrap() + .extensions_mut() + .insert(NvtxId(self.domain.range_start(a))); + } + } + + fn on_exit(&self, id: &Id, ctx: Context<'_, S>) { + let span = ctx.span(id).unwrap(); + let maybe_id = span.extensions_mut().remove::(); + if let Some(NvtxId(id)) = maybe_id { + self.domain.range_end(id) + } + } +} + +struct NvtxVisitor<'a, S> +where + S: tracing::Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + data: &'a mut NvtxData, + _traits: PhantomData, +} + +impl<'a, S> NvtxVisitor<'a, S> +where + S: tracing::Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + fn new(data: &'a mut NvtxData) -> NvtxVisitor<'a, S> { + NvtxVisitor { + data, + _traits: PhantomData, + } + } +} + +impl<'a, S> Visit for NvtxVisitor<'a, S> +where + S: tracing::Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + fn record_debug(&mut self, _field: &Field, _value: &dyn std::fmt::Debug) {} + + fn record_f64(&mut self, field: &Field, value: f64) { + if field.name() == "payload" { + self.data.payload = Some(Payload::Double(value)); + } + } + fn record_i64(&mut self, field: &Field, value: i64) { + if field.name() == "payload" { + self.data.payload = Some(Payload::Int64(value)); + } else if field.name() == "category" { + self.data.category = Some(value as u32); + } + } + fn record_u64(&mut self, field: &Field, value: u64) { + if field.name() == "payload" { + self.data.payload = Some(Payload::Uint64(value)); + } else if field.name() == "category" { + self.data.category = Some(value as u32); + } + } + fn record_bool(&mut self, field: &Field, value: bool) { + if field.name() == "payload" { + self.data.payload = Some(Payload::Int32(value as i32)); + } + } + fn record_str(&mut self, field: &Field, value: &str) { + if field.name() == "color" { + if let Ok([r, g, b]) = color_name::Color::val().by_string(value.to_string()) { + self.data.color = Some(Color::new(r, g, b, 255)); + } + } else if field.name() == "message" { + self.data.message = Some(value.to_string()); + } + } +}