From 99c3a1173754a88d2dcb2a22e729763834f0a611 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 29 Jul 2026 14:53:51 -0300 Subject: [PATCH 1/6] new opt --- crypto/math-cuda/kernels/fri.cu | 17 ++ crypto/math-cuda/kernels/inverse.cu | 35 +++ crypto/math-cuda/kernels/keccak.cu | 37 +++- crypto/math-cuda/kernels/ntt.cu | 57 +++++ crypto/math-cuda/src/device.rs | 49 +++++ crypto/math-cuda/src/fri.rs | 157 ++++++++------ crypto/math-cuda/src/inverse.rs | 69 ++++-- crypto/math-cuda/src/lde.rs | 120 ++++++++--- crypto/math-cuda/src/merkle.rs | 70 ++++++ crypto/stark/src/fri/fri_commitment.rs | 7 + crypto/stark/src/fri/mod.rs | 2 +- crypto/stark/src/gpu_lde.rs | 194 ++++++++++++++--- crypto/stark/src/logup_gpu.rs | 17 +- crypto/stark/src/lookup.rs | 23 +- crypto/stark/src/prover.rs | 287 +++++++++++++++++++------ 15 files changed, 919 insertions(+), 222 deletions(-) diff --git a/crypto/math-cuda/kernels/fri.cu b/crypto/math-cuda/kernels/fri.cu index 63d72cef1..bcc8f9e40 100644 --- a/crypto/math-cuda/kernels/fri.cu +++ b/crypto/math-cuda/kernels/fri.cu @@ -59,3 +59,20 @@ extern "C" __global__ void fri_update_twiddles( uint64_t old = tw_in[2 * j]; tw_out[j] = goldilocks::mul(old, old); } + +// Gather interleaved ext3 elements at arbitrary positions: one thread per +// query copies evals[positions[i]] (3 u64) into out[i]. Serves the FRI query +// phase's symmetric-eval reads off the resident layer buffers. +extern "C" __global__ void gather_ext3_at( + const uint64_t *evals, + const uint32_t *positions, + uint64_t q, + uint64_t *out +) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= q) return; + uint64_t p = positions[i]; + out[i * 3] = evals[p * 3]; + out[i * 3 + 1] = evals[p * 3 + 1]; + out[i * 3 + 2] = evals[p * 3 + 2]; +} diff --git a/crypto/math-cuda/kernels/inverse.cu b/crypto/math-cuda/kernels/inverse.cu index 4ee228d8c..65d54afc0 100644 --- a/crypto/math-cuda/kernels/inverse.cu +++ b/crypto/math-cuda/kernels/inverse.cu @@ -309,3 +309,38 @@ extern "C" __global__ void batch_inverse_combine_ext3( out_base[1] = res.b; out_base[2] = res.c; } + +// --------------------------------------------------------------------------- +// 7. invert_total_ext3 +// +// One-thread Fermat inversion of the scan total: out = src[n-1]^(p^3 - 2). +// Replaces the host round-trip (D2H + host Fermat + H2D + stream sync) so the +// whole batch inverse stays stream-ordered. The 192-bit exponent arrives as +// three little-endian u64 limbs. +// --------------------------------------------------------------------------- +extern "C" __global__ void invert_total_ext3( + const uint64_t *src, // 3 * n u64 (reads element n-1) + uint64_t n, + uint64_t e0, // exponent limbs, little-endian + uint64_t e1, + uint64_t e2, + uint64_t *out // 3 u64 +) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + const uint64_t *base = src + (n - 1) * 3; + ext3::Fe3 a = {base[0], base[1], base[2]}; + ext3::Fe3 r = ext3::one(); + uint64_t limbs[3] = {e0, e1, e2}; + for (int li = 2; li >= 0; --li) { + uint64_t bits = limbs[li]; + for (int b = 63; b >= 0; --b) { + r = ext3::mul(r, r); + if ((bits >> b) & 1) { + r = ext3::mul(r, a); + } + } + } + out[0] = r.a; + out[1] = r.b; + out[2] = r.c; +} diff --git a/crypto/math-cuda/kernels/keccak.cu b/crypto/math-cuda/kernels/keccak.cu index 7b62789f9..b026ff2b6 100644 --- a/crypto/math-cuda/kernels/keccak.cu +++ b/crypto/math-cuda/kernels/keccak.cu @@ -366,13 +366,11 @@ extern "C" __global__ void keccak_fri_leaves_ext3( // concatenation of two 32-byte siblings, identical to // `FieldElementVectorBackend::hash_new_parent` on host. // --------------------------------------------------------------------------- -extern "C" __global__ void keccak_merkle_level( +__device__ __forceinline__ void hash_merkle_parent( uint8_t *nodes, uint64_t parent_begin, // node index (counted in 32-byte nodes) - uint64_t n_pairs) { - uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= n_pairs) return; - + uint64_t n_pairs, + uint64_t tid) { uint64_t st[25]; #pragma unroll for (int i = 0; i < 25; ++i) st[i] = 0; @@ -393,6 +391,35 @@ extern "C" __global__ void keccak_merkle_level( finalize_keccak256(st, rate_pos, nodes + (parent_begin + tid) * 32); } +extern "C" __global__ void keccak_merkle_level( + uint8_t *nodes, + uint64_t parent_begin, // node index (counted in 32-byte nodes) + uint64_t n_pairs) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= n_pairs) return; + hash_merkle_parent(nodes, parent_begin, n_pairs, tid); +} + +// Build every remaining level (from `level_begin` up to the root) in ONE +// single-block launch: each level's pairs are grid-strided over the block, +// with a __syncthreads() barrier between levels. Replaces log2 launches of +// `keccak_merkle_level` for the small top levels of the tree, whose per-level +// work is dwarfed by launch overhead. +extern "C" __global__ void keccak_merkle_tail( + uint8_t *nodes, + uint64_t level_begin) { + uint64_t lb = level_begin; + while (lb != 0) { + uint64_t nb = lb / 2; + uint64_t n_pairs = lb - nb; + for (uint64_t tid = threadIdx.x; tid < n_pairs; tid += blockDim.x) { + hash_merkle_parent(nodes, nb, n_pairs, tid); + } + __syncthreads(); + lb = nb; + } +} + // Gather Merkle authentication paths for a batch of leaf positions, reading the // resident tree `nodes` (32-byte nodes; layout: inner nodes [0..leaves_len-1], // root at 0, leaves at [leaves_len-1..]). One thread per query walks leaf->root, diff --git a/crypto/math-cuda/kernels/ntt.cu b/crypto/math-cuda/kernels/ntt.cu index 13c1af688..35a4f7b20 100644 --- a/crypto/math-cuda/kernels/ntt.cu +++ b/crypto/math-cuda/kernels/ntt.cu @@ -411,3 +411,60 @@ extern "C" __global__ void matrix_transpose_strided( __syncthreads(); } } + +// First-8-levels fused DIT on row-major data: one block stages 256 consecutive +// rows x blockDim.x columns in shmem and runs levels 0..min(8,log_n) with +// __syncthreads between levels (row-major analog of ntt_dit_8_levels_batched +// with base_step == 0, whose twiddle math this reuses verbatim). Grid: +// x = column tiles, y = n/256 row blocks. Requires n >= 256. Shmem tile is +// padded (pitch = T+1) to break bank conflicts on the butterfly accesses. +extern "C" __global__ void ntt_dit_8_levels_row_major(uint64_t *data, + const uint64_t *tw, + uint64_t n, + uint64_t log_n, + uint64_t m) +{ + extern __shared__ uint64_t tile[]; + uint32_t T = blockDim.x; + uint32_t pitch = T + 1; + uint64_t col = (uint64_t)blockIdx.x * T + threadIdx.x; + bool live = col < m; + uint64_t row_base = (uint64_t)blockIdx.y * 256; + + for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { + if (live) tile[r * pitch + threadIdx.x] = data[(row_base + r) * m + col]; + } + __syncthreads(); + + uint32_t n_loc_steps = (uint32_t)min((uint64_t)8, log_n); + uint32_t remaining_high_bits = (uint32_t)(log_n - 1); + uint32_t high_mask = (1u << remaining_high_bits) - 1u; + + for (uint32_t loc_step = 0; loc_step < n_loc_steps; ++loc_step) { + for (uint32_t i = threadIdx.y; i < 128; i += blockDim.y) { + uint32_t half = 1u << loc_step; + uint32_t grp = i >> loc_step; + uint32_t grp_pos = i & (half - 1); + uint32_t idx1 = (grp << (loc_step + 1)) + grp_pos; + uint32_t idx2 = idx1 + half; + + uint32_t gs = loc_step; + uint32_t ggp = ((uint32_t)blockIdx.y << 7) + i; + ggp = (ggp & high_mask) + (ggp >> remaining_high_bits); + ggp = ggp & ((1u << gs) - 1u); + uint64_t factor = tw[(uint64_t)ggp * (n >> (gs + 1))]; + + if (live) { + uint64_t u = tile[idx1 * pitch + threadIdx.x]; + uint64_t v = mul(tile[idx2 * pitch + threadIdx.x], factor); + tile[idx1 * pitch + threadIdx.x] = add(u, v); + tile[idx2 * pitch + threadIdx.x] = sub(u, v); + } + } + __syncthreads(); + } + + for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { + if (live) data[(row_base + r) * m + col] = tile[r * pitch + threadIdx.x]; + } +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 8fd7f13de..292026401 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -186,6 +186,7 @@ pub struct Backend { // row-major NTT kernels pub bit_reverse_row_major: CudaFunction, pub ntt_dit_level_row_major: CudaFunction, + pub ntt_dit_8_levels_row_major: CudaFunction, pub pointwise_mul_row_major: CudaFunction, pub matrix_transpose_strided: CudaFunction, @@ -198,6 +199,7 @@ pub struct Backend { pub keccak_comp_poly_leaves_ext3: CudaFunction, pub keccak_fri_leaves_ext3: CudaFunction, pub keccak_merkle_level: CudaFunction, + pub keccak_merkle_tail: CudaFunction, pub merkle_gather_paths: CudaFunction, // barycentric.cubin @@ -214,6 +216,7 @@ pub struct Backend { // fri.cubin pub fri_fold_ext3: CudaFunction, + pub gather_ext3_at: CudaFunction, pub fri_update_twiddles: CudaFunction, // inverse.cubin @@ -223,6 +226,7 @@ pub struct Backend { pub block_inclusive_scan_rev_ext3: CudaFunction, pub apply_block_offsets_rev_ext3: CudaFunction, pub batch_inverse_combine_ext3: CudaFunction, + pub invert_total_ext3: CudaFunction, pub logup_fingerprint_ext3: CudaFunction, pub logup_term_ext3: CudaFunction, pub logup_row_sum_ext3: CudaFunction, @@ -412,6 +416,7 @@ impl Backend { scalar_mul_batched: ntt.load_function("scalar_mul_batched")?, bit_reverse_row_major: ntt.load_function("bit_reverse_row_major")?, ntt_dit_level_row_major: ntt.load_function("ntt_dit_level_row_major")?, + ntt_dit_8_levels_row_major: ntt.load_function("ntt_dit_8_levels_row_major")?, pointwise_mul_row_major: ntt.load_function("pointwise_mul_row_major")?, matrix_transpose_strided: ntt.load_function("matrix_transpose_strided")?, keccak256_leaves_base_row_major_row_pair: keccak @@ -425,6 +430,7 @@ impl Backend { keccak_comp_poly_leaves_ext3: keccak.load_function("keccak_comp_poly_leaves_ext3")?, keccak_fri_leaves_ext3: keccak.load_function("keccak_fri_leaves_ext3")?, keccak_merkle_level: keccak.load_function("keccak_merkle_level")?, + keccak_merkle_tail: keccak.load_function("keccak_merkle_tail")?, merkle_gather_paths: keccak.load_function("merkle_gather_paths")?, barycentric_base_batched: bary.load_function("barycentric_base_batched")?, barycentric_ext3_batched: bary.load_function("barycentric_ext3_batched")?, @@ -437,6 +443,7 @@ impl Backend { deep_composition_ext3_row: deep.load_function("deep_composition_ext3_row")?, bit_reverse_ext3_kernel: deep.load_function("bit_reverse_ext3_interleaved")?, fri_fold_ext3: fri.load_function("fri_fold_ext3")?, + gather_ext3_at: fri.load_function("gather_ext3_at")?, fri_update_twiddles: fri.load_function("fri_update_twiddles")?, compute_denoms_ext3: inverse.load_function("compute_denoms_ext3")?, block_inclusive_scan_fwd_ext3: inverse @@ -446,6 +453,7 @@ impl Backend { .load_function("block_inclusive_scan_rev_ext3")?, apply_block_offsets_rev_ext3: inverse.load_function("apply_block_offsets_rev_ext3")?, batch_inverse_combine_ext3: inverse.load_function("batch_inverse_combine_ext3")?, + invert_total_ext3: inverse.load_function("invert_total_ext3")?, logup_fingerprint_ext3: logup.load_function("logup_fingerprint_ext3")?, logup_term_ext3: logup.load_function("logup_term_ext3")?, logup_row_sum_ext3: logup.load_function("logup_row_sum_ext3")?, @@ -631,6 +639,47 @@ impl Drop for PendingD2H<'_> { /// stream-ordered on its own stream, so a `src` allocated on `stream` may be /// dropped after this call — the free queues behind the copy. Do NOT pass a /// `src` owned by a *different* stream and drop it before waiting. +/// Host→device copy staged through the pinned slot: one host memcpy into +/// pinned memory + one async DMA, instead of the driver's internal pageable +/// staging (small chunks; 2-3x slower for multi-hundred-MB traces and it +/// convoys under multi-thread load). Blocks until the DMA lands, so the slot +/// and `src_host` are both reusable on return. +pub fn htod_via( + stream: &Arc, + slot: &Mutex, + ctx: &CudaContext, + src_host: &[T], + dst: &mut cudarc::driver::CudaViewMut<'_, T>, +) -> Result<()> { + use cudarc::driver::DevicePtrMut; + let n_bytes = std::mem::size_of_val(src_host); + let u64_len = n_bytes.div_ceil(8); + let mut staging = slot.lock().unwrap(); + staging.ensure_capacity(u64_len, ctx)?; + ctx.bind_to_thread()?; + // SAFETY: the pinned allocation is stable while the lock is held and at + // least `n_bytes` long (`ensure_capacity`). The DMA reads it after the + // host memcpy (program order); `device_ptr_mut` orders the device write + // on `stream`. + unsafe { + std::ptr::copy_nonoverlapping( + src_host.as_ptr() as *const u8, + staging.ptr as *mut u8, + n_bytes, + ); + let (dst_ptr, _record) = dst.device_ptr_mut(stream); + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + dst_ptr, + staging.ptr as *const core::ffi::c_void, + n_bytes, + stream.cu_stream(), + ) + .result()?; + } + staging.record_event(stream)?; + staging.sync_event() +} + pub fn async_dtoh_via<'a, T: cudarc::driver::DeviceRepr>( stream: &Arc, slot: &'a Mutex, diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 8a477e1ee..8da80472d 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -40,23 +40,19 @@ fn check_fault_injection() -> Result<()> { Ok(()) } -/// Device-side state across FRI commit iterations. Owns two ext3 eval -/// buffers (flip-flopped as layer input / output) and the inv_twiddles -/// buffer. Freed when dropped. +/// Device-side state across FRI commit iterations. Owns the current fold +/// input (the previous layer's evals, Arc-shared with that layer's retained +/// `gpu_evals`) and the inv_twiddles buffer. Freed when dropped. pub struct FriCommitState { pub stream: Arc, - // Ping-pong evaluation buffers. Both sized `3 * n0` u64 at init. Each - // successive fold uses half the space. Cheap to pre-allocate vs. per- - // layer alloc. - evals_a: CudaSlice, - evals_b: CudaSlice, + /// Current fold input. Each fold allocates a fresh output buffer that is + /// both returned to the caller (kept resident for the query phase) and + /// becomes the next fold's input. + current: Arc>, /// Base-field inv_twiddles; `n0 / 2` u64 at init, halved each layer. inv_tw: CudaSlice, - /// Number of ext3 elements in the buffer currently acting as fold input - /// (`evals_a` or `evals_b`, selected by `a_is_input`). + /// Number of ext3 elements in `current`. pub current_n: usize, - /// Which buffer holds the current layer's input. Toggles each fold. - a_is_input: bool, } impl FriCommitState { @@ -71,20 +67,16 @@ impl FriCommitState { let be = backend()?; let stream = be.next_stream(); - // SAFETY: every byte of evals_a is overwritten by the H2D below. - // evals_b is written by the first fold before it is read. - let mut evals_a = unsafe { stream.alloc::(3 * n0) }?; - let evals_b = unsafe { stream.alloc::(3 * n0) }?; - stream.memcpy_htod(evals_host, &mut evals_a)?; + // SAFETY: every byte of evals is overwritten by the H2D below. + let mut evals = unsafe { stream.alloc::(3 * n0) }?; + stream.memcpy_htod(evals_host, &mut evals)?; let inv_tw = stream.clone_htod(inv_tw_host)?; Ok(Self { stream, - evals_a, - evals_b, + current: Arc::new(evals), inv_tw, current_n: n0, - a_is_input: true, }) } @@ -96,31 +88,32 @@ impl FriCommitState { assert_eq!(buf.len(), 3 * n); assert_eq!(inv_tw_host.len(), n / 2); - // SAFETY: evals_b is written by the first fold before it is read. - let evals_b = unsafe { stream.alloc::(3 * n) }?; let inv_tw = stream.clone_htod(inv_tw_host)?; Ok(Self { stream, - evals_a: buf, - evals_b, + current: Arc::new(buf), inv_tw, current_n: n, - a_is_input: true, }) } - /// Fold the current layer using `zeta`, run the row-pair Keccak leaves - /// + pair-hash Merkle tree kernels on the result, and D2H: - /// - the new root (32 bytes) - /// - the new layer's evals (3 * (current_n / 2) u64s) - /// - the new layer's Merkle tree nodes (standard layout, byte-packed) + /// Fold the current layer using `zeta`, run the row-pair Keccak leaves and + /// pair-hash Merkle tree kernels on the result, and return the layer's + /// evals — device-resident Arc, plus a host copy only when `want_host` — + /// with its resident Merkle tree (root D2H'd, 32 bytes). /// /// Also advances the internal twiddle factors for the next layer. + #[allow(clippy::type_complexity)] pub fn fold_and_commit_layer( &mut self, zeta_raw: [u64; 3], - ) -> Result<(Vec, crate::lde::GpuMerkleTree)> { + want_host: bool, + ) -> Result<( + Option>, + Arc>, + crate::lde::GpuMerkleTree, + )> { #[cfg(feature = "test-faults")] check_fault_injection()?; let be = backend()?; @@ -147,15 +140,11 @@ impl FriCommitState { }; let n_out_u64 = n_out as u64; - // Split the eval buffers into (input, output) based on a_is_input. - // Disjoint-field borrow is fine since evals_a and evals_b are - // separate fields. - let (input_evals, output_evals): (&CudaSlice, &mut CudaSlice) = if self.a_is_input - { - (&self.evals_a, &mut self.evals_b) - } else { - (&self.evals_b, &mut self.evals_a) - }; + // Fresh output buffer per layer: it is retained by the caller for the + // query phase and becomes the next fold's input. + // SAFETY: the fold kernel writes all 3 * n_out slots before any read. + let mut out = unsafe { self.stream.alloc::(3 * n_out) }?; + let input_evals: &CudaSlice = self.current.as_ref(); unsafe { self.stream .launch_builder(&be.fri_fold_ext3) @@ -163,7 +152,7 @@ impl FriCommitState { .arg(&n_out_u64) .arg(&self.inv_tw) .arg(&zeta_dev) - .arg(output_evals) + .arg(&mut out) .launch(cfg)?; } @@ -182,17 +171,10 @@ impl FriCommitState { block_dim: (128, 1, 1), shared_mem_bytes: 0, }; - // Leaves read from the layer's OUTPUT eval buffer (the buffer - // we just wrote to above). - let output_evals: &CudaSlice = if self.a_is_input { - &self.evals_b - } else { - &self.evals_a - }; unsafe { self.stream .launch_builder(&be.keccak_fri_leaves_ext3) - .arg(output_evals) + .arg(&out) .arg(&num_leaves_u64) .arg(&mut leaves_view) .launch(kcfg)?; @@ -225,39 +207,40 @@ impl FriCommitState { self.inv_tw = tw_out; } - // Sync and D2H. - self.stream.synchronize()?; - - // Layer evals: 3 * n_out u64 from the output buffer, staged through - // the per-worker pinned slab (async DMA) instead of a blocking - // pageable copy. The wait is deferred past the root copy below. + // Layer evals to host only when a host copy is wanted (fallback + // consumers), staged through the per-worker pinned slab (async DMA); + // the wait is deferred past the root copy below. let n_evals = 3 * n_out; - let pending = { - let output_evals: &CudaSlice = if self.a_is_input { - &self.evals_b - } else { - &self.evals_a - }; - crate::device::async_dtoh_via( + let pending = if want_host { + Some(crate::device::async_dtoh_via( &self.stream, be.pinned_staging(), &be.ctx, - output_evals, + &out, n_evals, - )? + )?) + } else { + None }; // Keep the layer tree resident on device; copy only the 32-byte root so // R4 query openings gather paths on device instead of copying the tree. - // This pageable copy drains the stream (including the evals DMA above), + // This pageable copy drains the stream (including any evals DMA above), // so the pending wait after it is instant — one block covers both. let mut root = [0u8; 32]; self.stream .memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; - let mut layer_evals = vec![0u64; n_evals]; - pending.wait_into_u64(&mut layer_evals)?; + let layer_evals = match pending { + Some(p) => { + let mut v = vec![0u64; n_evals]; + p.wait_into_u64(&mut v)?; + Some(v) + } + None => None, + }; - self.a_is_input = !self.a_is_input; + let out = Arc::new(out); + self.current = Arc::clone(&out); self.current_n = n_out; let tree = crate::lde::GpuMerkleTree { @@ -265,6 +248,42 @@ impl FriCommitState { leaves_len: num_leaves, root, }; - Ok((layer_evals, tree)) + Ok((layer_evals, out, tree)) + } +} + +/// Gather interleaved ext3 elements at `positions` from a resident evals +/// buffer — a small D2H of only the queried values (the FRI query phase's +/// `evaluation[index ^ 1]` reads). +pub fn gather_ext3_at( + evals: &CudaSlice, + positions: &[u32], + stream: &Arc, +) -> Result> { + let q = positions.len(); + if q == 0 { + return Ok(Vec::new()); + } + let be = backend()?; + let pos_dev = stream.clone_htod(positions)?; + // SAFETY: the gather kernel writes all 3 * q slots. + let mut out_dev = unsafe { stream.alloc::(3 * q) }?; + let cfg = LaunchConfig { + grid_dim: ((q as u32).div_ceil(128), 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + }; + let q_u64 = q as u64; + unsafe { + stream + .launch_builder(&be.gather_ext3_at) + .arg(evals) + .arg(&pos_dev) + .arg(&q_u64) + .arg(&mut out_dev) + .launch(cfg)?; } + let out = stream.clone_dtoh(&out_dev)?; + stream.synchronize()?; + Ok(out) } diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index a59c3950c..4dd49556c 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -78,12 +78,56 @@ pub fn batch_inverse_ext3(a: &[u64]) -> Result> { Ok(out) } +/// `p^3 - 2` as little-endian u64 limbs: the Fermat exponent for inversion in +/// the Goldilocks cubic extension (`|F_{p^3}^*| = p^3 - 1`). +const EXT3_FERMAT_EXP: [u64; 3] = ext3_fermat_exponent(); + +const fn ext3_fermat_exponent() -> [u64; 3] { + const P: u128 = 0xFFFF_FFFF_0000_0001; + let p2 = P * P; + let m0 = ((p2 as u64) as u128) * P; + let m1 = (p2 >> 64) * P + (m0 >> 64); + let l0 = m0 as u64; + // p^3 mod 2^64 ends in ...0001, so subtracting 2 never borrows. + assert!(l0 >= 2); + [l0 - 2, m1 as u64, (m1 >> 64) as u64] +} + +/// One-thread Fermat inversion of `src[n-1]` into `out[0..3]`, stream-ordered. +fn launch_invert_total( + stream: &Arc, + be: &crate::device::Backend, + src: &CudaSlice, + n: usize, + out: &mut CudaSlice, +) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + let n_u64 = n as u64; + let [e0, e1, e2] = EXT3_FERMAT_EXP; + unsafe { + stream + .launch_builder(&be.invert_total_ext3) + .arg(src) + .arg(&n_u64) + .arg(&e0) + .arg(&e1) + .arg(&e2) + .arg(&mut *out) + .launch(cfg)?; + } + Ok(()) +} + /// Device-input batch inverse. Allocates and returns a fresh `CudaSlice` /// of length `3 * n` holding the inverses. Requires `n >= 1`. /// -/// The caller's `stream` is used for every launch and synchronised at the -/// end (so the returned slice's data is committed before this function -/// returns). +/// Stream-ordered end to end: every launch (including the total's Fermat +/// inversion) goes on the caller's `stream`, so downstream same-stream +/// consumers need no synchronize. pub fn batch_inverse_ext3_dev( input: &CudaSlice, n: usize, @@ -101,13 +145,11 @@ pub fn batch_inverse_ext3_dev( )); } if n == 1 { - // Single element: D2H, host invert, H2D. Avoids running the - // scan + combine machinery for a degenerate case. - let host_view: Vec = stream.clone_dtoh(&input.slice(0..3))?; - stream.synchronize()?; - let inv = invert_ext3_host([host_view[0], host_view[1], host_view[2]])?; + // Single element: one-thread Fermat kernel, skipping the scan + + // combine machinery (and any host round-trip). + let be = backend()?; let mut out = unsafe { stream.alloc::(3) }?; - stream.memcpy_htod(&inv, &mut out)?; + launch_invert_total(stream, be, input, 1, &mut out)?; return Ok(out); } @@ -122,12 +164,11 @@ pub fn batch_inverse_ext3_dev( scan_into_fwd(stream, be, input, &mut prefix, n)?; scan_into_rev(stream, be, input, &mut suffix, n)?; - // total = prefix[n-1] = suffix[0]. Invert on host (one Fermat per batch). - let last_host: Vec = stream.clone_dtoh(&prefix.slice((n - 1) * 3..n * 3))?; - stream.synchronize()?; - let inv_total = invert_ext3_host([last_host[0], last_host[1], last_host[2]])?; + // total = prefix[n-1] = suffix[0]. One-thread Fermat inversion on device, + // keeping the whole batch inverse stream-ordered (the host round-trip here + // blocked the calling thread once per batch). let mut inv_total_dev = unsafe { stream.alloc::(3) }?; - stream.memcpy_htod(&inv_total, &mut inv_total_dev)?; + launch_invert_total(stream, be, &prefix, n, &mut inv_total_dev)?; // Combine: out[i] = prefix[i-1] * inv_total * suffix[i+1]. // SAFETY: the combine kernel writes every slot before any read. diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 5f13161aa..ff05e7312 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -262,9 +262,32 @@ fn run_row_major_ntt_body( log_n: u64, m: u64, ) -> Result<()> { + // Levels 0..8 fused in shmem (one DRAM pass instead of eight); the + // remaining high-stride levels keep one kernel per level. + let mut first_level = 0u64; + if n >= 256 { + let t: u32 = 8.min(m as u32).max(1); + let cfg = LaunchConfig { + grid_dim: ((m as u32).div_ceil(t), (n / 256) as u32, 1), + block_dim: (t, 128, 1), + shared_mem_bytes: 256 * (t + 1) * 8, + }; + unsafe { + stream + .launch_builder(&be.ntt_dit_8_levels_row_major) + .arg(&mut *buf) + .arg(tw) + .arg(&n) + .arg(&log_n) + .arg(&m) + .launch(cfg)?; + } + first_level = 8.min(log_n); + } + let col_tile: u32 = 32.min(m as u32); let row_tile: u32 = (256 / col_tile).max(1); - for level in 0..log_n { + for level in first_level..log_n { let cfg = LaunchConfig { grid_dim: ( (m as u32).div_ceil(col_tile), @@ -437,8 +460,15 @@ fn expand_row_major_on_stream( // Fill a zeroed lde_size*total_cols buffer; only the first n*total_cols rows // carry data, the remainder are already zero (zero-padding for LDE). Host // input uploads (H2D); device input copies in place (D2D, no PCIe upload). + // Big host traces go through the pinned staging slot: the driver's + // internal pageable staging is 2-3x slower and convoys across threads. + const PINNED_H2D_MIN_U64: usize = 1 << 20; let mut buf = stream.alloc_zeros::(lde_size * total_cols)?; match input { + InnerInput::Host(h) if h.len() >= PINNED_H2D_MIN_U64 => { + let mut dst = buf.slice_mut(0..n * total_cols); + crate::device::htod_via(stream, be.pinned_staging(), &be.ctx, h, &mut dst)?; + } InnerInput::Host(h) => stream.memcpy_htod(h, &mut buf.slice_mut(0..n * total_cols))?, InnerInput::Dev(d) => stream.memcpy_dtod(d, &mut buf.slice_mut(0..n * total_cols))?, } @@ -694,7 +724,7 @@ pub fn coset_lde_row_major_split_trees( weights: &[u64], split_col: usize, build_precomputed: bool, -) -> Result<(Option>, Vec, GpuLdeBase, Vec)> { +) -> Result<(Option>, GpuLdeBase, Vec)> { assert!(split_col > 0 && split_col < m, "split inside the row"); assert!(n.is_power_of_two(), "n must be a power of two"); assert_eq!(weights.len(), n, "weights length must match n"); @@ -727,7 +757,7 @@ pub fn coset_lde_row_major_split_trees( )?; // One subset tree per column range, built sequentially on the stream. - let build_subset_tree = |col_start: u64, col_end: u64| -> Result> { + let build_subset_tree_dev = |col_start: u64, col_end: u64| -> Result> { let mut nodes_dev = unsafe { stream.alloc::(nodes_bytes) }?; { let mut leaves_view = @@ -745,17 +775,32 @@ pub fn coset_lde_row_major_split_trees( )?; } crate::merkle::build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; - let mut nodes_host = vec![0u8; nodes_bytes]; - stream.memcpy_dtoh(&nodes_dev, &mut nodes_host)?; - Ok(nodes_host) + Ok(nodes_dev) }; + // Precomputed subset tree: full nodes to host (feeds the process-wide + // host tree cache keyed by root; built once per prove on cache miss). let precomputed_nodes = if build_precomputed { - Some(build_subset_tree(0, split_col as u64)?) + let nodes_dev = build_subset_tree_dev(0, split_col as u64)?; + let mut nodes_host = vec![0u8; nodes_bytes]; + stream.memcpy_dtoh(&nodes_dev, &mut nodes_host)?; + Some(nodes_host) } else { None }; - let mult_nodes = build_subset_tree(split_col as u64, cols_u64)?; + // Multiplicity subset tree: resident (per-epoch; the ~2x-leaves node + // download and host rebuild it used to pay are dropped — R4 openings + // gather paths on device). + let mult_tree = { + let nodes_dev = build_subset_tree_dev(split_col as u64, cols_u64)?; + let mut root = [0u8; 32]; + stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; + GpuMerkleTree { + nodes: Arc::new(nodes_dev), + leaves_len: num_leaves, + root, + } + }; // D2H the row-major LDE (preprocessed tables always keep the host copy — // they are excluded from the device-only gate). @@ -778,12 +823,12 @@ pub fn coset_lde_row_major_split_trees( buf: Arc::new(col_major_dev), m, lde_size, - tree: None, + tree: Some(mult_tree), ready: Some(Arc::new(ready)), trace_dev: trace_col_major.map(Arc::new), trace_rows: n, }; - Ok((precomputed_nodes, mult_nodes, handle, lde_out)) + Ok((precomputed_nodes, handle, lde_out)) } /// Row-major ext3 LDE + Keccak + Merkle, all on-device. @@ -1963,10 +2008,11 @@ pub fn coset_lde_batch_ext3_into( /// Batched ext3 coset LDE over columns ALREADY resident on device in slab /// layout (`3m` slabs of `lde_size` u64, first `n` of each filled, rest /// zero-padded), e.g. from the on-device degree-2 decomposition. Runs the -/// same butterfly pipeline as [`coset_lde_batch_ext3_into`], drains the -/// evaluations to `outputs` (interleaved ext3, `3*lde_size` u64 each), and -/// keeps the device buffer as a [`GpuLdeExt3`] handle (synchronized by the -/// drain, so `ready: None`). +/// same butterfly pipeline as [`coset_lde_batch_ext3_into`] and keeps the +/// device buffer as a [`GpuLdeExt3`] handle. With `outputs = Some(..)` the +/// evaluations are also drained to host (interleaved ext3, `3*lde_size` u64 +/// each; the drain synchronizes, so `ready: None`). With `None` nothing +/// leaves the device and the handle carries a `ready` event instead. pub fn coset_lde_batch_ext3_slabs_keep( stream: &Arc, mut buf: CudaSlice, @@ -1974,7 +2020,7 @@ pub fn coset_lde_batch_ext3_slabs_keep( n: usize, blowup_factor: usize, weights: &[u64], - outputs: &mut [&mut [u64]], + outputs: Option<&mut [&mut [u64]]>, ) -> Result { assert!(m > 0 && n.is_power_of_two(), "slab LDE shape"); assert_eq!(weights.len(), n, "weights length must match n"); @@ -1985,9 +2031,11 @@ pub fn coset_lde_batch_ext3_slabs_keep( let lde_size = n * blowup_factor; let mb = 3 * m; assert_eq!(buf.len(), mb * lde_size, "slab buffer shape"); - assert_eq!(outputs.len(), m, "outputs must match column count"); - for o in outputs.iter() { - assert_eq!(o.len(), 3 * lde_size, "each output must be 3*lde_size u64s"); + if let Some(outputs) = outputs.as_ref() { + assert_eq!(outputs.len(), m, "outputs must match column count"); + for o in outputs.iter() { + assert_eq!(o.len(), 3 * lde_size, "each output must be 3*lde_size u64s"); + } } assert_u32_domain(lde_size, "coset_lde_batch_ext3_slabs_keep lde_size"); let log_n = n.trailing_zeros() as u64; @@ -2049,22 +2097,38 @@ pub fn coset_lde_batch_ext3_slabs_keep( mb_u32, )?; - let pending = - crate::device::async_dtoh_via(stream, be.pinned_staging(), &be.ctx, &buf, mb * lde_size)?; - pending.wait_and_read(|bytes| { - // SAFETY: the pinned slab is u64-aligned by construction and the copy - // deposited exactly `mb * lde_size` u64s. - let pinned = - unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const u64, mb * lde_size) }; - unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); - })?; + let ready = match outputs { + Some(outputs) => { + let pending = crate::device::async_dtoh_via( + stream, + be.pinned_staging(), + &be.ctx, + &buf, + mb * lde_size, + )?; + pending.wait_and_read(|bytes| { + // SAFETY: the pinned slab is u64-aligned by construction and the + // copy deposited exactly `mb * lde_size` u64s. + let pinned = unsafe { + std::slice::from_raw_parts(bytes.as_ptr() as *const u64, mb * lde_size) + }; + unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); + })?; + None + } + None => { + let ready = be.take_event()?; + ready.event().record(stream)?; + Some(Arc::new(ready)) + } + }; Ok(GpuLdeExt3 { buf: Arc::new(buf), m, lde_size, tree: None, - ready: None, + ready, }) } diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index bfa756b13..1cb1c5b6f 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -158,10 +158,30 @@ pub(crate) fn build_inner_tree_levels( nodes_dev: &mut CudaSlice, leaves_len: usize, ) -> Result<()> { + // Once a level fits this many pairs, one single-block launch + // (`keccak_merkle_tail`) builds all remaining levels with barriers + // between them: the top ~11 levels of a big tree are each smaller than + // the per-launch overhead they used to pay. + const TAIL_MAX_PAIRS: u64 = 2048; let mut level_begin: u64 = (leaves_len - 1) as u64; while level_begin != 0 { let new_begin = level_begin / 2; let n_pairs = level_begin - new_begin; + if n_pairs <= TAIL_MAX_PAIRS { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (KECCAK_BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.keccak_merkle_tail) + .arg(&mut *nodes_dev) + .arg(&level_begin) + .launch(cfg)?; + } + return Ok(()); + } let cfg = keccak_launch_cfg(n_pairs); unsafe { stream @@ -456,6 +476,56 @@ fn build_comp_poly_tree_nodes_dev( Ok((nodes_dev, num_leaves, stream)) } +/// Build the composition Merkle tree straight from a device-resident slab +/// buffer (`3*m` slabs of `lde_size` u64s, component `k` of part `c` at +/// `(c*3 + k) * lde_size` — the [`crate::lde::GpuLdeExt3`] layout). No host +/// staging and no H2D: the leaves kernel reads `buf` in place on `stream`. +pub fn build_comp_poly_tree_from_slabs_dev( + stream: &Arc, + buf: &CudaSlice, + m: usize, + lde_size: usize, +) -> Result { + assert!(m > 0); + assert!(lde_size.is_power_of_two() && lde_size >= 2); + assert_eq!(buf.len(), 3 * m * lde_size, "slab buffer shape"); + let num_leaves = lde_size / 2; + let tight_total_nodes = 2 * num_leaves - 1; + let be = backend()?; + + let mut nodes_dev = unsafe { stream.alloc::(tight_total_nodes * 32) }?; + let leaves_offset_bytes = (num_leaves - 1) * 32; + { + let mut leaves_view = + nodes_dev.slice_mut(leaves_offset_bytes..leaves_offset_bytes + num_leaves * 32); + let col_stride_u64 = lde_size as u64; + let num_parts_u64 = m as u64; + let num_rows_u64 = lde_size as u64; + let log_num_rows = lde_size.trailing_zeros() as u64; + let cfg = keccak_launch_cfg(num_leaves as u64); + unsafe { + stream + .launch_builder(&be.keccak_comp_poly_leaves_ext3) + .arg(buf) + .arg(&col_stride_u64) + .arg(&num_parts_u64) + .arg(&num_rows_u64) + .arg(&log_num_rows) + .arg(&mut leaves_view) + .launch(cfg)?; + } + } + build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; + let mut root = [0u8; 32]; + stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; + stream.synchronize()?; + Ok(crate::lde::GpuMerkleTree { + nodes: Arc::new(nodes_dev), + leaves_len: num_leaves, + root, + }) +} + /// Build the comp poly Merkle tree on device and keep the nodes resident /// (returned as a [`crate::lde::GpuMerkleTree`] with its root), so R4 /// composition openings gather paths on device instead of copying the whole diff --git a/crypto/stark/src/fri/fri_commitment.rs b/crypto/stark/src/fri/fri_commitment.rs index 58c9eed77..1c199441d 100644 --- a/crypto/stark/src/fri/fri_commitment.rs +++ b/crypto/stark/src/fri/fri_commitment.rs @@ -18,6 +18,11 @@ where /// `merkle_tree` is a root only placeholder. `None` on the CPU path. #[cfg(feature = "cuda")] pub gpu_tree: Option, + /// The layer's evaluations kept resident on device (interleaved ext3, + /// `3 * len` u64). When `evaluation` is empty (device-only), the query + /// phase gathers `evaluation[index ^ 1]` from this buffer instead. + #[cfg(feature = "cuda")] + pub gpu_evals: Option>>, } impl FriLayer @@ -32,6 +37,8 @@ where merkle_tree, #[cfg(feature = "cuda")] gpu_tree: None, + #[cfg(feature = "cuda")] + gpu_evals: None, } } } diff --git a/crypto/stark/src/fri/mod.rs b/crypto/stark/src/fri/mod.rs index c3c16d123..1f53b51cf 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -150,7 +150,7 @@ where (final_poly_coeffs, fri_layer_list) } -pub fn query_phase( +pub fn query_phase( fri_layers: &[FriLayer>], iotas: &[usize], ) -> Vec> diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 2167fcb94..8c87a8a3d 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -572,13 +572,16 @@ where /// Fully device-resident degree-2 decomposition + half extension: takes the /// resident composition evals `H`, decomposes into H0/H1 on device, LDE-extends -/// both, drains the evaluations to host (R3/openings still read them) and -/// keeps the de-interleaved parts buffer as a `GpuLdeExt3` for R4 DEEP. +/// both and keeps the de-interleaved parts buffer as a `GpuLdeExt3` (commit +/// tree, R3 OOD, R4 DEEP and openings all read the handle). With `want_host` +/// the evaluations are also drained to host for the fallback consumers; +/// without it (device-only) the returned part Vecs are empty placeholders. /// `None` → the caller downloads `H` and runs the host decompose path. pub(crate) fn try_decompose_extend_d2_dev( h: &math_cuda::constraint_interp::GpuCompH, inv_2x: &std::sync::Arc>>, weights: &[FieldElement], + want_host: bool, ) -> Option<(Vec>>, math_cuda::lde::GpuLdeExt3)> where F: IsField + 'static, @@ -615,11 +618,26 @@ where GPU_EXTEND_HALVES_CALLS.fetch_add(1, Ordering::Relaxed); GPU_LDE_CALLS.fetch_add(6, Ordering::Relaxed); - let mut lde_h0 = vec![FieldElement::::zero(); lde_size]; - let mut lde_h1 = vec![FieldElement::::zero(); lde_size]; // SAFETY: F == Goldilocks (repr u64); ext3 outputs are [u64; 3] per element. let weights_u64: &[u64] = unsafe { from_raw_parts(weights.as_ptr() as *const u64, weights.len()) }; + + if !want_host { + let handle = math_cuda::lde::coset_lde_batch_ext3_slabs_keep( + &stream, + slabs, + 2, + n, + 2, + weights_u64, + None, + ) + .ok()?; + return Some((vec![Vec::new(), Vec::new()], handle)); + } + + let mut lde_h0 = vec![FieldElement::::zero(); lde_size]; + let mut lde_h1 = vec![FieldElement::::zero(); lde_size]; let ext3_len = lde_size .checked_mul(3) .expect("ext3 output length overflow"); @@ -634,7 +652,7 @@ where n, 2, weights_u64, - &mut outputs, + Some(&mut outputs), ) .ok()?; @@ -800,7 +818,7 @@ where GPU_LEAF_HASH_CALLS.fetch_add(1 + build_precomputed as u64, Ordering::Relaxed); GPU_MERKLE_TREE_CALLS.fetch_add(1 + build_precomputed as u64, Ordering::Relaxed); - let (pre_nodes, mult_nodes, handle, lde_u64) = math_cuda::lde::coset_lde_row_major_split_trees( + let (pre_nodes, handle, lde_u64) = math_cuda::lde::coset_lde_row_major_split_trees( raw, n, m, @@ -815,7 +833,15 @@ where Some(nodes) => Some(tree_from_node_bytes::(nodes)?), None => None, }; - let mult_tree = tree_from_node_bytes::(mult_nodes)?; + // Mult tree resident in the handle: the host tree is root only and R4 + // openings gather authentication paths on device. + let mult_tree = MerkleTree::::from_root( + handle + .tree + .as_ref() + .expect("split path always builds the mult tree") + .root, + ); // Transmute Vec → Vec> (zero-copy, E == GoldilocksField). let lde_out: Vec> = unsafe { @@ -1127,6 +1153,38 @@ where Some((host, dev_tree)) } +/// Device-resident variant of [`try_build_comp_poly_tree_gpu`]: hashes the +/// composition tree straight from the resident R2 parts handle, skipping the +/// host pack + H2D re-upload of data that is already on device. +pub(crate) fn try_build_comp_poly_tree_gpu_from_dev( + handle: &math_cuda::lde::GpuLdeExt3, +) -> Option<(MerkleTree, math_cuda::lde::GpuMerkleTree)> +where + E: IsField + 'static, + B: IsMerkleTreeBackend, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + if handle.m == 0 || !handle.lde_size.is_power_of_two() || handle.lde_size < gpu_lde_threshold() + { + return None; + } + let be = math_cuda::device::backend().ok()?; + let stream = be.next_stream(); + handle.wait_ready_on(&stream).ok()?; + let dev_tree = math_cuda::merkle::build_comp_poly_tree_from_slabs_dev( + &stream, + handle.buf.as_ref(), + handle.m, + handle.lde_size, + ) + .ok()?; + GPU_COMP_POLY_TREE_CALLS.fetch_add(1, Ordering::Relaxed); + let host = MerkleTree::::from_root(dev_tree.root); + Some((host, dev_tree)) +} + /// R3 GPU dispatch: batched strided barycentric OOD evaluation over the main /// (base-field) LDE columns kept on device from R1. Operates on the /// device-resident LDE in place; only the coset points and inv_denoms are @@ -1230,6 +1288,37 @@ pub(crate) fn try_barycentric_ext3_on_handle( inv_denoms_host: &[FieldElement], r3_ctx: Option<(&R3DevContext, usize)>, ) -> Option>> +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + try_barycentric_ext3_on_ext3_handle( + lde_trace.gpu_aux()?, + row_stride, + coset_points, + coset_offset_pow_n, + n_inv, + g_n_inv, + z_pow_n, + inv_denoms_host, + r3_ctx, + ) +} + +/// Same dispatch over an arbitrary resident ext3 handle (aux LDE or the R2 +/// composition parts). One column of OOD sums per handle column. +#[allow(clippy::too_many_arguments)] +pub(crate) fn try_barycentric_ext3_on_ext3_handle( + aux: &math_cuda::lde::GpuLdeExt3, + row_stride: usize, + coset_points: &[FieldElement], + coset_offset_pow_n: &FieldElement, + n_inv: &FieldElement, + g_n_inv: &FieldElement, + z_pow_n: &FieldElement, + inv_denoms_host: &[FieldElement], + r3_ctx: Option<(&R3DevContext, usize)>, +) -> Option>> where F: IsField + IsSubFieldOf + 'static, E: IsField + 'static, @@ -1240,7 +1329,6 @@ where if TypeId::of::() != TypeId::of::() { return None; } - let aux = lde_trace.gpu_aux()?; let num_cols = aux.m; if num_cols == 0 { return Some(Vec::new()); @@ -2137,6 +2225,7 @@ where Ok(s) => s, Err(_) => return None, }; + // Host-evals entry: the caller works with host copies, keep draining them. fri_commit_gpu_drive( state, transcript, @@ -2144,6 +2233,7 @@ where n0, blowup_log, final_poly_log_degree, + true, ) } @@ -2157,6 +2247,7 @@ pub(crate) fn try_fri_commit_gpu_from_dev( blowup_log: u32, final_poly_log_degree: u32, inv_twiddles: &[FieldElement], + want_host: bool, ) -> Option<( Vec>, Vec>>, @@ -2200,6 +2291,7 @@ where n0, blowup_log, final_poly_log_degree, + want_host, ) } @@ -2215,6 +2307,7 @@ fn fri_commit_gpu_drive( n0: usize, blowup_log: u32, final_poly_log_degree: u32, + want_host: bool, ) -> Option<( Vec>, Vec>>, @@ -2266,42 +2359,51 @@ where let zeta_ptr = &zeta as *const FieldElement as *const u64; let zeta_raw: [u64; 3] = unsafe { [*zeta_ptr, *zeta_ptr.add(1), *zeta_ptr.add(2)] }; - let (layer_evals_u64, dev_tree) = match state.fold_and_commit_layer(zeta_raw) { - Ok(v) => v, - Err(_) => { - *transcript = transcript_snapshot.clone(); - return None; - } - }; + let (layer_evals_u64, evals_dev, dev_tree) = + match state.fold_and_commit_layer(zeta_raw, want_host) { + Ok(v) => v, + Err(_) => { + *transcript = transcript_snapshot.clone(); + return None; + } + }; - // Build the FriLayer: ext3 evals and a root only host tree. The layer - // tree stays resident on device in `gpu_tree`; query openings gather - // paths from it via `gather_proofs_dev`. - let evaluation = u64_to_ext3_vec::(&layer_evals_u64); + // Build the FriLayer: a root only host tree, the tree and evals kept + // resident on device (`gpu_tree` / `gpu_evals`), and host evals only + // when a host copy was drained (fallback consumers). + let evaluation = layer_evals_u64 + .map(|v| u64_to_ext3_vec::(&v)) + .unwrap_or_default(); let root = dev_tree.root; let merkle_tree = MerkleTree::>::from_root(root); - let mut layer = FriLayer::new(&evaluation, merkle_tree); - layer.gpu_tree = Some(dev_tree); - fri_layer_list.push(layer); + fri_layer_list.push(FriLayer { + evaluation, + merkle_tree, + gpu_tree: Some(dev_tree), + gpu_evals: Some(evals_dev), + }); // >>>> Send commitment: [p_k] transcript.append_bytes(&root); } // Final (uncommitted) fold to the terminal codeword. n_out == terminal_len - // >= 2, so reuse fold_and_commit_layer and keep only its evaluations; the + // >= 2, so reuse fold_and_commit_layer and keep only its evaluations (the + // coefficient extraction below is host-side, so always drain them); the // Merkle root/nodes are discarded (the terminal layer is sent as coeffs). let zeta_final: FieldElement = transcript.sample_field_element(); let zeta_ptr = &zeta_final as *const FieldElement as *const u64; let zeta_raw: [u64; 3] = unsafe { [*zeta_ptr, *zeta_ptr.add(1), *zeta_ptr.add(2)] }; - let (terminal_evals_u64, _tree) = match state.fold_and_commit_layer(zeta_raw) { + let (terminal_evals_u64, _evals_dev, _tree) = match state.fold_and_commit_layer(zeta_raw, true) + { Ok(v) => v, Err(_) => { *transcript = transcript_snapshot; return None; } }; + let terminal_evals_u64 = terminal_evals_u64.expect("terminal fold drains to host"); debug_assert_eq!(terminal_evals_u64.len(), layout.terminal_len * 3); let terminal_codeword = u64_to_ext3_vec::(&terminal_evals_u64); @@ -2335,7 +2437,7 @@ pub(crate) fn try_fri_query_phase_gpu( iotas: &[usize], ) -> Option>> where - E: IsField, + E: IsField + 'static, FieldElement: AsBytes + Sync + Send, { if fri_layers.is_empty() { @@ -2376,6 +2478,30 @@ where ); } + // Symmetric evals per layer: read the host Vec when it was drained, + // otherwise a batched device gather off the resident layer evals + // (device-only, where no host copy exists). + let per_layer_syms: Vec>>> = fri_layers + .iter() + .enumerate() + .map(|(l, layer)| { + if !layer.evaluation.is_empty() { + return None; + } + let evals_dev = layer + .gpu_evals + .as_ref() + .expect("device-only FRI layer without resident evals"); + let positions: Vec = iotas.iter().map(|&iota| ((iota >> l) ^ 1) as u32).collect(); + let raw = math_cuda::fri::gather_ext3_at(evals_dev, &positions, &stream) + .expect("device FRI sym-eval gather failed; no host fallback"); + Some( + crate::constraint_ir::gpu_interp::ext3_u64_to_field::(&raw) + .expect("resident FRI evals are Goldilocks ext3"), + ) + }) + .collect(); + // Reassemble per-query decommitments, matching the host walk's order. let decommits = iotas .iter() @@ -2385,7 +2511,11 @@ where let mut layers_auth_paths = Vec::with_capacity(num_layers); let mut index = iota; for (l, layer) in fri_layers.iter().enumerate() { - layers_evaluations_sym.push(layer.evaluation[index ^ 1].clone()); + let sym = match &per_layer_syms[l] { + Some(v) => v[q].clone(), + None => layer.evaluation[index ^ 1].clone(), + }; + layers_evaluations_sym.push(sym); layers_auth_paths.push(per_layer_proofs[l][q].clone()); index >>= 1; } @@ -2458,25 +2588,31 @@ mod split_tree_tests { assert_eq!(mult_tree.root, cpu_mult_root, "multiplicity root"); // Openings must be byte-identical at scattered positions (pins the - // full node buffers, not just the roots). + // full node buffers, not just the roots). The mult tree is resident + // (host tree root only), so its paths come from the device gather — + // the exact production opening path. let num_leaves = n * blowup / 2; + let dev_tree = handle.tree.as_ref().expect("resident mult subset tree"); + let stream = math_cuda::device::backend().unwrap().next_stream(); for pos in [0usize, 1, 511, 12_345, num_leaves - 1] { assert_eq!( pre_tree.get_proof_by_pos(pos).unwrap().merkle_path, cpu_pre.get_proof_by_pos(pos).unwrap().merkle_path, "precomputed path at {pos}" ); + let dev_proofs = + gather_proofs_dev(dev_tree, &[pos], &stream).expect("device mult-tree path gather"); assert_eq!( - mult_tree.get_proof_by_pos(pos).unwrap().merkle_path, + dev_proofs[0].merkle_path, cpu_mult.get_proof_by_pos(pos).unwrap().merkle_path, "multiplicity path at {pos}" ); } + assert_eq!(mult_tree.root, dev_tree.root, "root-only host tree root"); // The handle must carry the column-major LDE for downstream rounds: // spot-check a few cells against the row-major host LDE. assert_eq!(handle.m, m); assert_eq!(handle.lde_size, n * blowup); - assert!(handle.tree.is_none(), "no device tree on the split path"); } } diff --git a/crypto/stark/src/logup_gpu.rs b/crypto/stark/src/logup_gpu.rs index 3fd49134d..9aed7c026 100644 --- a/crypto/stark/src/logup_gpu.rs +++ b/crypto/stark/src/logup_gpu.rs @@ -415,9 +415,10 @@ where /// straight to the aux LDE, no host round-trip) + the table contribution `L`. /// Returns `None` to fall back (non Goldilocks, below threshold, no GPU, GPU /// error). This is the residency path that avoids the term-column download. -pub fn try_build_aux_resident_gpu( +pub fn try_build_aux_resident_gpu<'a, F, E>( interactions: &[BusInteraction], - main_cols: &[Vec>], + num_cols: usize, + main_cols: impl FnOnce() -> &'a [Vec>], main_dev: Option<(&math_cuda::CudaSlice, usize)>, trace_len: usize, challenges: &[FieldElement], @@ -431,7 +432,7 @@ where { return None; } - if trace_len < GPU_LOGUP_MIN_ROWS || main_cols.is_empty() || interactions.is_empty() { + if trace_len < GPU_LOGUP_MIN_ROWS || num_cols == 0 || interactions.is_empty() { return None; } if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() { @@ -442,18 +443,18 @@ where return None; } - let num_cols = main_cols.len(); desc.assert_columns_in_bounds(num_cols); // Reuse the resident main trace from the R1 main LDE (column-major - // `[col*trace_len + row]`, same column order as `main_cols`) when it matches - // this table exactly; otherwise flatten + upload the host columns. The - // resident buffer skips the ~3 GB main re-upload. + // `[col*trace_len + row]`, same column order as the host columns) when it + // matches this table exactly; otherwise materialize + flatten + upload the + // host columns. The resident buffer skips both the host transpose and the + // ~3 GB main re-upload. let resident_main = main_dev.filter(|&(buf, rows)| rows == trace_len && buf.len() == num_cols * trace_len); let mut main_flat = Vec::new(); if resident_main.is_none() { main_flat = vec![0u64; num_cols * trace_len]; - for (c, col) in main_cols.iter().enumerate() { + for (c, col) in main_cols().iter().enumerate() { for (r, e) in col.iter().enumerate() { main_flat[c * trace_len + r] = unsafe { *(e.value() as *const _ as *const u64) }; } diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index d376ebd1f..698e89ba9 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -1155,8 +1155,10 @@ where return None; } - // Clone main columns once (shared across all interactions) - let main_segment_cols = trace.columns_main(); + // Host main columns, materialized lazily: the resident GPU aux path + // reads the device main in place and must not pay this transpose. + let main_cols_cell: std::cell::OnceCell>>> = + std::cell::OnceCell::new(); let trace_len = trace.num_rows(); let _table_name = self.name.as_deref().unwrap_or("UNKNOWN"); @@ -1188,7 +1190,12 @@ where if trace.resident_aux_ok() && let Some(ra) = crate::logup_gpu::try_build_aux_resident_gpu::( interactions, - &main_segment_cols, + trace.num_main_columns, + || { + main_cols_cell + .get_or_init(|| trace.columns_main()) + .as_slice() + }, resident_main.as_ref().map(|r| (r.buf.as_ref(), r.rows)), trace_len, challenges, @@ -1201,12 +1208,14 @@ where return Some(BusPublicInputs { table_contribution }); } + let main_segment_cols = main_cols_cell.get_or_init(|| trace.columns_main()); + // GPU aux build (Goldilocks + ext3 + above threshold) computes all term // columns on device, byte identical, and falls back to the CPU build. #[cfg(feature = "cuda")] let gpu_term_cols = crate::logup_gpu::try_build_term_columns_gpu::( interactions, - &main_segment_cols, + main_segment_cols, trace_len, challenges, ); @@ -1220,7 +1229,7 @@ where let build_pair = |i: usize| { compute_logup_term_column( &[&interactions[i * 2], &interactions[i * 2 + 1]], - &main_segment_cols, + main_segment_cols, trace_len, challenges, _table_name, @@ -1248,7 +1257,7 @@ where &interactions[num_interactions - 2], &interactions[num_interactions - 1], ], - &main_segment_cols, + main_segment_cols, trace_len, challenges, _table_name, @@ -1256,7 +1265,7 @@ where } else { compute_logup_term_column( &[&interactions[num_interactions - 1]], - &main_segment_cols, + main_segment_cols, trace_len, challenges, _table_name, diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 9a369b042..543ebf989 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1481,10 +1481,12 @@ pub trait IsStarkProver< let mut gpu_composition_parts: Option = None; // Fully device-resident d=2 path: H stays on device through decompose + - // half extension, the parts handle feeds R4 DEEP, and only the final - // evaluations are drained to host (for the commit tree and openings). - // Any miss falls through to the host path below (downloading H when - // the evaluation itself already ran on device). + // half extension, and the parts handle feeds the commit tree, R3 OOD, + // R4 DEEP and the openings. The evaluations are drained to host only + // while a host trace copy exists (fallback consumers); under + // device-only nothing leaves the device and the placeholders below + // stay empty. Any miss falls through to the host path (downloading H + // when the evaluation itself already ran on device). #[cfg(feature = "cuda")] let mut precomputed_parts: Option>>> = None; #[cfg(feature = "cuda")] @@ -1502,6 +1504,7 @@ pub trait IsStarkProver< &h_dev, twiddles.inv_2x(domain), &twiddles.composition(domain).weights, + !round_1_result.lde_trace.host_trace_empty(), ) { Some((parts, handle)) => { gpu_composition_parts = Some(handle); @@ -1612,18 +1615,28 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let t_sub = Instant::now(); - // GPU fast path for the comp-poly Merkle commit: row-pair Keccak - // leaves + device-side inner tree, both wrapping the host eval Vecs. - // GPU path keeps the composition tree resident on device (no whole tree - // copy) and returns a root only host tree. The device tree is threaded - // to R4 in `Round2.gpu_composition_tree`. + // GPU fast path for the comp-poly Merkle commit: hash straight from + // the resident parts handle when R2 kept one (no host pack + H2D + // re-upload); otherwise wrap the host eval Vecs. Either way the tree + // stays resident on device (no whole-tree copy), a root-only host tree + // is returned, and the device tree is threaded to R4 in + // `Round2.gpu_composition_tree`. #[cfg(feature = "cuda")] let (composition_poly_merkle_tree, composition_poly_root, gpu_composition_tree) = - match crate::gpu_lde::try_build_comp_poly_tree_gpu::< - FieldExtension, - BatchedMerkleTreeBackend, - >(&lde_composition_poly_parts_evaluations) - { + match gpu_composition_parts + .as_ref() + .and_then(|h| { + crate::gpu_lde::try_build_comp_poly_tree_gpu_from_dev::< + FieldExtension, + BatchedMerkleTreeBackend, + >(h) + }) + .or_else(|| { + crate::gpu_lde::try_build_comp_poly_tree_gpu::< + FieldExtension, + BatchedMerkleTreeBackend, + >(&lde_composition_poly_parts_evaluations) + }) { Some((host_tree, dev_tree)) => { let root = host_tree.root; (host_tree, root, Some(dev_tree)) @@ -1688,27 +1701,86 @@ pub trait IsStarkProver< // === Composition poly parts: barycentric evaluation at z^num_parts === let comp_z_pow_n = z_power.pow(domain_size); - let comp_inv_denoms = math::polynomial::barycentric_inv_denoms(&z_power, &dc.points); - let composition_poly_parts_ood_evaluation: Vec<_> = round_2_result - .lde_composition_poly_evaluations - .iter() - .map(|lde_evals| { - // Extract trace-size evaluations (stride = blowup_factor) - let evals: Vec> = (0..domain_size) - .map(|i| lde_evals[i * blowup_factor].clone()) - .collect(); - math::polynomial::interpolate_coset_eval_ext_with_g_n_inv( - &comp_z_pow_n, - &dc.offset_pow_n, - &dc.size_inv, - &dc.offset_pow_n_inv, - &dc.points, - &evals, - &comp_inv_denoms, - ) - }) - .collect(); + // GPU fast path: strided barycentric straight over the resident R2 + // parts handle (device inv_denoms for the single point z^P), skipping + // the host stride-extract and the sequential CPU fold per part. + #[cfg(feature = "cuda")] + let gpu_parts_ood: Option>> = + round_1_result + .lde_trace + .gpu_composition_parts() + .and_then(|parts_dev| { + let dispatch = |inv_host: &[FieldElement], + ctx: Option<(&crate::gpu_lde::R3DevContext, usize)>| { + crate::gpu_lde::try_barycentric_ext3_on_ext3_handle::( + parts_dev, + blowup_factor, + &dc.points, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &comp_z_pow_n, + inv_host, + ctx, + ) + }; + match crate::gpu_lde::try_prep_r3_dev_context::( + &dc.points, + std::slice::from_ref(&z_power), + round_1_result.lde_trace.bound_stream(), + ) { + Some(ctx) => dispatch(&[], Some((&ctx, 0))), + // Below the dev-context threshold (single eval point): + // host inv_denoms + the same strided kernel, mirroring the + // trace OOD's mixed arm. + None => { + let inv = + math::polynomial::barycentric_inv_denoms(&z_power, &dc.points); + dispatch(&inv, None) + } + } + }); + #[cfg(not(feature = "cuda"))] + let gpu_parts_ood: Option>> = None; + + let composition_poly_parts_ood_evaluation: Vec<_> = match gpu_parts_ood { + Some(v) => v, + None => { + // The host part evals are empty under device-only (the R2 + // drain is skipped); reaching this arm there is a mis-gate. + #[cfg(feature = "cuda")] + assert!( + round_2_result + .lde_composition_poly_evaluations + .first() + .is_none_or(|p| !p.is_empty()), + "R3 parts OOD fell back to the host part evals, but they are \ + device-only (empty)" + ); + let comp_inv_denoms = + math::polynomial::barycentric_inv_denoms(&z_power, &dc.points); + round_2_result + .lde_composition_poly_evaluations + .iter() + .map(|lde_evals| { + // Extract trace-size evaluations (stride = blowup_factor) + let evals: Vec> = (0..domain_size) + .map(|i| lde_evals[i * blowup_factor].clone()) + .collect(); + math::polynomial::interpolate_coset_eval_ext_with_g_n_inv( + &comp_z_pow_n, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &dc.points, + &evals, + &comp_inv_denoms, + ) + }) + .collect() + } + }; // === Trace polynomials: barycentric evaluation via LDE === let trace_ood_evaluations = crate::trace::get_trace_evaluations_from_lde( @@ -1812,6 +1884,7 @@ pub trait IsStarkProver< domain.blowup_factor.trailing_zeros(), air.options().fri_final_poly_log_degree as u32, domain.fri_inv_twiddles(), + !round_1_result.lde_trace.host_trace_empty(), ) }); #[cfg(not(feature = "cuda"))] @@ -2496,23 +2569,21 @@ pub trait IsStarkProver< // must succeed: there is no host tree to fall back to, so a gather error // is a hard abort. When the tree is not device resident the value is // `None` and the openings below walk the full host tree. + // For preprocessed tables the resident tree is the multiplicity subset + // tree (the host `main_commit.tree` is root only); values still come + // from the host LDE range gather below. #[cfg(feature = "cuda")] - let main_dev_proofs: Option>> = if is_preprocessed { - None - } else { - lde_trace - .gpu_main() - .and_then(|h| h.tree.as_ref()) - .map(|tree| { - let stream = lde_trace - .bound_stream() - .expect("bound stream for device-resident main-tree opening"); - // Row-pair leaves: one proof per query at position `challenge`. - crate::gpu_lde::gather_proofs_dev(tree, indexes_to_open, &stream).expect( - "device main-tree gather failed; resident tree has no host fallback", - ) - }) - }; + let main_dev_proofs: Option>> = lde_trace + .gpu_main() + .and_then(|h| h.tree.as_ref()) + .map(|tree| { + let stream = lde_trace + .bound_stream() + .expect("bound stream for device-resident main-tree opening"); + // Row-pair leaves: one proof per query at position `challenge`. + crate::gpu_lde::gather_proofs_dev(tree, indexes_to_open, &stream) + .expect("device main-tree gather failed; resident tree has no host fallback") + }); // Same for the aux trace tree, when it is device resident. #[cfg(feature = "cuda")] @@ -2554,8 +2625,10 @@ pub trait IsStarkProver< // *_dev_values.is_some()` on the Goldilocks path) and we never gather // rows for a tree that is not device resident. #[cfg(feature = "cuda")] - let main_dev_values: Option>> = - main_dev_proofs.as_ref().and_then(|_| { + let main_dev_values: Option>> = (!is_preprocessed) + .then_some(()) + .and(main_dev_proofs.as_ref()) + .and_then(|_| { lde_trace.gpu_main().and_then(|h| { Self::gather_query_rows_device( lde_trace, @@ -2595,12 +2668,65 @@ pub trait IsStarkProver< }) }); + // Composition part values off the resident R2 parts handle (one ext3 + // "column" per part), same row-pair gather as main/aux above. + #[cfg(feature = "cuda")] + let comp_num_parts = lde_trace + .gpu_composition_parts() + .map(|h| h.m) + .unwrap_or_else(|| round_2_result.lde_composition_poly_evaluations.len()); + #[cfg(feature = "cuda")] + let comp_dev_values: Option>> = + comp_dev_proofs.as_ref().and_then(|_| { + lde_trace.gpu_composition_parts().and_then(|h| { + Self::gather_query_rows_device( + lde_trace, + "composition", + |stream| { + math_cuda::barycentric::gather_rows_ext3_on_device( + h, + &query_rows, + stream, + ) + }, + |raw| { + crate::constraint_ir::gpu_interp::ext3_u64_to_field::( + raw, + ) + }, + ) + }) + }); + for (qi, index) in indexes_to_open.iter().enumerate() { #[cfg(not(feature = "cuda"))] let _ = qi; // For preprocessed tables, open the main split (multiplicities only); // for normal tables, open all main columns. let main_trace_opening = if is_preprocessed { + // Multiplicity subset: device proof (resident subset tree) + + // host range gather for the values. + #[cfg(feature = "cuda")] + { + match &main_dev_proofs { + Some(proofs) => Self::open_polys_with_proofs( + domain, + proofs[qi].clone(), + *index, + |row| { + lde_trace.gather_main_row_range( + row, + num_precomputed_cols, + total_cols, + ) + }, + ), + None => Self::open_polys_with(domain, &main_commit.tree, *index, |row| { + lde_trace.gather_main_row_range(row, num_precomputed_cols, total_cols) + }), + } + } + #[cfg(not(feature = "cuda"))] Self::open_polys_with(domain, &main_commit.tree, *index, |row| { lde_trace.gather_main_row_range(row, num_precomputed_cols, total_cols) }) @@ -2638,18 +2764,57 @@ pub trait IsStarkProver< let composition_openings = { #[cfg(feature = "cuda")] { - if let Some(proofs) = &comp_dev_proofs { - Self::open_composition_poly_with_proof( - proofs[qi].clone(), - &round_2_result.lde_composition_poly_evaluations, - *index, - ) - } else { - Self::open_composition_poly( + match (&comp_dev_proofs, &comp_dev_values) { + (Some(proofs), Some(vals)) => { + let (even, odd) = Self::device_row_pair(vals, qi, comp_num_parts); + // Cross-check against the host part evals while + // they are still resident (absent under full + // residency, where the gather is the only source). + if round_2_result + .lde_composition_poly_evaluations + .first() + .is_some_and(|p| !p.is_empty()) + { + let expected = Self::open_composition_poly_with_proof( + proofs[qi].clone(), + &round_2_result.lde_composition_poly_evaluations, + *index, + ); + assert_eq!( + even, expected.evaluations, + "device composition-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, expected.evaluations_sym, + "device composition-row gather mismatch (odd), query {qi}" + ); + } + PolynomialOpenings { + proof: proofs[qi].clone(), + evaluations: even, + evaluations_sym: odd, + } + } + (Some(proofs), None) => { + assert!( + round_2_result + .lde_composition_poly_evaluations + .first() + .is_none_or(|p| !p.is_empty()), + "R4 composition opening fell back to the host part evals, \ + but they are device-only (empty)" + ); + Self::open_composition_poly_with_proof( + proofs[qi].clone(), + &round_2_result.lde_composition_poly_evaluations, + *index, + ) + } + _ => Self::open_composition_poly( &round_2_result.composition_poly_merkle_tree, &round_2_result.lde_composition_poly_evaluations, *index, - ) + ), } } #[cfg(not(feature = "cuda"))] From ac3dee9e36f7821c74648f61161682804583fe9c Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 29 Jul 2026 15:43:09 -0300 Subject: [PATCH 2/6] fix(gpu): harden round-2 residency paths after review Grid-stride the fused row-major NTT past gridDim.y (lde >= 2^24 silently fell back to CPU), assert the device-only contract in the R2 composition commit and preprocessed opening fallbacks, validate htod_via bounds, retain FRI device evals only under device-only, and move the inverse fault-injection hook so every batch-inverse entry is covered. --- crypto/crypto/src/merkle_tree/merkle.rs | 13 ++++++ crypto/math-cuda/kernels/ntt.cu | 61 ++++++++++++++----------- crypto/math-cuda/src/device.rs | 33 ++++++++----- crypto/math-cuda/src/inverse.rs | 9 +++- crypto/math-cuda/src/lde.rs | 19 ++++---- crypto/stark/src/gpu_lde.rs | 26 ++++++----- crypto/stark/src/lookup.rs | 2 +- crypto/stark/src/prover.rs | 28 ++++++++++-- 8 files changed, 125 insertions(+), 66 deletions(-) diff --git a/crypto/crypto/src/merkle_tree/merkle.rs b/crypto/crypto/src/merkle_tree/merkle.rs index d53f06f10..e9ecd0bb5 100644 --- a/crypto/crypto/src/merkle_tree/merkle.rs +++ b/crypto/crypto/src/merkle_tree/merkle.rs @@ -172,6 +172,19 @@ where /// but no nodes. Used when paths are gathered from a device resident copy /// (GPU) instead of this host tree, so the host nodes are never built. /// [`get_proof_by_pos`](Self::get_proof_by_pos) must NOT be called on it. + /// True when this tree carries only its root (the nodes live elsewhere, + /// e.g. device-resident): openings must not walk this tree. + pub fn is_root_only(&self) -> bool { + #[cfg(feature = "disk-spill")] + { + self.nodes.is_empty() && self.mmap_backing.is_none() + } + #[cfg(not(feature = "disk-spill"))] + { + self.nodes.is_empty() + } + } + pub fn from_root(root: B::Node) -> Self { MerkleTree { root, diff --git a/crypto/math-cuda/kernels/ntt.cu b/crypto/math-cuda/kernels/ntt.cu index 35a4f7b20..1e6c83f5c 100644 --- a/crypto/math-cuda/kernels/ntt.cu +++ b/crypto/math-cuda/kernels/ntt.cu @@ -429,42 +429,49 @@ extern "C" __global__ void ntt_dit_8_levels_row_major(uint64_t *data, uint32_t pitch = T + 1; uint64_t col = (uint64_t)blockIdx.x * T + threadIdx.x; bool live = col < m; - uint64_t row_base = (uint64_t)blockIdx.y * 256; - - for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { - if (live) tile[r * pitch + threadIdx.x] = data[(row_base + r) * m + col]; - } - __syncthreads(); uint32_t n_loc_steps = (uint32_t)min((uint64_t)8, log_n); uint32_t remaining_high_bits = (uint32_t)(log_n - 1); uint32_t high_mask = (1u << remaining_high_bits) - 1u; - for (uint32_t loc_step = 0; loc_step < n_loc_steps; ++loc_step) { - for (uint32_t i = threadIdx.y; i < 128; i += blockDim.y) { - uint32_t half = 1u << loc_step; - uint32_t grp = i >> loc_step; - uint32_t grp_pos = i & (half - 1); - uint32_t idx1 = (grp << (loc_step + 1)) + grp_pos; - uint32_t idx2 = idx1 + half; + // Grid-stride over 256-row blocks: gridDim.y caps at 65535, so lde sizes + // >= 2^24 need more than one row block per y-slot. The trip count is + // uniform across the block, keeping every __syncthreads converged. + for (uint64_t rb = blockIdx.y; rb < (n >> 8); rb += gridDim.y) { + uint64_t row_base = rb * 256; - uint32_t gs = loc_step; - uint32_t ggp = ((uint32_t)blockIdx.y << 7) + i; - ggp = (ggp & high_mask) + (ggp >> remaining_high_bits); - ggp = ggp & ((1u << gs) - 1u); - uint64_t factor = tw[(uint64_t)ggp * (n >> (gs + 1))]; + for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { + if (live) tile[r * pitch + threadIdx.x] = data[(row_base + r) * m + col]; + } + __syncthreads(); - if (live) { - uint64_t u = tile[idx1 * pitch + threadIdx.x]; - uint64_t v = mul(tile[idx2 * pitch + threadIdx.x], factor); - tile[idx1 * pitch + threadIdx.x] = add(u, v); - tile[idx2 * pitch + threadIdx.x] = sub(u, v); + for (uint32_t loc_step = 0; loc_step < n_loc_steps; ++loc_step) { + for (uint32_t i = threadIdx.y; i < 128; i += blockDim.y) { + uint32_t half = 1u << loc_step; + uint32_t grp = i >> loc_step; + uint32_t grp_pos = i & (half - 1); + uint32_t idx1 = (grp << (loc_step + 1)) + grp_pos; + uint32_t idx2 = idx1 + half; + + uint32_t gs = loc_step; + uint32_t ggp = ((uint32_t)rb << 7) + i; + ggp = (ggp & high_mask) + (ggp >> remaining_high_bits); + ggp = ggp & ((1u << gs) - 1u); + uint64_t factor = tw[(uint64_t)ggp * (n >> (gs + 1))]; + + if (live) { + uint64_t u = tile[idx1 * pitch + threadIdx.x]; + uint64_t v = mul(tile[idx2 * pitch + threadIdx.x], factor); + tile[idx1 * pitch + threadIdx.x] = add(u, v); + tile[idx2 * pitch + threadIdx.x] = sub(u, v); + } } + __syncthreads(); } - __syncthreads(); - } - for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { - if (live) data[(row_base + r) * m + col] = tile[r * pitch + threadIdx.x]; + for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { + if (live) data[(row_base + r) * m + col] = tile[r * pitch + threadIdx.x]; + } + __syncthreads(); } } diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 292026401..c3dd4a2a9 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -611,7 +611,9 @@ pub fn backend() -> Result<&'static Backend> { /// /// Holding this value keeps the staging slot's mutex locked, which is what /// makes the whole scheme safe: no other caller (and no capacity growth) can -/// touch the slab while the DMA is in flight. +/// touch the slab while the DMA is in flight. Corollary: never call +/// `htod_via`/`async_dtoh_via` on the same slot from the thread holding a +/// live `PendingD2H` — the non-reentrant slot mutex self-deadlocks. pub struct PendingD2H<'a> { staging: std::sync::MutexGuard<'a, PinnedStaging>, n_bytes: usize, @@ -628,17 +630,6 @@ impl Drop for PendingD2H<'_> { } } -/// Enqueue an async D2H of `n_elems` of `src` into the pinned slab of `slot`, -/// without synchronizing the stream. Unlike `stream.memcpy_dtoh` into a plain -/// (pageable) slice — which the driver services synchronously — this returns -/// as soon as the copy is queued; the returned [`PendingD2H`] is awaited at -/// the point the host actually needs the bytes. -/// -/// SAFETY contract (upheld by construction for our callers): `src` must stay -/// alive until the copy completes. Dropping a `CudaSlice` frees it -/// stream-ordered on its own stream, so a `src` allocated on `stream` may be -/// dropped after this call — the free queues behind the copy. Do NOT pass a -/// `src` owned by a *different* stream and drop it before waiting. /// Host→device copy staged through the pinned slot: one host memcpy into /// pinned memory + one async DMA, instead of the driver's internal pageable /// staging (small chunks; 2-3x slower for multi-hundred-MB traces and it @@ -652,7 +643,14 @@ pub fn htod_via( dst: &mut cudarc::driver::CudaViewMut<'_, T>, ) -> Result<()> { use cudarc::driver::DevicePtrMut; + assert!( + dst.len() >= src_host.len(), + "htod_via: destination shorter than source" + ); let n_bytes = std::mem::size_of_val(src_host); + if n_bytes == 0 { + return Ok(()); + } let u64_len = n_bytes.div_ceil(8); let mut staging = slot.lock().unwrap(); staging.ensure_capacity(u64_len, ctx)?; @@ -680,6 +678,17 @@ pub fn htod_via( staging.sync_event() } +/// Enqueue an async D2H of `n_elems` of `src` into the pinned slab of `slot`, +/// without synchronizing the stream. Unlike `stream.memcpy_dtoh` into a plain +/// (pageable) slice — which the driver services synchronously — this returns +/// as soon as the copy is queued; the returned [`PendingD2H`] is awaited at +/// the point the host actually needs the bytes. +/// +/// SAFETY contract (upheld by construction for our callers): `src` must stay +/// alive until the copy completes. Dropping a `CudaSlice` frees it +/// stream-ordered on its own stream, so a `src` allocated on `stream` may be +/// dropped after this call — the free queues behind the copy. Do NOT pass a +/// `src` owned by a *different* stream and drop it before waiting. pub fn async_dtoh_via<'a, T: cudarc::driver::DeviceRepr>( stream: &Arc, slot: &'a Mutex, diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 4dd49556c..9931e3a8c 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -94,6 +94,11 @@ const fn ext3_fermat_exponent() -> [u64; 3] { } /// One-thread Fermat inversion of `src[n-1]` into `out[0..3]`, stream-ordered. +/// +/// Unlike the host Fermat this used to call, a zero total maps silently to +/// zero instead of `Err`. Unreachable with honest inputs (LogUp/barycentric +/// denominators are nonzero w.h.p.); callers must not rely on a zero-total +/// error. fn launch_invert_total( stream: &Arc, be: &crate::device::Backend, @@ -133,6 +138,8 @@ pub fn batch_inverse_ext3_dev( n: usize, stream: &Arc, ) -> Result> { + #[cfg(feature = "test-faults")] + check_inverse_fault_injection()?; 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 @@ -220,8 +227,6 @@ pub fn compute_and_invert_denoms_ext3_dev( sign: DenomSign, stream: &Arc, ) -> Result> { - #[cfg(feature = "test-faults")] - check_inverse_fault_injection()?; assert_eq!(z_scalars_host.len(), k_scalars * 3); assert!(n >= 1 && k_scalars >= 1); diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index ff05e7312..3d8bfa207 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -268,7 +268,7 @@ fn run_row_major_ntt_body( if n >= 256 { let t: u32 = 8.min(m as u32).max(1); let cfg = LaunchConfig { - grid_dim: ((m as u32).div_ceil(t), (n / 256) as u32, 1), + grid_dim: ((m as u32).div_ceil(t), ((n / 256) as u32).min(65535), 1), block_dim: (t, 128, 1), shared_mem_bytes: 256 * (t + 1) * 8, }; @@ -704,17 +704,16 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( /// `[split_col, m)` commit to separate trees over the same row-major LDE, /// mirroring the CPU `commit_rows_bit_reversed_subset` pair. /// -/// Both trees' complete node buffers are downloaded to host -/// (`(2*num_leaves - 1) * 32` bytes each, inner nodes first, root at offset 0, +/// The precomputed tree's complete node buffer is downloaded to host +/// (`(2*num_leaves - 1) * 32` bytes, inner nodes first, root at offset 0, /// leaves at the tail — the exact `MerkleTree::from_precomputed_nodes` -/// layout), because preprocessed-table openings walk host trees. The -/// precomputed tree is only built when `build_precomputed` is true (the -/// caller skips it on a process-cache hit). +/// layout) because it feeds the process-wide host tree cache; it is only +/// built when `build_precomputed` is true (the caller skips it on a cache +/// hit). The multiplicity tree stays resident in `handle.tree` — openings +/// gather its paths on device. /// -/// Returns `(precomputed_nodes, mult_nodes, handle, row_major_lde)`. The -/// handle carries the column-major LDE + trace snapshot for downstream GPU -/// rounds but NO device tree (`tree: None`) — openings for preprocessed -/// tables never gather from device. +/// Returns `(precomputed_nodes, handle, row_major_lde)`. The handle also +/// carries the column-major LDE + trace snapshot for downstream GPU rounds. #[allow(clippy::type_complexity)] pub fn coset_lde_row_major_split_trees( row_major: &[u64], diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 8c87a8a3d..98830fcc7 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -767,10 +767,11 @@ where /// one row-major GPU LDE of ALL columns plus TWO subset Merkle trees — the /// precomputed columns `[0, split_col)` and the multiplicity columns /// `[split_col, m)` — matching the CPU `commit_rows_bit_reversed_subset` -/// pair bit for bit. Trees come back as full HOST trees (openings for -/// preprocessed tables walk host trees); the handle keeps the column-major -/// LDE + trace snapshot device-resident for the downstream GPU rounds, with -/// no device tree. +/// pair bit for bit. The precomputed tree comes back as a full HOST tree +/// (it feeds the process-wide cache); the multiplicity tree stays resident +/// in the handle (root-only host tree, R4 openings gather paths on device). +/// The handle also keeps the column-major LDE + trace snapshot for the +/// downstream GPU rounds. /// /// `build_precomputed=false` skips the precomputed tree (process-cache hit); /// the first element is then `None`. @@ -1405,12 +1406,12 @@ pub fn gpu_fri_calls() -> u64 { /// Batch-invert dispatch counter (one per /// [`try_compute_and_invert_inv_denoms_dev`] call that actually built a -/// device handle). Fires at most twice per prove per table: once for R3 -/// OOD's `num_eval_points * trace_size` denominators and once for R4 -/// DEEP's `(1 + num_eval_points) * lde_size` denominators. R4 has two -/// chances at it (device-only DEEP, then the host DEEP arm), and both are -/// counted here, so a single failed dispatch does not necessarily lower the -/// total; R3's fallback is CPU-only, so a failure there does. +/// device handle). Fires up to three times per prove per table: R3 trace +/// OOD's `num_eval_points * trace_size` denominators, R3 parts OOD's single +/// point, and R4 DEEP's `(1 + num_eval_points) * lde_size` denominators. R4 +/// has two chances at it (device-only DEEP, then the host DEEP arm), and both +/// are counted here, so a single failed dispatch does not necessarily lower +/// the total; R3's fallbacks are CPU-only, so a failure there does. pub(crate) static GPU_BATCH_INVERT_CALLS: AtomicU64 = AtomicU64::new(0); pub fn gpu_batch_invert_calls() -> u64 { GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed) @@ -2376,11 +2377,14 @@ where .unwrap_or_default(); let root = dev_tree.root; let merkle_tree = MerkleTree::>::from_root(root); + // Retain the device evals only when no host copy exists (device-only): + // with a host copy the query phase reads it, and the retained buffer + // would be ~24 bytes/LDE-row of dead VRAM per table. fri_layer_list.push(FriLayer { evaluation, merkle_tree, gpu_tree: Some(dev_tree), - gpu_evals: Some(evals_dev), + gpu_evals: (!want_host).then_some(evals_dev), }); // >>>> Send commitment: [p_k] diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index 698e89ba9..ceda5417a 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -1286,7 +1286,7 @@ where let (per_bus_sums, per_bus_sender_sums, per_bus_receiver_sums) = compute_debug_bus_sums_batched( &self.auxiliary_trace_build_data.interactions, - &main_segment_cols, + main_segment_cols, trace_len, challenges, _table_name, diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 543ebf989..8612bb20a 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1642,6 +1642,14 @@ pub trait IsStarkProver< (host_tree, root, Some(dev_tree)) } None => { + // The host part evals are empty under device-only (the R2 + // drain is skipped); abort with the device-only contract's + // message instead of a misleading EmptyCommitment. + assert!( + !round_1_result.lde_trace.host_trace_empty(), + "R2 composition commit fell back to the host part evals, \ + but they are device-only (empty)" + ); let (tree, root) = crate::commitment::commit_bit_reversed( &lde_composition_poly_parts_evaluations, crate::commitment::ROWS_PER_LEAF, @@ -2721,9 +2729,23 @@ pub trait IsStarkProver< ) }, ), - None => Self::open_polys_with(domain, &main_commit.tree, *index, |row| { - lde_trace.gather_main_row_range(row, num_precomputed_cols, total_cols) - }), + None => { + // A root-only host tree means the nodes are + // device-resident: this arm would emit an empty + // path for query position 0 instead of failing. + assert!( + !main_commit.tree.is_root_only(), + "preprocessed opening fell back to the host tree, \ + but it is root-only (nodes device-resident)" + ); + Self::open_polys_with(domain, &main_commit.tree, *index, |row| { + lde_trace.gather_main_row_range( + row, + num_precomputed_cols, + total_cols, + ) + }) + } } } #[cfg(not(feature = "cuda"))] From 8a658846c143ea67f6c780f10ea27988d830fdf2 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 31 Jul 2026 16:44:27 -0300 Subject: [PATCH 3/6] =?UTF-8?q?fix(gpu):=20address=20round-2=20review=20?= =?UTF-8?q?=E2=80=94=20guarded=20zero-inverse,=20retargeted=20fault=20hook?= =?UTF-8?q?,=20merkle=20root-only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crypto/crypto/src/merkle_tree/merkle.rs | 15 ++++++++---- crypto/math-cuda/src/inverse.rs | 32 +++++++++++++++++++++---- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/crypto/crypto/src/merkle_tree/merkle.rs b/crypto/crypto/src/merkle_tree/merkle.rs index e9ecd0bb5..447654907 100644 --- a/crypto/crypto/src/merkle_tree/merkle.rs +++ b/crypto/crypto/src/merkle_tree/merkle.rs @@ -168,10 +168,6 @@ where }) } - /// Create a root only Merkle tree placeholder: stores the commitment root - /// but no nodes. Used when paths are gathered from a device resident copy - /// (GPU) instead of this host tree, so the host nodes are never built. - /// [`get_proof_by_pos`](Self::get_proof_by_pos) must NOT be called on it. /// True when this tree carries only its root (the nodes live elsewhere, /// e.g. device-resident): openings must not walk this tree. pub fn is_root_only(&self) -> bool { @@ -185,6 +181,10 @@ where } } + /// Create a root only Merkle tree placeholder: stores the commitment root + /// but no nodes. Used when paths are gathered from a device resident copy + /// (GPU) instead of this host tree, so the host nodes are never built. + /// [`get_proof_by_pos`](Self::get_proof_by_pos) must NOT be called on it. pub fn from_root(root: B::Node) -> Self { MerkleTree { root, @@ -266,7 +266,14 @@ where /// Returns a Merkle proof for the element/s at position pos /// For example, give me an inclusion proof for the 3rd element in the /// Merkle tree + /// + /// Returns `None` on a root-only tree ([`from_root`](Self::from_root)): + /// its nodes live elsewhere (e.g. device-resident), so a host path would + /// be a silently-empty bogus proof rather than an inclusion witness. pub fn get_proof_by_pos(&self, pos: usize) -> Option> { + if self.is_root_only() { + return None; + } let pos = pos + self.node_count() / 2; let Ok(merkle_path) = self.build_merkle_path(pos) else { return None; diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 9931e3a8c..833bf2906 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -97,8 +97,11 @@ const fn ext3_fermat_exponent() -> [u64; 3] { /// /// Unlike the host Fermat this used to call, a zero total maps silently to /// zero instead of `Err`. Unreachable with honest inputs (LogUp/barycentric -/// denominators are nonzero w.h.p.); callers must not rely on a zero-total -/// error. +/// denominators are nonzero w.h.p. under random Fiat-Shamir challenges); +/// callers must not rely on a zero-total error. Debug builds add a D2H+sync +/// invertibility guard (see below) that panics on a zero total so a +/// construction/kernel bug fails loudly in tests; release elides it to keep +/// the batch inverse fully stream-ordered (no per-batch host round-trip). fn launch_invert_total( stream: &Arc, be: &crate::device::Backend, @@ -124,6 +127,23 @@ fn launch_invert_total( .arg(&mut *out) .launch(cfg)?; } + // Debug-only invertibility guard. The Fermat kernel maps a zero total + // (some denominator was zero) silently to zero, so the batch would ship + // all-zero "inverses" instead of erroring. A valid inverse is never zero, + // so `out == 0` unambiguously flags a zero total. Gated off release: the + // D2H+sync would reintroduce the per-batch host block this path exists to + // avoid, and a zero total is unreachable with honest inputs — a hit here + // is a construction or kernel bug, which tests/CI are the place to catch. + #[cfg(debug_assertions)] + { + let mut host = [0u64; 3]; + stream.memcpy_dtoh(&out.slice(0..3), &mut host)?; + stream.synchronize()?; + assert_ne!( + host, [0u64; 3], + "batch inverse: zero total has no inverse (a denominator was zero)" + ); + } Ok(()) } @@ -138,8 +158,6 @@ pub fn batch_inverse_ext3_dev( n: usize, stream: &Arc, ) -> Result> { - #[cfg(feature = "test-faults")] - check_inverse_fault_injection()?; 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 @@ -227,6 +245,12 @@ pub fn compute_and_invert_denoms_ext3_dev( sign: DenomSign, stream: &Arc, ) -> Result> { + // Fault-injection hook lives here (not in the shared `batch_inverse_ext3_dev`) + // so `schedule_inverse_fault(N)` targets exactly the Nth R3/R4 denominator + // inversion the fallback test exercises — not the LogUp aux inverses that + // also route through `batch_inverse_ext3_dev` earlier in the prove. + #[cfg(feature = "test-faults")] + check_inverse_fault_injection()?; assert_eq!(z_scalars_host.len(), k_scalars * 3); assert!(n >= 1 && k_scalars >= 1); From 05fb8dd25b97e3097a5ac269c47fffdf0f92baa3 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 31 Jul 2026 17:17:00 -0300 Subject: [PATCH 4/6] perf(gpu): stage htod_via in fixed 64MB chunks to bound pinned footprint --- crypto/math-cuda/src/device.rs | 89 +++++++++++++++++++++++----------- 1 file changed, 62 insertions(+), 27 deletions(-) diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index c3dd4a2a9..b0620c715 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -630,11 +630,26 @@ impl Drop for PendingD2H<'_> { } } -/// Host→device copy staged through the pinned slot: one host memcpy into -/// pinned memory + one async DMA, instead of the driver's internal pageable -/// staging (small chunks; 2-3x slower for multi-hundred-MB traces and it -/// convoys under multi-thread load). Blocks until the DMA lands, so the slot -/// and `src_host` are both reusable on return. +/// Chunk size for [`htod_via`]'s staged upload — the upper bound a single H2D +/// puts on a staging slot's page-locked footprint. 64 MB is large enough to +/// amortize the per-chunk DMA launch + event sync, small enough to keep the +/// pinned slab independent of trace size. +const HTOD_CHUNK_BYTES: usize = 64 << 20; // 64 MB + +/// Host→device copy staged through the pinned slot, in fixed-size chunks: each +/// chunk is one host memcpy into pinned memory + one async DMA, instead of the +/// driver's internal pageable staging (small chunks; 2-3x slower for +/// multi-hundred-MB traces and it convoys under multi-thread load). Blocks +/// until the last DMA lands, so the slot and `src_host` are both reusable on +/// return. +/// +/// Chunking caps the slot's page-locked footprint at [`HTOD_CHUNK_BYTES`] +/// regardless of trace size. This matters on the device-only path +/// (`retain_host_lde = false`): there is no [`async_dtoh_via`] drain to size +/// the slot, so `htod_via` is its only writer — an uncapped copy would grow +/// the per-worker slab to a whole trace and, being grow-only, never shrink it. +/// The host-retaining path is unaffected: its later `async_dtoh_via` grows the +/// same slot to the full LDE anyway, and we simply reuse the first chunk of it. pub fn htod_via( stream: &Arc, slot: &Mutex, @@ -651,31 +666,51 @@ pub fn htod_via( if n_bytes == 0 { return Ok(()); } - let u64_len = n_bytes.div_ceil(8); + let elem_size = std::mem::size_of::(); + // Chunk in whole elements so a `T` never straddles a chunk boundary. + let chunk_elems = (HTOD_CHUNK_BYTES / elem_size.max(1)).max(1); + let mut staging = slot.lock().unwrap(); - staging.ensure_capacity(u64_len, ctx)?; + // Only ask for a chunk's worth of pinned memory (or the whole copy when + // smaller). If another path (`async_dtoh_via` on the host-retaining flow) + // already grew this slot larger, it stays larger — grow-only — and we just + // use the first chunk of it. + let want_u64 = (chunk_elems * elem_size).div_ceil(8).min(n_bytes.div_ceil(8)); + staging.ensure_capacity(want_u64, ctx)?; ctx.bind_to_thread()?; - // SAFETY: the pinned allocation is stable while the lock is held and at - // least `n_bytes` long (`ensure_capacity`). The DMA reads it after the - // host memcpy (program order); `device_ptr_mut` orders the device write - // on `stream`. - unsafe { - std::ptr::copy_nonoverlapping( - src_host.as_ptr() as *const u8, - staging.ptr as *mut u8, - n_bytes, - ); - let (dst_ptr, _record) = dst.device_ptr_mut(stream); - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( - dst_ptr, - staging.ptr as *const core::ffi::c_void, - n_bytes, - stream.cu_stream(), - ) - .result()?; + + // SAFETY: `device_ptr_mut` yields the destination base pointer and orders + // the device writes on `stream`; `dst.len() >= src_host.len()` (asserted), + // so every chunk's byte range stays within `dst`. + let (dst_base, _record) = dst.device_ptr_mut(stream); + let src = src_host.as_ptr() as *const u8; + let n_elems = src_host.len(); + let mut elem_off = 0usize; + while elem_off < n_elems { + let this_elems = (n_elems - elem_off).min(chunk_elems); + let this_bytes = this_elems * elem_size; + let byte_off = elem_off * elem_size; + // SAFETY: the pinned slab holds at least `chunk_elems * elem_size` + // bytes (or the whole copy when smaller). The previous chunk's DMA is + // synced below before this memcpy overwrites the slab, so the slab is + // never read (by an in-flight DMA) and written at the same time. + unsafe { + std::ptr::copy_nonoverlapping(src.add(byte_off), staging.ptr as *mut u8, this_bytes); + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + dst_base + byte_off as u64, + staging.ptr as *const core::ffi::c_void, + this_bytes, + stream.cu_stream(), + ) + .result()?; + } + // Single-buffered: wait for this chunk's DMA before the next memcpy + // reuses the slab. + staging.record_event(stream)?; + staging.sync_event()?; + elem_off += this_elems; } - staging.record_event(stream)?; - staging.sync_event() + Ok(()) } /// Enqueue an async D2H of `n_elems` of `src` into the pinned slab of `slot`, From 51536c2ae8a82b8ca7a7dd81f75ed5e4fd486073 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 31 Jul 2026 17:24:52 -0300 Subject: [PATCH 5/6] style: rustfmt htod_via chunk-size expression --- crypto/math-cuda/src/device.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index b0620c715..b7b837393 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -675,7 +675,9 @@ pub fn htod_via( // smaller). If another path (`async_dtoh_via` on the host-retaining flow) // already grew this slot larger, it stays larger — grow-only — and we just // use the first chunk of it. - let want_u64 = (chunk_elems * elem_size).div_ceil(8).min(n_bytes.div_ceil(8)); + let want_u64 = (chunk_elems * elem_size) + .div_ceil(8) + .min(n_bytes.div_ceil(8)); staging.ensure_capacity(want_u64, ctx)?; ctx.bind_to_thread()?; From db04e557e57d7dded4ad0cd96292d06ecdcd3642 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 31 Jul 2026 18:01:25 -0300 Subject: [PATCH 6/6] fix(gpu): gate R2 comp-tree host fallback on the parts, not host_trace_empty --- crypto/stark/src/prover.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 8612bb20a..d8cf4bf87 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1644,9 +1644,16 @@ pub trait IsStarkProver< None => { // The host part evals are empty under device-only (the R2 // drain is skipped); abort with the device-only contract's - // message instead of a misleading EmptyCommitment. + // message instead of a misleading EmptyCommitment. Gate on + // the parts the CPU fallback actually consumes, not on + // `host_trace_empty()`: the trace can stay device-resident + // while these parts were downloaded to the host anyway (the + // GPU decompose fell back to `decompose_and_extend_d2`), in + // which case this fallback is valid and must not panic. assert!( - !round_1_result.lde_trace.host_trace_empty(), + lde_composition_poly_parts_evaluations + .first() + .is_none_or(|p| !p.is_empty()), "R2 composition commit fell back to the host part evals, \ but they are device-only (empty)" );