Gpu opt 5090

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"]

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" }

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()

crypto/math-cuda/kernels/keccak.cu

@@ -474,3 +474,41 @@ extern "C" __global__ void keccak256_leaves_base_row_major_row_pair(
     }
     finalize_keccak256(st, rate_pos, hashed_leaves_out + tid * 32);
 }
+
+// Column-range variant of `keccak256_leaves_base_row_major_row_pair`: each leaf
+// hashes only columns `[col_start, col_end)` of the two bit-reversed rows,
+// while `m` remains the full row stride. Byte layout equals the CPU
+// `commit_rows_bit_reversed_subset(data, m, col_start, col_end)` — used for
+// preprocessed tables, whose precomputed and multiplicity column ranges commit
+// to separate Merkle trees over the same row-major LDE.
+extern "C" __global__ void keccak256_leaves_base_row_major_row_pair_range(
+    const uint64_t *data,
+    uint64_t m,
+    uint64_t col_start,
+    uint64_t col_end,
+    uint64_t num_rows,
+    uint64_t log_num_rows,
+    uint8_t *hashed_leaves_out)
+{
+    uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    uint64_t num_leaves = num_rows >> 1;
+    if (tid >= num_leaves) return;
+
+    uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows);
+    uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows);
+    const uint64_t *row_0 = data + br_0 * m;
+    const uint64_t *row_1 = data + br_1 * m;
+
+    uint64_t st[25];
+    #pragma unroll
+    for (int i = 0; i < 25; ++i) st[i] = 0;
+
+    uint32_t rate_pos = 0;
+    for (uint64_t c = col_start; c < col_end; ++c) {
+        absorb_lane(st, rate_pos, bswap64(goldilocks::canonical(row_0[c])));
+    }
+    for (uint64_t c = col_start; c < col_end; ++c) {
+        absorb_lane(st, rate_pos, bswap64(goldilocks::canonical(row_1[c])));
+    }
+    finalize_keccak256(st, rate_pos, hashed_leaves_out + tid * 32);
+}

crypto/math-cuda/src/barycentric.rs

@@ -45,9 +45,13 @@ pub fn barycentric_base(
     let be = backend()?;
     let stream = be.next_stream();
 
-    let cols_dev = stream.clone_htod(&columns[..num_cols * col_stride])?;
-    let points_dev = stream.clone_htod(coset_points)?;
-    let inv_dev = stream.clone_htod(inv_denoms_ext3)?;
+    let (cols_dev, points_dev, inv_dev) = {
+        (
+            stream.clone_htod(&columns[..num_cols * col_stride])?,
+            stream.clone_htod(coset_points)?,
+            stream.clone_htod(inv_denoms_ext3)?,
+        )
+    };
     let mut out_dev = stream.alloc_zeros::<u64>(3 * num_cols)?;
 
     let col_stride_u64 = col_stride as u64;
@@ -68,8 +72,10 @@ pub fn barycentric_base(
             .arg(&mut out_dev)
             .launch(cfg)?;
     }
-    let out = stream.clone_dtoh(&out_dev)?;
-    stream.synchronize()?;
+    let out = { stream.clone_dtoh(&out_dev)? };
+    {
+        stream.synchronize()?;
+    }
     Ok(out)
 }
 
@@ -100,9 +106,13 @@ pub fn barycentric_ext3(
     let be = backend()?;
     let stream = be.next_stream();
 
-    let cols_dev = stream.clone_htod(&columns[..num_cols * 3 * col_stride])?;
-    let points_dev = stream.clone_htod(coset_points)?;
-    let inv_dev = stream.clone_htod(inv_denoms_ext3)?;
+    let (cols_dev, points_dev, inv_dev) = {
+        (
+            stream.clone_htod(&columns[..num_cols * 3 * col_stride])?,
+            stream.clone_htod(coset_points)?,
+            stream.clone_htod(inv_denoms_ext3)?,
+        )
+    };
     let mut out_dev = stream.alloc_zeros::<u64>(3 * num_cols)?;
 
     let col_stride_u64 = col_stride as u64;
@@ -123,8 +133,10 @@ pub fn barycentric_ext3(
             .arg(&mut out_dev)
             .launch(cfg)?;
     }
-    let out = stream.clone_dtoh(&out_dev)?;
-    stream.synchronize()?;
+    let out = { stream.clone_dtoh(&out_dev)? };
+    {
+        stream.synchronize()?;
+    }
     Ok(out)
 }
 
@@ -149,9 +161,14 @@ pub fn barycentric_base_on_device(
 
     let be = backend()?;
     let stream = be.next_stream();
+    main_handle.wait_ready_on(&stream)?;
 
-    let points_dev = stream.clone_htod(coset_points)?;
-    let inv_dev = stream.clone_htod(inv_denoms_ext3)?;
+    let (points_dev, inv_dev) = {
+        (
+            stream.clone_htod(coset_points)?,
+            stream.clone_htod(inv_denoms_ext3)?,
+        )
+    };
     let mut out_dev = stream.alloc_zeros::<u64>(3 * num_cols)?;
 
     let col_stride_u64 = col_stride as u64;
@@ -174,8 +191,10 @@ pub fn barycentric_base_on_device(
             .arg(&mut out_dev)
             .launch(cfg)?;
     }
-    let out = stream.clone_dtoh(&out_dev)?;
-    stream.synchronize()?;
+    let out = { stream.clone_dtoh(&out_dev)? };
+    {
+        stream.synchronize()?;
+    }
     Ok(out)
 }
 
@@ -197,6 +216,7 @@ pub fn barycentric_base_on_device_with_dev_inv_denoms(
     inv_offset_u64: usize,
     n: usize,
 ) -> Result<Vec<u64>> {
+    main_handle.wait_ready_on(stream)?;
     assert!(coset_points_dev.len() >= n);
     let inv_end = inv_offset_u64
         .checked_add(3 * n)
@@ -233,8 +253,10 @@ pub fn barycentric_base_on_device_with_dev_inv_denoms(
             .arg(&mut out_dev)
             .launch(cfg)?;
     }
-    let out = stream.clone_dtoh(&out_dev)?;
-    stream.synchronize()?;
+    let out = { stream.clone_dtoh(&out_dev)? };
+    {
+        stream.synchronize()?;
+    }
     Ok(out)
 }
 
@@ -257,9 +279,14 @@ pub fn barycentric_ext3_on_device(
 
     let be = backend()?;
     let stream = be.next_stream();
+    aux_handle.wait_ready_on(&stream)?;
 
-    let points_dev = stream.clone_htod(coset_points)?;
-    let inv_dev = stream.clone_htod(inv_denoms_ext3)?;
+    let (points_dev, inv_dev) = {
+        (
+            stream.clone_htod(coset_points)?,
+            stream.clone_htod(inv_denoms_ext3)?,
+        )
+    };
     let mut out_dev = stream.alloc_zeros::<u64>(3 * num_cols)?;
 
     let col_stride_u64 = col_stride as u64;
@@ -282,8 +309,10 @@ pub fn barycentric_ext3_on_device(
             .arg(&mut out_dev)
             .launch(cfg)?;
     }
-    let out = stream.clone_dtoh(&out_dev)?;
-    stream.synchronize()?;
+    let out = { stream.clone_dtoh(&out_dev)? };
+    {
+        stream.synchronize()?;
+    }
     Ok(out)
 }
 
@@ -297,6 +326,7 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms(
     inv_offset_u64: usize,
     n: usize,
 ) -> Result<Vec<u64>> {
+    aux_handle.wait_ready_on(stream)?;
     assert!(coset_points_dev.len() >= n);
     let inv_end = inv_offset_u64
         .checked_add(3 * n)
@@ -333,8 +363,10 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms(
             .arg(&mut out_dev)
             .launch(cfg)?;
     }
-    let out = stream.clone_dtoh(&out_dev)?;
-    stream.synchronize()?;
+    let out = { stream.clone_dtoh(&out_dev)? };
+    {
+        stream.synchronize()?;
+    }
     Ok(out)
 }
 
@@ -347,12 +379,13 @@ pub fn gather_rows_base_on_device(
     rows: &[u32],
     stream: &Arc<CudaStream>,
 ) -> Result<Vec<u64>> {
+    main.wait_ready_on(stream)?;
     let num_cols = main.m;
     if num_cols == 0 || rows.is_empty() {
         return Ok(Vec::new());
     }
     let be = backend()?;
-    let rows_dev = stream.clone_htod(rows)?;
+    let rows_dev = { stream.clone_htod(rows)? };
     let mut out = stream.alloc_zeros::<u64>(rows.len() * num_cols)?;
     let col_stride = main.lde_size as u64;
     let num_cols_u64 = num_cols as u64;
@@ -373,8 +406,10 @@ pub fn gather_rows_base_on_device(
             .arg(&mut out)
             .launch(cfg)?;
     }
-    let host = stream.clone_dtoh(&out)?;
-    stream.synchronize()?;
+    let host = { stream.clone_dtoh(&out)? };
+    {
+        stream.synchronize()?;
+    }
     Ok(host)
 }
 
@@ -385,12 +420,13 @@ pub fn gather_rows_ext3_on_device(
     rows: &[u32],
     stream: &Arc<CudaStream>,
 ) -> Result<Vec<u64>> {
+    aux.wait_ready_on(stream)?;
     let num_cols = aux.m;
     if num_cols == 0 || rows.is_empty() {
         return Ok(Vec::new());
     }
     let be = backend()?;
-    let rows_dev = stream.clone_htod(rows)?;
+    let rows_dev = { stream.clone_htod(rows)? };
     let mut out = stream.alloc_zeros::<u64>(rows.len() * num_cols * 3)?;
     let col_stride = aux.lde_size as u64;
     let num_cols_u64 = num_cols as u64;
@@ -411,7 +447,9 @@ pub fn gather_rows_ext3_on_device(
             .arg(&mut out)
             .launch(cfg)?;
     }
-    let host = stream.clone_dtoh(&out)?;
-    stream.synchronize()?;
+    let host = { stream.clone_dtoh(&out)? };
+    {
+        stream.synchronize()?;
+    }
     Ok(host)
 }