Review fixes for the GPU LogUp aux build (#762)

.github/workflows/pr_main.yaml

@@ -127,7 +127,7 @@ jobs:
   test:
     name: Test
     if: always()
-    needs: [test-executor, test-cli, test-prover, test-disk-spill]
+    needs: [test-executor, test-cli, test-prover, test-disk-spill, test-stark-cuda-lib]
     runs-on: ubuntu-latest
     steps:
       - name: Check results
@@ -136,11 +136,13 @@ jobs:
           cli="${{ needs.test-cli.result }}"
           prover="${{ needs.test-prover.result }}"
           disk_spill="${{ needs.test-disk-spill.result }}"
+          stark_cuda_lib="${{ needs.test-stark-cuda-lib.result }}"
 
           echo "test-executor: $executor"
           echo "test-cli: $cli"
           echo "test-prover: $prover"
           echo "test-disk-spill: $disk_spill"
+          echo "test-stark-cuda-lib: $stark_cuda_lib"
 
           # Allow "success" or "skipped" (skipped on merge queue pushes)
           if [[ "$executor" != "success" && "$executor" != "skipped" ]]; then
@@ -155,6 +157,9 @@ jobs:
           if [[ "$disk_spill" != "success" && "$disk_spill" != "skipped" ]]; then
             exit 1
           fi
+          if [[ "$stark_cuda_lib" != "success" && "$stark_cuda_lib" != "skipped" ]]; then
+            exit 1
+          fi
 
   test-disk-spill:
     name: Disk-spill tests
@@ -215,6 +220,33 @@ jobs:
         run: |
           cargo test --release -p lambda-vm-prover --features disk-spill -- disk_spill count_table_lengths
 
+  test-stark-cuda-lib:
+    name: Stark cuda-feature lib tests
+    runs-on: ubuntu-latest
+    if: github.event_name != 'push' || github.actor != 'github-merge-queue[bot]'
+    steps:
+      - name: Checkout sources
+        uses: actions/checkout@v4
+
+      - name: Setup Rust Environment
+        uses: ./.github/actions/setup-rust
+
+      - name: Cache cargo build artifacts
+        uses: Swatinem/rust-cache@v2
+        with:
+          shared-key: "lambda-vm-stark-cuda-lib"
+          cache-all-crates: "true"
+
+      # The cuda feature gates the logup_gpu module, whose descriptor /
+      # CPU-mirror parity tests are pure CPU (they never touch the driver).
+      # Without nvcc the kernels build as empty PTX stubs, so these run on a
+      # plain runner; GPU-dependent tests are #[ignore] and stay skipped.
+      # Scoped to logup_gpu: the rest of the cuda-feature lib suite dispatches
+      # to the GPU and needs the CUDA driver.
+      - name: Run stark logup_gpu parity tests (no GPU required)
+        run: |
+          cargo test --release -p stark --features cuda --lib logup_gpu
+
   build-prover-tests:
     name: Build prover tests
     runs-on: ubuntu-latest

crypto/math-cuda/kernels/logup.cu

@@ -115,7 +115,7 @@ extern "C" __global__ void logup_term_ext3(
 // Accumulated column (K4): running sum of the term columns, on device.
 //   row_sum[i] = sum over all term columns of term[col][i]
 //   S = inclusive prefix scan of row_sum ;  L = S[n-1] ;  offset = L / N
-//   acc[i] = S[i] - (i+1) * offset          (matches build_accumulated_column)
+//   acc[i] = S[i] - (i+1) * offset          (matches build_accumulated_column_from_terms)
 // Additive 3-phase Hillis-Steele scan (mirrors inverse.cu, add not mul).
 // ===========================================================================
 

crypto/math-cuda/src/device.rs

@@ -265,9 +265,11 @@ impl Backend {
     fn init() -> Result<Self> {
         let ctx = CudaContext::new(0)?;
         // cudarc's default per-slice CudaEvent tracking adds two driver calls
-        // per alloc and serialises under the context lock. We never share
-        // slices across streams (every call scopes its own buffers and syncs
-        // before returning), so the tracking is pure overhead. Disable it.
+        // per alloc and serialises under the context lock. Slices are only
+        // shared across streams after the producing stream has been host-
+        // synchronised (e.g. the retained trace snapshot and the resident
+        // LogUp aux buffer; every producer syncs before its handle escapes),
+        // so the tracking is pure overhead. Disable it.
         unsafe { ctx.disable_event_tracking() };
 
         // Retain freed device memory in the stream ordered pool for reuse.

crypto/math-cuda/src/lde.rs

@@ -379,24 +379,26 @@ fn launch_row_to_col_major(
     Ok(dst)
 }
 
-/// Shared row-major LDE + Keccak + Merkle pipeline for the base and ext3 paths.
-///
-/// `total_cols` is the number of base-field columns in the row-major layout:
-/// `m` for base, `m * 3` for ext3. Because `Fp3 = [u64; 3]`, the three ext3
-/// components are just three adjacent base-field columns, so the same row-major
-/// NTT and Keccak kernels process all of them simultaneously — no de-interleave.
-///
-/// Single H2D, row-major NTT, single D2H — no CPU-side extract or transpose.
-/// Returns (merkle_nodes, column-major device buffer, row-major LDE Vec). The
-/// buffer is transposed to column-major (as required by the downstream GPU
-/// kernels DEEP/barycentric); callers wrap it in the appropriate LDE handle.
 /// Row-major LDE input: either a host slice (uploaded) or an already-resident
 /// device buffer (copied device-to-device, no PCIe upload).
 enum InnerInput<'a> {
     Host(&'a [u64]),
     Dev(&'a CudaSlice<u64>),
 }
 
+/// Shared row-major LDE + Keccak + Merkle pipeline for the base and ext3 paths.
+///
+/// `total_cols` is the number of base-field columns in the row-major layout:
+/// `m` for base, `m * 3` for ext3. Because `Fp3 = [u64; 3]`, the three ext3
+/// components are just three adjacent base-field columns, so the same row-major
+/// NTT and Keccak kernels process all of them simultaneously — no de-interleave.
+///
+/// Single H2D (or D2D), row-major NTT, single D2H — no CPU-side extract or
+/// transpose. Returns (merkle_nodes, column-major device buffer, row-major LDE
+/// Vec, optional trace-domain column-major snapshot — `Some` iff
+/// `retain_trace_col_major`). The buffer is transposed to column-major (as
+/// required by the downstream GPU kernels DEEP/barycentric); callers wrap it in
+/// the appropriate LDE handle.
 #[allow(clippy::type_complexity)]
 fn coset_lde_row_major_inner(
     input: InnerInput,

crypto/math-cuda/src/lib.rs

@@ -1,8 +1,9 @@
 //! GPU backend for the lambda-vm STARK prover.
 //!
-//! Primary entry point: [`lde::coset_lde_base`]. Everything else (`ntt`,
-//! element-wise arith) is either internal to the LDE pipeline or used by the
-//! parity test suite.
+//! Primary entry points: [`lde::coset_lde_base`] for the LDE pipeline and
+//! [`logup::logup_aux_resident`] for the device-resident LogUp aux build.
+//! Everything else (`ntt`, element-wise arith) is either internal to those
+//! pipelines or used by the parity test suite.
 
 pub mod barycentric;
 pub mod deep;

crypto/math-cuda/src/logup.rs

@@ -17,6 +17,8 @@ use crate::Result;
 use crate::device::backend;
 use crate::inverse::batch_inverse_ext3_dev;
 
+// Must match LOGUP_BLK in kernels/logup.cu: the block scan kernel assumes
+// exactly this many threads per block for its shared-memory array.
 const BLOCK_SIZE: u32 = 256;
 
 /// Flat LogUp descriptor for one table (CSR arrays, canonical Goldilocks). Built
@@ -41,12 +43,22 @@ pub struct LogupDescriptor<'a> {
     pub mult_term_col: &'a [u32],
 }
 
-fn cfg(total: usize) -> LaunchConfig {
-    LaunchConfig {
+fn cfg(total: usize) -> Result<LaunchConfig> {
+    // See `batch_inverse_ext3_dev` for the rationale: a u32 grid_dim is
+    // truncated past u32::MAX / BLOCK_SIZE, which would silently launch too
+    // few blocks and leave a tail of the (uninitialized) output unwritten.
+    // Runtime Err, not debug_assert, so release builds also route to the
+    // caller's CPU fallback.
+    if total > u32::MAX as usize / BLOCK_SIZE as usize {
+        return Err(cudarc::driver::DriverError(
+            cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE,
+        ));
+    }
+    Ok(LaunchConfig {
         grid_dim: ((total as u32).div_ceil(BLOCK_SIZE), 1, 1),
         block_dim: (BLOCK_SIZE, 1, 1),
         shared_mem_bytes: 0,
-    }
+    })
 }
 
 /// Fingerprint kernel over a device-resident main trace. Returns the ext3 fp
@@ -94,7 +106,7 @@ fn fingerprints_into_dev(
             .arg(&z1)
             .arg(&z2)
             .arg(&mut out)
-            .launch(cfg(total))?;
+            .launch(cfg(total)?)?;
     }
     Ok(out)
 }
@@ -168,7 +180,7 @@ pub fn logup_term_columns(
             .arg(&mult_term_coef)
             .arg(&mult_term_col)
             .arg(&mut out)
-            .launch(cfg(total))?;
+            .launch(cfg(total)?)?;
     }
     if timing {
         stream.synchronize()?;
@@ -266,9 +278,6 @@ impl PartialEq for ResidentAux {
 }
 impl Eq for ResidentAux {}
 
-/// Full aux build on device: fingerprints → invert → term columns → accumulate
-/// scan → assemble the row-major aux trace buffer, all resident. `inv_n` is
-/// `1/num_rows` embedded in ext3. The stream is synchronised before return.
 /// Main trace input for the resident aux build: either a host column-major
 /// buffer to upload, or an already-resident device buffer (from the R1 main
 /// LDE) to read in place. The device form skips the ~3 GB main re-upload.
@@ -278,6 +287,10 @@ pub enum ResidentMain<'a> {
     Dev(&'a CudaSlice<u64>),
 }
 
+/// Full aux build on device: fingerprints → invert → term columns → accumulate
+/// scan → assemble the row-major aux trace buffer, all resident. `inv_n` is
+/// `1/num_rows` embedded in ext3. Requires `num_rows >= 1`. The stream is
+/// synchronised before return.
 #[allow(clippy::too_many_arguments)]
 pub fn logup_aux_resident(
     main: ResidentMain,
@@ -288,6 +301,7 @@ pub fn logup_aux_resident(
     inv_n: [u64; 3],
     stream: &Arc<CudaStream>,
 ) -> Result<ResidentAux> {
+    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
     // time is attributed correctly. Off = no extra syncs, production path.
@@ -326,10 +340,13 @@ pub fn logup_aux_resident(
 
     // Term columns (committed + virtual), resident, layout [col][row].
     // num_out is always >= 1 (the accumulated column); num_committed = num_out - 1.
-    debug_assert!(
-        d.num_out_cols >= 1,
-        "logup descriptor has no output columns"
-    );
+    // Runtime Err, not debug_assert: in release a zero would wrap num_out - 1
+    // to usize::MAX and launch the assemble kernel with a bogus column count.
+    if d.num_out_cols == 0 {
+        return Err(cudarc::driver::DriverError(
+            cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE,
+        ));
+    }
     let num_out = d.num_out_cols;
     let mut terms = unsafe { stream.alloc::<u64>(num_out * num_rows * 3) }?;
     let out_col_offsets = stream.clone_htod(d.out_col_offsets)?;
@@ -356,7 +373,7 @@ pub fn logup_aux_resident(
             .arg(&mult_term_coef)
             .arg(&mult_term_col)
             .arg(&mut terms)
-            .launch(cfg(num_out * num_rows))?;
+            .launch(cfg(num_out * num_rows)?)?;
     }
     sync_if(timing)?;
     let t_term = std::time::Instant::now();
@@ -370,7 +387,7 @@ pub fn logup_aux_resident(
             .arg(&num_out_u32)
             .arg(&num_rows_u32)
             .arg(&mut row_sum)
-            .launch(cfg(num_rows))?;
+            .launch(cfg(num_rows)?)?;
     }
     scan_add_inplace(stream, be, &mut row_sum, num_rows)?; // row_sum now holds S
     let (i0, i1, i2) = (inv_n[0], inv_n[1], inv_n[2]);
@@ -385,7 +402,7 @@ pub fn logup_aux_resident(
             .arg(&i1)
             .arg(&i2)
             .arg(&mut accumulated)
-            .launch(cfg(num_rows))?;
+            .launch(cfg(num_rows)?)?;
     }
 
     // Assemble row-major aux buffer: committed (num_out-1) cols + accumulated.
@@ -401,7 +418,7 @@ pub fn logup_aux_resident(
             .arg(&accumulated)
             .arg(&num_rows_u32)
             .arg(&mut aux)
-            .launch(cfg(num_rows))?;
+            .launch(cfg(num_rows)?)?;
     }
     sync_if(timing)?;
     let t_accum_done = std::time::Instant::now();

crypto/stark/src/gpu_lde.rs

@@ -84,7 +84,9 @@ pub fn gpu_extend_halves_calls() -> u64 {
     GPU_EXTEND_HALVES_CALLS.load(Ordering::Relaxed)
 }
 
-/// LogUp aux-build term-column GPU path attempts (one per table that took it).
+/// Successful LogUp aux-build GPU dispatches (one per table that took either
+/// the resident or the term-column path; failed attempts fall back to CPU and
+/// are not counted).
 pub(crate) static GPU_LOGUP_CALLS: AtomicU64 = AtomicU64::new(0);
 pub fn gpu_logup_calls() -> u64 {
     GPU_LOGUP_CALLS.load(Ordering::Relaxed)
@@ -1108,12 +1110,10 @@ unsafe fn ext3_slice_to_u64<E: IsField>(col: &[FieldElement<E>]) -> &[u64] {
     unsafe { from_raw_parts(ptr, len) }
 }
 
-/// Convert ext3 evals (3*n u64s, interleaved) into a freshly allocated
-/// `Vec<FieldElement<E>>` of length `n`. Caller must have established
-/// `E == Ext3`.
 /// Like [`try_expand_leaf_and_tree_ext3_row_major_keep`] but the aux columns are
 /// already resident on device (from the GPU LogUp aux build) — no host upload.
-/// Consumes the resident buffer via the device-input LDE.
+/// The resident buffer is only borrowed: the device-input LDE copies it
+/// device-to-device into its own scratch, so `ra` stays valid afterwards.
 pub(crate) fn try_expand_leaf_and_tree_ext3_row_major_keep_dev<F, E, B>(
     ra: &math_cuda::logup::ResidentAux,
     blowup_factor: usize,
@@ -1150,6 +1150,10 @@ where
 
     let lde_out: Vec<FieldElement<E>> = unsafe {
         let mut v = std::mem::ManuallyDrop::new(lde_u64);
+        debug_assert!(
+            v.len() % 3 == 0 && v.capacity() % 3 == 0,
+            "lde_u64 len/capacity must be a multiple of 3 for Fp3 reinterpret"
+        );
         Vec::from_raw_parts(
             v.as_mut_ptr() as *mut FieldElement<E>,
             v.len() / 3,
@@ -1161,6 +1165,9 @@ where
     Some((tree, handle, lde_out))
 }
 
+/// Convert ext3 evals (3*n u64s, interleaved) into a freshly allocated
+/// `Vec<FieldElement<E>>` of length `n`. Caller must have established
+/// `E == Ext3`.
 pub(crate) fn u64_to_ext3_vec<E>(raw: &[u64]) -> Vec<FieldElement<E>>
 where
     E: IsField + 'static,