FRI early-termination review follow-ups (1/3): terminal-fold cleanup

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<u64>, &mut CudaSlice<u64>) = 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<u64> = 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]])
-    }
 }

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>,
     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*<w^blowup> = 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*<w^blowup> = 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<FieldElement<E>> = natural.into_iter().step_by(blowup).collect();
-    debug_assert_eq!(sub_coset.len(), keep);
+    let blowup = len / keep;
+    let sub_coset: Vec<FieldElement<E>> = (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::<F>(&sub_coset, terminal_offset)

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;