feat: add jolt-profiling crate

Cargo.toml

@@ -24,6 +24,7 @@ keywords = ["SNARK", "cryptography", "proofs"]
 
 [workspace]
 members = [
+  "crates/jolt-profiling",
   "crates/jolt-field",
     "jolt-core",
     "tracer",
@@ -235,6 +236,7 @@ sha3 = "0.10.8"
 blake2 = "0.10"
 blake3 = { version = "1.5.0" }
 light-poseidon = "0.4"
+digest = "0.10"
 jolt-optimizations = { git = "https://github.com/a16z/arkworks-algebra", branch = "dev/twist-shout" }
 dory = { package = "dory-pcs", version = "0.3.0", features = ["backends", "cache", "disk-persistence"] }
 
@@ -289,6 +291,8 @@ tracing-chrome = "0.7.1"
 tracing-subscriber = { version = "0.3.20", features = ["fmt", "env-filter"] }
 inferno = { version = "0.12.3" }
 allocative = { git = "https://github.com/facebookexperimental/allocative", rev = "85b773d85d526d068ce94724ff7a7b81203fc95e" }
+pprof = { version = "0.15", features = ["prost-codec", "flamegraph", "frame-pointer"] }
+prost = "0.14"
 
 # Parsing
 syn = { version = "2", features = ["full"] }

crates/jolt-profiling/Cargo.toml

@@ -0,0 +1,37 @@
+[package]
+name = "jolt-profiling"
+version = "0.1.0"
+authors = ["Jolt Contributors"]
+edition = "2021"
+description = "Profiling and tracing infrastructure for the Jolt proving system"
+license = "MIT OR Apache-2.0"
+repository = "https://github.com/a16z/jolt"
+keywords = ["profiling", "tracing", "performance"]
+categories = ["development-tools::profiling"]
+publish = false
+
+[lints]
+workspace = true
+
+[features]
+default = []
+monitor = ["dep:sysinfo"]
+pprof = ["dep:pprof", "dep:prost"]
+allocative = ["dep:inferno", "dep:allocative"]
+
+[dependencies]
+tracing.workspace = true
+tracing-chrome.workspace = true
+tracing-subscriber.workspace = true
+inferno = { workspace = true, optional = true }
+allocative = { workspace = true, optional = true }
+
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+memory-stats.workspace = true
+sysinfo = { workspace = true, optional = true }
+pprof = { workspace = true, optional = true }
+prost = { workspace = true, optional = true }
+
+[package.metadata.cargo-machete]
+# prost is required by pprof's prost-codec feature but not directly imported
+ignored = ["prost"]

crates/jolt-profiling/src/flamegraph.rs

@@ -0,0 +1,55 @@
+//! Heap flamegraph generation from `allocative`-instrumented data structures.
+
+use std::{fs::File, io::Cursor, path::Path};
+
+use allocative::{Allocative, FlameGraphBuilder};
+use inferno::flamegraph::Options;
+
+use crate::units::{format_memory_size, BYTES_PER_GIB, BYTES_PER_MIB};
+
+/// Logs the heap allocation size of an `Allocative`-instrumented value.
+pub fn print_data_structure_heap_usage<T: Allocative>(label: &str, data: &T) {
+    if tracing::enabled!(tracing::Level::DEBUG) {
+        let memory_gib = allocative::size_of_unique_allocated_data(data) as f64 / BYTES_PER_GIB;
+        tracing::debug!(
+            label = label,
+            usage = %format_memory_size(memory_gib),
+            "heap allocation size"
+        );
+    }
+}
+
+/// Renders a [`FlameGraphBuilder`] to an SVG flamegraph file.
+///
+/// Uses `inferno` for rendering with MiB units and flame-chart mode.
+/// Logs a warning and returns on I/O failure instead of panicking.
+pub fn write_flamegraph_svg<P: AsRef<Path>>(flamegraph: FlameGraphBuilder, path: P) {
+    let mut opts = Options::default();
+    opts.color_diffusion = true;
+    opts.count_name = String::from("MiB");
+    opts.factor = 1.0 / BYTES_PER_MIB;
+    opts.flame_chart = true;
+
+    let flamegraph_src = flamegraph.finish_and_write_flame_graph();
+    let input = Cursor::new(flamegraph_src);
+
+    let output = match File::create(path.as_ref()) {
+        Ok(f) => f,
+        Err(e) => {
+            tracing::warn!(
+                path = %path.as_ref().display(),
+                error = %e,
+                "failed to create flamegraph SVG file"
+            );
+            return;
+        }
+    };
+
+    if let Err(e) = inferno::flamegraph::from_reader(&mut opts, input, output) {
+        tracing::warn!(
+            path = %path.as_ref().display(),
+            error = %e,
+            "failed to render flamegraph SVG"
+        );
+    }
+}

crates/jolt-profiling/src/lib.rs

@@ -0,0 +1,97 @@
+//! Profiling and tracing infrastructure for the Jolt proving system.
+//!
+//! Provides a unified interface for performance analysis across all Jolt crates:
+//!
+//! - **Tracing subscriber setup** — configures `tracing-chrome` (Perfetto/Chrome JSON)
+//!   and `tracing-subscriber` (console output) for the host binary.
+//! - **Memory profiling** — tracks memory deltas across proving stages via `memory-stats`.
+//! - **System metrics monitoring** (`monitor` feature) — background thread sampling
+//!   CPU usage, memory, active cores, and thread count. Outputs structured counter events
+//!   compatible with the Perfetto postprocessing script.
+//! - **CPU profiling** (`pprof` feature) — scoped `pprof` guards that write `.pb`
+//!   flamegraph files on drop.
+//! - **Heap flamegraphs** (`allocative` feature) — generates SVG flamegraphs from
+//!   `allocative`-instrumented data structures.
+//!
+//! # Usage
+//!
+//! Individual crates add `tracing` as a dependency and instrument their functions with
+//! `#[tracing::instrument]`. The host binary (e.g. `jolt-zkvm` CLI) depends on
+//! `jolt-profiling` to configure the subscriber that captures those spans.
+//!
+//! ```no_run
+//! use jolt_profiling::{setup_tracing, TracingFormat};
+//!
+//! let _guards = setup_tracing(
+//!     &[TracingFormat::Chrome],
+//!     "my_benchmark_20260306",
+//! );
+//! // All tracing spans from any Jolt crate now flow to Perfetto JSON output.
+//! ```
+//!
+//! # Feature Flags
+//!
+//! | Flag | Description |
+//! |------|-------------|
+//! | `monitor` | Background system metrics sampling (CPU, memory, cores) |
+//! | `pprof` | Scoped CPU profiling via `pprof` with `.pb` output |
+//! | `allocative` | Heap flamegraph generation from `allocative`-instrumented types |
+//!
+//! # Dependency Position
+//!
+//! This is a leaf crate — imported by host binaries and benchmarks.
+//! Library crates depend only on `tracing` for instrumentation.
+
+pub mod setup;
+
+#[cfg(not(target_arch = "wasm32"))]
+pub mod memory;
+
+#[cfg(all(not(target_arch = "wasm32"), feature = "monitor"))]
+pub mod monitor;
+
+mod pprof_guard;
+
+#[cfg(feature = "allocative")]
+pub mod flamegraph;
+#[cfg(feature = "allocative")]
+pub use flamegraph::{print_data_structure_heap_usage, write_flamegraph_svg};
+
+mod units;
+
+pub use setup::{setup_tracing, TracingFormat, TracingGuards};
+pub use units::{format_memory_size, BYTES_PER_GIB, BYTES_PER_MIB};
+
+#[cfg(not(target_arch = "wasm32"))]
+pub use memory::{
+    end_memory_tracing_span, print_current_memory_usage, report_memory_usage,
+    start_memory_tracing_span,
+};
+
+#[cfg(target_arch = "wasm32")]
+pub fn start_memory_tracing_span(_label: &'static str) {}
+
+#[cfg(target_arch = "wasm32")]
+pub fn end_memory_tracing_span(_label: &'static str) {}
+
+#[cfg(target_arch = "wasm32")]
+pub fn report_memory_usage() {}
+
+#[cfg(target_arch = "wasm32")]
+pub fn print_current_memory_usage(_label: &str) {}
+
+#[cfg(all(not(target_arch = "wasm32"), feature = "monitor"))]
+pub use monitor::MetricsMonitor;
+
+#[cfg(all(target_arch = "wasm32", feature = "monitor"))]
+#[must_use = "monitor stops when dropped"]
+pub struct MetricsMonitor;
+
+#[cfg(all(target_arch = "wasm32", feature = "monitor"))]
+impl MetricsMonitor {
+    pub fn start(_interval_secs: f64) -> Self {
+        Self
+    }
+}
+
+pub use pprof_guard::PprofGuard;

crates/jolt-profiling/src/memory.rs

@@ -0,0 +1,122 @@
+//! Memory profiling utilities.
+//!
+//! Tracks physical memory deltas across labeled spans. Call
+//! [`start_memory_tracing_span`] before the section and
+//! [`end_memory_tracing_span`] after, then [`report_memory_usage`] to
+//! log all collected deltas.
+
+use memory_stats::memory_stats;
+use std::{
+    collections::BTreeMap,
+    sync::{LazyLock, Mutex},
+};
+
+use crate::units::{format_memory_size, BYTES_PER_GIB};
+
+static MEMORY_USAGE_MAP: LazyLock<Mutex<BTreeMap<&'static str, f64>>> =
+    LazyLock::new(|| Mutex::new(BTreeMap::new()));
+static MEMORY_DELTA_MAP: LazyLock<Mutex<BTreeMap<&'static str, f64>>> =
+    LazyLock::new(|| Mutex::new(BTreeMap::new()));
+
+/// Records the current physical memory usage at the start of a labeled span.
+///
+/// Logs a warning and returns without recording if memory stats are unavailable
+/// or if a span with the same label is already open.
+pub fn start_memory_tracing_span(label: &'static str) {
+    let Some(stats) = memory_stats() else {
+        tracing::warn!(
+            span = label,
+            "memory stats unavailable, skipping span start"
+        );
+        return;
+    };
+    let memory_gib = stats.physical_mem as f64 / BYTES_PER_GIB;
+    let mut map = MEMORY_USAGE_MAP.lock().unwrap_or_else(|e| e.into_inner());
+    if map.insert(label, memory_gib).is_some() {
+        tracing::warn!(span = label, "duplicate memory span label, overwriting");
+    }
+}
+
+/// Closes a labeled memory span and records the memory delta (in GiB).
+///
+/// Logs a warning and returns without recording if memory stats are unavailable
+/// or if no matching span was opened.
+pub fn end_memory_tracing_span(label: &'static str) {
+    let Some(stats) = memory_stats() else {
+        tracing::warn!(span = label, "memory stats unavailable, skipping span end");
+        return;
+    };
+    let memory_gib_end = stats.physical_mem as f64 / BYTES_PER_GIB;
+    let Some(memory_gib_start) = MEMORY_USAGE_MAP
+        .lock()
+        .unwrap_or_else(|e| e.into_inner())
+        .remove(label)
+    else {
+        tracing::warn!(span = label, "no open memory span, skipping span end");
+        return;
+    };
+
+    let delta = memory_gib_end - memory_gib_start;
+    let _ = MEMORY_DELTA_MAP
+        .lock()
+        .unwrap_or_else(|e| e.into_inner())
+        .insert(label, delta);
+}
+
+/// Logs all collected memory deltas and warns about any unclosed spans.
+pub fn report_memory_usage() {
+    let memory_usage_map = MEMORY_USAGE_MAP.lock().unwrap_or_else(|e| e.into_inner());
+    for label in memory_usage_map.keys() {
+        tracing::warn!(span = label, "unclosed memory tracing span");
+    }
+
+    let memory_delta_map = MEMORY_DELTA_MAP.lock().unwrap_or_else(|e| e.into_inner());
+    for (label, delta) in memory_delta_map.iter() {
+        tracing::info!(
+            span = label,
+            delta = %format_memory_size(*delta),
+            "memory delta"
+        );
+    }
+}
+
+/// Logs the current physical memory usage at the point of call.
+pub fn print_current_memory_usage(label: &str) {
+    if tracing::enabled!(tracing::Level::DEBUG) {
+        if let Some(usage) = memory_stats() {
+            let memory_gib = usage.physical_mem as f64 / BYTES_PER_GIB;
+            tracing::debug!(
+                label = label,
+                usage = %format_memory_size(memory_gib),
+                "current memory usage"
+            );
+        } else {
+            tracing::debug!(label = label, "memory stats unavailable");
+        }
+    }
+}
+
+#[cfg(test)]
+#[expect(clippy::unwrap_used)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn memory_span_start_end_records_delta() {
+        start_memory_tracing_span("test_span_lifecycle");
+        end_memory_tracing_span("test_span_lifecycle");
+        let map = MEMORY_DELTA_MAP.lock().unwrap();
+        assert!(map.contains_key("test_span_lifecycle"));
+    }
+
+    #[test]
+    fn duplicate_span_warns_without_panic() {
+        start_memory_tracing_span("test_span_dup");
+        start_memory_tracing_span("test_span_dup");
+    }
+
+    #[test]
+    fn end_without_start_warns_without_panic() {
+        end_memory_tracing_span("test_span_nonexistent");
+    }
+}