diff --git a/Cargo.lock b/Cargo.lock index 1f085845..7a5d984d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1748,6 +1748,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", + "tracing-texray", ] [[package]] @@ -1815,6 +1816,12 @@ dependencies = [ "cc", ] +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.2" @@ -1937,7 +1944,7 @@ dependencies = [ [[package]] name = "multi-stark" version = "0.1.0" -source = "git+https://github.com/argumentcomputer/multi-stark.git?rev=a8a15ea6aa2890f9f60f32a6e0e5e66afc1535ff#a8a15ea6aa2890f9f60f32a6e0e5e66afc1535ff" +source = "git+https://github.com/argumentcomputer/multi-stark.git?branch=sb%2Fstreaming-print#dbe77db42f71ac27ace80d54e69d7e33c9636d0e" dependencies = [ "bincode", "p3-air", @@ -1954,6 +1961,7 @@ dependencies = [ "p3-symmetric", "p3-util", "serde", + "tracing", ] [[package]] @@ -3144,6 +3152,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.52.0", +] + [[package]] name = "rustls" version = "0.23.40" @@ -3520,6 +3541,16 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" +[[package]] +name = "terminal_size" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" +dependencies = [ + "rustix", + "windows-sys 0.59.0", +] + [[package]] name = "thiserror" version = "2.0.18" @@ -3839,6 +3870,18 @@ dependencies = [ "tracing-log", ] +[[package]] +name = "tracing-texray" +version = "0.2.0" +source = "git+https://github.com/argumentcomputer/tracing-texray?branch=sb%2Fstreaming-print#8ce04e3422cd48e68ef47fab95dba7d06b8c368c" +dependencies = [ + "loom", + "parking_lot", + "terminal_size", + "tracing", + "tracing-subscriber", +] + [[package]] name = "transpose" version = "0.2.3" diff --git a/Cargo.toml b/Cargo.toml index ba93c5e8..88b2c907 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ itertools = "0.14.0" indexmap = { version = "2", features = ["rayon"] } lean-ffi = { git = "https://github.com/argumentcomputer/lean-ffi", rev = "cc98ebf67bf453ac3827cb767f78b13ea674dd6a" } mimalloc = { version = "0.1", default-features = false } -multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "a8a15ea6aa2890f9f60f32a6e0e5e66afc1535ff" } +multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", branch = "sb/streaming-print" } num-bigint = "0.4.6" rayon = "1" rustc-hash = "2" @@ -27,8 +27,9 @@ iroh = { version = "0.97", optional = true } iroh-base = { version = "0.97", optional = true } n0-error = { version = "0.1", optional = true } getrandom = { version = "0.3", optional = true } -tracing = { version = "0.1", optional = true } -tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +tracing-texray = { git = "https://github.com/argumentcomputer/tracing-texray", branch = "sb/streaming-print" } bincode = { version = "2.0.1", optional = true } serde = { version = "1.0.219", features = ["derive"], optional = true } @@ -41,7 +42,7 @@ quickcheck_macros = "1.0.0" default = [] parallel = ["multi-stark/parallel"] test-ffi = [] -net = ["bytes", "tokio", "iroh", "iroh-base", "n0-error", "getrandom", "tracing", "tracing-subscriber", "bincode", "serde" ] +net = ["bytes", "tokio", "iroh", "iroh-base", "n0-error", "getrandom", "bincode", "serde" ] [profile.dev] panic = "abort" diff --git a/Ix.lean b/Ix.lean index b1809a25..693e49ac 100644 --- a/Ix.lean +++ b/Ix.lean @@ -15,3 +15,4 @@ public import Ix.Claim public import Ix.Commit public import Ix.Benchmark.Bench public import Ix.Aiur +public import Ix.TracingTexray diff --git a/Ix/TracingTexray.lean b/Ix/TracingTexray.lean new file mode 100644 index 00000000..e63ebd3f --- /dev/null +++ b/Ix/TracingTexray.lean @@ -0,0 +1,44 @@ +/- + tracing-texray + streaming Lean API. + + The Rust prover emits `tracing` spans for each major Aiur proving stage + (`aiur/execute`, `aiur/witness`, `aiur/prove`) and each STARK proving + stage (`stark/...`). After `init` installs the subscriber: + * a one-line `[texray] : ── RAM Δ +X peak Y` is streamed + to stderr as each tracked span closes (when `streaming = true`); + * the full texray graph (with RAM block, Linux-only) prints when + `aiur/prove` exits. +-/ + +module + +public section + +namespace TracingTexray + +/-- Configuration for the tracing-texray subscriber. -/ +structure Settings where + /-- Comma-separated span-name prefixes to keep + (default `"aiur/,stark/"`; empty string disables filtering). -/ + namePrefixes : String := "aiur/,stark/" + /-- Sample VmRSS/VmHWM per span. Linux-only; zeros elsewhere. -/ + trackRam : Bool := true + /-- Emit per-span `[texray] : ── RAM Δ +X peak Y` lines + on stderr as each span closes. The texray graph prints either way + when the examined root span exits. -/ + streaming : Bool := true + +@[extern "rs_texray_init"] +private opaque initWith + (namePrefixes : @& String) + (trackRam : Bool) + (streaming : Bool) : IO Unit + +/-- Install the global tracing-texray subscriber with the given settings. + Idempotent: subsequent calls are no-ops once installed. -/ +def init (s : Settings := {}) : IO Unit := + initWith s.namePrefixes s.trackRam s.streaming + +end TracingTexray + +end diff --git a/src/aiur/synthesis.rs b/src/aiur/synthesis.rs index 6699c811..605ca954 100644 --- a/src/aiur/synthesis.rs +++ b/src/aiur/synthesis.rs @@ -109,6 +109,7 @@ impl AiurSystem { AiurSystem { system, key, toplevel } } + #[tracing::instrument(level = "info", skip_all, name = "aiur/prove")] pub fn prove( &self, fri_parameters: FriParameters, @@ -116,11 +117,17 @@ impl AiurSystem { input: &[G], io_buffer: &mut IOBuffer, ) -> (Vec, Proof) { + // Register the instrumented span for tracing-texray's tree dump on exit. + tracing_texray::examine_current(); + // Execute the Aiur bytecode. + let _g = tracing::info_span!("aiur/execute").entered(); let (query_record, output) = self.toplevel.execute(fun_idx, input.to_vec(), io_buffer); + drop(_g); // Build the `SystemWitness` + let _g = tracing::info_span!("aiur/witness").entered(); let functions = (0..self.toplevel.functions.len()).into_par_iter().filter_map(|idx| { if self.toplevel.functions[idx].constrained { @@ -152,6 +159,7 @@ impl AiurSystem { drop(query_record); // Early drop to free memory. let (traces, lookups) = witness_data.into_iter().unzip(); let witness = SystemWitness { traces, lookups }; + drop(_g); // Construct the claim. let mut claim = vec![function_channel(), G::from_usize(fun_idx)]; diff --git a/src/ffi.rs b/src/ffi.rs index 1cb987fb..646c40b2 100644 --- a/src/ffi.rs +++ b/src/ffi.rs @@ -14,6 +14,7 @@ pub mod byte_array; pub mod iroh; pub mod keccak; pub mod lean_env; +pub mod texray; pub mod unsigned; // Modular FFI structure diff --git a/src/ffi/texray.rs b/src/ffi/texray.rs new file mode 100644 index 00000000..85486129 --- /dev/null +++ b/src/ffi/texray.rs @@ -0,0 +1,58 @@ +//! FFI shim for installing the `tracing-texray` subscriber from Lean. +//! +//! With `streaming = true`, each tracked span (e.g. `aiur/execute`, +//! `aiur/witness`, `aiur/prove`) emits a one-line `[texray] : +//! ── RAM Δ +X peak Y` to stderr as it closes, and the full texray +//! graph prints when the examined root span exits. + +use lean_ffi::object::{LeanBorrowed, LeanIOResult, LeanOwned, LeanString}; +use tracing_subscriber::{ + Layer, Registry, filter::FilterFn, layer::SubscriberExt, + util::SubscriberInitExt, +}; + +/// Install the `tracing-texray` subscriber as the global default. +/// +/// Idempotent: subsequent calls are no-ops once the global subscriber has +/// been set. +/// +/// Parameters: +/// - `name_prefixes`: comma-separated list of span-name prefixes to render +/// (e.g. `"aiur/,stark/"`). Empty string disables filtering and renders +/// everything, including upstream library spans. +/// - `track_ram`: sample VmRSS/VmHWM per span. Linux-only; zeros elsewhere. +/// - `streaming`: emit per-span `[texray] : ── RAM Δ +X peak Y` +/// lines on stderr as each span closes. The texray graph prints either +/// way when the examined root span exits. +#[unsafe(no_mangle)] +extern "C" fn rs_texray_init( + name_prefixes: LeanString>, + track_ram: bool, + streaming: bool, +) -> LeanIOResult { + let prefixes: Vec = name_prefixes + .to_string() + .split(',') + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect(); + + let mut layer = tracing_texray::TeXRayLayer::new(); + if track_ram { + layer = layer.track_ram(); + } + if streaming { + layer = layer.streaming(); + } + + let filter = FilterFn::new(move |metadata| { + if !metadata.is_span() || prefixes.is_empty() { + return true; + } + let name = metadata.name(); + prefixes.iter().any(|p| name.starts_with(p.as_str())) + }); + + let _ = Registry::default().with(layer.with_filter(filter)).try_init(); + LeanIOResult::ok(LeanOwned::box_usize(0)) +}