From 71ee66941c5ec314a9fb96675264f195b2b40b25 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 6 Jul 2026 16:06:28 -0300 Subject: [PATCH 1/3] fix(stark): fall back to CPU when GPU FRI terminal_len == 1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The GPU final fold reuses `fold_and_commit_layer`, whose `assert!(n_out >= 2)` fires when the terminal codeword has length 1 (blowup_log + k == 0). That config only arises from a raw `ProofOptions` literal with `blowup_factor: 1` and `fri_final_poly_log_degree: 0` (every validated constructor rejects blowup 1), but when it does the assert aborts the prover mid-transcript instead of `try_fri_commit_gpu` returning None and letting the CPU fallback (which handles terminal_len == 1) produce the proof — violating the function's documented return-None-on-any-failure contract. Extend the early-return guard to also bail when terminal_len < 2, so the final fold below is always n_out >= 2. --- crypto/stark/src/gpu_lde.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 28e188d1b..ac11b63dc 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1641,9 +1641,12 @@ where 1usize << terminal_shift }; let total_folds = (n0 / terminal_len).trailing_zeros() as usize; - // The GPU path only runs above gpu_lde_threshold(); tiny clamped traces - // (total_folds == 0) are handled by the CPU fallback. - if total_folds == 0 { + // The GPU path only runs above gpu_lde_threshold(). Two cases fall back to + // the CPU path (which handles both correctly): tiny clamped traces + // (total_folds == 0), and terminal_len == 1 (blowup_log + k == 0), whose + // final fold would reach n_out == 1 and trip `fold_and_commit_layer`'s + // `n_out >= 2` assert. The final fold below is therefore always n_out >= 2. + if total_folds == 0 || terminal_len < 2 { return None; } let num_committed = total_folds - 1; From 7cd2a579cf2f7cf3e0a484cee05c25915a9664ea Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 6 Jul 2026 16:06:29 -0300 Subject: [PATCH 2/3] refactor(math-cuda): remove dead fold_final `fold_final` supported the old fold-to-a-constant terminal step (n_out == 1). After the FRI early-termination switch, the GPU final fold goes through `fold_and_commit_layer` and `fold_final` has no callers anywhere. Remove it, and fix the two stale references that still pointed at it (the fault-injection doc and `fold_and_commit_layer`'s assert message). --- crypto/math-cuda/src/fri.rs | 59 ++++--------------------------------- 1 file changed, 6 insertions(+), 53 deletions(-) diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index a2f96c07a..fb854d0a4 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -17,9 +17,9 @@ use crate::device::backend; use crate::merkle::build_inner_tree_levels; /// Test-only fault injection. When the `test-faults` feature is on, setting -/// this to a finite value forces the next `fold_and_commit_layer` / -/// `fold_final` call to return Err and decrement the counter. Tests use -/// this to exercise the CPU-fallback path in `try_fri_commit_gpu`. +/// this to a finite value forces the next `fold_and_commit_layer` call to +/// return Err and decrement the counter. Tests use this to exercise the +/// CPU-fallback path in `try_fri_commit_gpu`. #[cfg(feature = "test-faults")] pub static FAULT_FOLDS_REMAINING_UNTIL_ERR: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(-1); @@ -104,10 +104,11 @@ impl FriCommitState { let be = backend()?; let n_in = self.current_n; let n_out = n_in / 2; - // fold_final handles the n_out == 1 last layer (no Merkle commit). + // n_out == 1 (terminal_len < 2) never reaches this path: `try_fri_commit_gpu` + // filters it out and returns None so the CPU fallback handles it. assert!( n_out >= 2, - "fold_and_commit_layer requires n_out >= 2; use fold_final" + "fold_and_commit_layer requires n_out >= 2 (n_out == 1 falls back to the CPU path)" ); // Row-pair leaves: each leaf hashes two consecutive ext3 evals. @@ -231,52 +232,4 @@ impl FriCommitState { }; Ok((layer_evals, tree)) } - - /// Final fold, no Merkle commit. Returns the single ext3 output - /// element (the FRI last_value). - pub fn fold_final(&mut self, zeta_raw: [u64; 3]) -> Result<[u64; 3]> { - #[cfg(feature = "test-faults")] - check_fault_injection()?; - let be = backend()?; - let n_in = self.current_n; - let n_out = n_in / 2; - assert!(n_out >= 1); - - let zeta_dev = self.stream.clone_htod(&zeta_raw)?; - let cfg = LaunchConfig { - grid_dim: ((n_out as u32).div_ceil(128), 1, 1), - block_dim: (128, 1, 1), - shared_mem_bytes: 0, - }; - let n_out_u64 = n_out as u64; - - 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) - }; - unsafe { - self.stream - .launch_builder(&be.fri_fold_ext3) - .arg(input_evals) - .arg(&n_out_u64) - .arg(&self.inv_tw) - .arg(&zeta_dev) - .arg(output_evals) - .launch(cfg)?; - } - - self.stream.synchronize()?; - let out_first: Vec = if self.a_is_input { - let view = self.evals_b.slice(0..3); - self.stream.clone_dtoh(&view)? - } else { - let view = self.evals_a.slice(0..3); - self.stream.clone_dtoh(&view)? - }; - self.a_is_input = !self.a_is_input; - self.current_n = n_out; - Ok([out_first[0], out_first[1], out_first[2]]) - } } From e833b83ef67fd4ca4847ae7c174ceceda04ef310 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 6 Jul 2026 16:06:29 -0300 Subject: [PATCH 3/3] perf(stark): gather FRI terminal sub-coset via reverse_index MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `coeffs_from_terminal_codeword` cloned the whole terminal codeword and ran a full O(n) bit-reverse permute only to keep every blowup-th element. Since the codeword is already in bit-reversed order, gather the size-2^k sub-coset directly with reverse_index — no clone, no full permute, and only 2^k of the blowup*2^k evaluations are read. Behaviour is identical (verified by the terminal roundtrip and FRI early-termination tests). --- crypto/stark/src/fri/terminal.rs | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/crypto/stark/src/fri/terminal.rs b/crypto/stark/src/fri/terminal.rs index e3c663e49..f47875556 100644 --- a/crypto/stark/src/fri/terminal.rs +++ b/crypto/stark/src/fri/terminal.rs @@ -4,7 +4,7 @@ //! These are pure, self-contained helpers — no transcript, no FRI logic. //! They are used by the prover (`commit_phase_from_evaluations`) and verifier FRI step. -use math::fft::bit_reversing::in_place_bit_reverse_permute; +use math::fft::bit_reversing::{in_place_bit_reverse_permute, reverse_index}; use math::field::element::FieldElement; use math::field::traits::{IsFFTField, IsField, IsSubFieldOf}; use math::polynomial::Polynomial; @@ -33,16 +33,18 @@ where F: IsFFTField + IsSubFieldOf, E: IsField + Send + Sync, { - // Bit-reversed -> natural order. - let mut natural = codeword_bitrev.to_vec(); - in_place_bit_reverse_permute(&mut natural); - - // A degree-<2^k poly is determined by 2^k points: take the size-2^k sub-coset - // terminal_offset* = every `blowup`-th natural-order evaluation. + // A degree-<2^k poly is determined by 2^k points: the size-2^k sub-coset + // terminal_offset* = every `blowup`-th natural-order evaluation, + // i.e. natural-order index m*blowup for m in 0..2^k. The codeword is in + // bit-reversed order, so gather those points straight from it via + // reverse_index — no full-codeword clone or O(n) permute (only 2^k of the + // blowup*2^k evaluations are ever read). + let len = codeword_bitrev.len(); let keep = 1usize << final_poly_log_degree; - let blowup = natural.len() / keep; - let sub_coset: Vec> = natural.into_iter().step_by(blowup).collect(); - debug_assert_eq!(sub_coset.len(), keep); + let blowup = len / keep; + let sub_coset: Vec> = (0..keep) + .map(|m| codeword_bitrev[reverse_index(m * blowup, len as u64)].clone()) + .collect(); // Coset iFFT on the small domain -> the 2^k coefficients directly (no oversized trim). let poly = Polynomial::interpolate_offset_fft::(&sub_coset, terminal_offset)