From 84a087229d5c8b4f848a812e9f8ecc32ada6a75d Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Thu, 23 Jul 2026 15:52:00 -0300 Subject: [PATCH 01/12] add profiling --- Cargo.lock | 13 +- bin/cli/Cargo.toml | 2 + crypto/math-cuda/Cargo.toml | 6 + crypto/math-cuda/build.rs | 12 +- crypto/math-cuda/src/barycentric.rs | 8 + crypto/math-cuda/src/constraint_interp.rs | 4 + crypto/math-cuda/src/deep.rs | 5 + crypto/math-cuda/src/fri.rs | 2 + crypto/math-cuda/src/inverse.rs | 3 + crypto/math-cuda/src/lde.rs | 53 ++ crypto/math-cuda/src/lib.rs | 1 + crypto/math-cuda/src/logup.rs | 3 + crypto/math-cuda/src/merkle.rs | 12 + crypto/math-cuda/src/nvtx.rs | 276 +++++++++++ crypto/stark/Cargo.toml | 4 + crypto/stark/src/instruments.rs | 45 ++ prover/Cargo.toml | 2 + prover/src/continuation.rs | 41 ++ scripts/profiling/README.md | 145 ++++++ .../nsys_phase_busy.cpython-313.pyc | Bin 0 -> 18907 bytes .../__pycache__/nvml_sampler.cpython-313.pyc | Bin 0 -> 3539 bytes .../__pycache__/phase_table.cpython-313.pyc | Bin 0 -> 10820 bytes .../timeline_to_perfetto.cpython-313.pyc | Bin 0 -> 1947 bytes scripts/profiling/bench_mode.sh | 48 ++ scripts/profiling/capture_env.sh | 53 ++ scripts/profiling/flamegraphs.sh | 119 +++++ scripts/profiling/nsys_phase_busy.py | 386 +++++++++++++++ scripts/profiling/nvml_sampler.py | 71 +++ scripts/profiling/phase_table.py | 259 ++++++++++ scripts/profiling/run_profile.sh | 113 +++++ scripts/profiling/setup_machine.sh | 118 +++++ scripts/profiling/timeline_to_perfetto.py | 47 ++ thoughts/gpu-profiling/plan.md | 467 ++++++++++++++++++ 33 files changed, 2315 insertions(+), 3 deletions(-) create mode 100644 crypto/math-cuda/src/nvtx.rs create mode 100644 scripts/profiling/README.md create mode 100644 scripts/profiling/__pycache__/nsys_phase_busy.cpython-313.pyc create mode 100644 scripts/profiling/__pycache__/nvml_sampler.cpython-313.pyc create mode 100644 scripts/profiling/__pycache__/phase_table.cpython-313.pyc create mode 100644 scripts/profiling/__pycache__/timeline_to_perfetto.cpython-313.pyc create mode 100755 scripts/profiling/bench_mode.sh create mode 100755 scripts/profiling/capture_env.sh create mode 100755 scripts/profiling/flamegraphs.sh create mode 100755 scripts/profiling/nsys_phase_busy.py create mode 100755 scripts/profiling/nvml_sampler.py create mode 100755 scripts/profiling/phase_table.py create mode 100755 scripts/profiling/run_profile.sh create mode 100755 scripts/profiling/setup_machine.sh create mode 100755 scripts/profiling/timeline_to_perfetto.py create mode 100644 thoughts/gpu-profiling/plan.md diff --git a/Cargo.lock b/Cargo.lock index 74986dcc9..ac88f2ac5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -473,7 +473,7 @@ version = "0.19.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f071cd6a7b5d51607df76aa2d426aaabc7a74bc6bdb885b8afa63a880572ad9b" dependencies = [ - "libloading", + "libloading 0.9.0", ] [[package]] @@ -867,6 +867,16 @@ version = "0.2.178" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "libloading" version = "0.9.0" @@ -927,6 +937,7 @@ version = "0.1.0" dependencies = [ "crypto", "cudarc", + "libloading 0.8.9", "math", "rand 0.8.5", "rand_chacha 0.3.1", diff --git a/bin/cli/Cargo.toml b/bin/cli/Cargo.toml index e4fcdb7fd..b744714c3 100644 --- a/bin/cli/Cargo.toml +++ b/bin/cli/Cargo.toml @@ -19,3 +19,5 @@ env_logger = "0.11" jemalloc-stats = ["dep:tikv-jemalloc-ctl"] disk-spill = ["prover/disk-spill"] instruments = ["prover/instruments", "stark/instruments"] +# GPU profiling build (Nsight): CUDA prover + instruments spans + NVTX ranges. +nvtx = ["prover/nvtx", "instruments"] diff --git a/crypto/math-cuda/Cargo.toml b/crypto/math-cuda/Cargo.toml index 7a28498da..5a214f7bd 100644 --- a/crypto/math-cuda/Cargo.toml +++ b/crypto/math-cuda/Cargo.toml @@ -30,11 +30,17 @@ cudarc = { version = "0.19", default-features = false, features = [ ] } math = { path = "../math" } rayon = "1.7" +# NVTX range emission for Nsight timelines (dlopen'd at runtime, see src/nvtx.rs). +libloading = { version = "0.8", optional = true } [features] # Test-only fault injection in FriCommitState. Production builds leave this # off so the fault check is fully elided at compile time. test-faults = [] +# NVTX ranges around every GPU pipeline entry point, for Nsight Systems +# profiling. Zero-cost when disabled; when enabled but libnvToolsExt is +# absent at runtime, every call is a cheap no-op. +nvtx = ["dep:libloading"] [dev-dependencies] crypto = { path = "../crypto" } diff --git a/crypto/math-cuda/build.rs b/crypto/math-cuda/build.rs index 7bd7c04cc..fbd70eb5b 100644 --- a/crypto/math-cuda/build.rs +++ b/crypto/math-cuda/build.rs @@ -82,6 +82,7 @@ fn compile_kernel(src: &str, out_name: &str, have_nvcc: bool) { println!("cargo:rerun-if-env-changed=CUDA_HOME"); println!("cargo:rerun-if-env-changed=CUDA_PATH"); println!("cargo:rerun-if-env-changed=CUDARC_NVCC_ARCH"); + println!("cargo:rerun-if-env-changed=LAMBDA_VM_NVCC_LINEINFO"); // When nvcc is missing from PATH, emit an empty cubin stub so the crate // still compiles. include_bytes! in src/device.rs needs the file to exist @@ -115,8 +116,15 @@ fn compile_kernel(src: &str, out_name: &str, have_nvcc: bool) { .map(|a| to_real_arch(&a)) .unwrap_or_else(|_| detect_arch()); - let status = Command::new(nvcc_path()) - .args(["--cubin", "-O3", "-std=c++17", "-arch", &arch, "-o"]) + let mut cmd = Command::new(nvcc_path()); + cmd.args(["--cubin", "-O3", "-std=c++17", "-arch", &arch]); + // SASS→source line mapping for Nsight Compute. Unlike -G this does not + // change codegen, but keep it opt-in so production cubins stay byte-stable. + if env::var("LAMBDA_VM_NVCC_LINEINFO").is_ok_and(|v| v != "0" && !v.is_empty()) { + cmd.arg("-lineinfo"); + } + let status = cmd + .arg("-o") .arg(&out_path) .arg(&src_path) .status() diff --git a/crypto/math-cuda/src/barycentric.rs b/crypto/math-cuda/src/barycentric.rs index 8fdea14f1..5c0a5d379 100644 --- a/crypto/math-cuda/src/barycentric.rs +++ b/crypto/math-cuda/src/barycentric.rs @@ -29,6 +29,7 @@ pub fn barycentric_base( n: usize, num_cols: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("bary_base[n={n} cols={num_cols}]")); assert_eq!(coset_points.len(), n); assert_eq!(inv_denoms_ext3.len(), 3 * n); assert!(columns.len() >= num_cols * col_stride); @@ -84,6 +85,7 @@ pub fn barycentric_ext3( n: usize, num_cols: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("bary_ext3[n={n} cols={num_cols}]")); assert_eq!(coset_points.len(), n); assert_eq!(inv_denoms_ext3.len(), 3 * n); assert!(columns.len() >= num_cols * 3 * col_stride); @@ -139,6 +141,7 @@ pub fn barycentric_base_on_device( inv_denoms_ext3: &[u64], n: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("bary_base_dev[n={n} cols={}]", main_handle.m)); assert_eq!(coset_points.len(), n); assert_eq!(inv_denoms_ext3.len(), 3 * n); let num_cols = main_handle.m; @@ -197,6 +200,7 @@ pub fn barycentric_base_on_device_with_dev_inv_denoms( inv_offset_u64: usize, n: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("bary_base_dev2[n={n} cols={}]", main_handle.m)); assert!(coset_points_dev.len() >= n); let inv_end = inv_offset_u64 .checked_add(3 * n) @@ -247,6 +251,7 @@ pub fn barycentric_ext3_on_device( inv_denoms_ext3: &[u64], n: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("bary_ext3_dev[n={n} cols={}]", aux_handle.m)); assert_eq!(coset_points.len(), n); assert_eq!(inv_denoms_ext3.len(), 3 * n); let num_cols = aux_handle.m; @@ -297,6 +302,7 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms( inv_offset_u64: usize, n: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("bary_ext3_dev2[n={n} cols={}]", aux_handle.m)); assert!(coset_points_dev.len() >= n); let inv_end = inv_offset_u64 .checked_add(3 * n) @@ -347,6 +353,7 @@ pub fn gather_rows_base_on_device( rows: &[u32], stream: &Arc, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("gather_rows_base[q={}]", rows.len())); let num_cols = main.m; if num_cols == 0 || rows.is_empty() { return Ok(Vec::new()); @@ -385,6 +392,7 @@ pub fn gather_rows_ext3_on_device( rows: &[u32], stream: &Arc, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("gather_rows_ext3[q={}]", rows.len())); let num_cols = aux.m; if num_cols == 0 || rows.is_empty() { return Ok(Vec::new()); diff --git a/crypto/math-cuda/src/constraint_interp.rs b/crypto/math-cuda/src/constraint_interp.rs index ba11b4f64..fd2ea2477 100644 --- a/crypto/math-cuda/src/constraint_interp.rs +++ b/crypto/math-cuda/src/constraint_interp.rs @@ -53,6 +53,8 @@ pub fn eval_constraints_on_device( next_step: usize, num_rows: usize, ) -> Result> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("constraint_eval[rows={num_rows} nodes={num_nodes}]")); let num_roots = roots.len(); if num_rows == 0 || num_roots == 0 || num_nodes == 0 { return Ok(vec![0u64; num_roots * num_rows * 3]); @@ -169,6 +171,8 @@ pub fn eval_composition_on_device( num_rows: usize, accum: &CompositionAccum, ) -> Result> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("constraint_comp[rows={num_rows} nodes={num_nodes}]")); let num_roots = roots.len(); if num_rows == 0 { return Ok(Vec::new()); diff --git a/crypto/math-cuda/src/deep.rs b/crypto/math-cuda/src/deep.rs index 581fbc404..7b68917a6 100644 --- a/crypto/math-cuda/src/deep.rs +++ b/crypto/math-cuda/src/deep.rs @@ -41,6 +41,7 @@ pub fn deep_composition_ext3( row_stride: usize, domain_size: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("deep_comp[n={domain_size} parts={num_parts}]")); let be = backend()?; let stream = be.next_stream(); deep_composition_ext3_impl( @@ -86,6 +87,8 @@ pub fn deep_composition_ext3_with_dev_parts( row_stride: usize, domain_size: usize, ) -> Result> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("deep_comp_dev[n={domain_size} parts={num_parts}]")); let be = backend()?; let stream = be.next_stream(); deep_composition_ext3_impl( @@ -138,6 +141,8 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms( row_stride: usize, domain_size: usize, ) -> Result> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("deep_comp_dev2[n={domain_size} parts={num_parts}]")); assert_eq!(main_lde.m, num_main); assert_eq!(h_parts_dev.m, num_parts); assert_eq!(h_parts_dev.lde_size, main_lde.lde_size); diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index fb854d0a4..f2b1869f9 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -64,6 +64,7 @@ impl FriCommitState { /// initial inv_twiddles (base field, n0/2 u64). `n0` must be a power of /// two and >= 2. pub fn new(evals_host: &[u64], inv_tw_host: &[u64], n0: usize) -> Result { + let _nvtx = crate::nvtx::Range::fmt(|| format!("fri_init[n0={n0}]")); assert!(n0 >= 2 && n0.is_power_of_two()); assert_eq!(evals_host.len(), 3 * n0); assert_eq!(inv_tw_host.len(), n0 / 2); @@ -99,6 +100,7 @@ impl FriCommitState { &mut self, zeta_raw: [u64; 3], ) -> Result<(Vec, crate::lde::GpuMerkleTree)> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("fri_fold[n={}]", self.current_n)); #[cfg(feature = "test-faults")] check_fault_injection()?; let be = backend()?; diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 485e005f8..a4a479395 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -51,6 +51,7 @@ fn check_inverse_fault_injection() -> Result<()> { /// containing the inverses. Used by the parity-test suite; production /// callers should prefer `batch_inverse_ext3_dev` to avoid the D2H. pub fn batch_inverse_ext3(a: &[u64]) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("batch_inverse[n={}]", a.len() / 3)); assert!(a.len().is_multiple_of(3)); let n = a.len() / 3; if n == 0 { @@ -83,6 +84,7 @@ pub fn batch_inverse_ext3_dev( n: usize, stream: &Arc, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("batch_inverse_dev[n={n}]")); assert!(n >= 1, "batch_inverse_ext3_dev requires n >= 1"); // Runtime guard (not debug_assert): a u32 grid_dim is truncated past // u32::MAX / BLOCK_SIZE, which would silently launch too few blocks @@ -173,6 +175,7 @@ pub fn compute_and_invert_denoms_ext3_dev( sign: DenomSign, stream: &Arc, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("denoms_invert[n={n} k={k_scalars}]")); #[cfg(feature = "test-faults")] check_inverse_fault_injection()?; assert_eq!(z_scalars_host.len(), k_scalars * 3); diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 06c1a02cd..ba59b23ae 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -571,6 +571,8 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( weights: &[u64], retain_host_lde: bool, ) -> Result<(GpuLdeBase, Vec)> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("lde_tree_base[n={n} m={m} bf={blowup_factor}]")); let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner( InnerInput::Host(row_major), n, @@ -610,6 +612,8 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( weights: &[u64], retain_host_lde: bool, ) -> Result<(GpuLdeExt3, Vec)> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("lde_tree_ext3[n={n} m={m} bf={blowup_factor}]")); let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( InnerInput::Host(row_major), n, @@ -641,6 +645,8 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( weights: &[u64], retain_host_lde: bool, ) -> Result<(GpuLdeExt3, Vec)> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("lde_tree_ext3_dev[n={n} m={m} bf={blowup_factor}]")); let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( InnerInput::Dev(input_dev), n, @@ -709,6 +715,8 @@ pub struct GpuMerkleTree { } pub fn coset_lde_base(evals: &[u64], blowup_factor: usize, weights: &[u64]) -> Result> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("lde_base[n={} bf={blowup_factor}]", evals.len())); let n = evals.len(); // Empty input must short-circuit before the power-of-two assert // (is_power_of_two returns false for 0). @@ -795,6 +803,13 @@ pub fn coset_lde_batch_base( blowup_factor: usize, weights: &[u64], ) -> Result>> { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!( + "lde_batch_base[cols={} n={} bf={blowup_factor}]", + columns.len(), + columns.first().map_or(0, |c| c.len()) + ) + }); if columns.is_empty() { return Ok(Vec::new()); } @@ -950,6 +965,13 @@ pub fn coset_lde_batch_base_into( weights: &[u64], outputs: &mut [&mut [u64]], ) -> Result<()> { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!( + "lde_batch_base_into[cols={} n={} bf={blowup_factor}]", + columns.len(), + columns.first().map_or(0, |c| c.len()) + ) + }); if columns.is_empty() { return Ok(()); } @@ -1077,6 +1099,13 @@ pub fn coset_lde_batch_base_into_with_leaf_hash( outputs: &mut [&mut [u64]], hashed_leaves_out: &mut [u8], ) -> Result<()> { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!( + "lde_batch_base_hash[cols={} n={} bf={blowup_factor}]", + columns.len(), + columns.first().map_or(0, |c| c.len()) + ) + }); coset_lde_batch_base_into_with_merkle_tree_inner( columns, blowup_factor, @@ -1291,6 +1320,12 @@ pub fn evaluate_poly_coset_batch_ext3_into( weights: &[u64], outputs: &mut [&mut [u64]], ) -> Result<()> { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!( + "poly_lde_ext3[cols={} n={n} bf={blowup_factor}]", + coefs.len() + ) + }); evaluate_poly_coset_batch_ext3_into_inner( coefs, n, @@ -1313,6 +1348,12 @@ pub fn evaluate_poly_coset_batch_ext3_into_keep( weights: &[u64], outputs: &mut [&mut [u64]], ) -> Result { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!( + "poly_lde_ext3_keep[cols={} n={n} bf={blowup_factor}]", + coefs.len() + ) + }); let opt = evaluate_poly_coset_batch_ext3_into_inner( coefs, n, @@ -1481,6 +1522,12 @@ pub fn evaluate_poly_coset_batch_ext3_into_with_merkle_tree( outputs: &mut [&mut [u64]], merkle_nodes_out: &mut [u8], ) -> Result<()> { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!( + "poly_lde_ext3_tree[cols={} n={n} bf={blowup_factor}]", + coefs.len() + ) + }); evaluate_poly_coset_batch_ext3_into_inner( coefs, n, @@ -1519,6 +1566,12 @@ pub fn coset_lde_batch_ext3_into( weights: &[u64], outputs: &mut [&mut [u64]], ) -> Result<()> { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!( + "lde_batch_ext3[cols={} n={n} bf={blowup_factor}]", + columns.len() + ) + }); if columns.is_empty() { return Ok(()); } diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index 493480991..6b58d935b 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -15,6 +15,7 @@ pub mod lde; pub mod logup; pub mod merkle; pub mod ntt; +pub mod nvtx; // Re-exported for downstream crates so they can refer to CUDA primitive // types without depending on cudarc directly. diff --git a/crypto/math-cuda/src/logup.rs b/crypto/math-cuda/src/logup.rs index e9d120ee1..397ab3c3d 100644 --- a/crypto/math-cuda/src/logup.rs +++ b/crypto/math-cuda/src/logup.rs @@ -121,6 +121,7 @@ pub fn logup_fingerprints_dev( z: [u64; 3], stream: &Arc, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("logup_fingerprints[rows={num_rows}]")); let main_dev = stream.clone_htod(main_cols)?; let out = fingerprints_into_dev(&main_dev, num_rows, d, alpha_powers, z, stream)?; stream.synchronize()?; @@ -137,6 +138,7 @@ pub fn logup_term_columns( alpha_powers: &[u64], z: [u64; 3], ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("logup_terms[rows={num_rows}]")); let be = backend()?; let stream = be.next_stream(); let timing = std::env::var_os("LAMBDA_VM_LOGUP_TIMING").is_some(); @@ -301,6 +303,7 @@ pub fn logup_aux_resident( inv_n: [u64; 3], stream: &Arc, ) -> Result { + let _nvtx = crate::nvtx::Range::fmt(|| format!("logup_aux[rows={num_rows}]")); assert!(num_rows >= 1, "logup_aux_resident requires num_rows >= 1"); let be = backend()?; // Per-phase timing (env LAMBDA_VM_LOGUP_TIMING): sync between phases so wall diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index fb1125ea4..aeba53754 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -37,6 +37,9 @@ pub fn keccak_leaves_base( num_rows: usize, rows_per_leaf: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!("keccak_leaves_base[rows={num_rows} cols={num_cols} rpl={rows_per_leaf}]") + }); assert!(num_rows.is_power_of_two()); assert!(rows_per_leaf == 1 || rows_per_leaf == 2); assert!( @@ -87,6 +90,9 @@ pub fn keccak_leaves_ext3( num_rows: usize, rows_per_leaf: usize, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| { + format!("keccak_leaves_ext3[rows={num_rows} cols={num_cols} rpl={rows_per_leaf}]") + }); assert!(num_rows.is_power_of_two()); assert!(rows_per_leaf == 1 || rows_per_leaf == 2); assert!( @@ -283,6 +289,8 @@ pub(crate) fn launch_keccak_ext3_row_pair( /// /// `leaves_len` must be a power of two and >= 2. pub fn build_merkle_tree_on_device(hashed_leaves: &[u8]) -> Result> { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("merkle_tree[leaves={}]", hashed_leaves.len() / 32)); assert!(hashed_leaves.len().is_multiple_of(32)); let leaves_len = hashed_leaves.len() / 32; assert!(leaves_len >= 2, "tree needs at least two leaves"); @@ -330,6 +338,7 @@ pub fn gather_merkle_paths_dev( positions: &[u32], stream: &Arc, ) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("merkle_paths[q={}]", positions.len())); let num_queries = positions.len(); if num_queries == 0 { return Ok(Vec::new()); @@ -457,6 +466,8 @@ fn build_comp_poly_tree_nodes_dev( pub fn build_comp_poly_tree_from_evals_ext3_keep( parts_interleaved: &[&[u64]], ) -> Result { + let _nvtx = + crate::nvtx::Range::fmt(|| format!("comp_poly_tree[parts={}]", parts_interleaved.len())); let (nodes_dev, num_leaves, stream) = build_comp_poly_tree_nodes_dev(parts_interleaved)?; let mut root = [0u8; 32]; stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; @@ -475,6 +486,7 @@ pub fn build_comp_poly_tree_from_evals_ext3_keep( /// consecutive ext3 values; `num_leaves = evals.len() / 6`. Returns the /// `(2*num_leaves - 1) * 32`-byte node buffer in standard layout. pub fn build_fri_layer_tree_from_evals_ext3(evals: &[u64]) -> Result> { + let _nvtx = crate::nvtx::Range::fmt(|| format!("fri_layer_tree[n={}]", evals.len() / 3)); assert!( evals.len().is_multiple_of(6), "evals must hold whole pair-leaves" diff --git a/crypto/math-cuda/src/nvtx.rs b/crypto/math-cuda/src/nvtx.rs new file mode 100644 index 000000000..fa81f555c --- /dev/null +++ b/crypto/math-cuda/src/nvtx.rs @@ -0,0 +1,276 @@ +//! Minimal NVTX bindings so Nsight Systems timelines show named host-side +//! ranges (prover phases, kernel-pipeline entry points) instead of a wall of +//! anonymous CUDA API calls. +//! +//! Loading mirrors the crate's cudarc `dynamic-loading` philosophy: no +//! build-time or link-time dependency on the CUDA toolkit layout. At first use +//! we dlopen `libnvToolsExt.so` (NVTX v2, shipped with every CUDA toolkit and +//! honored by nsys/ncu); when it is absent every call is a cheap no-op, so a +//! `--features nvtx` binary runs unchanged on machines without the library. +//! Override the library path with `LAMBDA_VM_NVTX_LIB` if it lives somewhere +//! unusual. +//! +//! With the `nvtx` cargo feature *disabled* this module compiles to empty +//! inline stubs — the label closures passed to [`Range::fmt`] are never +//! evaluated and the whole thing vanishes. +//! +//! Semantics note: a [`Range`] measures the *host-side* extent of a call. For +//! the `_keep`/`_dev` entry points that enqueue async GPU work without +//! syncing, kernels execute after the range closes — that is fine: nsys +//! correlates each kernel to the range that launched it. + +pub use imp::*; + +#[cfg(feature = "nvtx")] +mod imp { + use std::ffi::{CString, c_char, c_int}; + use std::marker::PhantomData; + use std::path::PathBuf; + use std::sync::OnceLock; + + struct Api { + // Field order is drop order; the fn pointers are only valid while the + // library is loaded, and both live for the whole process anyway + // (static OnceLock). + range_push: unsafe extern "C" fn(*const c_char) -> c_int, + range_pop: unsafe extern "C" fn() -> c_int, + mark: unsafe extern "C" fn(*const c_char), + _lib: libloading::Library, + } + // SAFETY: the NVTX v2 API is thread-safe (push/pop stacks are per-thread) + // and the Library handle is only kept alive, never re-entered. + unsafe impl Send for Api {} + unsafe impl Sync for Api {} + + fn nvtx_lib_candidates() -> Vec { + let mut c = Vec::new(); + if let Some(p) = std::env::var_os("LAMBDA_VM_NVTX_LIB") { + c.push(PathBuf::from(p)); + } + c.push(PathBuf::from("libnvToolsExt.so.1")); + c.push(PathBuf::from("libnvToolsExt.so")); + let cuda_home = std::env::var_os("CUDA_HOME") + .or_else(|| std::env::var_os("CUDA_PATH")) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/usr/local/cuda")); + c.push(cuda_home.join("lib64").join("libnvToolsExt.so.1")); + c.push(cuda_home.join("lib64").join("libnvToolsExt.so")); + c + } + + fn api() -> Option<&'static Api> { + static API: OnceLock> = OnceLock::new(); + API.get_or_init(|| { + for path in nvtx_lib_candidates() { + // SAFETY: loading a shared library runs its initializers; + // libnvToolsExt is NVIDIA's stub dispatcher with no side + // effects beyond tool injection. + let Ok(lib) = (unsafe { libloading::Library::new(&path) }) else { + continue; + }; + // SAFETY: signatures match the NVTX v2 C API. + let syms = unsafe { + ( + lib.get:: c_int>( + b"nvtxRangePushA\0", + ) + .map(|s| *s), + lib.get:: c_int>(b"nvtxRangePop\0") + .map(|s| *s), + lib.get::(b"nvtxMarkA\0") + .map(|s| *s), + ) + }; + if let (Ok(range_push), Ok(range_pop), Ok(mark)) = syms { + return Some(Api { + range_push, + range_pop, + mark, + _lib: lib, + }); + } + } + None + }) + .as_ref() + } + + /// True when libnvToolsExt was found; use to skip label formatting work. + #[inline] + pub fn is_active() -> bool { + api().is_some() + } + + fn push_str(api: &Api, name: &str) { + // NVTX takes a NUL-terminated C string; a label containing NUL is a + // bug we don't care to surface here — fall back to a fixed name. + let c = CString::new(name).unwrap_or_else(|_| CString::new("invalid-label").unwrap()); + // SAFETY: `c` is a valid NUL-terminated string for the duration of the call. + unsafe { (api.range_push)(c.as_ptr()) }; + } + + /// Push a range on this thread's NVTX stack. Prefer [`Range`]; this raw + /// form exists for RAII guards that live in other crates (instruments). + #[inline] + pub fn range_push(name: &str) { + if let Some(api) = api() { + push_str(api, name); + } + } + + /// Pop this thread's innermost NVTX range. Must pair with [`range_push`]. + #[inline] + pub fn range_pop() { + if let Some(api) = api() { + // SAFETY: no arguments; unbalanced pops are handled by NVTX (no-op). + unsafe { (api.range_pop)() }; + } + } + + /// Instantaneous marker on the timeline. + #[inline] + pub fn mark(name: &str) { + if let Some(api) = api() { + let c = CString::new(name).unwrap_or_else(|_| CString::new("invalid-label").unwrap()); + // SAFETY: `c` is a valid NUL-terminated string for the duration of the call. + unsafe { (api.mark)(c.as_ptr()) }; + } + } + + /// RAII NVTX range: pushed on construction, popped on drop. `!Send` on + /// purpose — NVTX push/pop stacks are per-thread, so a guard must drop on + /// the thread that created it. + pub struct Range { + pushed: bool, + _not_send: PhantomData<*const ()>, + } + + impl Range { + #[inline] + pub fn new(name: &str) -> Range { + let pushed = api().map(|a| push_str(a, name)).is_some(); + Range { + pushed, + _not_send: PhantomData, + } + } + + /// Like [`Range::new`] but the label is only formatted when a + /// profiler-visible NVTX library is actually loaded. + #[inline] + pub fn fmt String>(label: F) -> Range { + if is_active() { + Range::new(&label()) + } else { + Range { + pushed: false, + _not_send: PhantomData, + } + } + } + } + + impl Drop for Range { + fn drop(&mut self) { + if self.pushed { + range_pop(); + } + } + } + + // --- CUDA profiler capture-range control ------------------------------- + // + // cuProfilerStart/Stop gate `nsys profile --capture-range=cudaProfilerApi`, + // letting a session capture one phase/epoch of a long prove instead of the + // whole run. Loaded from libcuda (already resident via cudarc) — separate + // dlopen so this module stays independent of cudarc's bound symbol set. + + struct ProfilerApi { + start: unsafe extern "C" fn() -> c_int, + stop: unsafe extern "C" fn() -> c_int, + _lib: libloading::Library, + } + unsafe impl Send for ProfilerApi {} + unsafe impl Sync for ProfilerApi {} + + fn profiler_api() -> Option<&'static ProfilerApi> { + static API: OnceLock> = OnceLock::new(); + API.get_or_init(|| { + for name in ["libcuda.so.1", "libcuda.so"] { + // SAFETY: libcuda is loaded by cudarc already; this bumps a refcount. + let Ok(lib) = (unsafe { libloading::Library::new(name) }) else { + continue; + }; + // SAFETY: signatures match the CUDA driver profiler API. + let syms = unsafe { + ( + lib.get:: c_int>(b"cuProfilerStart\0") + .map(|s| *s), + lib.get:: c_int>(b"cuProfilerStop\0") + .map(|s| *s), + ) + }; + if let (Ok(start), Ok(stop)) = syms { + return Some(ProfilerApi { + start, + stop, + _lib: lib, + }); + } + } + None + }) + .as_ref() + } + + /// Begin a profiler capture range (`nsys --capture-range=cudaProfilerApi`). + /// No-op without libcuda or outside a profiler session. + #[inline] + pub fn profiler_start() { + if let Some(api) = profiler_api() { + // SAFETY: no arguments; valid to call any time after libcuda loads. + unsafe { (api.start)() }; + } + } + + /// End a profiler capture range. Must pair with [`profiler_start`]. + #[inline] + pub fn profiler_stop() { + if let Some(api) = profiler_api() { + // SAFETY: no arguments; valid to call any time after libcuda loads. + unsafe { (api.stop)() }; + } + } +} + +#[cfg(not(feature = "nvtx"))] +mod imp { + /// No-op stub; see the `nvtx`-feature implementation above. + pub struct Range; + + impl Range { + #[inline(always)] + pub fn new(_name: &str) -> Range { + Range + } + #[inline(always)] + pub fn fmt String>(_label: F) -> Range { + Range + } + } + + #[inline(always)] + pub fn is_active() -> bool { + false + } + #[inline(always)] + pub fn range_push(_name: &str) {} + #[inline(always)] + pub fn range_pop() {} + #[inline(always)] + pub fn mark(_name: &str) {} + #[inline(always)] + pub fn profiler_start() {} + #[inline(always)] + pub fn profiler_stop() {} +} diff --git a/crypto/stark/Cargo.toml b/crypto/stark/Cargo.toml index 78d95be67..64e650aa7 100644 --- a/crypto/stark/Cargo.toml +++ b/crypto/stark/Cargo.toml @@ -56,6 +56,10 @@ debug-checks = [] # Enables v parallel = ["dep:rayon", "crypto/parallel"] cuda = ["dep:math-cuda"] test-cuda-faults = ["cuda", "math-cuda/test-faults"] +# NVTX ranges for Nsight Systems: math-cuda pipeline entry points plus every +# instruments span (prover phases) become named timeline ranges. Pulls in +# `instruments` so the span tree exists to mirror, and `cuda` for math-cuda. +nvtx = ["cuda", "instruments", "math-cuda/nvtx"] wasm = ["dep:wasm-bindgen", "dep:serde-wasm-bindgen", "dep:web-sys"] disk-spill = ["dep:memmap2", "dep:tempfile", "dep:libc", "crypto/disk-spill"] diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 96bf6ffae..06a5d2c70 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -43,9 +43,36 @@ pub struct SpanGuard { order: u32, start: Instant, start_ns: u128, + /// This span gates an nsys capture range (LAMBDA_VM_NSYS_CAPTURE_SPAN). + #[cfg(feature = "nvtx")] + capture: bool, +} + +/// Label of the span that brackets an `nsys --capture-range=cudaProfilerApi` +/// session, from `LAMBDA_VM_NSYS_CAPTURE_SPAN` (e.g. `rounds_2to4`, or +/// `proving` to capture one epoch of a continuations run). None = never. +#[cfg(feature = "nvtx")] +fn capture_span_label() -> Option<&'static str> { + static LABEL: OnceLock> = OnceLock::new(); + LABEL + .get_or_init(|| std::env::var("LAMBDA_VM_NSYS_CAPTURE_SPAN").ok()) + .as_deref() +} + +/// NVTX-only range with a runtime-formatted name (e.g. `epoch[i=3]`), for +/// callers that need per-instance identity on Nsight timelines. Instruments +/// spans require `'static` labels, so repeated phases (continuation epochs) +/// are told apart by instance order in the JSON timeline and by one of these +/// ranges in nsys. The closure only runs when a profiler-visible NVTX +/// library is loaded. +#[cfg(feature = "nvtx")] +pub fn nvtx_range_fmt String>(label: F) -> math_cuda::nvtx::Range { + math_cuda::nvtx::Range::fmt(label) } /// Open a wall-clock span; records elapsed time when the guard drops. +/// Under the `nvtx` feature the span is mirrored as an NVTX range so Nsight +/// timelines carry the same phase names as the instruments tree. pub fn span(label: &'static str) -> SpanGuard { let depth = SPAN_DEPTH.with(|d| { let v = d.get(); @@ -57,17 +84,35 @@ pub fn span(label: &'static str) -> SpanGuard { .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_nanos(); + #[cfg(feature = "nvtx")] + let capture = { + math_cuda::nvtx::range_push(label); + let capture = capture_span_label() == Some(label); + if capture { + math_cuda::nvtx::profiler_start(); + } + capture + }; SpanGuard { label, depth, order, start: Instant::now(), start_ns, + #[cfg(feature = "nvtx")] + capture, } } impl Drop for SpanGuard { fn drop(&mut self) { + #[cfg(feature = "nvtx")] + { + if self.capture { + math_cuda::nvtx::profiler_stop(); + } + math_cuda::nvtx::range_pop(); + } let wall = self.start.elapsed(); SPAN_DEPTH.with(|d| d.set(d.get().saturating_sub(1))); if let Ok(mut t) = TIMELINE.lock() { diff --git a/prover/Cargo.toml b/prover/Cargo.toml index 15344d138..e4b67d780 100644 --- a/prover/Cargo.toml +++ b/prover/Cargo.toml @@ -9,6 +9,8 @@ default = ["parallel"] parallel = ["stark/parallel", "math/parallel", "crypto/parallel", "dep:rayon"] cuda = ["stark/cuda"] test-cuda-faults = ["cuda", "stark/test-cuda-faults"] +# GPU profiling build: NVTX ranges on prover spans + math-cuda entry points. +nvtx = ["cuda", "instruments", "stark/nvtx"] debug-checks = ["stark/debug-checks"] instruments = ["stark/instruments"] profile-markers = ["stark/profile-markers"] diff --git a/prover/src/continuation.rs b/prover/src/continuation.rs index 169cd7278..e45edecdf 100644 --- a/prover/src/continuation.rs +++ b/prover/src/continuation.rs @@ -1009,6 +1009,11 @@ pub fn prove_continuation( )) })?; + #[cfg(feature = "instruments")] + stark::instruments::reset_timeline(); + #[cfg(feature = "instruments")] + let __root = stark::instruments::span("prove_continuation_total"); + let elf = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; let mut executor = Executor::new(&elf, private_inputs.to_vec()) .map_err(|e| Error::Execution(format!("{e}")))?; @@ -1047,6 +1052,14 @@ pub fn prove_continuation( local_to_global::MAX_EPOCHS ))); } + // Per-epoch identity on Nsight timelines (dynamic NVTX name); the + // instruments spans below carry static labels and are told apart by + // instance order (phase_table.py reports them per instance). + #[cfg(feature = "nvtx")] + let __epoch_nvtx = stark::instruments::nvtx_range_fmt(|| format!("epoch[i={index}]")); + #[cfg(feature = "instruments")] + let __epoch = stark::instruments::span("epoch"); + let register_init: Vec = if index == 0 { register::register_init_from_entry_point(elf.entry_point) } else { @@ -1060,6 +1073,8 @@ pub fn prove_continuation( }; // Run one epoch; `logs` is this epoch's chunk only (the executor clears it). + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("epoch_execute"); let logs = match executor .resume_with_limit(epoch_size) .map_err(|e| Error::Execution(format!("{e}")))? @@ -1067,6 +1082,8 @@ pub fn prove_continuation( Some(logs) => logs.to_vec(), None => break, }; + #[cfg(feature = "instruments")] + drop(__sp); let is_final = executor.pc() == 0; // Invariant: a non-final epoch ran the full `epoch_size` (a power of two), @@ -1079,6 +1096,8 @@ pub fn prove_continuation( } let label = local_to_global::epoch_label(index); + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("epoch_trace_build"); let traces = Traces::from_image_and_logs( &elf, &image, @@ -1093,12 +1112,18 @@ pub fn prove_continuation( )?; let boundary = local_to_global::epoch_boundary(&mut provenance, label, &traces.touched_memory_cells); + #[cfg(feature = "instruments")] + drop(__sp); let start = EpochStart { register_init: ®ister_init, label, }; + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("epoch_prove"); let epoch = prove_epoch(&elf, elf_bytes, &start, traces, is_final, &boundary, opts)?; + #[cfg(feature = "instruments")] + drop(__sp); prev_fini = Some(epoch.reg_fini.clone()); // Carry the image forward: this epoch's fini is the next epoch's init. @@ -1120,6 +1145,8 @@ pub fn prove_continuation( // SINGLE source of truth: the same page-base list drives the committed GLOBAL_MEMORY // tables and is shipped in the bundle, so the two can never diverge in set or order. let touched_page_bases = touched_page_bases(&all_boundaries); + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("prove_global"); let global = prove_global( &all_boundaries, elf_bytes, @@ -1128,6 +1155,20 @@ pub fn prove_continuation( num_private_input_pages, opts, )?; + #[cfg(feature = "instruments")] + { + drop(__sp); + // Same timeline output as the monolithic path (prover/src/lib.rs): + // print the wall-clock tree and honor LAMBDA_VM_TIMELINE_JSON. Without + // this, continuation runs record spans that are never drained. + drop(__root); + let spans = stark::instruments::take_timeline(); + print!("{}", stark::instruments::format_timeline(&spans)); + if let Ok(path) = std::env::var("LAMBDA_VM_TIMELINE_JSON") { + let _ = std::fs::write(&path, stark::instruments::timeline_json(&spans)); + println!("[timeline] wrote {path}"); + } + } Ok(ContinuationProof { epochs, diff --git a/scripts/profiling/README.md b/scripts/profiling/README.md new file mode 100644 index 000000000..b771672f0 --- /dev/null +++ b/scripts/profiling/README.md @@ -0,0 +1,145 @@ +# GPU profiling toolkit + +Tooling for profiling the CUDA prover on the dedicated RTX 5090 box. The full +methodology (what to measure, in what order, and how to read it) lives in +`thoughts/gpu-profiling/plan.md`; this directory is the executable part. + +## One-time machine setup + +```bash +scripts/profiling/setup_machine.sh # apt tooling, nsight, perf/eBPF perms — then REBOOT +``` + +Prereqs handled elsewhere: NVIDIA driver, and the build toolchain + guest +programs per `scripts/SERVER_SETUP.md` (`make compile-programs-asm`, etc.). +On Blackwell (RTX 5090, sm_120) CUDA ≥ 12.8 and nsys/ncu ≥ 2025.1 are hard +requirements. + +## Before every session + +```bash +sudo scripts/profiling/bench_mode.sh on # lock SM clocks, persistence, governor +make test-cuda-integration # sanity: every GPU counter fires +``` + +`bench_mode.sh off` restores defaults. Numbers taken with floating clocks are +not comparable across sessions. + +## The main entry point + +```bash +# 3 instrumented runs + phase table with GPU util per phase: +scripts/profiling/run_profile.sh executor/program_artifacts/asm/fib_iterative_372k.elf + +# the real workload, plus an nsys-traced run and the per-phase GPU-busy report: +scripts/profiling/run_profile.sh --nsys \ + executor/program_artifacts/rust/ethrex.elf \ + --private-input executor/tests/ethrex_5_transfers.bin + +# big continuation run, nsys capture limited to one `proving` span (one epoch): +LAMBDA_VM_NSYS_CAPTURE_SPAN=proving scripts/profiling/run_profile.sh --nsys --continuations \ + executor/program_artifacts/rust/ethrex.elf \ + --private-input executor/tests/ethrex_20_transfers.bin +``` + +Each invocation produces a self-contained bundle under +`reports/__/`: + +| file | what | +|---|---| +| `env.json` | driver, clocks, sha, env — the context that makes numbers comparable | +| `phase_table.md` | warm-run phase tree: median ms, % of total, jitter, GPU util per phase | +| `phase_table_cold.md` | run 1 alone (module load, twiddle build, mempool growth) | +| `trace_perfetto.json` | span tree for ui.perfetto.dev | +| `nsys_report.nsys-rep` | open in the Nsight Systems GUI (scp to a laptop) | +| `nsys_stats.txt` | stock nsys summaries (kernels, memcpy, API, NVTX) | +| `phase_busy.md` | **the ranking input**: per-phase GPU busy %, memcpy, top kernels | + +How to read `phase_busy.md`: a phase with low busy% is host-bound — fix +pipeline/overlap/syncs, don't touch its kernels. A phase with high busy% is +kernel-bound — take its top kernels to Nsight Compute. + +## Continuations: per-epoch data for parallelization + +`prove_continuation` is instrumented independently of the monolithic path +(`prover/src/continuation.rs`): a `prove_continuation_total` root, per-epoch +`epoch` → `epoch_execute` / `epoch_trace_build` / `epoch_prove` spans, a +`prove_global` span, and a dynamic `epoch[i=N]` NVTX range per epoch — and it +writes `LAMBDA_VM_TIMELINE_JSON` at the end (before this instrumentation, +continuation runs produced no timeline at all). + +With `run_profile.sh --continuations` the reports gain the epoch-level view +that drives parallelization decisions: + +- `phase_table.md` — per-epoch instances table: wall per epoch (uniformity — + can epochs be scheduled symmetrically?), gap→next, and **GPU util inside the + span vs inside the gap**. `epoch_execute`/`epoch_trace_build` rows with ~0 + GPU% while `epoch_prove` is high = the fraction of wall recoverable by + overlapping epoch N+1's CPU work with epoch N's GPU proving. +- `phase_busy.md` — "Epochs (continuations)" section from the `epoch[i=N]` + ranges: per-epoch GPU busy%, inter-epoch gaps and GPU work inside them + (≈0 → idle gap = pipelining headroom, quantified in ms). +- In the Nsight GUI, epochs appear as named `epoch[i=N]` ranges, so + cross-epoch overlap (or the lack of it) is visible at a glance; combine with + `LAMBDA_VM_NSYS_CAPTURE_SPAN=epoch_prove` to capture a single epoch's prove. + +## The `nvtx` cargo feature + +`run_profile.sh` builds the cli with `--features nvtx,jemalloc-stats`. The +`nvtx` feature (cli → prover → stark → math-cuda) does three things: + +1. every `instruments` span (prover phases) becomes an NVTX range, so Nsight + timelines carry the same names as the phase table; +2. every `math-cuda` public entry point pushes a range with its shape, e.g. + `lde_tree_base[n=1048576 m=48 bf=4]`; +3. `LAMBDA_VM_NSYS_CAPTURE_SPAN=