From 99c3a1173754a88d2dcb2a22e729763834f0a611 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 29 Jul 2026 14:53:51 -0300 Subject: [PATCH 1/8] 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/8] 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/8] =?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/8] 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/8] 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/8] 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)" ); From e75bcbed85110ad0347618075c94631f7287dadd Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 3 Aug 2026 14:50:48 -0300 Subject: [PATCH 7/8] =?UTF-8?q?chore(gpu):=20review=20follow-ups=20?= =?UTF-8?q?=E2=80=94=20gather=20bounds,=20release=20canaries,=20live=20zer?= =?UTF-8?q?o-total=20guard?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - gather_ext3_at asserts positions against the evals buffer host-side (same guard as gather_merkle_paths_dev). - The device-gather cross-checks keep query 0 as a release canary instead of paying every query; debug still checks all of them. - The batch-inverse zero-total guard also compiles under test-faults, so the GPU fallback suite (which runs --release) actually exercises it. - New htod_via round-trip test covering the 64 MB chunk loop and its partial tail. --- crypto/math-cuda/src/fri.rs | 8 +++++ crypto/math-cuda/src/inverse.rs | 15 +++++----- crypto/math-cuda/tests/htod_via.rs | 48 ++++++++++++++++++++++++++++++ crypto/stark/src/prover.rs | 17 +++++++---- 4 files changed, 75 insertions(+), 13 deletions(-) create mode 100644 crypto/math-cuda/tests/htod_via.rs diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 8da80472d..a3c29dedc 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -264,6 +264,14 @@ pub fn gather_ext3_at( if q == 0 { return Ok(Vec::new()); } + // Guard the kernel's device reads: a position past the evals buffer would + // be a silent out-of-bounds read. Positions are valid by construction; + // this catches a caller bug host-side before it becomes device garbage + // (matching `gather_merkle_paths_dev`). + assert!( + positions.iter().all(|&p| (p as usize) < evals.len() / 3), + "gather_ext3_at: position >= evals length" + ); let be = backend()?; let pos_dev = stream.clone_htod(positions)?; // SAFETY: the gather kernel writes all 3 * q slots. diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 833bf2906..1087e2ae4 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -127,14 +127,15 @@ 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 + // 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)] + // so `out == 0` unambiguously flags a zero total. Gated off plain release + // (the D2H+sync would reintroduce the per-batch host block this path + // exists to avoid); `test-faults` keeps it live in the GPU fallback + // suite, which runs --release — a hit is a construction or kernel bug, + // and that suite is where CI can actually catch it. + #[cfg(any(debug_assertions, feature = "test-faults"))] { let mut host = [0u64; 3]; stream.memcpy_dtoh(&out.slice(0..3), &mut host)?; diff --git a/crypto/math-cuda/tests/htod_via.rs b/crypto/math-cuda/tests/htod_via.rs new file mode 100644 index 000000000..6db1eb227 --- /dev/null +++ b/crypto/math-cuda/tests/htod_via.rs @@ -0,0 +1,48 @@ +//! Round-trip coverage for `htod_via`'s chunk loop: uploads larger than the +//! 64 MB pinned-staging chunk must arrive intact across every chunk boundary +//! (a stale slab or a bad byte offset would corrupt exactly one chunk). + +use math_cuda::device::{backend, htod_via}; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +fn roundtrip(n_u64: usize, seed: u64) { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let src: Vec = (0..n_u64).map(|_| rng.r#gen::()).collect(); + + let be = backend().expect("cuda backend"); + let stream = be.next_stream(); + let mut dst = stream.alloc_zeros::(n_u64).expect("device alloc"); + htod_via( + &stream, + be.pinned_staging(), + &be.ctx, + &src, + &mut dst.slice_mut(0..n_u64), + ) + .expect("htod_via"); + + let back = stream.clone_dtoh(&dst).expect("dtoh"); + stream.synchronize().expect("sync"); + assert_eq!(src.len(), back.len()); + // Compare in chunks so a failure names the offset instead of dumping 100M+ values. + for (i, (a, b)) in src.iter().zip(back.iter()).enumerate() { + assert_eq!( + a, b, + "htod_via round-trip mismatch at u64 offset {i} (n={n_u64})" + ); + } +} + +#[test] +fn htod_via_single_chunk_roundtrip() { + // Below the 64 MB chunk: single iteration of the loop. + roundtrip(1 << 20, 42); +} + +#[test] +fn htod_via_multi_chunk_roundtrip() { + // 3 full chunks + a partial tail: exercises slab reuse across iterations + // and the final short chunk. 64 MB chunk = 2^23 u64s. + roundtrip((3 << 23) + 12345, 43); +} diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index d8cf4bf87..2e5a59016 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -2521,8 +2521,10 @@ pub trait IsStarkProver< // Cross-check the device gather against the host LDE. Skipped under // device-only (host trace empty): the gather was proven bit-identical // while the host copy was resident, and there is nothing to check - // against. - if !lde_trace.host_trace_empty() { + // against. Release keeps query 0 as a canary (the GPU test suites run + // --release, and gather failure modes — stride/offset/layout — are + // systematic, so one query catches them); debug checks every query. + if (cfg!(debug_assertions) || qi == 0) && !lde_trace.host_trace_empty() { let domain_size = domain.lde_roots_of_unity_coset.len() as u64; let r_even = reverse_index(challenge * 2, domain_size); let r_odd = reverse_index(challenge * 2 + 1, domain_size); @@ -2799,10 +2801,13 @@ pub trait IsStarkProver< // 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()) + // Query 0 stays a release canary, same rationale + // as `open_trace_polys_device`. + if (cfg!(debug_assertions) || qi == 0) + && 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(), From 4789c87be45c1a76d7d667d48476330729fbbbbc Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:12:27 -0300 Subject: [PATCH 8/8] fix(gpu): drain htod_via on error; narrow the merkle-tail threshold (#892) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(gpu): drain htod_via on error, guard the R2 host-evaluator fallback Review follow-ups for the round-2 residency work, rebased onto e75bcbed — only the items that commit did not already cover. htod_via error path. Once a chunk's DMA is in flight, `record_event` / `sync_event` returning `Err` drops the staging `MutexGuard` with the device still reading the pinned slab, so the next locker's `ensure_capacity` can `cuMemFreeHost` it mid-copy. `async_dtoh_via` already guards this exact hazard and the file ships a `DrainOnErr` helper for it; `htod_via` was the one site not using it. R2 host-evaluator fallback. If the device decompose and the `H` download both fail under device-only, control reaches the host evaluator, which reads the intentionally-empty trace and panics with a bare out-of-bounds. Assert the device-only contract instead, matching the other fallback arms. Coverage. `batch_inverse_ext3_dev`'s `n == 1` branch is never exercised — `batch_inverse_n1` goes through the host-only short circuit in `batch_inverse_ext3`, as its own comment says. Add a direct device test. Docs. The preprocessed split-tree comment still claimed both trees come back as full host trees (the multiplicity tree is root-only + device resident), and `FriCommitState`'s doc claimed its input is always Arc-shared with a retained `gpu_evals` (only true on the device-only path). * perf(gpu): set the merkle-tail threshold to the block width TAIL_MAX_PAIRS = 2048 overshoots. The tail grid-strides a single 128-thread block on one SM, so a level of k pairs is k/128 SEQUENTIAL keccak-f1600s where the per-level launches it replaces spread them over k/128 parallel blocks. At 2048 the first four levels alone are 16+8+4+2 = 30 serial permutations against 4 parallel waves — order +100 us per large tree to save 4 launches worth order 10 us, and it sits on the critical path because the caller's 32-byte root memcpy_dtoh host-blocks on everything queued before it. At the block width the entry level is exactly one permutation per thread, so the tail still collapses the top levels into one launch but adds no serialization at all. --- crypto/math-cuda/src/device.rs | 23 ++++++++++++++++++++--- crypto/math-cuda/src/fri.rs | 6 ++++-- crypto/math-cuda/src/merkle.rs | 17 ++++++++++++++--- crypto/math-cuda/tests/batch_inverse.rs | 25 +++++++++++++++++++++++++ crypto/stark/src/prover.rs | 24 ++++++++++++++++++++---- 5 files changed, 83 insertions(+), 12 deletions(-) diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index b7b837393..8bc140f21 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -685,6 +685,15 @@ pub fn htod_via( // 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); + // Declared after the slot's MutexGuard so it drops FIRST: once a chunk's + // DMA is in flight, any `?`-return below must drain the stream before the + // guard releases the slot, or the next locker's `ensure_capacity` could + // `cuMemFreeHost` the slab while the device is still reading it. Same + // hazard `async_dtoh_via` guards against on its record-event failure. + let mut drain = DrainOnErr { + stream, + armed: false, + }; let src = src_host.as_ptr() as *const u8; let n_elems = src_host.len(); let mut elem_off = 0usize; @@ -698,18 +707,26 @@ pub fn htod_via( // 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( + let r = 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()?; + .result(); + // Armed even on failure: the driver may have enqueued the copy + // before reporting the error. + drain.armed = true; + r?; } // Single-buffered: wait for this chunk's DMA before the next memcpy - // reuses the slab. + // reuses the slab. Both calls can fail with the DMA still in flight, + // which is what `drain` covers. staging.record_event(stream)?; staging.sync_event()?; + // This chunk has landed; nothing is reading the slab until the next + // iteration re-arms. + drain.armed = false; elem_off += this_elems; } Ok(()) diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index a3c29dedc..533ff6e32 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -41,8 +41,10 @@ fn check_fault_injection() -> Result<()> { } /// 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. +/// input (the previous layer's evals) and the inv_twiddles buffer. The input +/// is an `Arc` because the caller may also retain it as that layer's +/// `gpu_evals` — it does so only on the device-only path, where no host copy +/// of the evals exists. Freed when the last holder drops. pub struct FriCommitState { pub stream: Arc, /// Current fold input. Each fold allocates a fresh output buffer that is diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index 1cb1c5b6f..c499df702 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -160,9 +160,20 @@ pub(crate) fn build_inner_tree_levels( ) -> 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; + // between them: the top levels of a big tree are each smaller than the + // per-launch overhead they used to pay. + // + // Set to the block width, so the entry level is exactly one permutation + // per thread and the tail adds NO serialization over the per-level + // launches it replaces. Going wider is not free: the tail grid-strides a + // single 128-thread block on one SM, so a level of `k` pairs costs + // `k / 128` *sequential* keccak-f1600s where separate launches would have + // spread them over `k / 128` parallel blocks. At 2048 the first four + // levels alone are 16+8+4+2 = 30 serial permutations against 4 parallel + // waves — order +100 us per large tree, to save 4 launches worth order + // 10 us. It stays on the critical path because the caller's 32-byte root + // `memcpy_dtoh` host-blocks on everything queued before it. + const TAIL_MAX_PAIRS: u64 = KECCAK_BLOCK_DIM as u64; let mut level_begin: u64 = (leaves_len - 1) as u64; while level_begin != 0 { let new_begin = level_begin / 2; diff --git a/crypto/math-cuda/tests/batch_inverse.rs b/crypto/math-cuda/tests/batch_inverse.rs index bc52f9fcb..087a0b082 100644 --- a/crypto/math-cuda/tests/batch_inverse.rs +++ b/crypto/math-cuda/tests/batch_inverse.rs @@ -72,6 +72,31 @@ fn batch_inverse_n1() { run(1, 1); } +/// `batch_inverse_ext3_dev`'s own `n == 1` branch, which the host entry point +/// above never reaches: `batch_inverse_ext3` short-circuits n==1 to +/// `invert_ext3_host`, so only a direct device call exercises the single +/// `invert_total_ext3` launch that serves this case. +#[test] +fn batch_inverse_dev_n1() { + let mut rng = ChaCha8Rng::seed_from_u64(7); + let x = rand_fp3_nonzero(&mut rng); + let expected = x.inv().expect("nonzero is invertible"); + + let be = math_cuda::device::backend().expect("cuda backend"); + let stream = be.next_stream(); + let input = stream.clone_htod(&ext3_to_u64s(&[x])).unwrap(); + + let out_dev = math_cuda::inverse::batch_inverse_ext3_dev(&input, 1, &stream).unwrap(); + let got = stream.clone_dtoh(&out_dev).unwrap(); + stream.synchronize().unwrap(); + + assert_eq!( + canon3(&got), + canon3(&ext3_to_u64s(&[expected])), + "device n==1 inverse" + ); +} + #[test] fn batch_inverse_single_block() { // All single-block sizes (no recursion). diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 2e5a59016..42142f770 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1027,10 +1027,12 @@ pub trait IsStarkProver< // Fused GPU split path for preprocessed tables (cuda only): one // row-major LDE of ALL columns plus two subset Merkle trees // (precomputed / multiplicity) built on device — leaves and levels are - // bit-identical to `commit_rows_bit_reversed_subset`, and the trees - // come back as full host trees so the preprocessed opening path and - // the process-wide precomputed-tree cache work unchanged. The handle - // keeps the LDE device-resident for the downstream GPU rounds. + // bit-identical to `commit_rows_bit_reversed_subset`. The precomputed + // tree comes back as a full host tree, so the process-wide + // precomputed-tree cache works unchanged; the multiplicity tree stays + // device-resident behind a root-only host tree and its opening paths + // are gathered on device. The handle keeps the LDE device-resident for + // the downstream GPU rounds. #[cfg(feature = "cuda")] if let Some((expected_precomputed_root, num_precomputed)) = precomputed { let (trace_slice, num_cols) = trace.main_data_row_major(); @@ -1528,6 +1530,20 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let t_sub = Instant::now(); + // Every arm below runs the HOST evaluator, which reads `get_main` / + // `get_aux`. Under device-only those buffers are intentionally empty, + // so landing here means the device decompose AND the `H` download both + // failed. Abort with the device-only contract's message rather than a + // bare index-out-of-bounds from somewhere inside the evaluator. + #[cfg(feature = "cuda")] + if precomputed_parts.is_none() { + assert!( + !round_1_result.lde_trace.host_trace_empty(), + "R2 composition fell back to the host evaluator, but the trace \ + is device-only (empty)" + ); + } + let lde_composition_poly_parts_evaluations = if let Some(parts) = precomputed_parts { parts } else if number_of_parts == 2 {