Feat/fri early termination

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/benches/profile_prover.rs

@@ -21,6 +21,7 @@ fn main() {
         fri_number_of_queries: 100,
         coset_offset: 3,
         grinding_factor: 0,
+        fri_final_poly_log_degree: 7,
     };
 
     let num_columns = 16;

crypto/stark/benches/prover_benchmark.rs

@@ -61,6 +61,7 @@ fn benchmark_proof_options() -> ProofOptions {
         fri_number_of_queries: 30,
         coset_offset: 3,
         grinding_factor: 0,
+        fri_final_poly_log_degree: 7,
     }
 }
 

crypto/stark/src/domain.rs

@@ -49,7 +49,6 @@ use super::traits::AIR;
 /// Full domain with pre-computed roots of unity. Used by the prover which needs
 /// all elements for FFT operations.
 pub struct Domain<F: IsFFTField> {
-    pub(crate) root_order: u32,
     pub(crate) lde_roots_of_unity_coset: Vec<FieldElement<F>>,
     pub(crate) trace_primitive_root: FieldElement<F>,
     pub(crate) trace_roots_of_unity: Vec<FieldElement<F>>,
@@ -88,7 +87,6 @@ impl<F: IsFFTField> Domain<F> {
         .unwrap();
 
         Self {
-            root_order,
             lde_roots_of_unity_coset,
             trace_primitive_root,
             trace_roots_of_unity,

crypto/stark/src/fri/mod.rs

@@ -1,6 +1,7 @@
 pub mod fri_commitment;
 pub mod fri_decommit;
 pub(crate) mod fri_functions;
+pub(crate) mod terminal;
 
 use crypto::fiat_shamir::is_transcript::IsStarkTranscript;
 use math::field::element::FieldElement;
@@ -16,25 +17,28 @@ use self::fri_functions::{
 };
 
 /// FRI commit phase from pre-computed bit-reversed evaluations, skipping the
-/// initial FFT. Use this when the caller already has the evaluation vector
-/// (e.g. from a fused LDE pipeline).
+/// initial FFT. Stops folding when the remaining codeword encodes a polynomial
+/// of degree < 2^`final_poly_log_degree` with blowup 2^`blowup_log`, and
+/// returns the coefficient vector of that terminal polynomial.
 ///
 /// The `T: Clone` and `F/E: 'static` bounds are required by the cuda GPU
 /// fast path (`try_fri_commit_gpu` snapshots the transcript and TypeId-
 /// checks the field types). They are present unconditionally (including
 /// in builds without the `cuda` feature) to keep one stable signature.
+#[allow(clippy::type_complexity)]
 pub fn commit_phase_from_evaluations<
     F: IsFFTField + IsSubFieldOf<E> + 'static,
-    E: IsField + 'static,
+    E: IsField + 'static + Send + Sync,
     T: IsStarkTranscript<E, F> + Clone,
 >(
-    number_layers: usize,
     mut evals: Vec<FieldElement<E>>,
     transcript: &mut T,
     coset_offset: &FieldElement<F>,
     domain_size: usize,
+    blowup_log: u32,
+    final_poly_log_degree: u32,
 ) -> (
-    FieldElement<E>,
+    Vec<FieldElement<E>>,
     Vec<FriLayer<E, FriLayerMerkleTreeBackend<E>>>,
 )
 where
@@ -50,27 +54,39 @@ where
     // had never been tried.
     #[cfg(feature = "cuda")]
     {
+        // Try the GPU early-termination FRI commit first. `try_fri_commit_gpu`
+        // drives the same commit phase on-device (Goldilocks + Ext3, above the
+        // LDE size threshold, and only when folding actually happens) and returns
+        // `Some` with the final-polynomial coefficients. It returns `None` on any
+        // precondition miss or cudarc error — restoring the transcript first — so
+        // the CPU path below then runs as if the GPU had never been tried.
         if let Some(result) = crate::gpu_lde::try_fri_commit_gpu::<F, E, T>(
-            number_layers,
             &evals,
             transcript,
             coset_offset,
             domain_size,
+            blowup_log,
+            final_poly_log_degree,
         ) {
             return result;
         }
     }
 
+    // Fold layout, shared with the GPU prover and the verifier — see `FriFoldLayout`.
+    let layout = crate::fri::terminal::FriFoldLayout::new(
+        evals.len().trailing_zeros(),
+        blowup_log,
+        final_poly_log_degree,
+    );
+    let num_committed = layout.num_committed;
+
     // Inverse twiddle factors for evaluation-form folding.
     let mut inv_twiddles = compute_coset_twiddles_inv(coset_offset, domain_size);
+    let mut fri_layer_list = Vec::with_capacity(num_committed);
 
-    // The loop commits `number_layers - 1` folded layers; the final fold below
-    // produces the (uncommitted) last value.
-    let num_committed_layers = number_layers.saturating_sub(1);
-    let mut fri_layer_list = Vec::with_capacity(num_committed_layers);
-
-    for _ in 0..num_committed_layers {
-        // <<<< Receive challenge 𝜁ₖ₋₁
+    // Commit `num_committed` folded layers to the transcript.
+    for _ in 0..num_committed {
+        // <<<< Receive challenge 𝜁ₖ
         let zeta = transcript.sample_field_element();
 
         // Fold evaluations in-place (no FFT needed).
@@ -93,21 +109,40 @@ where
         update_twiddles_in_place(&mut inv_twiddles);
     }
 
-    // <<<< Receive challenge: 𝜁ₙ₋₁
-    let zeta = transcript.sample_field_element();
-
-    // Final fold.
-    fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
-
-    let last_value = evals
-        .first()
-        .expect("FRI evals are non-empty after folding")
-        .clone();
-
-    // >>>> Send value: pₙ
-    transcript.append_field_element(&last_value);
+    // One final fold to reach the terminal codeword (size terminal_len), unless
+    // already there (total_folds == 0 means initial_len == terminal_len).
+    if layout.total_folds > 0 {
+        // <<<< Receive challenge: 𝜁_final
+        let zeta = transcript.sample_field_element();
+        fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
+    }
+    debug_assert_eq!(
+        evals.len(),
+        layout.terminal_len,
+        "terminal codeword size mismatch"
+    );
+
+    // Recover the low-degree polynomial coefficients from the terminal codeword
+    // and send them to the verifier.
+    //
+    // The coefficient count follows the *actual* terminal codeword via
+    // `layout.effective_k` (`min(k, trace_bits)`), not the requested
+    // `final_poly_log_degree`: for tiny inputs the codeword is clamped to the
+    // full LDE, so passing the raw `k` would over-pad with zeros and break the
+    // round-trip against the verifier's own `expected_k` reconstruction.
+    // The terminal coset offset is `coset_offset^(2^total_folds)` — the offset
+    // after `total_folds` squarings (matches the GPU prover and the verifier).
+    let terminal_offset = coset_offset.pow(1u64 << layout.total_folds);
+    let final_poly_coeffs = crate::fri::terminal::coeffs_from_terminal_codeword::<F, E>(
+        &evals,
+        &terminal_offset,
+        layout.effective_k,
+    );
+    for c in &final_poly_coeffs {
+        transcript.append_field_element(c);
+    }
 
-    (last_value, fri_layer_list)
+    (final_poly_coeffs, fri_layer_list)
 }
 
 pub fn query_phase<F: IsField>(