From a2a60d18c3456b158e2c5a6b2defe6425ad08c6b Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 18:01:20 -0300 Subject: [PATCH 01/26] feat(stark): mixed-height MMCS and batched FRI primitives --- .../backends/field_element_vector.rs | 136 +- crypto/crypto/src/merkle_tree/traits.rs | 71 + crypto/stark/src/fri/batched.rs | 1130 +++++++++++ crypto/stark/src/fri/mmcs.rs | 1805 +++++++++++++++++ crypto/stark/src/fri/mod.rs | 2 + crypto/stark/src/par.rs | 24 + 6 files changed, 3167 insertions(+), 1 deletion(-) create mode 100644 crypto/stark/src/fri/batched.rs create mode 100644 crypto/stark/src/fri/mmcs.rs diff --git a/crypto/crypto/src/merkle_tree/backends/field_element_vector.rs b/crypto/crypto/src/merkle_tree/backends/field_element_vector.rs index 6d0cc6491..a9f1f4b05 100644 --- a/crypto/crypto/src/merkle_tree/backends/field_element_vector.rs +++ b/crypto/crypto/src/merkle_tree/backends/field_element_vector.rs @@ -1,7 +1,7 @@ use core::marker::PhantomData; use crate::hash::poseidon::Poseidon; -use crate::merkle_tree::traits::IsMerkleTreeBackend; +use crate::merkle_tree::traits::{IsLeafHasher, IsMerkleTreeBackend, IsStreamingLeafBackend}; use alloc::vec::Vec; use digest::{Digest, Output}; use math::{ @@ -202,6 +202,140 @@ where } } +/// Exposes the streaming leaf routes to callers that reach this backend through +/// a commitment configuration rather than by name. Both bodies go through +/// [`hash_streamed`], which is where the absorbed byte layout is defined, so +/// they agree with `hash_data` by construction. +impl IsStreamingLeafBackend + for FieldElementVectorBackend +where + F: IsField, + FieldElement: AsBytes, + [u8; NUM_BYTES]: From>, + Vec>: Sync + Send, +{ + fn hash_bytes(data: &[u8]) -> [u8; NUM_BYTES] { + hash_streamed::(|sink| sink(data)) + } + + fn hash_data_from_slices(a: &[FieldElement], b: &[FieldElement]) -> [u8; NUM_BYTES] { + // A size threshold below which this streams straight through was tried + // and MEASURED NEUTRAL-TO-WORSE (963.56M vs 963.28M cycles on a blowup8 + // verify, with `verify_fri` unmoved to the cycle). The 6.9M `verify_fri` + // rise that this change costs is NOT the staging buffer — gating the + // buffer away does not recover it — so it is not worth a branch here. + // Do not re-add one without a measurement. + hash_streamed::(|sink| { + let mut stage = LeafStage::new(); + for element in a.iter().chain(b.iter()) { + element.stream_bytes(&mut |bytes| stage.push(bytes, sink)); + } + stage.flush(sink); + }) + } + + type LeafHasher = DigestLeafHasher; + + fn leaf_hasher() -> Self::LeafHasher { + DigestLeafHasher { + hasher: D::new(), + phantom: PhantomData, + } + } +} + +/// [`IsLeafHasher`] over the same digest the one-shot routes use. +/// +/// The split-invariance the trait demands is inherited rather than argued: +/// `hash_streamed` opens a fresh `D`, feeds it every element's `stream_bytes` +/// and finalizes, with no length prefix, padding or framing of its own — so +/// absorbing the same elements across several `update` calls presents `D` with +/// the identical byte stream. There is no place for a split to show. +/// +/// This is a PROVER-side construct: the guest verifier authenticates leaves it +/// receives whole, through `hash_data_from_slices`. +pub struct DigestLeafHasher { + hasher: D, + /// `fn() -> F` rather than `F`: the field is a type-level label here, never a + /// value, and the function-pointer form is unconditionally `Send`/`Sync`. The + /// bare `PhantomData` would make every leaf hasher's thread-safety hinge on + /// a marker type nobody ever moves. + phantom: PhantomData F>, +} + +impl IsLeafHasher for DigestLeafHasher +where + F: IsField, + FieldElement: AsBytes, + [u8; NUM_BYTES]: From>, +{ + type Node = [u8; NUM_BYTES]; + + fn update(&mut self, data: &[FieldElement]) { + for element in data { + element.stream_bytes(&mut |bytes| self.hasher.update(bytes)); + } + } + + fn finalize(self) -> [u8; NUM_BYTES] { + let mut result = [0u8; NUM_BYTES]; + result.copy_from_slice(&self.hasher.finalize()); + result + } +} + +/// Bytes of the leaf staging buffer. Large enough that the run reaching the +/// hasher is worth batching — 16 blocks. +const LEAF_STAGE_BYTES: usize = 1024; + +/// Coalesces a leaf's field elements into large aligned runs before they reach +/// the hasher. +/// +/// A leaf arrives one field element at a time — eight bytes per `stream_bytes` +/// call — so without staging the hasher only ever sees eight bytes at a time. +/// Coalescing presents it with fewer, larger `update` calls. +/// +/// **This cannot change any digest.** The same bytes reach the hasher in the +/// same order; only the call boundaries move, and the sponge is split-invariant +/// by construction, so it simply sees fewer, larger `update` calls. +#[repr(align(8))] +struct LeafStage { + buf: [u8; LEAF_STAGE_BYTES], + len: usize, +} + +impl LeafStage { + #[inline] + fn new() -> Self { + Self { + buf: [0u8; LEAF_STAGE_BYTES], + len: 0, + } + } + + #[inline] + fn push(&mut self, mut bytes: &[u8], sink: &mut dyn FnMut(&[u8])) { + while !bytes.is_empty() { + if self.len == LEAF_STAGE_BYTES { + sink(&self.buf[..LEAF_STAGE_BYTES]); + self.len = 0; + } + let take = (LEAF_STAGE_BYTES - self.len).min(bytes.len()); + self.buf[self.len..self.len + take].copy_from_slice(&bytes[..take]); + self.len += take; + bytes = &bytes[take..]; + } + } + + #[inline] + fn flush(&mut self, sink: &mut dyn FnMut(&[u8])) { + if self.len > 0 { + sink(&self.buf[..self.len]); + self.len = 0; + } + } +} + #[derive(Clone, Default)] pub struct BatchPoseidonTree { _poseidon: PhantomData

, diff --git a/crypto/crypto/src/merkle_tree/traits.rs b/crypto/crypto/src/merkle_tree/traits.rs index c09cff9d0..049bf5615 100644 --- a/crypto/crypto/src/merkle_tree/traits.rs +++ b/crypto/crypto/src/merkle_tree/traits.rs @@ -1,4 +1,7 @@ use alloc::vec::Vec; +use math::field::element::FieldElement; +use math::field::traits::IsField; +use math::traits::AsBytes; #[cfg(feature = "parallel")] use rayon::prelude::{IntoParallelRefIterator, ParallelIterator}; @@ -27,3 +30,71 @@ pub trait IsMerkleTreeBackend { /// It will be used in the construction of the Merkle tree. fn hash_new_parent(child_1: &Self::Node, child_2: &Self::Node) -> Self::Node; } + +/// A leaf backend that can hash a leaf without being handed one. +/// +/// [`IsMerkleTreeBackend::hash_data`] takes a `&Self::Data`, which for the +/// batched backends is a `Vec>`. Building one per leaf costs an +/// allocation per leaf — millions on a real trace — so the prover and verifier +/// never do: they serialize into a reused buffer, or hold two slices they want +/// hashed as if concatenated. These are the two shapes they use. +/// +/// Both must agree with `hash_data` on the bytes they absorb, so a leaf hashed +/// through either route is the leaf the tree was built from. That is the whole +/// contract, and it is why these live on a trait rather than staying inherent +/// methods on one concrete backend: a commitment configuration that names its +/// leaf backend generically still has to reach them. +pub trait IsStreamingLeafBackend: IsMerkleTreeBackend +where + F: IsField, + FieldElement: AsBytes, +{ + /// Hash a pre-serialized leaf buffer. Equals `hash_data` applied to the + /// elements `data` encodes, in that order. + fn hash_bytes(data: &[u8]) -> Self::Node; + + /// Hash `a ‖ b` without materializing the concatenation. Equals + /// `hash_data(&[a, b].concat())`. + fn hash_data_from_slices(a: &[FieldElement], b: &[FieldElement]) -> Self::Node; + + /// The incremental form of the same leaf hash. See [`IsLeafHasher`]. + /// + /// `Send` because there is one of these per leaf and the base layer of a real + /// epoch has millions: absorbing them is parallel across leaves, exactly as + /// the one-shot leaf hashing is. + type LeafHasher: IsLeafHasher + Send; + + /// A leaf hasher that has absorbed nothing yet. + fn leaf_hasher() -> Self::LeafHasher; +} + +/// One leaf's hash, absorbed in an arbitrary number of updates. +/// +/// [`IsStreamingLeafBackend::hash_data_from_slices`] covers the two-slice case, +/// which is every leaf the per-table trees hash. A mixed-height MMCS leaf is +/// different: it concatenates one row pair per matrix at that height, and a +/// prover that wants to produce those matrices ONE AT A TIME — absorbing each +/// into the leaves and dropping its buffer — cannot hand over all the slices at +/// once. This is the API that lets it, and the memory it costs is one hasher +/// state per leaf rather than one LDE per matrix. +/// +/// # Contract +/// +/// Splitting is free: for any partition of a leaf's elements into consecutive +/// chunks, updating with each chunk in order and finalizing must equal +/// [`IsMerkleTreeBackend::hash_data`] over the whole. A backend whose framing +/// depended on where the updates fell would produce leaves no verifier could +/// re-derive, since the verifier only ever sees the concatenation. +pub trait IsLeafHasher +where + F: IsField, + FieldElement: AsBytes, +{ + type Node; + + /// Absorb the next consecutive run of the leaf's elements. + fn update(&mut self, data: &[FieldElement]); + + /// Finish the leaf. + fn finalize(self) -> Self::Node; +} diff --git a/crypto/stark/src/fri/batched.rs b/crypto/stark/src/fri/batched.rs new file mode 100644 index 000000000..f9e27adf0 --- /dev/null +++ b/crypto/stark/src/fri/batched.rs @@ -0,0 +1,1130 @@ +//! Batched FRI: one FRI instance over an epoch's DEEP codewords instead of one +//! per table. +//! +//! Codewords are bucketed by height, mixed within a bucket with powers of a +//! single `alpha`, and then folded from the tallest bucket downward, each +//! shorter bucket being *injected* into the running codeword at the layer whose +//! length matches it. One set of query indices, drawn from the tallest domain, +//! tests the whole chain. +//! +//! # Termination +//! +//! Folding stops at the same terminal the unbatched +//! [`crate::fri::commit_phase_from_evaluations`] stops at — the codeword that +//! encodes a polynomial of degree `< 2^fri_final_poly_log_degree` — and sends +//! that polynomial's coefficients, rather than folding all the way down to a +//! scalar. [`BatchedFriLayout`] derives the fold count through the shared +//! [`FriFoldLayout`], with one batched-only floor: the terminal may not sit +//! above the SHORTEST injected codeword, or that codeword would never reach the +//! running word. So the early stop is `min(blowup_log + k, h_min)`. + +use crypto::fiat_shamir::is_transcript::{IsStarkTranscript, IsTranscript}; +use crypto::merkle_tree::merkle::MerkleTree; +use math::field::element::FieldElement; +use math::field::traits::{IsFFTField, IsField, IsSubFieldOf}; +use math::traits::AsBytes; +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +use crate::config::FriLayerMerkleTreeBackend; +use crate::fri::fri_commitment::FriLayer; +use crate::fri::fri_functions::{ + compute_coset_twiddles_inv, fold_evaluations_in_place, update_twiddles_in_place, +}; +use crate::fri::terminal::{FriFoldLayout, coeffs_from_terminal_codeword}; + +/// Accumulates DEEP codewords into per-height buckets as they are produced, +/// mixing the `i`-th absorbed codeword with `alpha^i`. +/// +/// The point of absorbing one codeword at a time is memory: a caller that +/// produces a table's quotient, absorbs it and drops it retains only one bucket +/// per distinct height (`O(2^h_max)` in total), where handing +/// [`combine_by_height`] a fully-materialized `Vec` of every table's codeword +/// retains `O(N_tables · 2^h)`. The result is identical either way — absorption +/// order defines the `alpha` powers, so the caller must absorb in the same +/// canonical per-epoch order the verifier assumes. +pub struct HeightCombiner { + buckets: Vec>>>, + alpha: FieldElement, + /// `alpha^i` for the next codeword to be absorbed. + next_power: FieldElement, +} + +impl HeightCombiner { + pub fn new(alpha: FieldElement) -> Self { + Self { + buckets: Vec::new(), + alpha, + next_power: FieldElement::one(), + } + } + + /// Absorb one codeword of length `2^height`, scaled by the next power of + /// `alpha`. + pub fn absorb(&mut self, codeword: &[FieldElement], height: usize) { + let expected_len = 1usize << height; + assert_eq!( + codeword.len(), + expected_len, + "codeword has length {} but height {height} expects {expected_len}", + codeword.len() + ); + + if self.buckets.len() <= height { + self.buckets.resize_with(height + 1, || None); + } + let scaled = &self.next_power; + // Data-parallel under `parallel`: the scale and the scale-accumulate + // are elementwise over up to 2^h_max elements, and this loop has no + // per-table overlap to hide behind — it was serial wall time once per + // absorbed table. Same arithmetic in both arms, identical result. + #[cfg(feature = "parallel")] + match &mut self.buckets[height] { + None => { + self.buckets[height] = Some( + codeword + .par_iter() + .map(|x| scaled * x) + .collect::>>(), + ); + } + Some(acc) => { + acc.par_iter_mut() + .zip(codeword.par_iter()) + .for_each(|(a, x)| { + *a = &*a + &(scaled * x); + }); + } + } + #[cfg(not(feature = "parallel"))] + match &mut self.buckets[height] { + None => { + self.buckets[height] = Some(codeword.iter().map(|x| scaled * x).collect()); + } + Some(acc) => { + for (a, x) in acc.iter_mut().zip(codeword.iter()) { + *a = &*a + &(scaled * x); + } + } + } + self.next_power = &self.next_power * &self.alpha; + } + + /// The per-height buckets. Index `h` is `Some(combined)` when at least one + /// codeword of height `h` was absorbed, `None` otherwise; the `Vec` is + /// `max_absorbed_height + 1` long, or empty if nothing was absorbed. + pub fn finish(self) -> Vec>>> { + self.buckets + } +} + +/// Combine DEEP polynomial codewords by their FRI height for batched FRI. +/// +/// Each element of `inputs` is a pair `(codeword, height)` where `height` is +/// the log₂ of the codeword length (i.e. `codeword.len() == 2^height`). +/// The global index `i` into `inputs` is used to derive the mixing power +/// `alpha^i` (index 0 → alpha^0 = 1, index 1 → alpha^1, …). +/// +/// Returns a `Vec` of length `max_height + 1`. Index `h` contains +/// `Some(combined)` where `combined[j] = Σ_{i : height_i == h} alpha^i * codeword_i[j]`, +/// or `None` when no input has height `h`. +/// +/// This is [`HeightCombiner`] with every codeword already materialized. Prefer +/// the combiner in the prover, where holding all of them at once is the whole +/// memory cost the batching is meant to remove. +pub fn combine_by_height( + inputs: &[(Vec>, usize)], + alpha: &FieldElement, +) -> Vec>>> +where + E: IsField, +{ + let mut combiner = HeightCombiner::new(alpha.clone()); + for (codeword, height) in inputs { + combiner.absorb(codeword, *height); + } + combiner.finish() +} + +/// How far a batched FRI instance folds, and what it sends at the end. +/// +/// Mirrors [`FriFoldLayout`] — same early stop, same terminal codeword, same +/// coefficient count — with the one difference batching forces: the terminal is +/// additionally floored at the SHORTEST injected codeword's height, since a +/// bucket below the terminal would never be folded into the running word. In a +/// real epoch the shortest table is normally well above `blowup_log + k`, so the +/// floor is inert and the layout is exactly the unbatched one. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BatchedFriLayout { + /// Folds from the tallest bucket down to the terminal codeword. + pub total_folds: u32, + /// Committed (Merkle-rooted) FRI layers. + pub num_committed: usize, + /// Terminal codeword length. + pub terminal_len: usize, + /// `log2` of the terminal polynomial's degree bound — the number of + /// coefficients sent is `2^effective_k`. + pub effective_k: u32, +} + +impl BatchedFriLayout { + /// Derive the layout from the epoch's codeword heights. + /// + /// * `h_max` / `h_min` — the tallest and shortest codeword heights present. + /// * `blowup_log` — log2 of the LDE blowup factor. + /// * `final_poly_log_degree` — the requested `fri_final_poly_log_degree`. + /// + /// Panics if `h_min < blowup_log` (a codeword shorter than the blowup is not + /// a Reed-Solomon word of any positive rate) or if `h_min > h_max`. + pub fn new(h_max: usize, h_min: usize, blowup_log: u32, final_poly_log_degree: u32) -> Self { + assert!(h_min <= h_max, "h_min {h_min} exceeds h_max {h_max}"); + assert!( + h_min as u32 >= blowup_log, + "codeword height {h_min} is below the blowup {blowup_log}" + ); + // Deriving at `h_min` is what applies the floor: `FriFoldLayout` clamps + // the terminal to its `lde_log` argument, so the terminal comes out at + // `min(blowup_log + k, h_min)`. Its terminal_len / effective_k are then + // exactly what the unbatched prover would send for that codeword. + let shortest = FriFoldLayout::new(h_min as u32, blowup_log, final_poly_log_degree); + let terminal_log = shortest.terminal_len.trailing_zeros(); + // The running codeword starts at h_max, not h_min, so the fold count is + // re-derived from where folding actually begins. + let total_folds = h_max as u32 - terminal_log; + Self { + total_folds, + num_committed: total_folds.saturating_sub(1) as usize, + terminal_len: shortest.terminal_len, + effective_k: shortest.effective_k, + } + } +} + +/// Which of an epoch's tables enter the ONE batched FRI instance, and which keep +/// a terminal-only instance of their own. +/// +/// # Why there are two classes +/// +/// A table whose own FRI would commit ZERO layers gains nothing from being +/// batched — there is no layer for the batch to share — while it pays the full +/// cost of being lifted to the tallest domain, which is where the proximity-gaps +/// term's `|D0|²` lives. At the measured epoch that is 13 of 28 legs carrying 92% +/// of the batch's width, so excluding them is a correction rather than a +/// compromise: it recovers ~3.6 bits of soundness AND removes work. +/// +/// A zero-layer table's FRI is degenerate in the useful sense — its terminal +/// codeword IS its deep-composition codeword — so its "own instance" is one +/// terminal polynomial and no layers at all. +/// +/// # ★ The index rule BETWEEN the classes — a hard precondition +/// +/// Both classes are opened at the SAME query indices, because the mixed-height +/// MMCS is unaffected by this split: it still commits every table, and the point +/// of one shared authentication path survives whole. What differs is the index +/// SPACE each class reads them in: +/// +/// ```text +/// batched class: iota, used directly (it is an index in the tallest domain) +/// standalone table: iota >> (h_max - h_t) +/// ``` +/// +/// This is the same reduction [`crate::fri::mmcs`]'s index-convention section +/// documents for a short round, and it fails the same silent way: prover and +/// verifier derive it from the shape, so a wrong shift is self-consistent — +/// honest proofs still verify while the short tables end up checked at positions +/// the FRI join never reaches. `each_instance_class_is_tamper_checked` is the +/// control, and it tampers a table of EACH class, because a control that only +/// touched the batched class would pass under any convention for the other. +/// +/// # Determinism +/// +/// The plan is a pure function of `(heights, blowup_log, final_poly_log_degree)`, +/// all of which the transcript has bound before any challenge that depends on it. +/// Prover and verifier therefore derive the SAME partition without it being sent, +/// which is why the split adds nothing to the wire and nothing to the shape +/// binding. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FriInstancePlan { + /// Table indices whose codewords are mixed into the batched instance, in + /// input order — the order that defines the `alpha` powers. + pub batched: Vec, + /// Table indices that keep a terminal-only instance, in input order. + pub standalone: Vec, + /// Tallest and shortest height WITHIN the batched class — the layout is + /// derived from these, not from the whole epoch. + pub h_max: usize, + pub h_min: usize, +} + +impl FriInstancePlan { + /// Partition an epoch's tables. `None` when `heights` is empty or carries a + /// height that cannot be a codeword length — both are proof-supplied, so both + /// are rejections rather than panics. + /// + /// The TALLEST table is always batched, even if it would classify as + /// standalone on its own. That keeps the batched class non-empty, so the + /// layout is always well defined; an epoch whose tallest table folds nothing + /// degenerates to a single terminal-only instance, which is what it should be. + pub fn new(heights: &[usize], blowup_log: u32, final_poly_log_degree: u32) -> Option { + if heights.is_empty() { + return None; + } + let &h_max_epoch = heights.iter().max()?; + if h_max_epoch == 0 || h_max_epoch >= u32::BITS as usize { + return None; + } + let tallest = heights.iter().position(|h| *h == h_max_epoch)?; + + let mut batched = Vec::with_capacity(heights.len()); + let mut standalone = Vec::new(); + for (t, &h) in heights.iter().enumerate() { + if h < blowup_log as usize { + return None; + } + let folds_a_layer = + FriFoldLayout::new(h as u32, blowup_log, final_poly_log_degree).num_committed > 0; + if folds_a_layer || t == tallest { + batched.push(t); + } else { + standalone.push(t); + } + } + + let h_max = batched.iter().map(|&t| heights[t]).max()?; + let h_min = batched.iter().map(|&t| heights[t]).min()?; + Some(Self { + batched, + standalone, + h_max, + h_min, + }) + } +} + +/// FRI commit phase over the bucketed output of [`combine_by_height`] / +/// [`HeightCombiner::finish`]. +/// +/// `combined[h]` is `Some(codeword)` when there are DEEP contributions at height +/// `h` (codeword length `2^h`), or `None` otherwise. +/// +/// Folding starts from the tallest bucket. After each fold to height `h`, the +/// bucket at `combined[h]` is injected into the running codeword with +/// coefficient `β²` (β being the fold challenge just used), before the layer is +/// committed. Termination follows [`BatchedFriLayout`]: the running codeword is +/// folded to the terminal length and the terminal polynomial's coefficients are +/// appended to the transcript, exactly as +/// [`crate::fri::commit_phase_from_evaluations`] does — not folded down to a +/// single scalar. +/// +/// Layer trees are built with `FriLayerMerkleTreeBackend`, the same commitment +/// backend the unbatched [`crate::fri::commit_phase_from_evaluations`] uses — so a +/// batched prover and the verifier that authenticates its openings through +/// `BatchedMerkleTreeBackend` agree on the hash by naming one backend, not by two +/// call sites coinciding. +#[allow(clippy::type_complexity)] +pub fn batched_commit_phase( + mut combined: Vec>>>, + transcript: &mut T, + coset_offset: &FieldElement, + blowup_log: u32, + final_poly_log_degree: u32, +) -> (Vec>, Vec>>) +where + F: IsFFTField + IsSubFieldOf + 'static, + E: IsField + 'static + Send + Sync, + T: IsStarkTranscript + Clone, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, +{ + let (h_min, h_max) = bucket_height_range(&combined) + .expect("batched_commit_phase: combined must have at least one Some entry"); + + // Take the starting codeword — NOT committed; it plays the role of layer 0. + let mut running = combined[h_max] + .take() + .expect("combined[h_max] is Some by construction"); + + let domain_size = 1usize << h_max; + debug_assert_eq!( + running.len(), + domain_size, + "starting codeword length must equal 2^h_max" + ); + + let layout = BatchedFriLayout::new(h_max, h_min, blowup_log, final_poly_log_degree); + + // Inverse twiddle factors for the initial domain size. + let mut inv_twiddles = compute_coset_twiddles_inv(coset_offset, domain_size); + + let mut fri_layer_list = Vec::with_capacity(layout.num_committed); + + for _ in 0..layout.num_committed { + // <<<< Receive challenge β + let beta = transcript.sample_field_element(); + + // Fold evaluations in-place; running halves in length. + fold_evaluations_in_place(&mut running, &beta, &inv_twiddles); + inject_bucket(&mut running, &mut combined, &beta); + + // Build the row-pair Merkle tree over the current running codeword. + let leaves: Vec<[FieldElement; 2]> = running + .chunks_exact(2) + .map(|chunk| [chunk[0].clone(), chunk[1].clone()]) + .collect(); + let merkle_tree = MerkleTree::>::build(&leaves) + .expect("FRI batched commit: Merkle tree construction must succeed"); + let root = merkle_tree.root; + fri_layer_list.push(FriLayer::new(&running, merkle_tree)); + + // >>>> Send commitment: append root to transcript. + transcript.append_bytes(&root); + + // Update twiddles for the next (halved) level. + update_twiddles_in_place(&mut inv_twiddles); + } + + // One final fold to reach the terminal codeword, unless already there. The + // bucket AT the terminal height is injected here: it is the last one that can + // still enter the running word, which is why the layout floors the terminal + // at the shortest height rather than at `blowup_log + k` alone. + if layout.total_folds > 0 { + let beta = transcript.sample_field_element(); + fold_evaluations_in_place(&mut running, &beta, &inv_twiddles); + inject_bucket(&mut running, &mut combined, &beta); + } + debug_assert_eq!( + running.len(), + layout.terminal_len, + "terminal codeword size mismatch" + ); + debug_assert!( + combined.iter().all(Option::is_none), + "every bucket must have been injected before the terminal" + ); + + // Recover the terminal polynomial's coefficients and send them, mirroring + // `commit_phase_from_evaluations`: the coefficient count follows + // `layout.effective_k` (the actual terminal), and the terminal coset offset + // is `coset_offset^(2^total_folds)`. + let terminal_offset = coset_offset.pow(1u64 << layout.total_folds); + let final_poly_coeffs = + coeffs_from_terminal_codeword::(&running, &terminal_offset, layout.effective_k); + for c in &final_poly_coeffs { + transcript.append_field_element(c); + } + + (final_poly_coeffs, fri_layer_list) +} + +/// The `(h_min, h_max)` of the occupied buckets, or `None` when none are. +fn bucket_height_range( + combined: &[Option>>], +) -> Option<(usize, usize)> { + let mut occupied = combined + .iter() + .enumerate() + .filter_map(|(h, slot)| slot.as_ref().map(|_| h)); + let first = occupied.next()?; + Some((first, occupied.next_back().unwrap_or(first))) +} + +/// `running += β² · combined[h]` for the running codeword's current height `h`, +/// consuming that bucket. A no-op when the bucket is empty. +fn inject_bucket( + running: &mut [FieldElement], + combined: &mut [Option>>], + beta: &FieldElement, +) { + let h = running.len().trailing_zeros() as usize; + let Some(bucket) = combined.get_mut(h).and_then(Option::take) else { + return; + }; + debug_assert_eq!( + bucket.len(), + running.len(), + "a bucket at height {h} must match the running codeword's length" + ); + let beta_sq = beta.square(); + for (val, contribution) in running.iter_mut().zip(bucket.iter()) { + *val = &*val + &(&beta_sq * contribution); + } +} + +/// Canonical, order-deterministic absorption of an epoch's table-SHAPE histogram +/// into the transcript. Single source of truth for the structural binding. +/// +/// The multiset of `lde_log_height`s across an epoch's tables fully determines +/// the fold order and injection points of the batched FRI (arity is uniformly +/// 2), so binding the heights binds the whole injection schedule. The widths are +/// bound alongside them because they are what makes the mixed-height MMCS leaf +/// parse unambiguous (see [`crate::fri::mmcs`]'s width-binding section) — the +/// verifier derives widths from the AIR set rather than the proof, so this is +/// defence in depth rather than the primary binding, and it costs one field per +/// table. +/// +/// Encoding (fixed-width, length-prefixed, order-preserving): +/// `u64::to_le_bytes(len)` followed by `u64::to_le_bytes(h)`, `u64::to_le_bytes(w)` +/// for each `(h, w)` pair, in the exact order given. Caller (prover and verifier +/// alike) must pass the shape in the same canonical per-epoch table order — this +/// function does not sort or deduplicate. +/// +/// Panics if `heights` and `widths` differ in length; both sides construct them +/// from the same table list. +pub fn absorb_shape_histogram(transcript: &mut T, heights: &[usize], widths: &[usize]) +where + E: IsField, + T: IsTranscript, +{ + assert_eq!( + heights.len(), + widths.len(), + "the shape histogram needs one width per height" + ); + transcript.append_bytes(&(heights.len() as u64).to_le_bytes()); + for (h, w) in heights.iter().zip(widths.iter()) { + transcript.append_bytes(&(*h as u64).to_le_bytes()); + transcript.append_bytes(&(*w as u64).to_le_bytes()); + } +} + +/// Challenges derived from replaying the shared batched round-4 transcript +/// sequence. See [`derive_batched_fri_challenges`]. +#[derive(Debug, Clone)] +pub struct BatchedFriChallenges { + /// Sampled once after the shape histogram (and, at the call site, after all + /// per-table OOD evaluations have been absorbed). + pub alpha: FieldElement, + /// One per committed layer, plus one for the final fold when there is one: + /// `betas.len() == layout.num_committed + (layout.total_folds > 0) as usize`. + pub betas: Vec>, + /// The layout the betas and the terminal were derived under. + pub layout: BatchedFriLayout, + /// Transcript state right before the grinding nonce bytes are appended. + /// All-zero when `grinding_factor == 0` or `nonce` is `None`. + pub grinding_seed: [u8; 32], + /// One `sample_u64(2^(h_max - 1))` draw per query — a row-PAIR index in the + /// tallest domain. A round whose own `h_max` is lower must reduce these; see + /// [`crate::fri::mmcs`]'s index-convention section. + pub iotas: Vec, + /// Which tables the batched instance carries and which keep a terminal-only + /// instance of their own. Derived from the shape, never sent. + pub plan: FriInstancePlan, +} + +/// Replays the shared batched round-4 transcript sequence (shape histogram, +/// alpha, per-layer beta/root, final beta, terminal coefficients, grinding, query +/// iotas) and returns the derived challenges. The one routine the prover and the +/// verifier both call, so they provably derive identical challenges. +/// +/// `standalone_coeffs[t]` is table `t`'s terminal-only polynomial, `Some` +/// exactly for the standalone class — presence is checked against the derived +/// plan and every coefficient is ABSORBED, right after `α` and before the +/// first `ζ`. That absorb is load-bearing: the standalone check evaluates the +/// sent polynomial at the query indices drawn BELOW, so a polynomial that +/// were not bound here could be chosen after the indices are known, and each +/// query's proximity test would bind nothing until the queries saturate the +/// table's domain. The unbatched path absorbs its terminal before sampling +/// queries for the same reason; this keeps the batched path's binding equal. +/// +/// Returns `None` when the proof's layer-root count disagrees with the layout the +/// epoch's shape implies, when the terminal coefficient count is wrong, or when +/// a standalone polynomial is present for the wrong class — all prover-supplied, +/// all rejections, not panics. +#[allow(clippy::too_many_arguments)] +pub fn derive_batched_fri_challenges( + transcript: &mut T, + heights: &[usize], + widths: &[usize], + layer_roots: &[[u8; 32]], + final_poly_coeffs: &[FieldElement], + standalone_coeffs: &[Option<&[FieldElement]>], + blowup_log: u32, + final_poly_log_degree: u32, + grinding_factor: u8, + nonce: Option, + num_queries: usize, +) -> Option> +where + E: IsField, + T: IsTranscript, +{ + // The partition is derived, not sent: it is a pure function of the shape the + // histogram below binds, so both sides reach the same one. `None` on any + // height that cannot be a codeword length — heights come from proof-supplied + // trace lengths, so a bogus one is a rejection, never a panic on the + // verifier's path. + let plan = FriInstancePlan::new(heights, blowup_log, final_poly_log_degree)?; + let (h_max, h_min) = (plan.h_max, plan.h_min); + let layout = BatchedFriLayout::new(h_max, h_min, blowup_log, final_poly_log_degree); + if layer_roots.len() != layout.num_committed + || final_poly_coeffs.len() != 1usize << layout.effective_k + || standalone_coeffs.len() != heights.len() + { + return None; + } + + absorb_shape_histogram(transcript, heights, widths); + + let alpha = transcript.sample_field_element(); + + // The standalone class's terminal polynomials, bound before any query can + // depend on them — per table ascending, each coefficient in order. The + // length pin (`2^(h_t − blowup_log)`, exactly) stays with + // `verify_epoch_commitments`. + for (table, coeffs) in standalone_coeffs.iter().enumerate() { + if coeffs.is_some() != plan.standalone.contains(&table) { + return None; + } + if let Some(coeffs) = coeffs { + for c in coeffs.iter() { + transcript.append_field_element(c); + } + } + } + + let mut betas = Vec::with_capacity(layout.num_committed + 1); + for root in layer_roots { + let beta = transcript.sample_field_element(); + transcript.append_bytes(root); + betas.push(beta); + } + + if layout.total_folds > 0 { + betas.push(transcript.sample_field_element()); + } + for c in final_poly_coeffs { + transcript.append_field_element(c); + } + + let mut grinding_seed = [0u8; 32]; + if grinding_factor > 0 + && let Some(nonce_value) = nonce + { + grinding_seed = transcript.state(); + transcript.append_bytes(&nonce_value.to_be_bytes()); + } + + let iotas = (0..num_queries) + .map(|_| transcript.sample_u64(1u64 << (h_max - 1)) as usize) + .collect(); + + Some(BatchedFriChallenges { + alpha, + betas, + layout, + grinding_seed, + iotas, + plan, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::fri::commit_phase_from_evaluations; + use crate::fri::fri_functions::{compute_coset_twiddles_inv, fold_evaluations_in_place}; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use crypto::fiat_shamir::is_transcript::IsTranscript; + use math::field::element::FieldElement; + use math::field::goldilocks::GoldilocksField; + + type FE = FieldElement; + type Transcript = DefaultTranscript; + + #[test] + fn combine_by_height_two_height3_one_height2() { + // Three codewords: indices 0, 1 have height 3 (length 8); + // index 2 has height 2 (length 4). + let cw0: Vec = (1u64..=8).map(FE::from).collect(); + let cw1: Vec = (10u64..=17).map(FE::from).collect(); + let cw2: Vec = (100u64..=103).map(FE::from).collect(); + + let alpha = FE::from(7u64); + + let inputs: Vec<(Vec, usize)> = + vec![(cw0.clone(), 3), (cw1.clone(), 3), (cw2.clone(), 2)]; + + let out = combine_by_height(&inputs, &alpha); + + // Output vec length = max_height + 1 = 4 (indices 0..=3 only). + assert_eq!(out.len(), 4, "output length should be max_height+1 = 4"); + + // Heights 0 and 1 have no inputs. + assert!(out[0].is_none(), "height 0 should be None"); + assert!(out[1].is_none(), "height 1 should be None"); + + // Height 3: combined[j] = alpha^0 * cw0[j] + alpha^1 * cw1[j] + let alpha0 = FE::one(); + let alpha1 = alpha; + let expected3: Vec = cw0 + .iter() + .zip(cw1.iter()) + .map(|(a, b)| &(&alpha0 * a) + &(&alpha1 * b)) + .collect(); + + let got3 = out[3].as_ref().expect("height 3 should be Some"); + assert_eq!( + got3.len(), + 8, + "height-3 combined codeword should have length 8" + ); + assert_eq!(got3, &expected3, "height-3 combined values mismatch"); + + // Height 2: combined[j] = alpha^2 * cw2[j] + let alpha2 = &alpha * α + let expected2: Vec = cw2.iter().map(|x| &alpha2 * x).collect(); + + let got2 = out[2].as_ref().expect("height 2 should be Some"); + assert_eq!( + got2.len(), + 4, + "height-2 combined codeword should have length 4" + ); + assert_eq!(got2, &expected2, "height-2 combined values mismatch"); + } + + /// Absorbing codewords one at a time — the shape a prover uses so it never + /// holds every table's quotient at once — must land on the same buckets as + /// handing them all over materialized. + #[test] + fn streaming_absorption_matches_materialized_combine() { + let inputs: Vec<(Vec, usize)> = vec![ + ((1u64..=16).map(FE::from).collect(), 4), + ((50u64..=57).map(FE::from).collect(), 3), + ((90u64..=105).map(FE::from).collect(), 4), + ((200u64..=203).map(FE::from).collect(), 2), + ((300u64..=307).map(FE::from).collect(), 3), + ]; + let alpha = FE::from(11u64); + + let eager = combine_by_height(&inputs, &alpha); + + let mut combiner = HeightCombiner::new(alpha); + for (codeword, height) in &inputs { + combiner.absorb(codeword, *height); + } + assert_eq!( + combiner.finish(), + eager, + "streaming absorption must equal the materialized combine" + ); + } + + /// After the first fold in `batched_commit_phase`, the committed layer[0] + /// evaluation must equal `fold(combined[4], β₀) + β₀² · combined[3]`. + #[test] + fn batched_commit_phase_first_layer_matches_manual_fold_and_inject() { + // Build synthetic codewords for h=4 (len 16) and h=3 (len 8). + let data_h4: Vec = (1u64..=16).map(FE::from).collect(); + let data_h3: Vec = (101u64..=108).map(FE::from).collect(); + + // combined = [None, None, None, Some(data_h3), Some(data_h4)] + let combined: Vec>> = vec![ + None, + None, + None, + Some(data_h3.clone()), + Some(data_h4.clone()), + ]; + + let coset_offset = FE::from(3u64); + let (blowup_log, k) = (1u32, 1u32); + + // Create transcript; clone before mutating so we can replay independently. + let mut transcript = Transcript::new(b"batched_fri_test"); + let mut transcript_check = transcript.clone(); + + let (_coeffs, layers) = batched_commit_phase::<_, _, _>( + combined, + &mut transcript, + &coset_offset, + blowup_log, + k, + ); + + // Terminal at min(blowup_log + k, h_min) = min(2, 3) = 2, so folds run + // 4 -> 2: two folds, one committed layer. + let layout = BatchedFriLayout::new(4, 3, blowup_log, k); + assert_eq!(layout.total_folds, 2); + assert_eq!( + layers.len(), + layout.num_committed, + "committed layers must follow the layout" + ); + + // --- Independent recomputation of layer[0] --- + let beta_0 = transcript_check.sample_field_element(); + + let inv_twiddles_h4 = compute_coset_twiddles_inv::(&coset_offset, 16); + let mut expected = data_h4.clone(); + fold_evaluations_in_place(&mut expected, &beta_0, &inv_twiddles_h4); + // expected now has length 8 (height 3) + + // Inject combined[3]: expected[j] += beta_0² · data_h3[j] + let beta_0_sq = beta_0.square(); + for (j, val) in data_h3.iter().enumerate() { + expected[j] = &expected[j] + &(&beta_0_sq * val); + } + + assert_eq!( + layers[0].evaluation, expected, + "layer[0] evaluation does not match manual fold+inject" + ); + } + + /// ★ M-12: the batched commit phase must terminate where the unbatched one + /// does. With a single bucket the two are the same protocol, so they must + /// agree on the committed-layer count, the terminal coefficients, and the + /// resulting transcript state — pinning that batching did not silently switch + /// to folding all the way to a scalar (which for this input would commit + /// `h_max - 1 = 9` layers instead of 4). + #[test] + fn single_bucket_terminal_matches_the_unbatched_commit_phase() { + let h = 10usize; + let (blowup_log, k) = (1u32, 5u32); + let coset_offset = FE::from(3u64); + let evals: Vec = (0..(1u64 << h)).map(|i| FE::from(i * 7 + 1)).collect(); + let inv_twiddles = compute_coset_twiddles_inv::(&coset_offset, 1 << h); + + let mut t_unbatched = Transcript::new(b"terminal_parity"); + let (unbatched_coeffs, unbatched_layers) = commit_phase_from_evaluations::< + GoldilocksField, + GoldilocksField, + Transcript, + >( + evals.clone(), + &mut t_unbatched, + &coset_offset, + 1 << h, + blowup_log, + k, + &inv_twiddles, + ); + + let mut combined: Vec>> = vec![None; h + 1]; + combined[h] = Some(evals); + let mut t_batched = Transcript::new(b"terminal_parity"); + let (batched_coeffs, batched_layers) = batched_commit_phase::<_, _, _>( + combined, + &mut t_batched, + &coset_offset, + blowup_log, + k, + ); + + // total_folds = 10 - (1 + 5) = 4, so 3 committed layers — not h_max-1 = 9. + assert_eq!(unbatched_layers.len(), 3); + assert_eq!( + batched_layers.len(), + unbatched_layers.len(), + "batched and unbatched must commit the same number of layers" + ); + assert_eq!( + batched_coeffs.len(), + 1usize << k, + "the terminal polynomial must carry 2^k coefficients" + ); + assert_eq!( + batched_coeffs, unbatched_coeffs, + "batched and unbatched must send the same terminal polynomial" + ); + for (b, u) in batched_layers.iter().zip(unbatched_layers.iter()) { + assert_eq!(b.merkle_tree.root, u.merkle_tree.root); + } + assert_eq!( + t_batched.state(), + t_unbatched.state(), + "the two commit phases must leave the transcript in the same state" + ); + } + + /// The batched-only floor: the terminal may not sit above the shortest + /// injected codeword, or that bucket would never enter the running word. + #[test] + fn terminal_is_floored_at_the_shortest_codeword() { + let (blowup_log, k) = (1u32, 5u32); + + // Shortest codeword above blowup_log + k = 6: the floor is inert and the + // layout is the unbatched one for h_max. + let inert = BatchedFriLayout::new(10, 8, blowup_log, k); + assert_eq!(inert.total_folds, 4, "10 -> 6"); + assert_eq!(inert.effective_k, k); + + // Shortest codeword BELOW blowup_log + k: folding must continue down to + // it, and the terminal polynomial shrinks accordingly. + let floored = BatchedFriLayout::new(10, 4, blowup_log, k); + assert_eq!(floored.total_folds, 6, "10 -> 4"); + assert_eq!(floored.effective_k, 3, "terminal_log 4 - blowup_log 1"); + + // And the commit phase really does consume that low bucket. + let coset_offset = FE::from(3u64); + let mut combined: Vec>> = vec![None; 8]; + combined[7] = Some((0..128u64).map(|i| FE::from(i + 1)).collect()); + combined[4] = Some((0..16u64).map(|i| FE::from(i * 3 + 5)).collect()); + let mut transcript = Transcript::new(b"floor_test"); + let (coeffs, layers) = batched_commit_phase::<_, _, _>( + combined, + &mut transcript, + &coset_offset, + blowup_log, + k, + ); + let layout = BatchedFriLayout::new(7, 4, blowup_log, k); + assert_eq!(layers.len(), layout.num_committed); + assert_eq!(coeffs.len(), 1usize << layout.effective_k); + } + + /// The prover, by hand, runs exactly the round-4 sequence; the shared replay + /// routine must reproduce byte-identical outputs from the same start state. + #[test] + fn batched_round4_prover_inline_matches_verifier_replay() { + let heights: Vec = vec![10, 10, 8, 8, 8, 7]; + let widths: Vec = vec![3, 5, 2, 2, 9, 1]; + let (blowup_log, k) = (1u32, 5u32); + // total_folds = 10 - 6 = 4 -> 3 committed layers, 4 betas. + let layout = BatchedFriLayout::new(10, 7, blowup_log, k); + assert_eq!((layout.num_committed, layout.total_folds), (3, 4)); + + let layer_roots: Vec<[u8; 32]> = (0u8..3).map(|i| [i; 32]).collect(); + let final_poly_coeffs: Vec = (0..(1u64 << layout.effective_k)).map(FE::from).collect(); + // Height 7 folds no layer at these parameters, so table 5 is standalone + // and its terminal polynomial is part of the round-4 sequence. + let standalone_terminal: Vec = (0..(1u64 << (7 - blowup_log))).map(FE::from).collect(); + + let grinding_factor: u8 = 4; + let num_queries = 3; + + let seed_transcript = Transcript::new(b"batched_round4_test"); + let mut transcript_a = seed_transcript.clone(); + let mut transcript_b = seed_transcript.clone(); + + // --- Clone A: prover-inline sequence, by hand --- + absorb_shape_histogram(&mut transcript_a, &heights, &widths); + let alpha_a = transcript_a.sample_field_element(); + for c in &standalone_terminal { + transcript_a.append_field_element(c); + } + + let mut betas_a = Vec::with_capacity(layer_roots.len() + 1); + for root in &layer_roots { + let beta = transcript_a.sample_field_element(); + transcript_a.append_bytes(root); + betas_a.push(beta); + } + betas_a.push(transcript_a.sample_field_element()); + for c in &final_poly_coeffs { + transcript_a.append_field_element(c); + } + assert_eq!( + betas_a.len(), + layout.total_folds as usize, + "one beta per fold, matching batched_commit_phase" + ); + + let grinding_seed_a = transcript_a.state(); + // Test-only: derive a real PoW nonce so the grinding step is exercised + // identically by both sides (the nonce search itself is not under test). + let nonce = crate::grinding::generate_nonce(&grinding_seed_a, grinding_factor) + .expect("a valid grinding nonce exists for this small grinding_factor"); + transcript_a.append_bytes(&nonce.to_be_bytes()); + + let iotas_a: Vec = (0..num_queries) + .map(|_| transcript_a.sample_u64(1u64 << 9) as usize) + .collect(); + + // --- Clone B: shared replay routine --- + let standalone: Vec> = vec![ + None, + None, + None, + None, + None, + Some(standalone_terminal.as_slice()), + ]; + let result = derive_batched_fri_challenges( + &mut transcript_b, + &heights, + &widths, + &layer_roots, + &final_poly_coeffs, + &standalone, + blowup_log, + k, + grinding_factor, + Some(nonce), + num_queries, + ) + .expect("a well-formed layer-root and coefficient count"); + + assert_eq!(result.alpha, alpha_a, "alpha mismatch"); + assert_eq!(result.betas, betas_a, "beta vector mismatch"); + assert_eq!(result.layout, layout, "layout mismatch"); + assert_eq!( + result.grinding_seed, grinding_seed_a, + "grinding seed mismatch" + ); + assert_eq!(result.iotas, iotas_a, "iotas mismatch"); + assert!( + result.iotas.iter().all(|&i| i < 1usize << 9), + "iotas must be row-pair indices in the tallest domain" + ); + } + + /// A layer-root or coefficient count that disagrees with the shape's layout is + /// prover-supplied, so it is a rejection rather than a panic. + #[test] + fn derive_rejects_a_layer_count_that_contradicts_the_shape() { + let heights: Vec = vec![10, 8]; + let widths: Vec = vec![2, 3]; + let (blowup_log, k) = (1u32, 5u32); + let layout = BatchedFriLayout::new(10, 8, blowup_log, k); + let coeffs: Vec = vec![FE::one(); 1usize << layout.effective_k]; + let roots: Vec<[u8; 32]> = vec![[0u8; 32]; layout.num_committed]; + + let no_standalone: Vec> = vec![None; heights.len()]; + let mut ok = Transcript::new(b"reject"); + assert!( + derive_batched_fri_challenges( + &mut ok, + &heights, + &widths, + &roots, + &coeffs, + &no_standalone, + blowup_log, + k, + 0, + None, + 1 + ) + .is_some() + ); + + let mut too_few = Transcript::new(b"reject"); + assert!( + derive_batched_fri_challenges( + &mut too_few, + &heights, + &widths, + &roots[..roots.len() - 1], + &coeffs, + &no_standalone, + blowup_log, + k, + 0, + None, + 1 + ) + .is_none(), + "one fewer layer root than the shape implies must be rejected" + ); + + let mut bad_coeffs = Transcript::new(b"reject"); + assert!( + derive_batched_fri_challenges( + &mut bad_coeffs, + &heights, + &widths, + &roots, + &coeffs[..coeffs.len() - 1], + &no_standalone, + blowup_log, + k, + 0, + None, + 1 + ) + .is_none(), + "a short terminal polynomial must be rejected" + ); + } + + /// `heights` comes from proof-supplied trace lengths, so every out-of-range + /// value is a rejection rather than a shift overflow or a layout assert. + #[test] + fn derive_rejects_out_of_range_heights_without_panicking() { + let widths = vec![2usize, 3]; + let (blowup_log, k) = (1u32, 5u32); + let coeffs: Vec = vec![FE::one(); 1usize << k]; + let roots: Vec<[u8; 32]> = vec![[0u8; 32]; 3]; + + let derive = |heights: &[usize]| { + let no_standalone: Vec> = vec![None; heights.len()]; + derive_batched_fri_challenges( + &mut Transcript::new(b"range"), + heights, + &widths, + &roots, + &coeffs, + &no_standalone, + blowup_log, + k, + 0, + None, + 1, + ) + .is_some() + }; + + assert!(derive(&[10, 8]), "a well-formed shape is accepted"); + assert!(!derive(&[0, 0]), "a zero height must be rejected"); + assert!( + !derive(&[10, 0]), + "a height below the blowup must be rejected" + ); + assert!( + !derive(&[u32::BITS as usize, 8]), + "a height at the shift width must be rejected" + ); + assert!( + !derive(&[usize::MAX, 8]), + "an absurd height must be rejected, not wrapped by the u32 cast" + ); + let empty: [usize; 0] = []; + assert!(!derive(&empty), "an empty epoch must be rejected"); + } + + /// Tampering the shape histogram (without changing anything else) must change + /// the derived batching challenge α — the structural binding that protects the + /// fold/injection schedule. Heights and widths are both bound (M-13a), so a + /// change to either alone must move α. + #[test] + fn absorb_shape_histogram_binds_heights_and_widths_into_alpha() { + let heights: Vec = vec![10, 10, 8, 8, 8, 5]; + let widths: Vec = vec![4, 4, 2, 2, 2, 1]; + + let alpha_of = |h: &[usize], w: &[usize]| { + let mut t = Transcript::new(b"histogram_binding_test"); + absorb_shape_histogram(&mut t, h, w); + t.sample_field_element() + }; + + let base = alpha_of(&heights, &widths); + + let mut other_height = heights.clone(); + other_height[5] = 6; + assert_ne!( + base, + alpha_of(&other_height, &widths), + "different height histograms must yield different alpha" + ); + + let mut other_width = widths.clone(); + other_width[5] = 2; + assert_ne!( + base, + alpha_of(&heights, &other_width), + "different width histograms must yield different alpha" + ); + + // The length prefix plus fixed-width fields make the encoding injective: + // swapping a (height, width) pair between tables also moves alpha. + let swapped_h = vec![10, 10, 8, 8, 5, 8]; + let swapped_w = vec![4, 4, 2, 2, 1, 2]; + assert_ne!( + base, + alpha_of(&swapped_h, &swapped_w), + "table order must be bound, not just the multiset" + ); + } +} diff --git a/crypto/stark/src/fri/mmcs.rs b/crypto/stark/src/fri/mmcs.rs new file mode 100644 index 000000000..d940dd4e9 --- /dev/null +++ b/crypto/stark/src/fri/mmcs.rs @@ -0,0 +1,1805 @@ +//! Mixed-height, row-pair MMCS (Merkle Mixed Commitment Scheme). +//! +//! Commits ALL of an epoch's matrices (one per table, of possibly different +//! heights) into ONE mixed-height Merkle tree, so a single query opens ONE +//! authentication path that covers every table's row at that query — the +//! proof-size / opening-path win of the unified-shard design (SP1 / OpenVM / +//! Plonky3). Mirrors Plonky3's `MerkleTreeMmcs`, adapted to the concrete keccak +//! commitment backends and to the row-pair `(x, -x)` leaf layout (#735). +//! +//! This is a standalone primitive: the prover and verifier do not build epoch +//! commitments with it yet. The leaf and injection layout documented below is +//! the single source of truth for whoever wires it in. +//! +//! # Inputs +//! +//! [`MixedMmcs::commit`] reads matrices through a [`LeafSource`], which reports +//! each matrix's `(log_height, width)` and serves its rows on demand: +//! - `log_height`: `log2` of the row count; the matrix has `2^log_height` rows. +//! - `width`: number of committed columns. +//! - rows are addressed by **bit-reversed** LDE position (the same layout the +//! per-table trace commit produces internally). +//! +//! # Row-pair leaves +//! +//! Leaf `k` of a matrix groups LDE positions `2k` and `2k+1` (the FRI fold pair +//! `x` and `-x`), all `width` columns batched. A matrix of `log_height h` has +//! `2^(h-1)` leaves. In [`MixedMmcs::open_batch`] / [`PolynomialOpenings`]: +//! `evaluations` = row `2k`, `evaluations_sym` = row `2k+1`. +//! +//! # Tree layout (the soundness-relevant contract) +//! +//! Let `h_max = max(log_height)`. The base digest layer (layer 0) has +//! `N0 = 2^(h_max-1)` nodes. Layer `i` has `N0 >> i` nodes; the root is the sole +//! node of layer `h_max-1`. A matrix of `log_height h` is *injected* at layer +//! index `i = h_max - h` (so the tallest matrices, `h == h_max`, populate the +//! base layer; shorter matrices enter where the layer width matches their leaf +//! count `2^(h-1)`). +//! +//! Hashing (`H = >::hash_data` over a `Vec` of field +//! elements; `C = >::hash_new_parent`, the 2-input +//! compression — the same two functions, on the same backend, that the existing +//! per-table tree uses): +//! +//! - **Base layer** node `k` (`k in [0, N0)`): +//! `layer0[k] = H( CONCAT_{m : h_m == h_max} (row_m(2k) || row_m(2k+1)) )` +//! where matrices of height `h_max` are concatenated in INPUT order. +//! - **Climb** from layer `i` to layer `i+1` (`j in [0, N_{i+1})`): +//! `parent = C(layer_i[2j], layer_i[2j+1])`. Let `inject_h = h_max - 1 - i`. If +//! any matrix has `h_m == inject_h`, then +//! `layer_{i+1}[j] = C( parent, H( CONCAT_{m : h_m == inject_h} (row_m(2j) || row_m(2j+1)) ) )` +//! (injecting matrices concatenated in INPUT order); otherwise +//! `layer_{i+1}[j] = parent`. +//! - `root = layer_{h_max-1}[0]`. +//! +//! Because the leaf and parent hashes come from `BatchedMerkleTreeBackend` — +//! the backend the per-table row-pair tree already commits with — a single-matrix `MixedMmcs` is +//! byte-identical to that tree by construction, not by coincidence. There is no +//! second encoding of a leaf to keep in step. +//! +//! # Query opening +//! +//! For query `iota in [0, N0)`, matrix `m` is opened at leaf +//! `k_m = iota >> (h_max - h_m)` (`= iota >> i_m`). The shared authentication +//! path holds, for each level `level in [0, h_max-1)`, the sibling +//! `layer_level[(iota >> level) ^ 1]`. ONE path authenticates all matrices. +//! The per-matrix [`PolynomialOpenings::proof`] fields are empty; the single +//! [`MixedOpening::proof`] is the authenticator. +//! +//! # ★ Index convention — a HARD PRECONDITION on the caller +//! +//! `iota` is a leaf index **in THIS tree**: it must be drawn from +//! `[0, 2^(h_max-1))` where `h_max` is *this MMCS's* tallest matrix. +//! [`MixedMmcs::verify_batch`] walks the path with `(iota >> level) & 1`, i.e. it +//! consumes the **low** `h_max - 1` bits, while a shorter matrix inside the tree +//! is located by `iota >> (h_max - h_m)`, i.e. by the **high** bits. Both are +//! consistent only when the two `h_max` agree. +//! +//! A caller that batches several rounds under one shared FRI query index must +//! therefore reduce a global index before calling in: +//! +//! ```text +//! iota_round = iota_fri >> (h_max_fri - h_max_round) +//! ``` +//! +//! Passing the un-reduced `iota_fri` to a round whose `h_max` is below the FRI's +//! is not a loud error — prover and verifier share this routine, so a wrong +//! convention is self-consistent: honest proofs still verify and the failure is +//! that short matrices end up authenticated at positions the FRI join never +//! checks. [`MixedMmcs::verify_batch`] rejects an `iota` outside `[0, 2^(h_max-1))` +//! to turn most of that class of misuse into a rejection rather than a silent +//! mis-binding, but the reduction remains the caller's obligation: an index that +//! happens to land in range is accepted at the wrong leaf. +//! `short_round_low_bit_convention_is_exercised` is the control on this. +//! +//! # Width binding (soundness) +//! +//! [`MixedMmcs::verify_batch`] takes per-matrix `widths` alongside `heights`. +//! Within a height group the leaf hash is over the FLAT concatenation of every +//! matrix's opened row pair (`A.eval ‖ A.eval_sym ‖ B.eval ‖ B.eval_sym ‖ …`), +//! which does NOT by itself record where each matrix's columns end. Fixing +//! `widths[m]` (matrix `m`'s column count) makes those boundaries unambiguous: +//! without it a prover could shift a boundary — e.g. lengthen one matrix's +//! `evaluations` by one element and shorten its `evaluations_sym` by one — +//! leaving the flat bytes (and therefore the group hash) identical while feeding +//! a corrupted row downstream. Consumers MUST pass the committed public +//! per-table column counts, in the same INPUT order as `heights`, derived from +//! the AIR set rather than read out of the proof. +//! +//! `heights` and `widths` must ALSO be bound into the Fiat-Shamir transcript by +//! the consumer, before any challenge that depends on the epoch's shape — see +//! [`crate::fri::batched::absorb_shape_histogram`], which is the canonical +//! encoding of that binding. +//! +//! # Determinism +//! +//! The tree is a pure function of `(matrices, input order)`. Grouping within a +//! height (base batching and injection) follows INPUT order; the prover and +//! verifier MUST pass matrices and `heights` in the same per-epoch order. +//! +//! # Memory: what the caller may drop, and when +//! +//! The MMCS owns no evaluations. It stores the digest layers +//! (`O(2^(h_max-1))` nodes) plus each matrix's `(log_height, width)`; rows are +//! pulled through [`LeafSource`] both at commit and at open time. Two properties +//! follow, and `commit_reads_each_height_group_in_one_contiguous_phase` is the +//! control on the second: +//! +//! - `commit` reads matrix `m`'s rows **only while building level +//! `h_max - h_m`**, and levels are built in descending height order. A caller +//! may therefore produce a height group's LDEs, commit, and drop them before +//! the next group is needed. +//! - Within one height group the leaf is a single `hash_data` over the group's +//! concatenated rows, so `commit` reads every matrix of that height at every +//! leaf: their access windows overlap, and a caller serving them from in-RAM +//! LDE buffers holds the whole group at once. Since the tallest group is most +//! of a real epoch's tables, that is `O(N)` resident at the base layer. +//! +//! A `LeafSource` serving rows from disk, device memory or recomputation is the +//! escape for that base-group residency: it lets a caller stream a matrix in, +//! hash it, and drop it without holding the whole group in RAM at once. +//! +//! (A streaming, incremental-leaf-hasher builder that keeps one hasher per leaf +//! and absorbs matrices as they arrive is planned but not part of this phase.) + +use core::marker::PhantomData; + +use crypto::merkle_tree::proof::Proof; +use crypto::merkle_tree::traits::{IsLeafHasher, IsMerkleTreeBackend, IsStreamingLeafBackend}; +use math::fft::bit_reversing::reverse_index; +use math::field::element::FieldElement; +use math::field::traits::IsField; +use math::traits::AsBytes; + +use crate::config::{BatchedMerkleTreeBackend, Commitment}; +use crate::proof::stark::PolynomialOpenings; + +/// On-demand supplier of committed matrix rows, so [`MixedMmcs`] builds its +/// digests and serves openings WITHOUT owning a copy of the (large) LDE buffers. +/// Both [`MixedMmcs::commit`] and [`MixedMmcs::open_batch`] read every leaf +/// through this trait, so the root and opened rows are byte-identical to those a +/// matrix-owning MMCS would produce — the prover keeps only the LDE buffers it +/// already retains for DEEP, and each MMCS stores just digests. +/// +/// Rows are addressed in each matrix's committed row-pair layout: `append_row(m, +/// r, out)` appends matrix `m`'s row at **bit-reversed** LDE position `r` (its +/// `width(m)` committed columns, in column order). This is the same `r`-indexing +/// the module's "Tree layout" section uses; an implementor holding the +/// natural-order LDE maps `r` to `reverse_index(r, 2^log_height(m))`. +pub trait LeafSource { + /// Number of committed matrices, in canonical input order. + fn num_matrices(&self) -> usize; + /// `log2` of matrix `m`'s row count. Row-pair leaves require `>= 1`. + fn log_height(&self, m: usize) -> usize; + /// Matrix `m`'s committed column count. + fn width(&self, m: usize) -> usize; + /// Append matrix `m`'s bit-reversed LDE row `bitrev_row` (its `width(m)` + /// committed columns) to `out`. `bitrev_row in [0, 2^log_height(m))`. + fn append_row(&self, m: usize, bitrev_row: usize, out: &mut Vec>); +} + +/// One committed matrix borrowed from a retained LDE buffer. Resolves each +/// bit-reversed row on demand (mapping through `reverse_index`) so the MMCS owns +/// no copy of the evaluations. See [`LeafSource`]. +pub enum BorrowedMatrix<'a, E: IsField> { + /// A `stride`-wide, row-major, NATURAL-order LDE buffer (the main / aux LDE + /// retained in `Round1::lde_trace`). This matrix occupies columns + /// `[col_start, col_start + width)`; its bit-reversed row `r` lives at + /// natural-order row `reverse_index(r, 2^log_height)`. + RowMajorNatural { + data: &'a [FieldElement], + stride: usize, + col_start: usize, + width: usize, + log_height: usize, + }, + /// Column-major NATURAL-order columns (the composition-poly LDE retained in + /// `Round2::lde_composition_poly_evaluations`): `cols[c][nat]` is column `c` + /// at natural-order row `nat`. Every committed column is used. + ColMajorNatural { + cols: &'a [Vec>], + log_height: usize, + }, +} + +impl BorrowedMatrix<'_, E> { + fn log_height(&self) -> usize { + match self { + BorrowedMatrix::RowMajorNatural { log_height, .. } + | BorrowedMatrix::ColMajorNatural { log_height, .. } => *log_height, + } + } + + fn width(&self) -> usize { + match self { + BorrowedMatrix::RowMajorNatural { width, .. } => *width, + BorrowedMatrix::ColMajorNatural { cols, .. } => cols.len(), + } + } + + fn append_row(&self, bitrev_row: usize, out: &mut Vec>) { + match self { + BorrowedMatrix::RowMajorNatural { + data, + stride, + col_start, + width, + log_height, + } => { + let nat = reverse_index(bitrev_row, 1u64 << log_height); + let base = nat * stride + col_start; + out.extend_from_slice(&data[base..base + width]); + } + BorrowedMatrix::ColMajorNatural { cols, log_height } => { + let nat = reverse_index(bitrev_row, 1u64 << log_height); + for col in cols.iter() { + out.push(col[nat].clone()); + } + } + } + } +} + +impl LeafSource for Vec> { + fn num_matrices(&self) -> usize { + self.len() + } + fn log_height(&self, m: usize) -> usize { + self[m].log_height() + } + fn width(&self, m: usize) -> usize { + self[m].width() + } + fn append_row(&self, m: usize, bitrev_row: usize, out: &mut Vec>) { + self[m].append_row(bitrev_row, out); + } +} + +/// A committed mixed-height, row-pair MMCS under the concrete keccak commitment +/// backends. Stores ONLY the digest layers (to serve the shared authentication path) +/// plus each matrix's `(log_height, width)` (to locate leaves). The row DATA is +/// served on demand by the caller's [`LeafSource`] — the MMCS never owns a copy +/// of the LDE. +pub struct MixedMmcs { + root: Commitment, + /// `layers[0]` is the base digest layer; `layers[h_max-1] == [root]`. + layers: Vec>, + /// Per committed matrix, in input order: `(log_height, width)`. + dims: Vec<(usize, usize)>, + h_max: usize, + _marker: PhantomData, +} + +/// The opening of ALL matrices at one query index, authenticated by a single +/// shared Merkle path. +#[derive( + Debug, + Clone, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] +#[serde(bound = "")] +pub struct MixedOpening { + /// The one authentication path covering every matrix's row at the query. + pub proof: Proof, + /// Per-matrix row pair (in the same INPUT order as `commit`). Each entry's + /// own `proof` is empty — [`MixedOpening::proof`] is the authenticator. + pub per_matrix: Vec>, +} + +/// Hash the row pair `(row(2*leaf), row(2*leaf+1))` of every matrix whose index +/// is in `group` (in the given order), all columns batched, into one digest. +/// Rows are pulled from `source` — the MMCS owns no copy. +fn hash_group_leaf(source: &S, group: &[usize], leaf: usize) -> Commitment +where + E: IsField + 'static, + S: LeafSource, + FieldElement: AsBytes + Sync + Send, +{ + let mut buf: Vec> = Vec::new(); + for &m in group { + source.append_row(m, 2 * leaf, &mut buf); + source.append_row(m, 2 * leaf + 1, &mut buf); + } + as IsMerkleTreeBackend>::hash_data(&buf) +} + +/// Verifier-side analogue of [`hash_group_leaf`]: hash the opened row pairs of a +/// group of openings (in the given order) into one digest. +fn hash_group_openings(group: &[&PolynomialOpenings]) -> Commitment +where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + let mut buf: Vec> = Vec::new(); + for o in group { + buf.extend_from_slice(&o.evaluations); + buf.extend_from_slice(&o.evaluations_sym); + } + as IsMerkleTreeBackend>::hash_data(&buf) +} + +#[inline] +fn compress(left: &Commitment, right: &Commitment) -> Commitment +where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + as IsMerkleTreeBackend>::hash_new_parent(left, right) +} + +impl MixedMmcs +where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + /// Commit the matrices supplied by `source` into one mixed-height row-pair + /// tree, storing only the digest layers. See the module docs for the exact + /// leaf/injection layout. `source` provides each matrix's dimensions and its + /// bit-reversed rows on demand; no copy of the evaluations is retained. + /// + /// Leaf hashing (the base layer and each injected climb layer) is parallel + /// across leaves via [`crate::par::par_map_collect`]; the per-level output is + /// index-ordered, so the root and layers are byte-identical to a sequential + /// build. `S: Sync` lets leaf closures read `source` from worker threads. + /// + /// Levels are built in descending height order and matrix `m` is read only + /// while its own level is built, so the caller may release a height group's + /// buffers once the next level starts — see the module's memory section. + pub fn commit + Sync>(source: &S) -> Self { + let num_matrices = source.num_matrices(); + assert!( + num_matrices > 0, + "MixedMmcs::commit requires at least one matrix" + ); + + let dims: Vec<(usize, usize)> = (0..num_matrices) + .map(|m| { + let log_height = source.log_height(m); + assert!( + log_height >= 1, + "log_height must be >= 1 (row-pair leaves need at least 2 rows)" + ); + (log_height, source.width(m)) + }) + .collect(); + + let h_max = dims + .iter() + .map(|(log_height, _)| *log_height) + .max() + .expect("dims is non-empty"); + + // Per-height group leaf digests, built in descending height order — the + // order that makes the memory claim in the module header true. Index `h` + // is `Some` exactly when some matrix has that height. + let mut group_digests: Vec>> = vec![None; h_max + 1]; + for h in (1..=h_max).rev() { + let group: Vec = (0..num_matrices).filter(|&m| dims[m].0 == h).collect(); + if group.is_empty() { + continue; + } + // 2^(h-1) independent group-leaf hashes; at `h == h_max` that is the + // bulk of the tree's hashing (half of all nodes). Parallel across + // leaves. + group_digests[h] = Some(crate::par::par_map_collect(0..1usize << (h - 1), |k| { + hash_group_leaf::(source, &group, k) + })); + } + + Self::from_group_digests(dims, h_max, group_digests) + } + + /// Build the tree from each height group's already-hashed leaf digests. + /// + /// The single climb implementation. [`Self::commit`] reaches it having hashed + /// every group leaf in one pass; a future streaming builder would reach it + /// having hashed them incrementally, matrix by matrix. That the two produce + /// the same tree is therefore a property of calling one function, not a + /// coincidence two code paths have to be shown to share. + fn from_group_digests( + dims: Vec<(usize, usize)>, + h_max: usize, + mut group_digests: Vec>>, + ) -> Self { + let mut layers: Vec> = Vec::with_capacity(h_max); + layers.push( + group_digests[h_max] + .take() + .expect("the tallest height group is occupied by construction"), + ); + + // Climb, compressing pairs and injecting shorter matrices where the layer + // width matches their leaf count. Each level's nodes are independent + // (they read only the previous, already-materialized layer), so parallel + // across nodes; levels stay sequential. + let mut i = 0usize; + while layers[i].len() > 1 { + let next_len = layers[i].len() / 2; + let injected = group_digests[h_max - 1 - i].take(); + + let cur = &layers[i]; + let next: Vec = crate::par::par_map_collect(0..next_len, |j| { + let parent = compress::(&cur[2 * j], &cur[2 * j + 1]); + match &injected { + Some(digests) => compress::(&parent, &digests[j]), + None => parent, + } + }); + layers.push(next); + i += 1; + } + + let root = layers.last().expect("at least the base layer exists")[0]; + + MixedMmcs { + root, + layers, + dims, + h_max, + _marker: PhantomData, + } + } + + /// The committed root. + pub fn root(&self) -> Commitment { + self.root + } + + /// `log2` of the tallest committed matrix. The query index this MMCS accepts + /// lives in `[0, 2^(h_max-1))` — see the module's index-convention section. + pub fn h_max(&self) -> usize { + self.h_max + } + + /// Per committed matrix, in input order: `(log_height, width)`. The verifier + /// is expected to rebuild these from the AIR set rather than read them here; + /// this accessor exists so a prover can bind the shape it actually committed. + pub fn dims(&self) -> &[(usize, usize)] { + &self.dims + } + + /// The leaf of matrix `m` that query `iota` opens: `iota >> (h_max - h_m)`. + /// `None` when `m` is not a committed matrix or `iota` is out of this tree's + /// index space. + /// + /// Exposed alongside [`Self::auth_path`] so a prover can assemble a + /// [`MixedOpening`] ONE MATRIX AT A TIME. [`Self::open_batch`] wants a + /// `LeafSource` describing the whole round, which means every matrix's rows + /// readable at once — the same `O(N)` residency a streaming commit path + /// exists to keep out of the commit. Query indices are only known after the + /// FRI, so without these two the win would be given back at opening time. + pub fn row_pair_leaf(&self, iota: usize, m: usize) -> Option { + if iota >= 1usize << (self.h_max - 1) { + return None; + } + let (log_height, _) = *self.dims.get(m)?; + Some(iota >> (self.h_max - log_height)) + } + + /// The shared authentication path for `iota`, reading no matrix rows at all. + /// `None` when `iota` is outside this tree's index space. + pub fn auth_path(&self, iota: usize) -> Option> { + if iota >= 1usize << (self.h_max - 1) { + return None; + } + let mut merkle_path = Vec::with_capacity(self.h_max - 1); + for level in 0..(self.h_max - 1) { + merkle_path.push(self.layers[level][(iota >> level) ^ 1]); + } + Some(Proof { merkle_path }) + } + + /// Open all matrices at query `iota in [0, 2^(h_max-1))`, returning each + /// matrix's row pair plus one shared authentication path. Row data is served + /// by `source`, which MUST describe the same matrices (same order and + /// dimensions) as the one passed to [`Self::commit`]. + pub fn open_batch>(&self, iota: usize, source: &S) -> MixedOpening { + let n0 = 1usize << (self.h_max - 1); + assert!(iota < n0, "iota {iota} out of range (n0 = {n0})"); + debug_assert_eq!( + source.num_matrices(), + self.dims.len(), + "leaf source matrix count must match the committed tree" + ); + + let per_matrix: Vec> = (0..self.dims.len()) + .map(|m| { + let (log_height, width) = self.dims[m]; + debug_assert_eq!(source.log_height(m), log_height); + debug_assert_eq!(source.width(m), width); + let k = iota >> (self.h_max - log_height); + let mut evaluations = Vec::with_capacity(width); + source.append_row(m, 2 * k, &mut evaluations); + let mut evaluations_sym = Vec::with_capacity(width); + source.append_row(m, 2 * k + 1, &mut evaluations_sym); + PolynomialOpenings { + proof: Proof { + merkle_path: Vec::new(), + }, + evaluations, + evaluations_sym, + } + }) + .collect(); + + let mut merkle_path = Vec::with_capacity(self.h_max - 1); + for level in 0..(self.h_max - 1) { + let sibling = (iota >> level) ^ 1; + merkle_path.push(self.layers[level][sibling]); + } + + MixedOpening { + proof: Proof { merkle_path }, + per_matrix, + } + } + + /// Verify a batched opening at `iota` against `root`. `heights[m]` is the + /// `log_height` of matrix `m` and `widths[m]` its column count, both in the + /// SAME order as `opening.per_matrix`, and both supplied by the verifier from + /// the AIR set rather than read out of the proof. + /// + /// `widths` binds each matrix's boundary inside the per-height-group leaf + /// hash (see the module `# Width binding` section): the group leaf hashes the + /// FLAT concatenation of every matrix's `evaluations ‖ evaluations_sym`, so + /// without fixed widths a prover could shift a matrix boundary while keeping + /// the flat bytes — and thus the hash — identical. Pinning `widths` makes the + /// boundaries unambiguous and closes that forgery. + /// + /// `iota` must already be reduced to this tree's index space — see the + /// module's index-convention section. Out-of-range indices are rejected here, + /// but that check is a backstop, not a substitute for the reduction. + /// + /// Returns `false` on every malformed input; it never panics, so a verifier + /// can call it on adversarial data. + pub fn verify_batch( + root: &Commitment, + iota: usize, + opening: &MixedOpening, + heights: &[usize], + widths: &[usize], + ) -> bool { + if opening.per_matrix.len() != heights.len() + || heights.len() != widths.len() + || heights.is_empty() + { + return false; + } + // Bind per-matrix boundaries: every opened matrix must present exactly + // `widths[m]` columns in BOTH rows of its pair. A boundary shift keeps the + // flat per-group concatenation identical but changes these lengths. + for (o, w) in opening.per_matrix.iter().zip(widths.iter()) { + if o.evaluations.len() != *w || o.evaluations_sym.len() != *w { + return false; + } + } + let Some(&h_max) = heights.iter().max() else { + return false; + }; + // Honest heights are >= 1 (row-pair leaves need >= 2 rows) and far below + // the shift width; guard both ends rather than trust the proof's shape. + if h_max == 0 || h_max >= usize::BITS as usize { + return false; + } + // Only the low `h_max - 1` bits of `iota` are consumed (one per level), so + // an index from a taller domain would authenticate the short matrices at a + // position nothing else checks. Reject it instead. + if iota >= 1usize << (h_max - 1) { + return false; + } + if opening.proof.merkle_path.len() != h_max - 1 { + return false; + } + + // Base node: batch all tallest matrices' opened row pairs (input order). + let base_group: Vec<&PolynomialOpenings> = opening + .per_matrix + .iter() + .zip(heights.iter()) + .filter(|(_, h)| **h == h_max) + .map(|(o, _)| o) + .collect(); + let mut acc = hash_group_openings::(&base_group); + + for level in 0..(h_max - 1) { + let sibling = &opening.proof.merkle_path[level]; + let bit = (iota >> level) & 1; + let mut parent = if bit == 0 { + compress::(&acc, sibling) + } else { + compress::(sibling, &acc) + }; + + // Inject matrices whose leaf count matches this (halved) layer, in + // INPUT order — mirroring `commit`'s climb exactly. + let inject_h = h_max - 1 - level; + let inject_group: Vec<&PolynomialOpenings> = opening + .per_matrix + .iter() + .zip(heights.iter()) + .filter(|(_, h)| **h == inject_h) + .map(|(o, _)| o) + .collect(); + if !inject_group.is_empty() { + let inj = hash_group_openings::(&inject_group); + parent = compress::(&parent, &inj); + } + acc = parent; + } + + &acc == root + } +} + +/// One leaf hasher of the batched (keccak) leaf backend. +type LeafHasherOf = as IsStreamingLeafBackend>::LeafHasher; + +/// Builds a [`MixedMmcs`] by absorbing matrices ONE AT A TIME, so a prover never +/// has to hold a height group's LDE buffers simultaneously. +/// +/// # Why this exists +/// +/// [`MixedMmcs::commit`] reads matrix `m` only while building level +/// `h_max - h_m`, so a caller may drop a height group before the next is needed. +/// That is not enough for the group that matters. Within one height the leaf is a +/// single hash over the concatenation of every matrix's row pair, so `commit` +/// needs them all readable at once — and the tallest group is most of an epoch's +/// tables. A caller serving those rows from full in-RAM LDE buffers is back to +/// `O(N)` at the base layer, which is the whole memory win given back. +/// +/// This builder inverts the loop: it keeps one incremental leaf hasher per leaf +/// ([`IsLeafHasher`]) and absorbs matrices into them as they arrive, so the +/// caller produces one matrix's LDE, absorbs it, and drops it. Retained state is +/// `O(leaves × hasher_state)` — bounded by the epoch's tallest height and +/// independent of how many matrices there are or how wide they get. +/// +/// # Contract +/// +/// The shape is declared up front and matrices arrive in that order: the leaf +/// concatenation binds input order (see the module's determinism section), and a +/// builder that let matrices arrive out of order would commit a different tree +/// than [`MixedMmcs::commit`] over the same input. The resulting tree IS that +/// tree — both finish through one climb — which is what makes the two +/// interchangeable rather than merely tested to agree. +pub struct StreamingMmcsBuilder +where + FieldElement: AsBytes + Sync + Send, +{ + dims: Vec<(usize, usize)>, + h_max: usize, + /// Indexed by height: the in-progress leaf hashers of that height group, + /// present from construction until the group's last matrix is absorbed. + pending: Vec>>>, + /// Indexed by height: the group's finalized leaf digests. + group_digests: Vec>>, + /// Matrices of each height still to arrive. A height reaching zero is what + /// releases that group's hashers. + remaining: Vec, + next: usize, +} + +impl StreamingMmcsBuilder +where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + /// Declare the epoch's shape: `(log_height, width)` per matrix, in the order + /// the matrices will be absorbed and in the order the verifier will present + /// their openings. + pub fn new(dims: &[(usize, usize)]) -> Self { + assert!( + !dims.is_empty(), + "StreamingMmcsBuilder requires at least one matrix" + ); + assert!( + dims.iter().all(|(log_height, _)| *log_height >= 1), + "log_height must be >= 1 (row-pair leaves need at least 2 rows)" + ); + let h_max = dims + .iter() + .map(|(log_height, _)| *log_height) + .max() + .expect("dims is non-empty"); + + let mut remaining = vec![0usize; h_max + 1]; + for (log_height, _) in dims { + remaining[*log_height] += 1; + } + + let pending = (0..=h_max) + .map(|h| { + (remaining[h] > 0).then(|| { + (0..1usize << (h - 1)) + .map(|_| as IsStreamingLeafBackend>::leaf_hasher()) + .collect() + }) + }) + .collect(); + + Self { + dims: dims.to_vec(), + h_max, + pending, + group_digests: vec![None; h_max + 1], + remaining, + next: 0, + } + } + + /// Absorb the next declared matrix, reading its rows from `source` at index + /// `m`. The caller may drop that matrix's buffers as soon as this returns. + /// + /// Panics when the arriving matrix's shape disagrees with what was declared — + /// a prover-side programming error, not proof data. + pub fn absorb + Sync>(&mut self, source: &S, m: usize) { + let index = self.next; + assert!( + index < self.dims.len(), + "absorbed more matrices ({}) than were declared ({})", + index + 1, + self.dims.len() + ); + let (log_height, width) = self.dims[index]; + assert_eq!( + (source.log_height(m), source.width(m)), + (log_height, width), + "matrix {index} arrived with a shape the builder was not declared for" + ); + + let hashers = self.pending[log_height] + .as_mut() + .expect("a height with matrices outstanding still holds its hashers"); + // One update per leaf, parallel across leaves — the same shape, and the + // same cost, as `commit`'s one-shot group hash. + crate::par::par_for_each_mut_indexed(hashers, |leaf, hasher| { + let mut row_pair = Vec::with_capacity(2 * width); + source.append_row(m, 2 * leaf, &mut row_pair); + source.append_row(m, 2 * leaf + 1, &mut row_pair); + hasher.update(&row_pair); + }); + + self.next += 1; + self.remaining[log_height] -= 1; + if self.remaining[log_height] == 0 { + let hashers = self.pending[log_height] + .take() + .expect("the group was present a moment ago"); + self.group_digests[log_height] = + Some(hashers.into_iter().map(IsLeafHasher::finalize).collect()); + } + } + + /// Finish the tree. Panics if a declared matrix never arrived — the digests + /// would silently commit to a leaf that absorbed less than it claims. + pub fn finish(self) -> MixedMmcs { + assert_eq!( + self.next, + self.dims.len(), + "{} of {} declared matrices were absorbed", + self.next, + self.dims.len() + ); + MixedMmcs::from_group_digests(self.dims, self.h_max, self.group_digests) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::commitment::commit_bit_reversed; + use math::field::element::FieldElement; + use math::field::goldilocks::GoldilocksField; + use std::sync::Mutex; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + type FE = FieldElement; + type Mmcs = MixedMmcs; + + /// Reference [`LeafSource`] owning bit-reversed row-major matrices. Every + /// test commits/opens through this, so the byte-parity assertion against + /// `commit_bit_reversed` pins the tree contract; `borrowed_sources_match_ + /// owned_reference` cross-checks it against the borrowed (natural-order) + /// sources a prover would use. + struct OwnedMatrices { + /// Each entry: `(bit-reversed row-major data, log_height, width)`. + mats: Vec<(Vec>, usize, usize)>, + } + + impl LeafSource for OwnedMatrices { + fn num_matrices(&self) -> usize { + self.mats.len() + } + fn log_height(&self, m: usize) -> usize { + self.mats[m].1 + } + fn width(&self, m: usize) -> usize { + self.mats[m].2 + } + fn append_row(&self, m: usize, bitrev_row: usize, out: &mut Vec>) { + let (data, _log_height, width) = &self.mats[m]; + out.extend_from_slice(&data[bitrev_row * width..(bitrev_row + 1) * width]); + } + } + + fn owned(mats: Vec<(Vec, usize, usize)>) -> OwnedMatrices { + OwnedMatrices { mats } + } + + /// Build a row-major, bit-reversed flat vec from column-major natural-order + /// `columns`, matching the layout the existing trace commit consumes: row `j` + /// of the output = `[col_0[br(j)], ..., col_{w-1}[br(j)]]` with + /// `br = reverse_index(., num_rows)`. + fn row_major_bit_reversed(columns: &[Vec], num_rows: usize) -> Vec { + let width = columns.len(); + let mut out = vec![FE::from(0u64); num_rows * width]; + for (r, chunk) in out.chunks_exact_mut(width).enumerate() { + let br = reverse_index(r, num_rows as u64); + for (c, col) in columns.iter().enumerate() { + chunk[c] = col[br]; + } + } + out + } + + /// Build a row-major flat vec in NATURAL order (no bit reversal): row `r` = + /// `[col_0[r], ..., col_{w-1}[r]]`. This is the layout the prover's + /// `BorrowedMatrix::RowMajorNatural` reads (the retained main/aux LDE buffer). + fn row_major_natural(columns: &[Vec], num_rows: usize) -> Vec { + let width = columns.len(); + let mut out = vec![FE::from(0u64); num_rows * width]; + for (r, chunk) in out.chunks_exact_mut(width).enumerate() { + for (c, col) in columns.iter().enumerate() { + chunk[c] = col[r]; + } + } + out + } + + fn make_columns(width: usize, num_rows: usize, seed: u64) -> Vec> { + (0..width) + .map(|c| { + (0..num_rows) + .map(|r| { + FE::from(seed.wrapping_mul(31) + (c as u64) * 1009 + (r as u64) * 7 + 1) + }) + .collect() + }) + .collect() + } + + #[test] + fn single_matrix_commit_open_verify_and_tamper() { + let log_height = 2usize; + let num_rows = 1usize << log_height; + let width = 3usize; + let columns = make_columns(width, num_rows, 5); + let data = row_major_bit_reversed(&columns, num_rows); + + let src = owned(vec![(data.clone(), log_height, width)]); + let mmcs = Mmcs::commit(&src); + let heights = [log_height]; + let widths = [width]; + let n0 = 1usize << (log_height - 1); + + for iota in 0..n0 { + let opening = mmcs.open_batch(iota, &src); + assert_eq!(opening.per_matrix.len(), 1); + let k = iota; + let row_2k = data[(2 * k) * width..(2 * k + 1) * width].to_vec(); + let row_2k1 = data[(2 * k + 1) * width..(2 * k + 2) * width].to_vec(); + assert_eq!(opening.per_matrix[0].evaluations, row_2k); + assert_eq!(opening.per_matrix[0].evaluations_sym, row_2k1); + assert!(Mmcs::verify_batch( + &mmcs.root(), + iota, + &opening, + &heights, + &widths + )); + } + + let mut opening = mmcs.open_batch(0, &src); + opening.per_matrix[0].evaluations[0] = + &opening.per_matrix[0].evaluations[0] + &FE::from(1u64); + assert!(!Mmcs::verify_batch( + &mmcs.root(), + 0, + &opening, + &heights, + &widths + )); + } + + /// ★ The backward-compatibility statement: a single-matrix MMCS IS the + /// existing per-table row-pair tree. It holds by construction — both go + /// through `BatchedMerkleTreeBackend`'s `hash_data` / `hash_new_parent` — + /// and this pins that no second leaf encoding crept in. + /// + /// Both sides use the SAME concrete keccak backend, which is what makes the + /// comparison meaningful: `commit_bit_reversed` commits through + /// `BatchedMerkleTreeBackend`, so a hash difference here would be a layout + /// difference, not a hash-configuration mismatch. + #[test] + fn single_matrix_root_matches_existing_row_pair_tree() { + let log_height = 3usize; + let num_rows = 1usize << log_height; + let width = 4usize; + let columns = make_columns(width, num_rows, 9); + + let (_, existing_root) = + commit_bit_reversed(&columns, 2).expect("non-empty columns build a tree"); + + let data = row_major_bit_reversed(&columns, num_rows); + let mmcs = Mmcs::commit(&owned(vec![(data, log_height, width)])); + + assert_eq!(mmcs.root(), existing_root); + } + + #[test] + fn mixed_height_open_positions_verify_and_tamper() { + // Three matrices, log_heights {5, 5, 3}, widths {2, 1, 4}. + let (ha, hb, hc) = (5usize, 5usize, 3usize); + let (wa, wb, wc) = (2usize, 1usize, 4usize); + let a = row_major_bit_reversed(&make_columns(wa, 1 << ha, 1), 1 << ha); + let b = row_major_bit_reversed(&make_columns(wb, 1 << hb, 2), 1 << hb); + let c = row_major_bit_reversed(&make_columns(wc, 1 << hc, 3), 1 << hc); + + let src = owned(vec![ + (a.clone(), ha, wa), + (b.clone(), hb, wb), + (c.clone(), hc, wc), + ]); + let mmcs = Mmcs::commit(&src); + let heights = [ha, hb, hc]; + let widths = [wa, wb, wc]; + let h_max = 5usize; + let n0 = 1usize << (h_max - 1); // 16 + + let row = |data: &[FE], w: usize, r: usize| data[r * w..(r + 1) * w].to_vec(); + + for iota in [0usize, 1, 2, 3, 7, 8, 13, n0 - 1] { + let opening = mmcs.open_batch(iota, &src); + assert_eq!(opening.per_matrix.len(), 3); + + // Tall matrices open at k = iota >> 0 = iota. + assert_eq!(opening.per_matrix[0].evaluations, row(&a, wa, 2 * iota)); + assert_eq!( + opening.per_matrix[0].evaluations_sym, + row(&a, wa, 2 * iota + 1) + ); + assert_eq!(opening.per_matrix[1].evaluations, row(&b, wb, 2 * iota)); + + // Height-3 matrix opens at k = iota >> (5 - 3) = iota >> 2. + let kc = iota >> (h_max - hc); + assert_eq!(opening.per_matrix[2].evaluations, row(&c, wc, 2 * kc)); + assert_eq!( + opening.per_matrix[2].evaluations_sym, + row(&c, wc, 2 * kc + 1) + ); + + assert!( + Mmcs::verify_batch(&mmcs.root(), iota, &opening, &heights, &widths), + "honest opening at iota={iota} must verify" + ); + } + + // Tamper the height-3 matrix's opened row -> rejection (proves the short + // matrix is bound by the shared path via injection). + let iota = 6usize; + let mut opening = mmcs.open_batch(iota, &src); + opening.per_matrix[2].evaluations[0] = + &opening.per_matrix[2].evaluations[0] + &FE::from(1u64); + assert!( + !Mmcs::verify_batch(&mmcs.root(), iota, &opening, &heights, &widths), + "tampered height-3 row must be rejected" + ); + + // Tamper a tall-matrix row too -> rejection. + let mut opening2 = mmcs.open_batch(iota, &src); + opening2.per_matrix[0].evaluations[0] = + &opening2.per_matrix[0].evaluations[0] + &FE::from(1u64); + assert!( + !Mmcs::verify_batch(&mmcs.root(), iota, &opening2, &heights, &widths), + "tampered tall-matrix row must be rejected" + ); + } + + /// Vector test: hand-compute the root for `{log_height 2, log_height 1}` + /// matrices per the documented layout and assert equality. Pins the + /// leaf/injection contract, plus determinism. + #[test] + fn vector_root_layout_contract_and_determinism() { + // A: log_height 2 (4 rows), width 2 ; B: log_height 1 (2 rows), width 3. + let a_data = row_major_bit_reversed(&make_columns(2, 4, 3), 4); + let b_data = row_major_bit_reversed(&make_columns(3, 2, 8), 2); + + let src = owned(vec![(a_data.clone(), 2, 2), (b_data.clone(), 1, 3)]); + let mmcs = Mmcs::commit(&src); + + // Hand recomputation via the backend primitives, in the documented order. + let arow = |r: usize| a_data[r * 2..(r + 1) * 2].to_vec(); + let brow = |r: usize| b_data[r * 3..(r + 1) * 3].to_vec(); + let h = |v: Vec| { + as IsMerkleTreeBackend>::hash_data(&v) + }; + + // Base layer (matrix A only): leaf k = H(A.row(2k) || A.row(2k+1)). + let mut leaf0 = arow(0); + leaf0.extend(arow(1)); + let mut leaf1 = arow(2); + leaf1.extend(arow(3)); + let l00 = h(leaf0); + let l01 = h(leaf1); + + // Climb to layer 1 (root): compress the base pair, then inject B (h=1). + let parent = compress::(&l00, &l01); + let mut binj = brow(0); + binj.extend(brow(1)); + let inj = h(binj); + let expected_root = compress::(&parent, &inj); + + assert_eq!( + mmcs.root(), + expected_root, + "root must match the hand-computed mixed-height layout" + ); + + // Determinism: a second commit over the same inputs yields the same root. + let mmcs2 = Mmcs::commit(&owned(vec![(a_data, 2, 2), (b_data, 1, 3)])); + assert_eq!(mmcs.root(), mmcs2.root(), "commit must be deterministic"); + + for iota in 0..2usize { + let opening = mmcs.open_batch(iota, &src); + // heights {2, 1}, widths {2, 3}. + assert!(Mmcs::verify_batch( + &mmcs.root(), + iota, + &opening, + &[2, 1], + &[2, 3] + )); + } + } + + /// Two SAME-HEIGHT matrices share one base-group leaf, whose hash is over the + /// FLAT concatenation `A.eval ‖ A.eval_sym ‖ B.eval ‖ B.eval_sym`. A malicious + /// prover can shift the A|A_sym boundary (move one element from A's + /// `evaluations_sym` into A's `evaluations`) leaving that flat concatenation — + /// and hence the leaf hash — byte-identical, so a width-blind `verify_batch` + /// would accept it. The per-matrix width binding rejects the shift. + #[test] + fn boundary_shift_forgery_rejected() { + let h = 2usize; + let num_rows = 1usize << h; + let (wa, wb) = (2usize, 1usize); // wA >= 2 so we can steal one column. + let a = row_major_bit_reversed(&make_columns(wa, num_rows, 11), num_rows); + let b = row_major_bit_reversed(&make_columns(wb, num_rows, 22), num_rows); + + let src = owned(vec![(a, h, wa), (b, h, wb)]); + let mmcs = Mmcs::commit(&src); + let heights = [h, h]; + let widths = [wa, wb]; + + let iota = 0usize; + let opening = mmcs.open_batch(iota, &src); + assert!( + Mmcs::verify_batch(&mmcs.root(), iota, &opening, &heights, &widths), + "honest opening must verify" + ); + + // Forge: lengthen A.evaluations by one element taken from A.evaluations_sym. + let mut forged = mmcs.open_batch(iota, &src); + let moved = forged.per_matrix[0].evaluations_sym.remove(0); + forged.per_matrix[0].evaluations.push(moved); + + // The FLAT per-group concatenation is byte-identical to the honest one, so + // the group leaf hash is UNCHANGED — the rejection must come from the width + // check, not from a differing hash. + let flat = |o: &MixedOpening| -> Vec { + let mut v = Vec::new(); + for m in &o.per_matrix { + v.extend_from_slice(&m.evaluations); + v.extend_from_slice(&m.evaluations_sym); + } + v + }; + assert_eq!( + flat(&opening), + flat(&forged), + "the flat concatenation must be byte-identical (boundary-only shift)" + ); + + assert!( + !Mmcs::verify_batch(&mmcs.root(), iota, &forged, &heights, &widths), + "boundary-shift forgery must be rejected by the width binding" + ); + } + + /// Extension-field (Fp3) coverage: the aux and composition matrices an epoch + /// batches are cubic-extension. Byte-parity cross-check of a single Fp3 matrix + /// against the existing per-table row-pair tree, plus an open/verify/tamper + /// roundtrip over the extension path. + #[test] + fn single_matrix_fp3_root_matches_existing_row_pair_tree() { + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Fp3; + type F3 = FieldElement; + + let log_height = 3usize; + let num_rows = 1usize << log_height; + let width = 3usize; + + // Populate ALL three components so the 24-byte extension serialization is + // exercised (not just the embedded-base subset). + let columns: Vec> = (0..width) + .map(|c| { + (0..num_rows) + .map(|r| { + F3::new([ + FE::from((c as u64) * 7 + r as u64 + 1), + FE::from((r as u64) * 13 + 2), + FE::from((c as u64) * 5 + (r as u64) * 3 + 4), + ]) + }) + .collect() + }) + .collect(); + + let (_, existing_root) = + commit_bit_reversed(&columns, 2).expect("non-empty columns build a tree"); + + // Row-major bit-reversed equivalent of the same column-major data. + let mut data = vec![F3::zero(); num_rows * width]; + for (r, chunk) in data.chunks_exact_mut(width).enumerate() { + let br = reverse_index(r, num_rows as u64); + for (c, col) in columns.iter().enumerate() { + chunk[c] = col[br]; + } + } + + let src = OwnedMatrices { + mats: vec![(data, log_height, width)], + }; + let mmcs = MixedMmcs::::commit(&src); + assert_eq!( + mmcs.root(), + existing_root, + "Fp3 single-matrix root must match the existing row-pair tree" + ); + + let heights = [log_height]; + let widths = [width]; + for iota in 0..(1usize << (log_height - 1)) { + let opening = mmcs.open_batch(iota, &src); + assert!(MixedMmcs::::verify_batch( + &mmcs.root(), + iota, + &opening, + &heights, + &widths + )); + } + + let mut opening = mmcs.open_batch(0, &src); + opening.per_matrix[0].evaluations[0] = &opening.per_matrix[0].evaluations[0] + &F3::one(); + assert!(!MixedMmcs::::verify_batch( + &mmcs.root(), + 0, + &opening, + &heights, + &widths + )); + } + + /// Equivalence (the soundness contract a batched prover relies on): the + /// digest-only MMCS built from borrowed, NATURAL-order leaf sources yields the + /// SAME root and the SAME opened rows as the reference owning source over the + /// bit-reversed data — for the row-major (main / aux) layout, the column-major + /// (composition) layout, AND a main-split column sub-range (`col_start > 0`). + /// Only the leaf-byte source changes; nothing the verifier sees does. + #[test] + fn borrowed_sources_match_owned_reference() { + // Mixed heights {5, 5, 3}; the height-3 matrix exercises injection. + let specs = [(5usize, 3usize, 100u64), (5, 1, 200), (3, 4, 300)]; + + // Column-major natural-order columns per matrix. + let cols: Vec>> = specs + .iter() + .map(|&(lh, w, seed)| make_columns(w, 1 << lh, seed)) + .collect(); + + // Reference: owned, bit-reversed row-major. + let owned_src = owned( + specs + .iter() + .zip(cols.iter()) + .map(|(&(lh, w, _), c)| (row_major_bit_reversed(c, 1 << lh), lh, w)) + .collect(), + ); + + // Borrowed row-major NATURAL (the retained main / aux LDE buffer). + let rm_natural: Vec> = specs + .iter() + .zip(cols.iter()) + .map(|(&(lh, _, _), c)| row_major_natural(c, 1 << lh)) + .collect(); + let rm_src: Vec> = specs + .iter() + .zip(rm_natural.iter()) + .map(|(&(lh, w, _), data)| BorrowedMatrix::RowMajorNatural { + data: data.as_slice(), + stride: w, + col_start: 0, + width: w, + log_height: lh, + }) + .collect(); + + // Borrowed column-major NATURAL (the retained composition-poly LDE). + let cm_src: Vec> = specs + .iter() + .zip(cols.iter()) + .map(|(&(lh, _, _), c)| BorrowedMatrix::ColMajorNatural { + cols: c.as_slice(), + log_height: lh, + }) + .collect(); + + let owned_mmcs = Mmcs::commit(&owned_src); + let rm_mmcs = Mmcs::commit(&rm_src); + let cm_mmcs = Mmcs::commit(&cm_src); + assert_eq!( + owned_mmcs.root(), + rm_mmcs.root(), + "row-major natural root must match the owned reference" + ); + assert_eq!( + owned_mmcs.root(), + cm_mmcs.root(), + "column-major natural root must match the owned reference" + ); + + let n0 = 1usize << (5 - 1); + for iota in 0..n0 { + let o = owned_mmcs.open_batch(iota, &owned_src); + let rm = rm_mmcs.open_batch(iota, &rm_src); + let cm = cm_mmcs.open_batch(iota, &cm_src); + assert_eq!(o.proof.merkle_path, rm.proof.merkle_path); + assert_eq!(o.proof.merkle_path, cm.proof.merkle_path); + for i in 0..specs.len() { + assert_eq!(o.per_matrix[i].evaluations, rm.per_matrix[i].evaluations); + assert_eq!( + o.per_matrix[i].evaluations_sym, + rm.per_matrix[i].evaluations_sym + ); + assert_eq!(o.per_matrix[i].evaluations, cm.per_matrix[i].evaluations); + assert_eq!( + o.per_matrix[i].evaluations_sym, + cm.per_matrix[i].evaluations_sym + ); + } + } + + // Main-split sub-range: a RowMajorNatural over a wider buffer with a + // leading prefix (`col_start = prefix`) must match an owned matrix built + // over ONLY the committed trailing columns. + let (lh, prefix, w) = (4usize, 2usize, 3usize); + let num_rows = 1usize << lh; + let full = make_columns(prefix + w, num_rows, 42); + let full_natural = row_major_natural(&full, num_rows); + let sub_cols: Vec> = full[prefix..].to_vec(); + let sub_owned = owned(vec![(row_major_bit_reversed(&sub_cols, num_rows), lh, w)]); + let split_src: Vec> = + vec![BorrowedMatrix::RowMajorNatural { + data: full_natural.as_slice(), + stride: prefix + w, + col_start: prefix, + width: w, + log_height: lh, + }]; + let sub_owned_mmcs = Mmcs::commit(&sub_owned); + let split_mmcs = Mmcs::commit(&split_src); + assert_eq!( + sub_owned_mmcs.root(), + split_mmcs.root(), + "main-split (col_start>0) root must match the owned sub-range" + ); + for iota in 0..(1usize << (lh - 1)) { + let a = sub_owned_mmcs.open_batch(iota, &sub_owned); + let b = split_mmcs.open_batch(iota, &split_src); + assert_eq!(a.per_matrix[0].evaluations, b.per_matrix[0].evaluations); + assert_eq!( + a.per_matrix[0].evaluations_sym, + b.per_matrix[0].evaluations_sym + ); + } + } + + /// ★ The index-convention control (the module's "HARD PRECONDITION" section). + /// + /// A round whose tallest matrix is SHORTER than the FRI's tallest is the case + /// where the two index conventions disagree: `verify_batch` consumes the LOW + /// `h_max_round - 1` bits of whatever index it is handed, while a matrix + /// inside the tree is located by the HIGH bits of the FRI index. This asserts + /// three things about that case: + /// + /// 1. honest-path control — the correctly reduced index verifies; + /// 2. a tampered row of a SHORT (injected) matrix is rejected, so the low-bits + /// walk really does authenticate the short matrices at the reduced index; + /// 3. handing the un-reduced FRI index straight in is rejected — the misuse is + /// detectable, not silently accepted at some other leaf. + /// + /// A tamper control on the tallest matrix alone would pass under either + /// convention and catch none of this. + #[test] + fn short_round_low_bit_convention_is_exercised() { + // A hypothetical FRI over a 2^6 domain: iota_fri in [0, 2^5). + let h_max_fri = 6usize; + // This round's matrices are shorter: heights {4, 2}. + let (h_tall, h_short) = (4usize, 2usize); + let (w_tall, w_short) = (3usize, 2usize); + let tall = row_major_bit_reversed(&make_columns(w_tall, 1 << h_tall, 77), 1 << h_tall); + let short = row_major_bit_reversed(&make_columns(w_short, 1 << h_short, 88), 1 << h_short); + + let src = owned(vec![(tall, h_tall, w_tall), (short, h_short, w_short)]); + let mmcs = Mmcs::commit(&src); + let heights = [h_tall, h_short]; + let widths = [w_tall, w_short]; + assert_eq!(mmcs.h_max(), h_tall, "the round's h_max is below the FRI's"); + + // The reduction the caller owes: iota_round = iota_fri >> (h_fri - h_round). + let shift = h_max_fri - h_tall; + // Pick a FRI index whose low bits differ from the reduced index's, so the + // two conventions genuinely disagree here. + let iota_fri = 0b10110usize; + let iota_round = iota_fri >> shift; + assert_ne!( + iota_fri & ((1 << (h_tall - 1)) - 1), + iota_round, + "the test index must distinguish the low-bit and high-bit conventions" + ); + + // (1) Honest-path control at the reduced index. + let opening = mmcs.open_batch(iota_round, &src); + assert!( + Mmcs::verify_batch(&mmcs.root(), iota_round, &opening, &heights, &widths), + "the correctly reduced index must verify" + ); + + // (2) Tamper the SHORT (injected) matrix — the matrix a tall-only control + // would never touch, and the one the disagreeing conventions move. + let mut tampered = mmcs.open_batch(iota_round, &src); + tampered.per_matrix[1].evaluations[0] = + &tampered.per_matrix[1].evaluations[0] + &FE::from(1u64); + assert!( + !Mmcs::verify_batch(&mmcs.root(), iota_round, &tampered, &heights, &widths), + "a tampered SHORT-matrix row must be rejected at the reduced index" + ); + + // (3) The misuse: hand the un-reduced FRI index in. It is out of this + // tree's range, so the range guard rejects it rather than walking to some + // unrelated leaf. + assert!( + iota_fri >= 1usize << (h_tall - 1), + "the un-reduced index is outside this round's leaf range" + ); + assert!( + !Mmcs::verify_batch(&mmcs.root(), iota_fri, &opening, &heights, &widths), + "an un-reduced FRI index must be rejected, not accepted at another leaf" + ); + + // And an in-range index that is simply the wrong leaf is rejected too, so + // the guard is not the only thing standing between the two conventions. + let wrong_but_in_range = iota_fri & ((1 << (h_tall - 1)) - 1); + assert!( + !Mmcs::verify_batch( + &mmcs.root(), + wrong_but_in_range, + &opening, + &heights, + &widths + ), + "an opening replayed at the wrong in-range leaf must be rejected" + ); + } + + /// The malformed-input surface of `verify_batch`: every shape error returns + /// `false` rather than panicking, since a verifier calls this on proof data. + #[test] + fn verify_batch_rejects_malformed_shapes_without_panicking() { + let h = 3usize; + let w = 2usize; + let data = row_major_bit_reversed(&make_columns(w, 1 << h, 4), 1 << h); + let src = owned(vec![(data, h, w)]); + let mmcs = Mmcs::commit(&src); + let root = mmcs.root(); + let opening = mmcs.open_batch(1, &src); + + assert!(Mmcs::verify_batch(&root, 1, &opening, &[h], &[w])); + // Mismatched metadata lengths. + assert!(!Mmcs::verify_batch(&root, 1, &opening, &[h, h], &[w])); + assert!(!Mmcs::verify_batch(&root, 1, &opening, &[h], &[w, w])); + // Empty metadata. + assert!(!Mmcs::verify_batch(&root, 1, &opening, &[], &[])); + // A height that would overflow the level shift. + assert!(!Mmcs::verify_batch( + &root, + 1, + &opening, + &[usize::BITS as usize], + &[w] + )); + // An index past this tree's leaf count. + assert!(!Mmcs::verify_batch( + &root, + 1usize << (h - 1), + &opening, + &[h], + &[w] + )); + // A path of the wrong length. + let mut short_path = opening.clone(); + short_path.proof.merkle_path.pop(); + assert!(!Mmcs::verify_batch(&root, 1, &short_path, &[h], &[w])); + } + + /// Wraps a source and records, per matrix, the first and last global access + /// sequence number, plus a residency model the caller drives. `Mutex` / + /// atomics (not `Cell`) because both `commit` and the streaming builder read + /// the source from rayon workers. + struct Tracing<'a, E: IsField> { + inner: &'a OwnedMatrices, + clock: AtomicUsize, + window: Mutex>, + /// The residency model: which matrices the caller says it is holding. + resident: Vec, + live: AtomicUsize, + peak: AtomicUsize, + /// Rows served for a matrix the caller had already dropped. Any nonzero + /// count means the access pattern does not fit the residency policy. + reads_while_dropped: AtomicUsize, + } + + impl<'a, E: IsField> Tracing<'a, E> { + fn new(inner: &'a OwnedMatrices) -> Self { + let n = inner.num_matrices(); + Self { + inner, + clock: AtomicUsize::new(0), + window: Mutex::new(vec![(usize::MAX, 0); n]), + resident: (0..n).map(|_| AtomicBool::new(false)).collect(), + live: AtomicUsize::new(0), + peak: AtomicUsize::new(0), + reads_while_dropped: AtomicUsize::new(0), + } + } + + /// Declare every matrix held for the whole build — the only policy + /// `MixedMmcs::commit` can be served under. + fn materialize_all(&self) { + for m in 0..self.inner.num_matrices() { + self.materialize(m); + } + } + + fn materialize(&self, m: usize) { + if !self.resident[m].swap(true, Ordering::SeqCst) { + let live = self.live.fetch_add(1, Ordering::SeqCst) + 1; + self.peak.fetch_max(live, Ordering::SeqCst); + } + } + + // Retained for a future streaming-commit phase's residency tests; the + // one kept `commit` test drives residency via `materialize_all` only. + #[allow(dead_code)] + fn drop_matrix(&self, m: usize) { + if self.resident[m].swap(false, Ordering::SeqCst) { + self.live.fetch_sub(1, Ordering::SeqCst); + } + } + + fn windows(self) -> (Vec<(usize, usize)>, usize, usize) { + let peak = self.peak.load(Ordering::SeqCst); + let dropped_reads = self.reads_while_dropped.load(Ordering::SeqCst); + let windows = self.window.into_inner().expect("uncontended after commit"); + (windows, peak, dropped_reads) + } + } + + impl LeafSource for Tracing<'_, E> { + fn num_matrices(&self) -> usize { + self.inner.num_matrices() + } + fn log_height(&self, m: usize) -> usize { + self.inner.log_height(m) + } + fn width(&self, m: usize) -> usize { + self.inner.width(m) + } + fn append_row(&self, m: usize, bitrev_row: usize, out: &mut Vec>) { + if !self.resident[m].load(Ordering::SeqCst) { + self.reads_while_dropped.fetch_add(1, Ordering::SeqCst); + } + let t = self.clock.fetch_add(1, Ordering::SeqCst); + let mut w = self.window.lock().expect("no test thread panics here"); + w[m].0 = w[m].0.min(t); + w[m].1 = w[m].1.max(t); + drop(w); + self.inner.append_row(m, bitrev_row, out); + } + } + + fn residency_fixture() -> ([(usize, usize, u64); 4], OwnedMatrices) { + let specs = [(5usize, 2usize, 1u64), (5, 3, 2), (3, 1, 3), (2, 4, 4)]; + let inner = owned( + specs + .iter() + .map(|&(lh, w, seed)| { + ( + row_major_bit_reversed(&make_columns(w, 1 << lh, seed), 1 << lh), + lh, + w, + ) + }) + .collect(), + ); + (specs, inner) + } + + /// The streaming builder is not a second implementation of the tree: it + /// finishes through the same climb `commit` does. This pins the consequence — + /// same root, same layers, same openings — so a future change that forked the + /// two would fail here rather than at a verifier three modules away. + #[test] + fn streaming_builder_commits_the_same_tree_as_commit() { + let (specs, inner) = residency_fixture(); + let dims: Vec<(usize, usize)> = specs.iter().map(|&(lh, w, _)| (lh, w)).collect(); + + let mut builder = StreamingMmcsBuilder::::new(&dims); + for m in 0..dims.len() { + builder.absorb(&inner, m); + } + let streamed = builder.finish(); + let reference = Mmcs::commit(&inner); + + assert_eq!( + streamed.root(), + reference.root(), + "the streamed root must equal the one-shot root" + ); + assert_eq!(streamed.h_max(), reference.h_max()); + assert_eq!(streamed.dims(), reference.dims()); + + let heights: Vec = specs.iter().map(|&(lh, _, _)| lh).collect(); + let widths: Vec = specs.iter().map(|&(_, w, _)| w).collect(); + for iota in 0..1usize << (streamed.h_max() - 1) { + let opening = streamed.open_batch(iota, &inner); + assert!( + Mmcs::verify_batch(&streamed.root(), iota, &opening, &heights, &widths), + "an opening of the streamed tree must verify at iota {iota}" + ); + assert_eq!( + opening.proof.merkle_path, + reference.open_batch(iota, &inner).proof.merkle_path, + "the authentication path at iota {iota} must be the same path" + ); + } + } + + /// ★ The acceptance test for the batched commit's memory claim. + /// + /// `commit`'s contract is per height GROUP: it reads a group inside one + /// contiguous phase, so a caller may drop the group before the next. That is + /// not enough. Within the tallest group the leaf is one hash over every + /// matrix's concatenated row pair, so `commit` reads all of them at every + /// leaf — their access windows OVERLAP, and a caller has to hold the whole + /// group. On a real epoch the tallest group is most of the tables, so that is + /// `O(N)` resident at the base layer: the memory batching exists to remove, + /// given back. + /// + /// The streaming builder's windows are pairwise disjoint across ALL matrices, + /// same-height ones included, so the residency policy "materialize, absorb, + /// drop" serves it with exactly ONE matrix live. Both halves are traced here; + /// the second is the property the batched R1 / aux / parts commits must be + /// built on, and the first is what makes it a real difference rather than a + /// restatement. + #[test] + fn streaming_builder_serves_the_base_group_without_holding_it() { + let (specs, inner) = residency_fixture(); + let dims: Vec<(usize, usize)> = specs.iter().map(|&(lh, w, _)| (lh, w)).collect(); + let base_group: Vec = (0..specs.len()).filter(|&m| specs[m].0 == 5).collect(); + assert!( + base_group.len() > 1, + "the fixture must batch more than one matrix at the tallest height" + ); + + // --- What `commit` requires: the whole group resident at once. --- + let tracing = Tracing::new(&inner); + tracing.materialize_all(); + let commit_root = Mmcs::commit(&tracing).root(); + let (commit_windows, commit_peak, commit_dropped_reads) = tracing.windows(); + assert_eq!(commit_dropped_reads, 0, "the control held everything"); + assert_eq!( + commit_peak, + specs.len(), + "serving `commit` needs every matrix resident" + ); + for (i, &m) in base_group.iter().enumerate() { + for &n in &base_group[i + 1..] { + let (fm, lm) = commit_windows[m]; + let (fn_, ln) = commit_windows[n]; + assert!( + fm <= ln && fn_ <= lm, + "matrices {m} and {n} share the base height, so `commit` must \ + read them in OVERLAPPING windows [{fm},{lm}] / [{fn_},{ln}] — \ + if this ever stops holding, the escape below is no longer the \ + thing that buys the memory" + ); + } + } + + // --- What the streaming builder requires: one matrix at a time. --- + let tracing = Tracing::new(&inner); + let mut builder = StreamingMmcsBuilder::::new(&dims); + for m in 0..dims.len() { + tracing.materialize(m); + builder.absorb(&tracing, m); + tracing.drop_matrix(m); + } + let streamed_root = builder.finish().root(); + let (streamed_windows, streamed_peak, streamed_dropped_reads) = tracing.windows(); + + assert_eq!( + streamed_root, commit_root, + "the escape must not change what is committed" + ); + assert_eq!( + streamed_dropped_reads, 0, + "no row may be read after the caller dropped its matrix" + ); + assert_eq!( + streamed_peak, + 1, + "the base height group must be served with ONE matrix resident, not \ + {} — this is the batched commit's whole memory claim", + specs.len() + ); + for (m, &(first, last)) in streamed_windows.iter().enumerate() { + assert!(first <= last, "matrix {m} was never read"); + for (n, &(fn_, ln)) in streamed_windows.iter().enumerate().skip(m + 1) { + assert!( + last < fn_ || ln < first, + "matrices {m} and {n} were read in overlapping windows \ + [{first},{last}] / [{fn_},{ln}] — the builder must finish one \ + matrix before the next is needed, at EVERY height" + ); + } + } + } + + /// A declared matrix that never arrives would leave its group's leaves having + /// absorbed less than the shape says, committing a tree no verifier rebuilds. + /// The builder refuses rather than producing it. + #[test] + #[should_panic(expected = "of 4 declared matrices were absorbed")] + fn finishing_with_a_matrix_missing_panics() { + let (specs, inner) = residency_fixture(); + let dims: Vec<(usize, usize)> = specs.iter().map(|&(lh, w, _)| (lh, w)).collect(); + let mut builder = StreamingMmcsBuilder::::new(&dims); + for m in 0..dims.len() - 1 { + builder.absorb(&inner, m); + } + builder.finish(); + } + + /// The incremental leaf hasher's whole contract: where the updates fall must + /// not show. Checked at every split point of a leaf, and for the extension + /// field the aux and composition matrices actually use — a framing bug that + /// only appeared at an element boundary would slip past a base-field check. + #[test] + fn leaf_hasher_splits_anywhere_and_matches_hash_data() { + use crypto::merkle_tree::traits::IsLeafHasher; + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Fp3; + + fn check(leaf: Vec>) + where + FieldElement: AsBytes + Sync + Send, + { + let expected = + as IsMerkleTreeBackend>::hash_data(&leaf); + for split in 0..=leaf.len() { + let mut hasher = + as IsStreamingLeafBackend>::leaf_hasher(); + hasher.update(&leaf[..split]); + hasher.update(&leaf[split..]); + assert_eq!( + hasher.finalize(), + expected, + "splitting the leaf at {split} changed the digest" + ); + } + // Three updates, so an implementation that only ever saw two would not + // pass by accident. + let mut hasher = + as IsStreamingLeafBackend>::leaf_hasher(); + for element in &leaf { + hasher.update(core::slice::from_ref(element)); + } + assert_eq!(hasher.finalize(), expected, "element-at-a-time must agree"); + } + + check::((1u64..=9).map(FE::from).collect()); + check::( + (1u64..=9) + .map(|i| { + FieldElement::::new([FE::from(i), FE::from(i * 7 + 1), FE::from(i * 13)]) + }) + .collect(), + ); + } + + /// The memory contract from the module's "what the caller may drop" section, + /// made falsifiable: `commit` reads each height group's rows inside ONE + /// contiguous window of the build, and the windows run in descending height + /// order. A rewrite that materialized every matrix up front, or that revisited + /// a group after moving on, would fail here. + #[test] + fn commit_reads_each_height_group_in_one_contiguous_phase() { + // Heights {5, 5, 3, 2}: two groups sharing the base layer, two injected. + let specs = [(5usize, 2usize, 1u64), (5, 3, 2), (3, 1, 3), (2, 4, 4)]; + let inner = owned( + specs + .iter() + .map(|&(lh, w, seed)| { + ( + row_major_bit_reversed(&make_columns(w, 1 << lh, seed), 1 << lh), + lh, + w, + ) + }) + .collect(), + ); + let tracing = Tracing::new(&inner); + tracing.materialize_all(); + + let traced_root = Mmcs::commit(&tracing).root(); + assert_eq!( + traced_root, + Mmcs::commit(&inner).root(), + "tracing must not change what is committed" + ); + + let (windows, _peak, _dropped) = tracing.windows(); + for (m, (first, last)) in windows.iter().enumerate() { + assert!(*first <= *last, "matrix {m} was never read"); + } + + // Same-height matrices share a window; different heights must not overlap, + // and taller groups must come first. + for (m, &(fm, lm)) in windows.iter().enumerate() { + for (n, &(fn_, ln)) in windows.iter().enumerate() { + if specs[m].0 <= specs[n].0 { + continue; + } + assert!( + lm < fn_ || ln < fm, + "matrices {m} (h={}) and {n} (h={}) were read in overlapping \ + windows [{fm},{lm}] / [{fn_},{ln}] — a height group must be \ + readable and then droppable", + specs[m].0, + specs[n].0 + ); + assert!( + lm < fn_, + "the taller matrix {m} (h={}) must be read before the shorter \ + {n} (h={})", + specs[m].0, + specs[n].0 + ); + } + } + } +} diff --git a/crypto/stark/src/fri/mod.rs b/crypto/stark/src/fri/mod.rs index 1f53b51cf..cee278a16 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -1,6 +1,8 @@ +pub mod batched; pub mod fri_commitment; pub mod fri_decommit; pub(crate) mod fri_functions; +pub mod mmcs; pub(crate) mod terminal; use crypto::fiat_shamir::is_transcript::IsStarkTranscript; diff --git a/crypto/stark/src/par.rs b/crypto/stark/src/par.rs index cee693e3f..5ad4accbd 100644 --- a/crypto/stark/src/par.rs +++ b/crypto/stark/src/par.rs @@ -92,3 +92,27 @@ pub(crate) fn par_try_for_each_mut( slice.iter_mut().try_for_each(f) } } + +/// Run `f(i, &mut item)` for each element of `slice` with its index. Parallel +/// when `feature = "parallel"`, sequential otherwise. +#[cfg_attr(not(feature = "parallel"), allow(dead_code))] +pub(crate) fn par_for_each_mut_indexed( + slice: &mut [T], + f: impl Fn(usize, &mut T) + Sync + Send, +) { + #[cfg(feature = "parallel")] + { + use rayon::prelude::*; + slice + .par_iter_mut() + .enumerate() + .for_each(|(i, item)| f(i, item)); + } + #[cfg(not(feature = "parallel"))] + { + slice + .iter_mut() + .enumerate() + .for_each(|(i, item)| f(i, item)); + } +} From b3101713de760535244e1c8773a1b855625a3b2e Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 18:01:55 -0300 Subject: [PATCH 02/26] refactor(stark): plain-data prover/verifier helpers for batched reuse --- crypto/math/src/fft/bit_reversing.rs | 5 +- crypto/stark/src/prover.rs | 456 +++++++++++++++++++++++++-- crypto/stark/src/trace.rs | 6 +- crypto/stark/src/verifier.rs | 57 +++- 4 files changed, 481 insertions(+), 43 deletions(-) diff --git a/crypto/math/src/fft/bit_reversing.rs b/crypto/math/src/fft/bit_reversing.rs index 8e830888b..c5fc69ce6 100644 --- a/crypto/math/src/fft/bit_reversing.rs +++ b/crypto/math/src/fft/bit_reversing.rs @@ -40,10 +40,7 @@ pub fn reverse_index(i: usize, size: u64) -> usize { /// `debug_assert!`): a non-power-of-two `n` would break the disjointness the /// parallel path relies on, turning a bad caller's input into a data race. #[cfg(feature = "alloc")] -pub(crate) fn in_place_bit_reverse_permute_row_major( - buf: &mut [E], - num_cols: usize, -) { +pub fn in_place_bit_reverse_permute_row_major(buf: &mut [E], num_cols: usize) { if num_cols == 0 || buf.is_empty() { return; } diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index d31ea09a2..a8dc48726 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -188,7 +188,7 @@ fn precomputed_tree_cache() CACHE.get_or_init(|| Mutex::new(std::collections::HashMap::new())) } -fn precomputed_tree_cache_get( +pub(crate) fn precomputed_tree_cache_get( root: &Commitment, ) -> Option>> where @@ -201,7 +201,7 @@ where .and_then(|any| any.downcast::>().ok()) } -fn precomputed_tree_cache_put( +pub(crate) fn precomputed_tree_cache_put( root: Commitment, tree: Arc>, ) where @@ -397,6 +397,9 @@ pub(crate) struct LdeTwiddles { /// `two_half_fwd` size-`n·blowup` forward. two_half_inv: TwoHalfTwiddles, two_half_fwd: TwoHalfTwiddles, + /// Size-`n` FORWARD set, built lazily — only the batched phase-4 coset + /// evaluation wants it (a full LDE's forward set is size `n·blowup`). + two_half_fwd_n: OnceLock>, coset_weights: Vec>, /// Composition half-extension cache, initialized only when the degree-2 /// decomposition path actually runs on CPU. @@ -477,12 +480,21 @@ impl LdeTwiddles { .expect("valid inverse two-half twiddles"), two_half_fwd: TwoHalfTwiddles::::new(lde_size.trailing_zeros() as usize, false) .expect("valid forward two-half twiddles"), + two_half_fwd_n: OnceLock::new(), coset_weights, composition: OnceLock::new(), inv_2x: OnceLock::new(), } } + /// The size-`n` forward set for the phase-4 coset evaluation, built once. + pub(crate) fn fwd_n(&self, domain_size: usize) -> &TwoHalfTwiddles { + self.two_half_fwd_n.get_or_init(|| { + TwoHalfTwiddles::::new(domain_size.trailing_zeros() as usize, false) + .expect("valid size-n forward two-half twiddles") + }) + } + fn composition(&self, domain: &Domain) -> &CompositionLdeTwiddles { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; let half_size = lde_size / 2; @@ -537,7 +549,10 @@ fn domain_twiddle_cache() -> &'static std::sync::Mutex< CACHE.get_or_init(Default::default) } -fn domain_and_twiddles(air: &A, trace_length: usize) -> (Arc>, Arc>) +pub(crate) fn domain_and_twiddles( + air: &A, + trace_length: usize, +) -> (Arc>, Arc>) where F: IsFFTField + 'static, FieldElement: Send + Sync, @@ -808,12 +823,33 @@ where pub(crate) gpu_composition_tree: Option, } +/// The composition-polynomial parts, before any commitment is taken over them. +/// +/// Returned by [`IsStarkProver::compute_composition_parts`], which round 2 and +/// the batched prover share. The device handle rides along rather than being +/// installed on the `Round1` inside, because the two callers install it at +/// different points: round 2 folds it into the table's own LDE session, the +/// batched prover keeps every table's parts alive only until they have been +/// absorbed into the epoch's MMCS. +pub(crate) struct CompositionParts +where + FieldElement: AsBytes + Sync + Send, +{ + pub(crate) parts: Vec>>, + #[cfg(feature = "cuda")] + pub(crate) gpu_parts: Option, + #[cfg(feature = "instruments")] + pub(crate) constraints_dur: Duration, + #[cfg(feature = "instruments")] + pub(crate) fft_dur: Duration, +} + /// A container for the results of the third round of the STARK Prove protocol. pub(crate) struct Round3 { /// Evaluations of the trace polynomials, main and auxiliary, at the out-of-domain challenge. - trace_ood_evaluations: Table, + pub(crate) trace_ood_evaluations: Table, /// Evaluations of the composition polynomial parts at the out-of-domain challenge. - composition_poly_parts_ood_evaluation: Vec>, + pub(crate) composition_poly_parts_ood_evaluation: Vec>, } /// A container for the results of the fourth round of the STARK Prove protocol. @@ -1387,6 +1423,129 @@ pub trait IsStarkProver< Ok((commit, (main_data, total_cols))) } + /// Expand a table's main trace to its coset LDE, row-major, without + /// building any Merkle tree. + /// + /// The Round-1 CPU commit and the `ResidencyMode::RecomputeLde` recompute + /// both go through here, which is what makes the recomputed buffer + /// bit-identical to the one the tree was built from — identical by + /// construction rather than by argument. The twiddles are process-cached, + /// so the second call re-runs the NTT over the same inputs. + fn expand_main_lde_row_major( + trace: &TraceTable, + domain: &Domain, + twiddles: &LdeTwiddles, + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + ) -> (Vec>, usize) { + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let (trace_data, total_cols) = trace.main_data_row_major(); + + let mut main_data: Vec> = Vec::with_capacity(lde_size * total_cols); + main_data.extend_from_slice(trace_data); + + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + trace.main_table.advise_drop_cache(); + } + + Polynomial::>::coset_lde_full_expand_row_major::( + &mut main_data, + total_cols, + domain.blowup_factor, + &twiddles.coset_weights, + &twiddles.two_half_inv, + &twiddles.two_half_fwd, + ) + .expect("row-major coset LDE expansion"); + + (main_data, total_cols) + } + + /// The MAIN trace's size-`n` coset evaluation, row-major — the stride- + /// `blowup` subsample of [`Self::expand_main_lde_row_major`]'s output, + /// computed directly: iFFT(n) → coset weights → FFT(n), about 37% of a + /// full expansion's work and a quarter of its bytes. Values are + /// bit-identical to the subsample (exact modular arithmetic; both compute + /// the same DFT), which is what the batched phase 4 reads and ALL it + /// reads. + fn expand_main_coset_eval_row_major( + trace: &TraceTable, + domain: &Domain, + twiddles: &LdeTwiddles, + ) -> (Vec>, usize) { + let (trace_data, total_cols) = trace.main_data_row_major(); + let mut main_data: Vec> = Vec::with_capacity(trace_data.len()); + main_data.extend_from_slice(trace_data); + Polynomial::>::coset_lde_full_expand_row_major::( + &mut main_data, + total_cols, + 1, + &twiddles.coset_weights, + &twiddles.two_half_inv, + twiddles.fwd_n(domain.interpolation_domain_size), + ) + .expect("row-major coset evaluation"); + (main_data, total_cols) + } + + /// The AUX counterpart of [`Self::expand_main_coset_eval_row_major`]. + fn expand_aux_coset_eval_row_major( + trace: &TraceTable, + domain: &Domain, + twiddles: &LdeTwiddles, + ) -> (Vec>, usize) { + let (trace_data, total_cols) = trace.aux_data_row_major(); + let mut aux_data: Vec> = Vec::with_capacity(trace_data.len()); + aux_data.extend_from_slice(trace_data); + Polynomial::>::coset_lde_full_expand_row_major::( + &mut aux_data, + total_cols, + 1, + &twiddles.coset_weights, + &twiddles.two_half_inv, + twiddles.fwd_n(domain.interpolation_domain_size), + ) + .expect("row-major aux coset evaluation"); + (aux_data, total_cols) + } + + /// Expand a table's auxiliary trace to its coset LDE, row-major, without + /// building any Merkle tree — the aux counterpart of + /// [`Self::expand_main_lde_row_major`], and extracted for the same reason: + /// the batched prover rebuilds this buffer once per phase instead of + /// retaining it, and a second expansion written elsewhere would be a second + /// encoding of the committed one. + fn expand_aux_lde_row_major( + trace: &TraceTable, + domain: &Domain, + twiddles: &LdeTwiddles, + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + ) -> (Vec>, usize) { + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let (trace_data, total_cols) = trace.aux_data_row_major(); + + let mut aux_data: Vec> = + Vec::with_capacity(lde_size * total_cols); + aux_data.extend_from_slice(trace_data); + + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + trace.aux_table.advise_drop_cache(); + } + + Polynomial::>::coset_lde_full_expand_row_major::( + &mut aux_data, + total_cols, + domain.blowup_factor, + &twiddles.coset_weights, + &twiddles.two_half_inv, + &twiddles.two_half_fwd, + ) + .expect("row-major aux coset LDE expansion"); + + (aux_data, total_cols) + } + /// Spill a committed Merkle tree to disk when `storage_mode` is `Disk`, /// tagging any I/O error with `label`. No-op otherwise. Shared by every commit /// site (main / preprocessed split / aux). @@ -1626,6 +1785,249 @@ pub trait IsStarkProver< .expect("coset extension") } + /// The evaluations of the composition-polynomial parts over the LDE domain, + /// and nothing else — no commitment. + /// + /// This is the half of round 2 that the batched path shares with the + /// per-table one. Round 2 commits each table's parts to its own Merkle tree; + /// [`crate::batched::prover::multi_prove_batched`] streams every table's + /// parts into one mixed-height MMCS instead. Both need the same parts, and + /// the arm selection (`number_of_parts` 1 / 2 / d>2, the device paths and + /// their fallbacks) is intricate enough that a second copy would drift — so + /// there is one function, and the commitment is what differs. + #[allow(clippy::too_many_arguments)] + fn compute_composition_parts( + air: &dyn AIR, + pub_inputs: &PI, + domain: &Domain, + twiddles: &LdeTwiddles, + // `&mut` for the device-only recovery below: when the host evaluator is + // reached on a device-resident trace, the LDEs are downloaded back into + // the host buffers in place rather than aborting the table. + lde_trace: &mut LDETraceTable, + rap_challenges: &[FieldElement], + bus_public_inputs: Option<&BusPublicInputs>, + transition_coefficients: &[FieldElement], + boundary_coefficients: &[FieldElement], + ) -> Result, ProvingError> + where + FieldElement: AsBytes, + FieldElement: AsBytes, + { + // Compute the evaluations of the composition polynomial on the LDE domain. + let trace_length = domain.interpolation_domain_size; + let evaluator = ConstraintEvaluator::new( + air, + pub_inputs, + rap_challenges, + bus_public_inputs, + trace_length, + ); + let number_of_parts = air.composition_poly_degree_bound(trace_length) / trace_length; + + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + #[cfg(feature = "cuda")] + let mut gpu_composition_parts: Option = None; + + // Fully device-resident d=2 path: H stays on device through decompose + + // 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; + // A downloaded `H` awaiting the host decompose: produced under the + // lock below, consumed after it — the host iFFT + LDEs are pure CPU + // work and must not serialize other tables' device windows. + #[cfg(feature = "cuda")] + let mut downloaded_h: Option>> = None; + #[cfg(feature = "cuda")] + if number_of_parts == 2 && !crate::gpu_lde::gpu_force_downgrade() { + // Serializing this window across tables (device constraint eval + + // decompose, where H is born) empirically eliminates a transient + // whole-buffer H corruption seen under concurrent R2 windows on + // VRAM pressure. What the guard orders is submission: a + // device-only table's window is enqueue-only, so its kernels may + // still overlap another table's on device. The commit, the host + // decompose of a downloaded `H` and every host arm run outside + // the lock. The force-downgrade test hook skips this fast path so + // every device-only table exercises the host recovery below. + let _r2_serial_guard = crate::gpu_lde::r2_serialize_guard(); + if let Some(h_dev) = evaluator.evaluate_dev( + air, + lde_trace, + domain, + transition_coefficients, + boundary_coefficients, + rap_challenges, + ) { + match crate::gpu_lde::try_decompose_extend_d2_dev::( + &h_dev, + twiddles.inv_2x(domain), + &twiddles.composition(domain).weights, + !lde_trace.host_trace_empty(), + ) { + Some((parts, handle)) => { + gpu_composition_parts = Some(handle); + precomputed_parts = Some(parts); + } + None => { + downloaded_h = + crate::gpu_lde::download_comp_h_to_field::(&h_dev); + } + } + } + } + #[cfg(feature = "cuda")] + if let Some(h) = downloaded_h.take() { + precomputed_parts = Some(Self::decompose_and_extend_d2(&h, domain, twiddles)); + } + #[cfg(not(feature = "cuda"))] + let precomputed_parts: Option>>> = None; + + #[cfg(feature = "instruments")] + let constraints_dur = t_sub.elapsed(); + #[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. The gate is a static predicate and cannot mirror every + // dynamic decline, so recover rather than abort: download the resident + // LDEs into the host buffers (which also clears the device-only flag) + // and let the host arms run — slower for this table, never wrong. The + // assert is left for the case where the handles themselves cannot + // serve the data, so that failure carries 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() && lde_trace.host_trace_empty() { + let recovered = crate::gpu_lde::materialize_lde_trace_host(lde_trace); + if recovered { + // Rare by design; the name tells which condition the gate is + // missing so it can be mirrored as an optimization. + eprintln!( + "[gpu] device-only downgrade: table={} n={} num_parts={} \ + (device R2 path declined; continuing on host)", + air.name(), + trace_length, + number_of_parts, + ); + } + assert!( + recovered, + "R2 composition fell back to the host evaluator on a device-only \ + trace and the resident handles could not be downloaded: \ + table={} n={} num_parts={} main_cols={} aux_cols={}", + air.name(), + trace_length, + number_of_parts, + lde_trace.num_main_cols(), + lde_trace.num_aux_cols(), + ); + } + + let lde_composition_poly_parts_evaluations = if let Some(parts) = precomputed_parts { + parts + } else if number_of_parts == 2 { + // Direct quotient decomposition: avoid full-size iFFT by algebraically + // splitting H(x) = H₀(x²) + x·H₁(x²) using: + // H₀(x²) = (H(x) + H(-x)) / 2 + // H₁(x²) = (H(x) - H(-x)) / (2x) + // On the LDE coset {g·ω^i}, we have -g·ω^i = g·ω^{i+N} since ω^N = -1. + let constraint_evaluations = evaluator.evaluate( + air, + lde_trace, + domain, + transition_coefficients, + boundary_coefficients, + rap_challenges, + ); + Self::decompose_and_extend_d2(&constraint_evaluations, domain, twiddles) + } else if number_of_parts == 1 { + // Degree bound equals trace length: constraint evals are the LDE directly. + vec![evaluator.evaluate( + air, + lde_trace, + domain, + transition_coefficients, + boundary_coefficients, + rap_challenges, + )] + } else { + // Fallback for any future AIR with d > 2. + let constraint_evaluations = evaluator.evaluate( + air, + lde_trace, + domain, + transition_coefficients, + boundary_coefficients, + rap_challenges, + ); + let composition_poly = + Polynomial::interpolate_offset_fft(&constraint_evaluations, &domain.coset_offset)?; + let composition_poly_parts = composition_poly.break_in_parts(number_of_parts); + + let cpu_eval = || -> Result>>, ProvingError> { + composition_poly_parts + .iter() + .map(|part| { + evaluate_polynomial_on_lde_domain( + part, + domain.blowup_factor, + domain.interpolation_domain_size, + &domain.coset_offset, + ) + .map_err(ProvingError::from) + }) + .collect() + }; + + // GPU fast path: batched ext3 LDE for all parts in one call. + // `_keep` variant retains the de-interleaved device buffer as a + // `GpuLdeExt3` handle stored on Round2 so R4 DEEP can skip the + // `num_parts * 3 * lde_size * 8` byte H2D. + #[cfg(feature = "cuda")] + { + let parts_slices: Vec<&[FieldElement]> = composition_poly_parts + .iter() + .map(|p| p.coefficients.as_slice()) + .collect(); + match crate::gpu_lde::try_evaluate_parts_on_lde_gpu_keep::( + &parts_slices, + domain.blowup_factor, + domain.interpolation_domain_size, + &domain.coset_offset, + ) { + Some((evals, handle)) => { + gpu_composition_parts = Some(handle); + evals + } + None => cpu_eval()?, + } + } + #[cfg(not(feature = "cuda"))] + cpu_eval()? + }; + + #[cfg(feature = "instruments")] + let fft_dur = t_sub.elapsed(); + + Ok(CompositionParts { + parts: lde_composition_poly_parts_evaluations, + #[cfg(feature = "cuda")] + gpu_parts: gpu_composition_parts, + #[cfg(feature = "instruments")] + constraints_dur, + #[cfg(feature = "instruments")] + fft_dur, + }) + } + /// Returns the result of the second round of the STARK Prove protocol. fn round_2_compute_composition_polynomial( air: &dyn AIR, @@ -1937,15 +2339,18 @@ pub trait IsStarkProver< fn round_3_evaluate_polynomials_in_out_of_domain_element( air: &dyn AIR, domain: &Domain, - round_1_result: &mut Round1, - round_2_result: &mut Round2, + // Both `&mut` for the device-only recoveries below: the parts OOD arm + // repopulates the host part evals from the resident handle, and the + // trace OOD reads through a trace that may have to be materialized. + lde_trace: &mut LDETraceTable, + composition_parts: &mut [Vec>], z: &FieldElement, ) -> Round3 where FieldElement: AsBytes, FieldElement: AsBytes, { - let num_parts = round_2_result.lde_composition_poly_evaluations.len(); + let num_parts = composition_parts.len(); let z_power = z.pow(num_parts); let domain_size = domain.interpolation_domain_size; let blowup_factor = domain.blowup_factor; @@ -1961,8 +2366,7 @@ pub trait IsStarkProver< // the host stride-extract and the sequential CPU fold per part. #[cfg(feature = "cuda")] let gpu_parts_ood: Option>> = - round_1_result - .lde_trace + lde_trace .gpu_composition_parts() .and_then(|parts_dev| { let dispatch = |inv_host: &[FieldElement], @@ -1982,7 +2386,7 @@ pub trait IsStarkProver< match crate::gpu_lde::try_prep_r3_dev_context::( &dc.points, std::slice::from_ref(&z_power), - round_1_result.lde_trace.bound_stream(), + lde_trace.bound_stream(), ) { Some(ctx) => dispatch(&[], Some((&ctx, 0))), // Below the dev-context threshold (single eval point): @@ -2008,8 +2412,8 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] { let recovered = crate::gpu_lde::materialize_composition_parts_host( - &round_1_result.lde_trace, - &mut round_2_result.lde_composition_poly_evaluations, + lde_trace, + composition_parts, ); assert!( recovered, @@ -2020,8 +2424,7 @@ pub trait IsStarkProver< } let comp_inv_denoms = math::polynomial::barycentric_inv_denoms(&z_power, &dc.points); - round_2_result - .lde_composition_poly_evaluations + composition_parts .iter() .map(|lde_evals| { // Extract trace-size evaluations (stride = blowup_factor) @@ -2044,7 +2447,7 @@ pub trait IsStarkProver< // === Trace polynomials: barycentric evaluation via LDE === let trace_ood_evaluations = crate::trace::get_trace_evaluations_from_lde( - &mut round_1_result.lde_trace, + lde_trace, domain, z, &air.context().transition_offsets, @@ -2173,7 +2576,7 @@ pub trait IsStarkProver< let t_sub = Instant::now(); let deep_evals = Self::compute_deep_composition_poly_evaluations( &mut round_1_result.lde_trace, - round_2_result, + &mut round_2_result.lde_composition_poly_evaluations, round_3_result, z, domain, @@ -2334,8 +2737,11 @@ pub trait IsStarkProver< #[allow(clippy::too_many_arguments)] fn compute_deep_composition_poly_evaluations( + // Both `&mut` for the device-only recovery in the host DEEP loop: the + // trace and the part evals are downloaded back in place there rather + // than aborting the table. lde_trace: &mut LDETraceTable, - round_2_result: &mut Round2, + composition_parts: &mut [Vec>], round_3_result: &Round3, z: &FieldElement, domain: &Domain, @@ -2347,7 +2753,7 @@ pub trait IsStarkProver< FieldElement: AsBytes, FieldElement: AsBytes, { - let num_parts = round_2_result.lde_composition_poly_evaluations.len(); + let num_parts = composition_parts.len(); let z_power = z.pow(num_parts); // pole for H terms // Number of evaluation points per trace column (= transition_offsets.len() * step_size) @@ -2396,7 +2802,7 @@ pub trait IsStarkProver< crate::gpu_lde::try_deep_composition_gpu::( lde_trace, lde_trace.gpu_composition_parts(), - &round_2_result.lde_composition_poly_evaluations, + composition_parts, h_ood, &trace_ood_columns, composition_poly_gammas, @@ -2432,7 +2838,7 @@ pub trait IsStarkProver< crate::gpu_lde::try_deep_composition_gpu::( lde_trace, lde_trace.gpu_composition_parts(), - &round_2_result.lde_composition_poly_evaluations, + composition_parts, h_ood, &trace_ood_columns, composition_poly_gammas, @@ -2464,7 +2870,7 @@ pub trait IsStarkProver< } let parts_recovered = crate::gpu_lde::materialize_composition_parts_host( lde_trace, - &mut round_2_result.lde_composition_poly_evaluations, + composition_parts, ); assert!( parts_recovered, @@ -2516,7 +2922,7 @@ pub trait IsStarkProver< // H terms for j in 0..num_parts { - let h_j_val = &round_2_result.lde_composition_poly_evaluations[j][i]; + let h_j_val = &composition_parts[j][i]; let h_j_ood = &h_ood[j]; result += &composition_poly_gammas[j] * (h_j_val - h_j_ood) * &inv_h[i]; } @@ -4435,8 +4841,8 @@ pub trait IsStarkProver< let round_3_result = Self::round_3_evaluate_polynomials_in_out_of_domain_element( air, domain, - round_1_result, - &mut round_2_result, + &mut round_1_result.lde_trace, + &mut round_2_result.lde_composition_poly_evaluations, &z, ); #[cfg(feature = "instruments")] diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index f953faac8..348f91d65 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -841,7 +841,11 @@ where E: IsField + 'static, { let n = domain.interpolation_domain_size; - let bf = domain.blowup_factor; + // The read stride is the TABLE's own blowup, not the domain's: the per-table + // path hands an LDE at the domain blowup (so these agree), and the batched + // phase 4 hands a blowup-1 table that IS the stride subsample already — + // same values, a quarter the buffer. + let bf = lde_trace.blowup_factor; let num_main_cols = lde_trace.num_main_cols(); let num_aux_cols = lde_trace.num_aux_cols(); let table_width = num_main_cols + num_aux_cols; diff --git a/crypto/stark/src/verifier.rs b/crypto/stark/src/verifier.rs index ca6f15152..8d869a302 100644 --- a/crypto/stark/src/verifier.rs +++ b/crypto/stark/src/verifier.rs @@ -263,9 +263,17 @@ pub trait IsStarkVerifier< }) } + /// The three proof-derived inputs are passed as plain data rather than read + /// off a `StarkProofView`, because the batched epoch verifier + /// ([`crate::batched::verifier`]) has to run this identical check against a + /// proof that has no such view. One constraint check, two callers. + #[allow(clippy::too_many_arguments)] fn step_2_verify_claimed_composition_polynomial( air: &dyn AIR, - proof: StarkProofView<'_, Field, FieldExtension, PI>, + trace_length: usize, + bus_table_contribution: Option>, + ood_current_row: &[FieldElement], + composition_parts_ood: &[FieldElement], public_inputs: &PI, domain: &VerifierDomain, challenges: &Challenges, @@ -280,11 +288,10 @@ pub trait IsStarkVerifier< crate::profile_markers::step_marker::< { crate::profile_markers::STEP_VERIFY_CLAIMED_COMPOSITION_POLYNOMIAL }, >(); - let trace_length = proof.trace_length(); // Owned `BusPublicInputs` (just the table contribution L — one field // element) reconstructed for the AIR boundary call. - let bus_public_inputs = proof - .bus_table_contribution() + let bus_public_inputs = bus_table_contribution + .clone() .map(BusPublicInputs::from_contribution); let boundary_constraints = air.boundary_constraints( @@ -309,8 +316,7 @@ pub trait IsStarkVerifier< .collect(); let main_trace_width = air.trace_layout().0; - let trace_ood_evaluations = proof.trace_ood_evaluations(); - let ood_row = trace_ood_evaluations.get_row(0); + let ood_row = ood_current_row; let (boundary_c_i_evaluations_num, mut boundary_c_i_evaluations_den): ( Vec>, @@ -352,8 +358,8 @@ pub trait IsStarkVerifier< // aux count; reject instead of underflowing. The current-row block keeps // the full trace width even under g·z pruning, so this still yields the // main width. - let num_main_trace_columns = match trace_ood_evaluations - .width() + let num_main_trace_columns = match ood_current_row + .len() .checked_sub(air.num_auxiliary_rap_columns()) { Some(n) => n, @@ -370,7 +376,7 @@ pub trait IsStarkVerifier< Vec::new() }; - let logup_table_offset = match proof.bus_table_contribution() { + let logup_table_offset = match bus_table_contribution { Some(contribution) => { let n = FieldElement::::from(trace_length as u64); match n.inv() { @@ -418,8 +424,7 @@ pub trait IsStarkVerifier< let composition_poly_ood_evaluation = &boundary_quotient_ood_evaluation + transition_c_i_evaluations_sum; - let composition_poly_claimed_ood_evaluation = proof - .composition_poly_parts_ood_evaluation() + let composition_poly_claimed_ood_evaluation = composition_parts_ood .iter() .rev() .fold(FieldElement::zero(), |acc, coeff| { @@ -836,12 +841,36 @@ pub trait IsStarkVerifier< /// elsewhere), not from `proof.trace_ood_evaluations()` which now carries /// only the current-row block. Pruned positions are zero in both the grid /// and `trace_term_coeffs`, so next rows sum only the window columns. + /// The per-proof invariant terms, from a `StarkProofView`. + /// + /// A thin delegate to [`Self::query_invariant_deep_terms_from_parts`], kept + /// at this name and signature because it is public API with callers outside + /// this crate (`prover/src/lfm/`'s constraint and join oracles). The batched + /// epoch verifier calls the slice-taking form instead — one OOD walk, two + /// entry points, rather than a second copy of it. fn compute_query_invariant_deep_terms( challenges: &Challenges, proof: StarkProofView<'_, Field, FieldExtension, PI>, ood_full: &Table, next_row_cols: &[usize], step_size: usize, + ) -> Option> { + Self::query_invariant_deep_terms_from_parts( + challenges, + proof.composition_poly_parts_ood_evaluation(), + ood_full, + next_row_cols, + step_size, + ) + } + + /// `composition_parts_ood` is the proof's `Hᵢ(z^N)` values. + fn query_invariant_deep_terms_from_parts( + challenges: &Challenges, + composition_parts_ood: &[FieldElement], + ood_full: &Table, + next_row_cols: &[usize], + step_size: usize, ) -> Option> { let ood_evaluations_table_height = ood_full.height; let ood_evaluations_table_width = ood_full.width; @@ -874,7 +903,6 @@ pub trait IsStarkVerifier< ood_row_sum.push(sum); } - let composition_parts_ood = proof.composition_poly_parts_ood_evaluation(); let number_of_parts = composition_parts_ood.len(); let z_pow = challenges.z.pow(number_of_parts); @@ -1710,7 +1738,10 @@ pub trait IsStarkVerifier< if !Self::step_2_verify_claimed_composition_polynomial( air, - proof, + proof.trace_length(), + proof.bus_table_contribution(), + proof.trace_ood_evaluations().get_row(0), + proof.composition_poly_parts_ood_evaluation(), public_inputs, &domain, &challenges, From 627e3bf4f2173cecd4f820382645f290f0887c36 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 18:02:58 -0300 Subject: [PATCH 03/26] feat(stark): batched MMCS prover/verifier (multi_prove_batched) + tests --- crypto/stark/src/batched/mod.rs | 27 + crypto/stark/src/batched/proof.rs | 204 ++ crypto/stark/src/batched/prover.rs | 1208 ++++++++++++ crypto/stark/src/batched/round4.rs | 892 +++++++++ crypto/stark/src/batched/shape.rs | 381 ++++ crypto/stark/src/batched/verifier.rs | 1068 +++++++++++ crypto/stark/src/lib.rs | 2 + crypto/stark/src/residency_mode.rs | 36 + .../src/tests/batched_mmcs_soundness_tests.rs | 1640 +++++++++++++++++ .../stark/src/tests/batched_prover_tests.rs | 1221 ++++++++++++ crypto/stark/src/tests/mod.rs | 2 + 11 files changed, 6681 insertions(+) create mode 100644 crypto/stark/src/batched/mod.rs create mode 100644 crypto/stark/src/batched/proof.rs create mode 100644 crypto/stark/src/batched/prover.rs create mode 100644 crypto/stark/src/batched/round4.rs create mode 100644 crypto/stark/src/batched/shape.rs create mode 100644 crypto/stark/src/batched/verifier.rs create mode 100644 crypto/stark/src/residency_mode.rs create mode 100644 crypto/stark/src/tests/batched_mmcs_soundness_tests.rs create mode 100644 crypto/stark/src/tests/batched_prover_tests.rs diff --git a/crypto/stark/src/batched/mod.rs b/crypto/stark/src/batched/mod.rs new file mode 100644 index 000000000..1260dd962 --- /dev/null +++ b/crypto/stark/src/batched/mod.rs @@ -0,0 +1,27 @@ +//! The batched-commitment proving path: one mixed-height MMCS per round and one +//! FRI instance per epoch, instead of one tree and one FRI instance per table. +//! +//! This is an OPT-IN path. The per-table prover and verifier +//! ([`crate::prover::IsStarkProver::multi_prove`], +//! [`crate::verifier::IsStarkVerifier::multi_verify`]) are untouched and produce +//! byte-identical proofs; nothing here is reachable from them. +//! +//! The primitives live one level down — [`crate::fri::mmcs`] (the mixed-height +//! tree) and [`crate::fri::batched`] (height combination, the batched commit +//! phase, and the shared challenge derivation). This module is the wiring: it +//! fixes the transcript sequence, the query-index convention and the per-query +//! fold-with-injection recursion that the prover and the verifier must agree on. +//! +//! - [`shape`] — which table contributes which matrix to which round, derived +//! from the AIR set on both sides and never read from a proof. +//! - [`round4`] — the round-4 transcript sequence and the per-query FRI check. +//! - [`proof`] — what a batched epoch proof carries. +//! - [`prover`] — the phase architecture the barriers force. +//! - [`verifier`] — the transcript replay, and ⛔ only the commitment half of a +//! verification. Read its header before assuming otherwise. + +pub mod proof; +pub mod prover; +pub mod round4; +pub mod shape; +pub mod verifier; diff --git a/crypto/stark/src/batched/proof.rs b/crypto/stark/src/batched/proof.rs new file mode 100644 index 000000000..62341e346 --- /dev/null +++ b/crypto/stark/src/batched/proof.rs @@ -0,0 +1,204 @@ +//! What a batched epoch proof carries. +//! +//! These types live here and NOT in [`crate::proof`] on purpose: the per-table +//! `StarkProof` / `MultiProof` rkyv layouts are the production wire format, and +//! the batched path is opt-in. Keeping its types in this module makes the +//! default path byte-identical by construction rather than by test — +//! `git diff ..HEAD -- crypto/stark/src/proof/` stays empty. +//! +//! # What is NOT here, and why +//! +//! No per-table Merkle roots, no per-table FRI layer roots, no per-table query +//! list. One epoch commits four mixed-height MMCS roots (preprocessed, main, +//! aux, composition parts) and runs ONE FRI instance, so a query costs one +//! authentication path per round instead of one per table per round. That is the +//! proof-size win; everything per-table that survives is data the verifier +//! cannot derive — OOD evaluations, bus sums, public inputs. + +use math::field::element::FieldElement; +use math::field::traits::IsField; + +use crate::config::Commitment; +use crate::fri::fri_decommit::FriDecommitment; +use crate::fri::mmcs::MixedOpening; +use crate::lookup::BusPublicInputs; +use crate::proof::stark::PolynomialOpenings; +use crate::table::Table; + +/// The per-table data a batched epoch proof still has to carry. +#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] +pub struct BatchedTableData { + /// This table's interpolation-domain size. The verifier derives the table's + /// height — and therefore every index reduction — from this, so it is bound + /// into the round-4 shape histogram before any challenge depends on it. + pub trace_length: usize, + /// tⱼ(z·gᵏ): the current-row block (all columns). + pub trace_ood_evaluations: Table, + /// tⱼ(z·gᵏ): the pruned next-row block (masked columns only). + pub trace_ood_next_evaluations: Table, + /// Hᵢ(z^N). + pub composition_poly_parts_ood_evaluation: Vec>, + /// LogUp bus sums, when the table has a RAP. + pub bus_public_inputs: Option>, + /// Public inputs for the boundary constraints. + pub public_inputs: PI, + /// A table excluded from the batched FRI class keeps a terminal-only + /// instance of its own: its DEEP codeword IS its terminal codeword, sent as + /// the `2^(h - blowup_log)` coefficients of the polynomial it evaluates. + /// `None` for a table in the batched class. + /// + /// See [`crate::fri::batched::FriInstancePlan`]: the partition is derived + /// from the shape by both sides and is never sent, so this field's presence + /// is checked against the derived plan rather than trusted. + pub standalone_final_poly_coeffs: Option>>, +} + +/// One query's openings: one authentication path per batched round, plus the +/// FRI layer decommitment. +#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] +pub struct BatchedQueryOpening { + /// Preprocessed openings, ONE PER PREPROCESSED TABLE in AIR order — each a + /// standard row-pair opening against that table's own precomputed tree. + /// + /// ★ Deliberately NOT a round of the mixed MMCS (this is #768's + /// arrangement, kept for the same reason): the per-table precomputed trees + /// are exactly the ones `air.precomputed_commitment()` pins, so the + /// verifier absorbs and compares roots it already owns — the per-table + /// path's critical soundness check, verbatim — and a recursive verifier + /// binds each root with the provenance machinery that already exists + /// (interned constant / derived in-machine / ELF-attested). A fused + /// mixed-height prep root has no in-machine binding story: its provenance + /// classes are mixed into one digest, which is the M-8 blocker this layout + /// dissolves. Empty when the epoch has no preprocessed table. + pub prep: Vec>, + /// Main round — always present; every table contributes a matrix. + pub main: MixedOpening, + /// Auxiliary round. `None` when no table has a RAP. + pub aux: Option>, + /// Composition-parts round — always present. + pub parts: MixedOpening, + /// The carved table's main row pair, a standard row-pair opening against + /// [`BatchedMultiProof::carved_main_root`] at the REDUCED index + /// (`reduce_iota_to_round(iota, h_max, h_carved)`). Present iff the epoch + /// is carved ([`crate::batched::shape::CarvedMain`]); the verifier rejects + /// a stray or missing one. + pub carved_main: Option>, + /// The batched FRI instance's per-layer openings for this query. + pub fri: FriDecommitment, +} + +/// One epoch, one proof. +#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] +pub struct BatchedMultiProof { + pub tables: Vec>, + /// ★ There is deliberately NO `prep_root` here. Preprocessed matrices are + /// committed per table and their roots are `air.precomputed_commitment()` + /// — absorbed by both sides FROM THE AIR SET, never from the proof, + /// exactly as the per-table path's Phase A does. The proof carries only + /// the per-query openings ([`BatchedQueryOpening::prep`]). + pub main_root: Commitment, + /// The carved table's standalone main-tree root — PROOF-CARRIED (absorbed + /// from the proof, like `main_root`), unlike the preprocessed roots, which + /// absorb from the AIR set. It is absorbed after the preprocessed roots and + /// before `main_root`, so every challenge is drawn after it. Present iff + /// the epoch is carved; whether the epoch IS carved is verifier-owned + /// configuration, never read from the proof. For a continuation epoch this + /// root is the L2G commitment `verify_l2g_commitment_binding_view` compares + /// against the global proof — byte-identical to the per-table L2G tree. + pub carved_main_root: Option, + pub aux_root: Option, + pub parts_root: Commitment, + /// The batched FRI instance's committed layer roots. + pub fri_layer_roots: Vec, + /// The batched FRI instance's terminal polynomial. + pub fri_final_poly_coeffs: Vec>, + pub nonce: Option, + pub queries: Vec>, +} + +/// What the batched prove cost in residency and in recomputation. +/// +/// Returned rather than logged because it is the number the campaign's +/// projection is missing (MMCS-PLAN §1.1 prices the commitment work but not the +/// LDE rebuilds the phase barriers force). A test can assert on it, which is +/// what keeps "the batched builder does not hold every table's LDE" falsifiable +/// at the PROVER level instead of only at the primitive's. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct BatchedProveStats { + /// Highest number of bytes of MAIN and AUX LDE simultaneously alive across + /// the whole prove, counted as the buffers are created and dropped. + /// + /// ★ This is the number the acceptance test asserts on, and the reason it is + /// reported separately from the parts. Under `ResidencyMode::RecomputeLde` + /// it must be flat in the table count — bounded by the widest single table, + /// not by the epoch. If the streaming builder were bypassed, or if any phase + /// quietly retained what it read, this would grow with `N` instead, which is + /// precisely the failure MMCS-PLAN §3.3 warns gives the win back inside the + /// same commit. + pub peak_trace_lde_bytes: usize, + /// Bytes of composition parts held at the peak. These are `O(N)` BY DESIGN — + /// recomputing them is a second constraint evaluation — so they are counted + /// apart from the trace LDEs rather than allowed to mask their behaviour. + pub retained_parts_bytes: usize, + /// The two above at the moment either was highest. Reported for budgeting; + /// the falsifiable claim lives in `peak_trace_lde_bytes`. + pub peak_lde_bytes: usize, + /// How many times a main LDE was expanded from a trace. One per table is the + /// floor (the commit itself); every phase barrier that follows costs another + /// forward NTT per table under `RecomputeLde`. + pub main_lde_expansions: usize, + /// Same for the auxiliary LDE. + pub aux_lde_expansions: usize, + /// Size-`n` coset evaluations (main), phase 4's cheap materialization + /// under `RecomputeLde` — about 37% of a full expansion's work and a + /// quarter of its bytes, counted apart so the full-expansion budget + /// stays an honest number. + pub main_coset_evals: usize, + /// Same for aux. + pub aux_coset_evals: usize, + /// How many times a table's composition parts were computed. Recomputing + /// these means re-running constraint evaluation, so the batched prover + /// retains them instead; this counter exists to make that visible if it ever + /// stops being true. + pub parts_computations: usize, + /// Wall clock per phase, indices 0..6 = phases 1..6 (main commit, aux + /// commit, composition parts, OOD, DEEP+FRI, openings). A latency + /// breakdown of the whole prove: the six entries plus the pre-phase + /// setup sum to the call's wall time. Returned in the stats — not logged — + /// for the same reason the residency numbers are: the A/B harness prints + /// the struct, so every box run carries its own phase profile. + pub phase_wall: [core::time::Duration; 6], + /// Wall clock spent inside LDE expansions (main and aux, all phases) — + /// the recompute traffic itself, separated from what the phases do with + /// the buffers. Under `RecomputeLde` this is the price of the residency + /// mode; under `Retain` it is the floor (one main + one aux per table). + pub lde_expansion_wall: core::time::Duration, +} + +/// Running account of live LDE bytes, so [`BatchedProveStats`] reports what the +/// prover actually held rather than what its comments claim. +#[derive(Debug, Default)] +pub(crate) struct ResidencyLedger { + live: usize, + peak: usize, +} + +impl ResidencyLedger { + pub(crate) fn alloc(&mut self, bytes: usize) { + self.live += bytes; + self.peak = self.peak.max(self.live); + } + + pub(crate) fn free(&mut self, bytes: usize) { + self.live = self.live.saturating_sub(bytes); + } + + pub(crate) fn peak(&self) -> usize { + self.peak + } +} + +/// Bytes a row-major LDE buffer of `len` field elements occupies. +pub(crate) fn lde_bytes(len: usize) -> usize { + len * core::mem::size_of::>() +} diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs new file mode 100644 index 000000000..b32cfa0c6 --- /dev/null +++ b/crypto/stark/src/batched/prover.rs @@ -0,0 +1,1208 @@ +//! The batched prover: four mixed-height MMCS roots and ONE FRI instance per +//! epoch. +//! +//! # The phase architecture, and why it is not `multi_prove` with a different +//! commit call +//! +//! `multi_prove` forks the transcript per table after the LogUp challenges and +//! then runs aux-build → aux-commit → rounds 2-4 FUSED per table, so a table +//! never waits on another. Batching cannot keep that: a batched root cannot be +//! absorbed until every contributing matrix exists, so each batched commitment +//! is a phase BARRIER. What survives of the fork is nothing — every challenge +//! here is drawn from the one shared transcript, in a fixed table order, and the +//! verifier replays that order exactly. +//! +//! ```text +//! shape histogram <- bound BEFORE the first root +//! per table: main LDE -> per-table prep tree + main MMCS builder [barrier] +//! per-table prep roots (from the AIR set), main_root +//! LogUp challenges +//! per table: aux trace + aux LDE -> aux MMCS builder [barrier] +//! aux_root +//! per table: bus contribution +//! per table: beta_t, composition parts -> parts builder [barrier] +//! parts_root +//! per table: z_t, OOD evaluations +//! per table: gamma_t +//! ONE batched FRI: alpha, (beta, layer root)*, terminal, grinding, iotas +//! openings, one table at a time +//! ``` +//! +//! # ★ The cost the plan does not price: the barriers force LDE rebuilds +//! +//! MMCS-PLAN §3.3 makes one memory argument — stream the tree build so a height +//! group's LDEs are not simultaneously resident — and [`StreamingMmcsBuilder`] +//! delivers it. But the tree build is not the only consumer of a table's LDE. +//! Constraint evaluation (round 2), the OOD evaluations (round 3), the DEEP +//! codeword (round 4) and the query openings all read it, and a barrier sits +//! between every pair of those: `beta` cannot be drawn before `aux_root` is +//! absorbed, `z` cannot be drawn before `parts_root` is, `alpha` cannot be drawn +//! before every table's OOD values are, and the query indices do not exist until +//! the FRI is over. +//! +//! So a table's main and aux LDEs are needed in FIVE phases that cannot be +//! merged, and the prover must either hold them (`O(N)`, which is what batching +//! was supposed to remove) or rebuild them (one forward NTT each, per phase). +//! [`ResidencyMode`] selects, exactly as it does in `multi_prove`, and +//! [`BatchedProveStats`] reports what it cost — `main_lde_expansions` and +//! `aux_lde_expansions` are the honest budget, not an estimate. +//! +//! The composition parts are the exception and are ALWAYS retained: recomputing +//! them means re-running constraint evaluation, which is the dominant cost of a +//! prove. `parts_computations` stays at one per table, and the counter is there +//! so that stops being silent if it ever changes. +//! +//! # What is deliberately absent +//! +//! No device paths. The GPU mixed-height MMCS exists (`crypto/math-cuda/`) and +//! is box-gated; wiring it in is a separate step, and a batched prover that +//! silently fell back between host and device arms would make the residency +//! numbers above unreproducible. Under `--features cuda` this path compiles and +//! runs on the host. + +use math::fft::bit_reversing::in_place_bit_reverse_permute_row_major; +use math::field::element::FieldElement; +use math::field::traits::{IsFFTField, IsField, IsSubFieldOf}; +use math::traits::AsBytes; + +use crate::batched::proof::{ + BatchedMultiProof, BatchedProveStats, BatchedQueryOpening, BatchedTableData, ResidencyLedger, + lde_bytes, +}; +use crate::batched::round4::commit_batched_fri; +use crate::batched::shape::{EpochShape, RoundShape, ShapeError}; +use crate::config::BatchedMerkleTreeBackend; +use crate::domain::Domain; +use crate::fri::batched::HeightCombiner; +use crate::fri::mmcs::{BorrowedMatrix, LeafSource, MixedMmcs, MixedOpening, StreamingMmcsBuilder}; +use crate::fri::terminal::coeffs_from_terminal_codeword; +use crate::lookup::{BusPublicInputs, LOGUP_NUM_CHALLENGES}; +use crate::proof::stark::PolynomialOpenings; +use crate::prover::{IsStarkProver, ProvingError, domain_and_twiddles}; +use crate::residency_mode::ResidencyMode; +#[cfg(feature = "disk-spill")] +use crate::storage_mode::StorageMode; +use crate::trace::{LDETraceTable, TraceTable}; +use crate::traits::AIR; +use crypto::merkle_tree::merkle::MerkleTree; + +use crypto::fiat_shamir::is_transcript::IsStarkTranscript; + +impl From for ProvingError { + fn from(e: ShapeError) -> Self { + ProvingError::WrongParameter(format!("batched epoch shape: {e}")) + } +} + +/// The AIR, its trace and its public inputs, as `multi_prove` takes them. +pub type BatchedAirTracePair<'a, Field, FieldExtension, PI> = ( + &'a dyn AIR, + &'a mut TraceTable, + &'a PI, +); + +/// A retention slot for one table's main LDE: `Some` while the buffer is being +/// held between phases, `None` while it is out on loan or was dropped. +type MainSlots<'a, Field> = &'a mut [Option<(Vec>, usize)>]; +/// The same for the auxiliary LDE. +type AuxSlots<'a, FieldExtension> = &'a mut [Option<(Vec>, usize)>]; +/// A preprocessed table's own row-pair tree — `None` for a table with no +/// preprocessed columns. +type PrepTreeSlot = Option>>; + +/// A table's LDE buffers, alive only for as long as the current phase needs +/// them, and accounted for while they are. +struct LdePair { + main: (Vec>, usize), + aux: (Vec>, usize), + bytes: usize, +} + +/// Prove one epoch with batched commitments. +/// +/// Preprocessed matrices are committed per table (the trees +/// `air.precomputed_commitment()` pins), so the per-table path's stale-constant +/// guard runs here unconditionally: a built prep tree that disagrees with the +/// AIR's own root fails the prove with the same error the per-table prover +/// raises. +#[allow(clippy::too_many_arguments)] +pub fn multi_prove_batched( + air_trace_pairs: Vec>, + transcript: &mut (impl IsStarkTranscript + Clone + Send), + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + residency: ResidencyMode, +) -> Result< + ( + BatchedMultiProof, + BatchedProveStats, + ), + ProvingError, +> +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + Copy + 'static, + FieldExtension: IsField + Send + Sync + Copy + 'static, + FieldElement: AsBytes + math::traits::ByteConversion + Sync + Send, + FieldElement: AsBytes + math::traits::ByteConversion + Sync + Send, + PI: Send + Sync + Clone, + P: IsStarkProver + ?Sized, + ::BaseType: math::spill_safe::SpillSafe, + ::BaseType: math::spill_safe::SpillSafe, +{ + multi_prove_batched_carved::( + air_trace_pairs, + transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + residency, + None, + ) +} + +/// As [`multi_prove_batched`], with one table's main matrix carved into a +/// standalone row-pair tree ([`crate::batched::shape::CarvedMain`]). +/// +/// The carved tree is built by `commit_rows_bit_reversed_subset` over the FULL +/// committed-main range of the same LDE expansion — the identical call the +/// per-table prover makes for a non-preprocessed table — so the carved root is +/// byte-identical to the root a per-table prove of the same trace commits. +/// The root is absorbed after the preprocessed roots and before `main_root`, +/// ahead of every challenge draw. +#[allow(clippy::too_many_arguments)] +pub fn multi_prove_batched_carved( + mut air_trace_pairs: Vec>, + transcript: &mut (impl IsStarkTranscript + Clone + Send), + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, + residency: ResidencyMode, + carved_main: Option, +) -> Result< + ( + BatchedMultiProof, + BatchedProveStats, + ), + ProvingError, +> +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + Copy + 'static, + FieldExtension: IsField + Send + Sync + Copy + 'static, + FieldElement: AsBytes + math::traits::ByteConversion + Sync + Send, + FieldElement: AsBytes + math::traits::ByteConversion + Sync + Send, + PI: Send + Sync + Clone, + P: IsStarkProver + ?Sized, + // The same two bounds `multi_prove` carries: under `disk-spill` the aux + // trace is spilled through an mmap backing, which only a field whose + // `BaseType` is plain data can be laid out in. + ::BaseType: math::spill_safe::SpillSafe, + ::BaseType: math::spill_safe::SpillSafe, +{ + let num_tables = air_trace_pairs.len(); + let mut stats = BatchedProveStats::default(); + // Two accounts, because they behave differently on purpose: the trace LDEs + // must stay flat in the table count, the retained parts must not. + let mut ledger = ResidencyLedger::default(); + let mut parts_ledger = ResidencyLedger::default(); + + // ===================================================================== + // Phase 0 — domains, shape, and the shape binding + // ===================================================================== + let mut domains = Vec::with_capacity(num_tables); + let mut twiddles = Vec::with_capacity(num_tables); + for (air, trace, _) in &*air_trace_pairs { + let (domain, tw) = domain_and_twiddles(*air, trace.num_rows()); + domains.push(domain); + twiddles.push(tw); + } + + let airs: Vec<&dyn AIR> = + air_trace_pairs.iter().map(|(air, _, _)| *air).collect(); + let trace_lengths: Vec = domains + .iter() + .map(|d| d.interpolation_domain_size) + .collect(); + let (shape, params) = EpochShape::derive_carved(&airs, &trace_lengths, carved_main)?; + let h_max = shape.h_max(); + let coset_offset = FieldElement::::from(params.coset_offset); + + // ★ Addendum A's recommendation S, adopted. `commit_batched_fri` binds the + // shape again in round 4, which is where the batched FRI's own challenges + // need it; binding it HERE, before the first root, is what turns "no + // rounds-1-3 challenge is shape-exploitable" from a collision-resistance + // argument into a transcript-ordering one. Two field-sized absorptions per + // table, and every later challenge inherits the binding. + crate::fri::batched::absorb_shape_histogram::( + transcript, + &shape.heights, + &shape.total_widths(), + ); + + // ===================================================================== + // Phase 1 — the preprocessed and main rounds, one main LDE pass per table + // ===================================================================== + let t_phase = std::time::Instant::now(); + // Both builders are fed from the SAME expansion: a preprocessed table's + // precomputed columns and its multiplicity columns are two column ranges of + // one row-major main LDE, exactly as `commit_main_trace` splits them. + // ★ Per-table preprocessed trees — #768's arrangement, kept for the same + // reason (see `BatchedQueryOpening::prep`): each preprocessed table keeps + // its OWN row-pair tree, the one `air.precomputed_commitment()` pins, and + // both sides absorb that root FROM THE AIR SET, never from the proof — + // the per-table path's critical soundness check, verbatim. The trees are + // process-cached by root, so continuation epochs stop re-committing the + // execution-independent tables (DECODE, BITWISE, ...), exactly as the + // per-table prover does. + let mut prep_trees: Vec>> = + (0..num_tables).map(|_| None).collect(); + let mut main_builder = StreamingMmcsBuilder::::new(&shape.main.dims); + let mut retained_main: Vec>, usize)>> = + (0..num_tables).map(|_| None).collect(); + // The carved table's standalone main tree and its root. Built inside the + // phase-1 loop from the same expansion every other table commits from; the + // root is absorbed AFTER the loop (after every preprocessed root) and + // before `main_root`. + let mut carved_tree: Option>> = None; + let mut carved_root: Option = None; + + for table in 0..num_tables { + let (air, trace, _) = &air_trace_pairs[table]; + let (main_data, total_cols) = P::expand_main_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + stats.main_lde_expansions += 1; + let bytes = lde_bytes::(main_data.len()); + ledger.alloc(bytes); + + let height = shape.heights[table]; + let is_carved = shape.carved_main.map(|c| c.table) == Some(table); + let num_precomputed = if is_carved { + // `derive_carved` rejects a preprocessed carved table, so the + // carved matrix is the full main range. + 0 + } else { + total_cols - matrix_width(&shape.main, table) + }; + + if num_precomputed > 0 { + // The root every verifier will absorb is the AIR's own; building a + // tree that disagrees with it is a stale constant or a wrong LDE, + // and the per-table path's error is the honest name for both. + let expected = air.precomputed_commitment(); + let tree = + match crate::prover::precomputed_tree_cache_get::(&expected) { + Some(tree) => tree, + None => { + let (tree, root) = P::commit_rows_bit_reversed_subset::( + &main_data, + total_cols, + 0, + num_precomputed, + ) + .ok_or(ProvingError::PrecomputedCommitmentMismatch)?; + if root != expected { + return Err(ProvingError::PrecomputedCommitmentMismatch); + } + let tree = std::sync::Arc::new(tree); + crate::prover::precomputed_tree_cache_put( + expected, + std::sync::Arc::clone(&tree), + ); + tree + } + }; + transcript.append_bytes(&expected); + prep_trees[table] = Some(tree); + } + if is_carved { + // The carve: the identical committer call the per-table prover + // makes for a non-preprocessed table (`commit_rows_bit_reversed` = + // the subset call over the full range), on the identical + // expansion — the root is byte-identical to the per-table tree's. + let (tree, root) = + P::commit_rows_bit_reversed_subset::(&main_data, total_cols, 0, total_cols) + .ok_or_else(|| { + ProvingError::WrongParameter( + "the carved table's main matrix has no committable rows".to_string(), + ) + })?; + carved_tree = Some(tree); + carved_root = Some(root); + } else { + let src = vec![BorrowedMatrix::RowMajorNatural { + data: &main_data, + stride: total_cols, + col_start: num_precomputed, + width: total_cols - num_precomputed, + log_height: height, + }]; + main_builder.absorb(&src, 0); + } + + // The root is what Fiat-Shamir needs; the buffer is not. Under + // `RecomputeLde` it dies here and every later phase rebuilds it. + match residency { + ResidencyMode::Retain => retained_main[table] = Some((main_data, total_cols)), + ResidencyMode::RecomputeLde => { + drop(main_data); + ledger.free(bytes); + } + } + } + + // The carved root's transcript slot: after every preprocessed root, before + // `main_root` — so every challenge (LogUp, beta, z, gamma, alpha, iotas) is + // drawn after it. Proof-carried on the verifier's side, absorbed here from + // the tree just built. + if let Some(root) = carved_root.as_ref() { + transcript.append_bytes(root); + } + + let main_mmcs = main_builder.finish(); + let main_root = main_mmcs.root(); + transcript.append_bytes(&main_root); + + // ===================================================================== + // Phase 2 — LogUp challenges, then the auxiliary round + // ===================================================================== + stats.phase_wall[0] = t_phase.elapsed(); + let t_phase = std::time::Instant::now(); + let needs_lookup = airs.iter().any(|air| air.has_aux_trace()); + let lookup_challenges: Vec> = if needs_lookup { + (0..LOGUP_NUM_CHALLENGES) + .map(|_| transcript.sample_field_element()) + .collect() + } else { + Vec::new() + }; + + // The aux round expands its LDE from the host trace columns + // (`expand_aux_lde_row_major` below); the device-resident aux build + // returns the columns device-side only and leaves the host trace + // unwritten, so it is disabled for every table here — the same switch + // the per-table prover throws under disk-spill and `RecomputeLde`. + #[cfg(feature = "cuda")] + for (_, trace, _) in air_trace_pairs.iter_mut() { + trace.set_resident_aux_ok(false); + } + + let mut bus_public_inputs: Vec>> = + (0..num_tables).map(|_| None).collect(); + let mut aux_builder = (!shape.aux.is_empty()) + .then(|| StreamingMmcsBuilder::::new(&shape.aux.dims)); + let mut retained_aux: Vec>, usize)>> = + (0..num_tables).map(|_| None).collect(); + + for table in 0..num_tables { + let (air, trace, _) = &mut air_trace_pairs[table]; + if !air.has_aux_trace() { + continue; + } + bus_public_inputs[table] = air.build_auxiliary_trace(trace, &lookup_challenges); + + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + trace + .spill_aux_to_disk() + .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; + } + + let Some(builder) = aux_builder.as_mut() else { + continue; + }; + let (aux_data, aux_cols) = P::expand_aux_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + stats.aux_lde_expansions += 1; + let bytes = lde_bytes::(aux_data.len()); + ledger.alloc(bytes); + let src = vec![BorrowedMatrix::RowMajorNatural { + data: &aux_data, + stride: aux_cols, + col_start: 0, + width: aux_cols, + log_height: shape.heights[table], + }]; + builder.absorb(&src, 0); + match residency { + ResidencyMode::Retain => retained_aux[table] = Some((aux_data, aux_cols)), + ResidencyMode::RecomputeLde => { + drop(aux_data); + ledger.free(bytes); + } + } + } + + let aux_mmcs = aux_builder.map(StreamingMmcsBuilder::finish); + let aux_root = aux_mmcs.as_ref().map(MixedMmcs::root); + if let Some(root) = aux_root { + transcript.append_bytes(&root); + } + + // ===================================================================== + // Phase 3 — bus contributions, beta per table, the composition-parts round + // ===================================================================== + stats.phase_wall[1] = t_phase.elapsed(); + let t_phase = std::time::Instant::now(); + for bpi in bus_public_inputs.iter().flatten() { + transcript.append_field_element(&bpi.table_contribution); + } + + let mut parts_builder = StreamingMmcsBuilder::::new(&shape.parts.dims); + let mut retained_parts: Vec>>> = + (0..num_tables).map(|_| Vec::new()).collect(); + + for table in 0..num_tables { + let beta: FieldElement = transcript.sample_field_element(); + let (air, _, pub_inputs) = &air_trace_pairs[table]; + let domain = &domains[table]; + + let num_transition_constraints = air.context().num_transition_constraints; + let num_boundary_constraints = air + .boundary_constraints( + pub_inputs, + &lookup_challenges, + bus_public_inputs[table].as_ref(), + domain.interpolation_domain_size, + ) + .constraints + .len(); + let mut coefficients: Vec> = + core::iter::successors(Some(FieldElement::one()), |x| Some(x * &beta)) + .take(num_boundary_constraints + num_transition_constraints) + .collect(); + let transition_coefficients: Vec<_> = + coefficients.drain(..num_transition_constraints).collect(); + let boundary_coefficients = coefficients; + + let ldes = materialize_ldes::( + table, + &air_trace_pairs, + &domains, + &twiddles, + &shape, + &mut retained_main, + &mut retained_aux, + &mut stats, + &mut ledger, + residency, + #[cfg(feature = "disk-spill")] + storage_mode, + ); + let (mut lde_trace, carried_bytes) = + lde_trace_take(ldes, air.step_size(), domain.blowup_factor); + + let computed = P::compute_composition_parts( + *air, + pub_inputs, + domain, + &twiddles[table], + &mut lde_trace, + &lookup_challenges, + bus_public_inputs[table].as_ref(), + &transition_coefficients, + &boundary_coefficients, + )?; + stats.parts_computations += 1; + let parts = computed.parts; + + let parts_bytes: usize = parts + .iter() + .map(|p| lde_bytes::(p.len())) + .sum(); + parts_ledger.alloc(parts_bytes); + let src = vec![BorrowedMatrix::ColMajorNatural { + cols: &parts, + log_height: shape.heights[table], + }]; + parts_builder.absorb(&src, 0); + + // Parts are RETAINED: rebuilding them is a second constraint evaluation. + retained_parts[table] = parts; + release_ldes( + ldes_from_trace(lde_trace, carried_bytes), + &mut retained_main, + &mut retained_aux, + table, + &mut ledger, + residency, + ); + } + + let parts_mmcs = parts_builder.finish(); + let parts_root = parts_mmcs.root(); + transcript.append_bytes(&parts_root); + + // ===================================================================== + // Phase 4 — z per table, OOD evaluations + // ===================================================================== + stats.phase_wall[2] = t_phase.elapsed(); + let t_phase = std::time::Instant::now(); + let mut zs = Vec::with_capacity(num_tables); + let mut round3s = Vec::with_capacity(num_tables); + let mut ood_blocks = Vec::with_capacity(num_tables); + + for table in 0..num_tables { + let (air, _, _) = &air_trace_pairs[table]; + let domain = &domains[table]; + // `sample_z_ood_with_domain_params` rather than `sample_z_ood`: the + // verifier has the trace length and the blowup but not the domain + // vectors, so naming the routine both sides can reach is what makes the + // two agree by construction instead of by two call sites coinciding. + let z = transcript.sample_z_ood_with_domain_params( + domain.interpolation_domain_size, + domain.interpolation_domain_size * domain.blowup_factor, + &coset_offset, + ); + + // Phase 4 reads the trace ONLY at stride `blowup` — the size-`n` + // coset evaluation. Under `Retain` the full LDE is already on hand + // and the strided read is free; under `RecomputeLde` a full 4n + // expansion here would be paid just to subsample it, so the + // recompute arm materializes the n-sized evaluation directly + // (bit-identical values, ~37% of the work, a quarter of the bytes) + // and hands `round_3` a blowup-1 table, whose OWN stride the trace + // reads follow. + let round3 = if retained_main[table].is_some() { + let ldes = materialize_ldes::( + table, + &air_trace_pairs, + &domains, + &twiddles, + &shape, + &mut retained_main, + &mut retained_aux, + &mut stats, + &mut ledger, + residency, + #[cfg(feature = "disk-spill")] + storage_mode, + ); + let (mut lde_trace, carried_bytes) = + lde_trace_take(ldes, air.step_size(), domain.blowup_factor); + let round3 = P::round_3_evaluate_polynomials_in_out_of_domain_element( + *air, + domain, + &mut lde_trace, + &mut retained_parts[table], + &z, + ); + release_ldes( + ldes_from_trace(lde_trace, carried_bytes), + &mut retained_main, + &mut retained_aux, + table, + &mut ledger, + residency, + ); + round3 + } else { + let (_, trace, _) = &air_trace_pairs[table]; + let t_expand = std::time::Instant::now(); + let main = P::expand_main_coset_eval_row_major(trace, domain, &twiddles[table]); + let aux = if matrix_index(&shape.aux, table).is_some() { + let aux = P::expand_aux_coset_eval_row_major(trace, domain, &twiddles[table]); + stats.aux_coset_evals += 1; + aux + } else { + (Vec::new(), 0) + }; + stats.lde_expansion_wall += t_expand.elapsed(); + stats.main_coset_evals += 1; + let bytes = lde_bytes::(main.0.len()) + lde_bytes::(aux.0.len()); + ledger.alloc(bytes); + let mut lde_trace = + LDETraceTable::from_row_major(main.0, main.1, aux.0, aux.1, air.step_size(), 1); + let round3 = P::round_3_evaluate_polynomials_in_out_of_domain_element( + *air, + domain, + &mut lde_trace, + &mut retained_parts[table], + &z, + ); + drop(lde_trace); + ledger.free(bytes); + round3 + }; + + let (block0, block1) = P::ood_layout(*air).split_full(&round3.trace_ood_evaluations); + for block in [&block0, &block1] { + for col in block.columns().iter() { + for elem in col.iter() { + transcript.append_field_element(elem); + } + } + } + for element in round3.composition_poly_parts_ood_evaluation.iter() { + transcript.append_field_element(element); + } + + zs.push(z); + ood_blocks.push((block0, block1)); + round3s.push(round3); + } + + // ===================================================================== + // Phase 5 — gamma per table, then ONE batched FRI + // ===================================================================== + stats.phase_wall[3] = t_phase.elapsed(); + let t_phase = std::time::Instant::now(); + let gammas: Vec> = (0..num_tables) + .map(|_| transcript.sample_field_element()) + .collect(); + + let commit = { + let air_trace_pairs = &air_trace_pairs; + let domains = &domains; + let twiddles = &twiddles; + let shape = &shape; + // `&mut`: the DEEP host loop repopulates a table's part evals from the + // resident handle when the device-only gate left them empty. + let retained_parts = &mut retained_parts; + let round3s = &round3s; + let zs = &zs; + let gammas = &gammas; + let retained_main = &mut retained_main; + let retained_aux = &mut retained_aux; + let stats = &mut stats; + let ledger = &mut ledger; + let coset_offset_ref = &coset_offset; + + commit_batched_fri::( + transcript, + &shape.heights, + &shape.total_widths(), + move |alpha, plan| { + // The standalone class's terminal polynomials, handed back so + // `commit_batched_fri` binds them into the transcript and the + // wire carries the very coefficients that were bound. + let mut standalone_coeffs: Vec>>> = + (0..num_tables).map(|_| None).collect(); + let mut combiner = HeightCombiner::new(*alpha); + // Ascending table order, which is also `plan.batched`'s order — + // absorption order is what defines the alpha powers, so the two + // must not be allowed to drift apart. + for table in 0..num_tables { + let (air, _, _) = &air_trace_pairs[table]; + let domain = &domains[table]; + let ldes = materialize_ldes::( + table, + air_trace_pairs, + domains, + twiddles, + shape, + retained_main, + retained_aux, + stats, + ledger, + residency, + #[cfg(feature = "disk-spill")] + storage_mode, + ); + let (mut lde_trace, carried_bytes) = + lde_trace_take(ldes, air.step_size(), domain.blowup_factor); + let mut deep = deep_codeword::( + *air, + domain, + &mut lde_trace, + &mut retained_parts[table], + &round3s[table], + &zs[table], + &gammas[table], + ); + release_ldes( + ldes_from_trace(lde_trace, carried_bytes), + retained_main, + retained_aux, + table, + ledger, + residency, + ); + // Row-major variant at one column = the parallel path; the + // serial swap loop was pure wall time, 27 times per epoch. + in_place_bit_reverse_permute_row_major(&mut deep, 1); + + if plan.batched.contains(&table) { + combiner.absorb(&deep, shape.heights[table]); + } else { + // A standalone table's terminal codeword IS this + // codeword; the proof carries the polynomial it + // evaluates, at its own degree bound. + let log_degree = (shape.heights[table] as u32) - params.blowup_log; + standalone_coeffs[table] = Some(coeffs_from_terminal_codeword::< + Field, + FieldExtension, + >( + &deep, coset_offset_ref, log_degree + )); + } + } + (combiner.finish(), standalone_coeffs) + }, + &coset_offset, + params.blowup_log, + params.final_poly_log_degree, + params.grinding_factor, + params.num_queries, + ) + }; + + // ===================================================================== + // Phase 6 — openings, one table at a time + // ===================================================================== + stats.phase_wall[4] = t_phase.elapsed(); + let t_phase = std::time::Instant::now(); + let iotas = commit.iotas.clone(); + let fri_decommitments = crate::fri::query_phase::(&commit.layers, &iotas); + + // Per-query, per-prep-table standard openings (prep-table order = + // `shape.prep.tables`, which is AIR order). + let mut prep_openings: Vec>> = + (0..iotas.len()).map(|_| Vec::new()).collect(); + let mut main_openings = empty_openings::(&iotas, shape.main.tables.len()); + let mut aux_openings = empty_openings::(&iotas, shape.aux.tables.len()); + let mut parts_openings = empty_openings::(&iotas, shape.parts.tables.len()); + let mut carved_openings: Vec>> = + (0..iotas.len()).map(|_| None).collect(); + + // ★ Each round is read in ITS OWN index space, and the reduction happens + // exactly once, here. Doing it inside the read would be wrong twice over: a + // round shorter than the FRI would be asked for a leaf it does not have + // (the prep round's `h_max` is below the FRI's whenever the tallest + // preprocessed table is not the tallest table), and a round that reduced + // again on the way out would land somewhere else entirely. + let main_iotas = reduced_iotas(&iotas, h_max, main_mmcs.h_max()); + let aux_iotas = aux_mmcs + .as_ref() + .map(|mmcs| reduced_iotas(&iotas, h_max, mmcs.h_max())); + let parts_iotas = reduced_iotas(&iotas, h_max, parts_mmcs.h_max()); + + for table in 0..num_tables { + let (air, _, _) = &air_trace_pairs[table]; + let ldes = materialize_ldes::( + table, + &air_trace_pairs, + &domains, + &twiddles, + &shape, + &mut retained_main, + &mut retained_aux, + &mut stats, + &mut ledger, + residency, + #[cfg(feature = "disk-spill")] + storage_mode, + ); + let _ = air; + let height = shape.heights[table]; + let (main_data, total_cols) = &ldes.main; + let is_carved = shape.carved_main.map(|c| c.table) == Some(table); + let num_precomputed = if is_carved { + 0 + } else { + total_cols - matrix_width(&shape.main, table) + }; + + if is_carved { + let tree = carved_tree + .as_ref() + .expect("the carved tree was built in phase 1"); + // The carved tree lives in the TABLE's own index space, exactly + // like a preprocessed tree: reduce the shared FRI index once. + let table_iotas = reduced_iotas(&iotas, h_max, height); + for (q, &idx) in table_iotas.iter().enumerate() { + carved_openings[q] = Some(P::open_polys_with(&domains[table], tree, idx, |row| { + main_data[row * total_cols..(row + 1) * total_cols].to_vec() + })); + } + } + + if let Some(tree) = prep_trees[table].as_ref() { + // The per-table tree lives in the TABLE's own index space; reduce + // the shared FRI index by the height difference once, here. + let table_iotas = reduced_iotas(&iotas, h_max, height); + for (q, &idx) in table_iotas.iter().enumerate() { + prep_openings[q].push(P::open_polys_with(&domains[table], tree, idx, |row| { + main_data[row * total_cols..row * total_cols + num_precomputed].to_vec() + })); + } + } + if let Some(m) = matrix_index(&shape.main, table) { + let src = vec![BorrowedMatrix::RowMajorNatural { + data: main_data, + stride: *total_cols, + col_start: num_precomputed, + width: total_cols - num_precomputed, + log_height: height, + }]; + fill_openings(&main_mmcs, m, &src, &main_iotas, &mut main_openings); + } + if let (Some(mmcs), Some(m)) = (aux_mmcs.as_ref(), matrix_index(&shape.aux, table)) { + let (aux_data, aux_cols) = &ldes.aux; + let src = vec![BorrowedMatrix::RowMajorNatural { + data: aux_data, + stride: *aux_cols, + col_start: 0, + width: *aux_cols, + log_height: height, + }]; + let indices = aux_iotas.as_ref().expect("the aux MMCS exists here"); + fill_openings(mmcs, m, &src, indices, &mut aux_openings); + } + if let Some(m) = matrix_index(&shape.parts, table) { + let src = vec![BorrowedMatrix::ColMajorNatural { + cols: &retained_parts[table], + log_height: height, + }]; + fill_openings(&parts_mmcs, m, &src, &parts_iotas, &mut parts_openings); + } + + release_ldes( + ldes, + &mut retained_main, + &mut retained_aux, + table, + &mut ledger, + residency, + ); + } + + let queries = (0..iotas.len()) + .map(|q| BatchedQueryOpening { + prep: std::mem::take(&mut prep_openings[q]), + main: assemble(&main_mmcs, main_iotas[q], &mut main_openings, q) + .expect("the main round was opened at these very indices"), + aux: aux_mmcs.as_ref().map(|mmcs| { + let indices = aux_iotas.as_ref().expect("the aux MMCS exists here"); + assemble(mmcs, indices[q], &mut aux_openings, q) + .expect("the aux round was opened at these very indices") + }), + parts: assemble(&parts_mmcs, parts_iotas[q], &mut parts_openings, q) + .expect("the parts round was opened at these very indices"), + carved_main: carved_openings[q].take(), + fri: fri_decommitments[q].clone(), + }) + .collect(); + + let tables = (0..num_tables) + .map(|table| { + let (block0, block1) = ood_blocks[table].clone(); + BatchedTableData { + trace_length: trace_lengths[table], + trace_ood_evaluations: block0, + trace_ood_next_evaluations: block1, + composition_poly_parts_ood_evaluation: round3s[table] + .composition_poly_parts_ood_evaluation + .clone(), + bus_public_inputs: bus_public_inputs[table].clone(), + public_inputs: air_trace_pairs[table].2.clone(), + standalone_final_poly_coeffs: commit.standalone_coeffs[table].clone(), + } + }) + .collect(); + + stats.phase_wall[5] = t_phase.elapsed(); + stats.peak_trace_lde_bytes = ledger.peak(); + stats.retained_parts_bytes = parts_ledger.peak(); + stats.peak_lde_bytes = stats.peak_trace_lde_bytes + stats.retained_parts_bytes; + + Ok(( + BatchedMultiProof { + tables, + main_root, + carved_main_root: carved_root, + aux_root, + parts_root, + fri_layer_roots: commit.layer_roots, + fri_final_poly_coeffs: commit.final_poly_coeffs, + nonce: commit.nonce, + queries, + }, + stats, + )) +} + +/// Matrix index of `table` inside `round`, or `None` when it does not +/// contribute one. +fn matrix_index(round: &RoundShape, table: usize) -> Option { + round.tables.iter().position(|&t| t == table) +} + +/// The width `table` contributes to `round`. Zero when it contributes nothing. +fn matrix_width(round: &RoundShape, table: usize) -> usize { + matrix_index(round, table).map_or(0, |m| round.dims[m].1) +} + +#[allow(clippy::type_complexity)] +fn empty_openings( + iotas: &[usize], + num_matrices: usize, +) -> Vec>>> { + iotas + .iter() + .map(|_| (0..num_matrices).map(|_| None).collect()) + .collect() +} + +/// Read one matrix's row pair at every query, so a table's openings are +/// harvested while its LDE is alive and never after. +fn fill_openings( + mmcs: &MixedMmcs, + matrix: usize, + source: &S, + iotas: &[usize], + out: &mut [Vec>>], +) where + E: IsField + 'static, + S: LeafSource, + FieldElement: AsBytes + Sync + Send, +{ + // `iotas` are in THIS round's index space already (see `reduced_iotas`). + // Passing the FRI's raw indices here does not corrupt anything quietly: a + // shorter round rejects them as out of range and produces no opening at + // all, which is what `the_preprocessed_round_is_committed_and_authenticates` + // caught the first time this was written the other way round. + for (q, &iota) in iotas.iter().enumerate() { + let Some(leaf) = mmcs.row_pair_leaf(iota, matrix) else { + continue; + }; + let mut evaluations = Vec::new(); + source.append_row(0, 2 * leaf, &mut evaluations); + let mut evaluations_sym = Vec::new(); + source.append_row(0, 2 * leaf + 1, &mut evaluations_sym); + out[q][matrix] = Some(PolynomialOpenings { + proof: crypto::merkle_tree::proof::Proof { + merkle_path: Vec::new(), + }, + evaluations, + evaluations_sym, + }); + } +} + +/// Reduce every FRI query index into one round's index space. +/// +/// `h_max_round <= h_max_fri` always holds for a round of this epoch — a round +/// commits a subset of the epoch's tables, so its tallest matrix cannot exceed +/// the epoch's — which is why this is infallible here and +/// `reduce_iota_to_round` returns an `Option` on the verifier's path, where the +/// heights are proof-supplied. +fn reduced_iotas(iotas: &[usize], h_max_fri: usize, h_max_round: usize) -> Vec { + iotas + .iter() + .map(|&iota| { + crate::batched::round4::reduce_iota_to_round(iota, h_max_fri, h_max_round) + .expect("a round of this epoch is never taller than the epoch") + }) + .collect() +} + +/// Turn one query's per-matrix rows into a [`MixedOpening`] by attaching the +/// round's shared authentication path. +/// +/// `iota` is already in this round's index space — see [`reduced_iotas`]. +fn assemble( + mmcs: &MixedMmcs, + iota: usize, + openings: &mut [Vec>>], + query: usize, +) -> Option> +where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + let proof = mmcs.auth_path(iota)?; + let per_matrix = openings[query] + .iter_mut() + .map(|slot| slot.take()) + .collect::>>()?; + Some(MixedOpening { proof, per_matrix }) +} + +/// Build (or take back) a table's main and aux LDEs for the phase about to read +/// them. +#[allow(clippy::too_many_arguments)] +fn materialize_ldes( + table: usize, + air_trace_pairs: &[BatchedAirTracePair<'_, Field, FieldExtension, PI>], + domains: &[std::sync::Arc>], + twiddles: &[std::sync::Arc>], + shape: &EpochShape, + retained_main: MainSlots<'_, Field>, + retained_aux: AuxSlots<'_, FieldExtension>, + stats: &mut BatchedProveStats, + ledger: &mut ResidencyLedger, + residency: ResidencyMode, + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, +) -> LdePair +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + Copy + 'static, + FieldExtension: IsField + Send + Sync + Copy + 'static, + FieldElement: AsBytes + math::traits::ByteConversion + Sync + Send, + FieldElement: AsBytes + math::traits::ByteConversion + Sync + Send, + P: IsStarkProver + ?Sized, +{ + let (_, trace, _) = &air_trace_pairs[table]; + let mut bytes = 0usize; + + let main = match retained_main[table].take() { + Some(lde) => lde, + None => { + let t_expand = std::time::Instant::now(); + let lde = P::expand_main_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + stats.lde_expansion_wall += t_expand.elapsed(); + stats.main_lde_expansions += 1; + let b = lde_bytes::(lde.0.len()); + ledger.alloc(b); + bytes += b; + lde + } + }; + + let aux = if matrix_index(&shape.aux, table).is_some() { + match retained_aux[table].take() { + Some(lde) => lde, + None => { + let t_expand = std::time::Instant::now(); + let lde = P::expand_aux_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + stats.lde_expansion_wall += t_expand.elapsed(); + stats.aux_lde_expansions += 1; + let b = lde_bytes::(lde.0.len()); + ledger.alloc(b); + bytes += b; + lde + } + } + } else { + (Vec::new(), 0) + }; + + let _ = residency; + LdePair { main, aux, bytes } +} + +/// Give a table's LDEs back to the retention slots, or drop them. +fn release_ldes( + ldes: LdePair, + retained_main: MainSlots<'_, Field>, + retained_aux: AuxSlots<'_, FieldExtension>, + table: usize, + ledger: &mut ResidencyLedger, + residency: ResidencyMode, +) { + match residency { + ResidencyMode::Retain => { + retained_main[table] = Some(ldes.main); + if ldes.aux.1 > 0 { + retained_aux[table] = Some(ldes.aux); + } + } + ResidencyMode::RecomputeLde => { + drop(ldes.main); + drop(ldes.aux); + ledger.free(ldes.bytes); + } + } +} + +/// Move a table's LDE buffers into the trace view the phase reads — no copy. +/// The phases never mutate the buffers on the host path (the one bulk writer, +/// the cuda `set_host_data`, only FILLS deliberately-empty buffers), so the +/// same allocation flows phase → view → [`ldes_from_trace`] → retention, and +/// the transient double-residency the old clone created — one table's whole +/// main+aux LDE, invisible to the ledger — is gone. +fn lde_trace_take( + ldes: LdePair, + step_size: usize, + blowup_factor: usize, +) -> (LDETraceTable, usize) +where + Field: IsFFTField + IsSubFieldOf, + FieldExtension: IsField, +{ + let LdePair { main, aux, bytes } = ldes; + ( + LDETraceTable::from_row_major(main.0, main.1, aux.0, aux.1, step_size, blowup_factor), + bytes, + ) +} + +/// Take the buffers back out of the trace view for release or retention — +/// the inverse of [`lde_trace_take`], carrying the byte account through. +fn ldes_from_trace( + lde_trace: LDETraceTable, + bytes: usize, +) -> LdePair +where + Field: IsFFTField + IsSubFieldOf, + FieldExtension: IsField, +{ + LdePair { + main: (lde_trace.main_data, lde_trace.num_main_cols), + aux: (lde_trace.aux_data, lde_trace.num_aux_cols), + bytes, + } +} + +/// One table's DEEP composition codeword, in NATURAL order. +#[allow(clippy::too_many_arguments)] +fn deep_codeword( + air: &dyn AIR, + domain: &Domain, + // `&mut` to match `compute_deep_composition_poly_evaluations`, whose host + // loop downloads the resident trace and part evals in place when the + // device-only gate left them empty. + lde_trace: &mut LDETraceTable, + composition_parts: &mut [Vec>], + round3: &crate::prover::Round3, + z: &FieldElement, + gamma: &FieldElement, +) -> Vec> +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + Copy + 'static, + FieldExtension: IsField + Send + Sync + Copy + 'static, + FieldElement: AsBytes + math::traits::ByteConversion + Sync + Send, + FieldElement: AsBytes + math::traits::ByteConversion + Sync + Send, + P: IsStarkProver + ?Sized, +{ + let n_terms_composition_poly = composition_parts.len(); + let layout = P::ood_layout(air); + let num_terms_trace = layout.num_surviving(); + + let mut deep_composition_coefficients: Vec> = + core::iter::successors(Some(FieldElement::one()), |x| Some(x * gamma)) + .take(n_terms_composition_poly + num_terms_trace) + .collect(); + let trace_term_powers: Vec<_> = deep_composition_coefficients + .drain(..num_terms_trace) + .collect(); + let trace_term_coeffs = layout.build_trace_term_coeffs(&trace_term_powers); + let gammas = deep_composition_coefficients; + + P::compute_deep_composition_poly_evaluations( + lde_trace, + composition_parts, + round3, + z, + domain, + &domain.trace_primitive_root, + &gammas, + &trace_term_coeffs, + ) +} diff --git a/crypto/stark/src/batched/round4.rs b/crypto/stark/src/batched/round4.rs new file mode 100644 index 000000000..1f9bd03f5 --- /dev/null +++ b/crypto/stark/src/batched/round4.rs @@ -0,0 +1,892 @@ +//! Round 4 of the batched path: ONE FRI instance over the epoch's height-combined +//! DEEP codewords. +//! +//! # The transcript sequence, and why it has one owner +//! +//! ```text +//! shape histogram → α → standalone terminals → (β, layer root)* → β_final → terminal coeffs → grinding → iotas +//! ``` +//! +//! [`commit_batched_fri`] walks it on the prover's side; +//! [`crate::fri::batched::derive_batched_fri_challenges`] walks it on the +//! verifier's. The two are pinned to each other by +//! `prover_commit_matches_verifier_derivation`, not by review of two call sites. +//! α is sampled AFTER the shape is absorbed and BEFORE any codeword is combined, +//! which is why this function takes a `combine` closure rather than the codewords: +//! the prover cannot mix with α until the transcript has produced it, and the +//! closure is where a caller streams table by table (see +//! [`crate::fri::batched::HeightCombiner`]). +//! +//! # ★ TWO instance classes, and the index rule between them +//! +//! Not every table belongs in the batch. A table whose own FRI commits ZERO +//! layers gains nothing from being batched — there is no layer for the batch to +//! share — while it pays the full lift to the tallest domain, which is where the +//! proximity-gaps term's `|D0|^2` lives. At the measured epoch that is 13 of 28 +//! legs carrying 92% of the batch's width. [`FriInstancePlan`] partitions them, +//! and the excluded tables keep a terminal-only instance +//! ([`verify_standalone_fri_query`]) that costs one polynomial and no layers. +//! +//! The MMCS is untouched by this split — it still commits every table, so the +//! one-shared-authentication-path win survives whole. What differs is the index +//! SPACE: the batched class reads `iota` directly, a standalone table at height +//! `h` reads `iota >> (h_max - h)`. Both classes need a tamper control, since a +//! control that only touched the batched one would pass under any convention for +//! the other. +//! +//! # Query indices and the injection convention +//! +//! One `iota` per query, drawn from `[0, 2^(h_max-1))` — a row-PAIR index in the +//! TALLEST codeword's domain. Every shorter object is located by shifting it +//! down, which is what makes "one index, shared across all tables" true rather +//! than aspirational: +//! +//! - a matrix of height `h` in a round whose own tallest matrix is `h_max_round` +//! is opened at MMCS leaf `iota >> (h_max_fri - h)` — but note that +//! [`crate::fri::mmcs::MixedMmcs::verify_batch`] wants an index in ITS OWN +//! space, so a round whose `h_max_round` is below the FRI's must first reduce +//! (see that module's index-convention section, and [`reduce_iota_to_round`]). +//! - the codeword bucket at height `h` is read at position +//! [`injection_position`], which is exactly one of the two rows of the pair the +//! MMCS opened. That coincidence is not luck: both are the same row-pair +//! layout, which is why a single opening serves both the authentication and the +//! FRI join. +//! +//! # What "injection" costs the verifier +//! +//! The prover's [`crate::fri::batched::batched_commit_phase`] folds, then adds +//! `β² · bucket_h` to the running codeword before committing the layer. So the +//! verifier's per-query recursion adds the same term to the value it computed by +//! folding — and only to that value. The symmetric value at each layer comes from +//! the proof and is Merkle-authenticated against the layer root, so it already +//! carries its own injection; re-adding one would double it. + +use crypto::fiat_shamir::is_transcript::{IsStarkTranscript, IsTranscript}; +use crypto::merkle_tree::proof::verify_merkle_path; +use math::field::element::FieldElement; +use math::field::traits::{IsFFTField, IsField, IsSubFieldOf}; +use math::traits::AsBytes; + +use crate::config::{BatchedMerkleTreeBackend, Commitment, FriLayerMerkleTreeBackend}; +use crate::fri::batched::{ + BatchedFriLayout, FriInstancePlan, absorb_shape_histogram, batched_commit_phase, + derive_batched_fri_challenges, +}; +use crate::fri::fri_commitment::FriLayer; +use crate::fri::fri_decommit::FriDecommitment; +use crate::grinding; + +/// What the prover produced in the batched round 4, plus the challenges it drew +/// on the way. The layers are kept so the caller can run the query phase over +/// them; everything else is what goes on the wire. +pub struct BatchedFriCommit +where + FieldElement: AsBytes + Sync + Send, +{ + pub layers: Vec>>, + pub layer_roots: Vec, + pub final_poly_coeffs: Vec>, + pub layout: BatchedFriLayout, + /// The grinding nonce, `None` when `grinding_factor == 0`. + pub nonce: Option, + /// Row-pair indices in the tallest domain, one per query. + pub iotas: Vec, + /// The mixing challenge the codewords were combined with. Kept because the + /// query phase needs it to rebuild each table's contribution. + pub alpha: FieldElement, + /// Which tables this instance carries, and which keep a terminal-only + /// instance of their own. See [`FriInstancePlan`]. + pub plan: FriInstancePlan, + /// Per table: the standalone class's terminal polynomial, `Some` exactly + /// for `plan.standalone`. Produced by `combine`, ABSORBED here (right + /// after α, before the first ζ — see `derive_batched_fri_challenges` for + /// why that absorb is load-bearing), and returned so the caller puts the + /// very coefficients the transcript bound onto the wire. + pub standalone_coeffs: Vec>>>, +} + +/// Prover side of the batched round-4 sequence. +/// +/// `heights[t]` is `log2` of table `t`'s LDE length and `widths[t]` its committed +/// column count, both in the epoch's canonical table order — the same order the +/// verifier rebuilds from the AIR set, and the same order `combine` must absorb +/// codewords in, since absorption order is what defines the α powers. +/// +/// `combine` receives α and returns the per-height buckets (see +/// [`crate::fri::batched::HeightCombiner::finish`]) TOGETHER WITH the +/// standalone class's terminal polynomials, per table (`Some` exactly for +/// `plan.standalone`). It is a closure rather than a materialized `Vec` so a +/// caller can produce one table's DEEP codeword, absorb it and drop it: +/// holding all of them at once is the memory cost batching exists to remove. +#[allow(clippy::too_many_arguments)] +pub fn commit_batched_fri( + transcript: &mut T, + heights: &[usize], + widths: &[usize], + combine: C, + coset_offset: &FieldElement, + blowup_log: u32, + final_poly_log_degree: u32, + grinding_factor: u8, + num_queries: usize, +) -> BatchedFriCommit +where + F: IsFFTField + IsSubFieldOf + 'static, + E: IsField + 'static + Send + Sync, + T: IsStarkTranscript + Clone, + C: FnOnce( + &FieldElement, + &FriInstancePlan, + ) -> ( + Vec>>>, + Vec>>>, + ), + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, +{ + // Derived from the shape, exactly as the verifier derives it — the partition + // is never sent. The tables whose own FRI commits no layer are left out of the + // batch: they gain nothing from it and pay the full lift to the tallest + // domain, which is where the proximity-gaps term's `|D0|^2` lives. + let plan = FriInstancePlan::new(heights, blowup_log, final_poly_log_degree) + .expect("commit_batched_fri: the epoch's shape is the prover's own"); + let h_max = plan.h_max; + + absorb_shape_histogram::(transcript, heights, widths); + let alpha = transcript.sample_field_element(); + + let (combined, standalone_coeffs) = combine(&alpha, &plan); + + // Bind the standalone class's terminal polynomials BEFORE the first ζ — + // the same walk `derive_batched_fri_challenges` replays, and the reason it + // does (its doc): a polynomial not bound here could be chosen after the + // query indices are known. + for (table, coeffs) in standalone_coeffs.iter().enumerate() { + assert_eq!( + coeffs.is_some(), + plan.standalone.contains(&table), + "the standalone terminals exist for exactly the standalone class" + ); + if let Some(coeffs) = coeffs { + for c in coeffs.iter() { + transcript.append_field_element(c); + } + } + } + + let (final_poly_coeffs, layers) = batched_commit_phase::( + combined, + transcript, + coset_offset, + blowup_log, + final_poly_log_degree, + ); + let layer_roots: Vec = layers.iter().map(|layer| layer.merkle_tree.root).collect(); + + // Grinding runs on the CONFIGURATION's transcript hash, not a hard-wired + // one — the same rule the unbatched `prover.rs` follows. `H` names both the + // commitment family and the Fiat-Shamir hash, so a batched proof committed + // with BLAKE3 grinds with BLAKE3 and one committed with keccak grinds with + // keccak, without either side being told twice. + let nonce = (grinding_factor > 0).then(|| { + let value = grinding::generate_nonce(&transcript.state(), grinding_factor) + .expect("nonce not found"); + transcript.append_bytes(&value.to_be_bytes()); + value + }); + + let iotas = (0..num_queries) + .map(|_| transcript.sample_u64(1u64 << (h_max - 1)) as usize) + .collect(); + + BatchedFriCommit { + layers, + layer_roots, + final_poly_coeffs, + layout: BatchedFriLayout::new(plan.h_max, plan.h_min, blowup_log, final_poly_log_degree), + nonce, + iotas, + alpha, + plan, + standalone_coeffs, + } +} + +/// Verify one query against a STANDALONE table's terminal-only instance. +/// +/// A table whose own FRI commits no layer has a terminal codeword that IS its +/// deep-composition codeword, so there is nothing to fold and nothing to +/// authenticate: the check is that the value the query opened is the value the +/// sent terminal polynomial encodes at that position. +/// +/// ★ `iota` is the SHARED batched query index and is reduced here — the two +/// instance classes read the same index in different spaces (see +/// [`FriInstancePlan`]). `deep` is the table's own deep-composition pair at its +/// reduced row pair, which the caller reconstructs from authenticated openings. +/// +/// Returns `false` on every malformed input; it never panics. +pub fn verify_standalone_fri_query( + iota: usize, + h_max_fri: usize, + h_table: usize, + deep: (&FieldElement, &FieldElement), + terminal_codeword: &[FieldElement], +) -> bool +where + E: IsField + 'static, +{ + let Some(reduced) = reduce_iota_to_round(iota, h_max_fri, h_table) else { + return false; + }; + terminal_codeword + .get(reduced * 2) + .is_some_and(|t| deep.0 == t) + && terminal_codeword + .get(reduced * 2 + 1) + .is_some_and(|t| deep.1 == t) +} + +/// Position, inside the codeword of height `h`, that query `iota` reads. +/// +/// `iota` is a row-pair index in the tallest domain (height `h_max`); the layer +/// whose codeword has height `h` is reached after `h_max - h` folds, and the +/// query's position there is `iota >> (h_max - h - 1)`. Both rows of the pair a +/// height-`h` MMCS opening returns — leaf `iota >> (h_max - h)`, i.e. LDE rows +/// `2k` and `2k+1` — are candidates, and the low bit of this position picks +/// between them; see [`injected_value_at_query`]. +/// +/// Not defined at `h == h_max`: the tallest codeword is the FRI's layer 0, which +/// the query reads as a PAIR (`2·iota`, `2·iota+1`) rather than at one position. +#[inline] +pub fn injection_position(iota: usize, h_max: usize, h: usize) -> usize { + debug_assert!( + h < h_max, + "the tallest codeword is read as a pair, not at a position" + ); + iota >> (h_max - h - 1) +} + +/// The value a height-`h` matrix contributes to its injection layer, chosen from +/// the row pair its MMCS opening returned. +/// +/// `evaluation` is the opening's row `2k` and `evaluation_sym` its row `2k+1`, +/// with `k = iota >> (h_max - h)`. The pair straddles the injection position, so +/// the choice is exactly that position's low bit. +#[inline] +pub fn injected_value_at_query<'a, E: IsField>( + iota: usize, + h_max: usize, + h: usize, + evaluation: &'a FieldElement, + evaluation_sym: &'a FieldElement, +) -> &'a FieldElement { + if injection_position(iota, h_max, h) & 1 == 0 { + evaluation + } else { + evaluation_sym + } +} + +/// Reduce a FRI query index to the index space of a round whose tallest matrix +/// is shorter than the FRI's. +/// +/// [`crate::fri::mmcs::MixedMmcs::verify_batch`] walks its path with the LOW bits +/// of the index it is given, while it locates a short matrix inside the tree by +/// the HIGH bits — consistent only when the index comes from that tree's own +/// `h_max`. The batched preprocessed round is the case that breaks it (its +/// tallest matrix sits below the FRI's), so every caller reduces here rather than +/// each writing the shift out. Returns `None` when the round claims to be TALLER +/// than the FRI, which no honest shape can be. +#[inline] +pub fn reduce_iota_to_round(iota: usize, h_max_fri: usize, h_max_round: usize) -> Option { + (h_max_round <= h_max_fri).then(|| iota >> (h_max_fri - h_max_round)) +} + +/// Verify one query of the batched FRI: the fold-with-injection recursion, every +/// committed layer's opening, and the terminal check. +/// +/// `p0` is the query's pair of values in the tallest codeword — the α-mixed DEEP +/// evaluations of the tables at height `h_max`, at LDE positions `2·iota` and +/// `2·iota + 1`. `bucket_at_height[h]` is `Some(v)` when at least one table has +/// height `h < h_max`, with `v` that height group's α-mixed value at +/// [`injection_position`]; `None` when no table sits at `h`. Both are the +/// caller's to reconstruct from authenticated openings — this function does no +/// authentication of trace data, only of FRI layers. +/// +/// Returns `false` on every malformed input; it never panics. +#[allow(clippy::too_many_arguments)] +pub fn verify_batched_fri_query( + layer_roots: &[Commitment], + betas: &[FieldElement], + layout: &BatchedFriLayout, + h_max: usize, + iota: usize, + decommitment: &FriDecommitment, + evaluation_point_inv: &FieldElement, + p0: (&FieldElement, &FieldElement), + bucket_at_height: &[Option>], + terminal_codeword: &[FieldElement], +) -> bool +where + F: IsFFTField + IsSubFieldOf + 'static, + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, +{ + // The decommitment vectors are prover-supplied and are NOT bound into the + // transcript, so their lengths are pinned here before anything zips them — + // the same reason `step_3_verify_fri` pins them in the unbatched path. A + // short vector would make the fold loop run fewer rounds and accept the query + // without ever reaching the terminal. + if layer_roots.len() != layout.num_committed + || decommitment.layers_auth_paths.len() != layout.num_committed + || decommitment.layers_evaluations_sym.len() != layout.num_committed + || betas.len() != layout.num_committed + usize::from(layout.total_folds > 0) + { + return false; + } + if h_max == 0 || h_max >= usize::BITS as usize || iota >= 1usize << (h_max - 1) { + return false; + } + if bucket_at_height.len() < h_max { + return false; + } + + // No-fold case: the codeword never folds, so the terminal IS the tallest + // codeword and the query's two points sit at `2·iota` and `2·iota + 1`. No + // bucket can exist below `h_max` here — `h_min == h_max` is what makes + // `total_folds` zero — so there is nothing to inject. + if layout.total_folds == 0 { + return terminal_codeword.get(iota * 2).is_some_and(|t| p0.0 == t) + && terminal_codeword + .get(iota * 2 + 1) + .is_some_and(|t| p0.1 == t); + } + + // First fold: layer 0 (the tallest codeword) is not committed, so this fold + // consumes `p0` rather than an authenticated opening. Then the height just + // below joins, exactly as `batched_commit_phase` does before it commits. + let mut point_inv = evaluation_point_inv.clone(); + let mut v = (p0.0 + p0.1) + &point_inv * &betas[0] * (p0.0 - p0.1); + let mut index = iota; + inject(&mut v, &betas[0], bucket_at_height, h_max - 1); + + let mut openings_ok = true; + for i in 0..layout.num_committed { + let evaluation_sym = &decommitment.layers_evaluations_sym[i]; + openings_ok &= verify_layer_opening::( + &layer_roots[i], + decommitment.layers_auth_paths[i].merkle_path.as_slice(), + &v, + evaluation_sym, + index, + ); + + point_inv = point_inv.square(); + v = (&v + evaluation_sym) + &point_inv * &betas[i + 1] * (&v - evaluation_sym); + index >>= 1; + // The injection height descends with the running codeword. `checked_sub` + // rather than `h_max - 2 - i`: `layout`'s fields are only consistent with + // `h_max` when the layout was DERIVED from the same heights, and this + // function is on the verifier's path, where an overflow panic is not a + // rejection. An inconsistent layout simply injects nothing and fails at + // the terminal. + if let Some(height) = (h_max - 1).checked_sub(i + 1) { + inject(&mut v, &betas[i + 1], bucket_at_height, height); + } + } + + // `v` is now the query's value in the terminal codeword and `index` its + // position there. `.get` fails closed on an out-of-range index. + openings_ok & terminal_codeword.get(index).is_some_and(|t| &v == t) +} + +/// `running += β² · bucket_h` for the height the running codeword has just +/// reached. A no-op when no table sits at that height, and when the height is +/// below the terminal (`bucket_at_height` is indexed by height, so a fold that +/// runs past index 0 has nothing to read). +fn inject( + value: &mut FieldElement, + beta: &FieldElement, + bucket_at_height: &[Option>], + height: usize, +) { + if let Some(Some(contribution)) = bucket_at_height.get(height) { + *value = &*value + &(beta.square() * contribution); + } +} + +/// Authenticate a committed FRI layer's row pair against its root. `index` is the +/// query's position in that layer; the leaf is the pair at `index >> 1`, ordered +/// by `index`'s low bit — the same convention the unbatched +/// `verify_fri_layer_openings` uses, and the same one `query_phase` opens with. +fn verify_layer_opening( + root: &Commitment, + auth_path: &[Commitment], + evaluation: &FieldElement, + evaluation_sym: &FieldElement, + index: usize, +) -> bool +where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + let leaf = if index % 2 == 1 { + vec![evaluation_sym.clone(), evaluation.clone()] + } else { + vec![evaluation.clone(), evaluation_sym.clone()] + }; + verify_merkle_path::>(auth_path, root, index >> 1, &leaf) +} + +/// Replay the batched round-4 transcript sequence and return the challenges, +/// or `None` when the proof's shape contradicts the epoch's. +/// +/// A thin alias for [`derive_batched_fri_challenges`], re-exported here so the +/// verifier reaches the sequence through the same module the prover's +/// [`commit_batched_fri`] lives in — the two are one protocol, and splitting them +/// across modules is how they drift. +#[allow(clippy::too_many_arguments)] +pub fn replay_batched_fri( + transcript: &mut T, + heights: &[usize], + widths: &[usize], + layer_roots: &[Commitment], + final_poly_coeffs: &[FieldElement], + standalone_coeffs: &[Option<&[FieldElement]>], + blowup_log: u32, + final_poly_log_degree: u32, + grinding_factor: u8, + nonce: Option, + num_queries: usize, +) -> Option> +where + E: IsField, + T: IsTranscript, +{ + derive_batched_fri_challenges( + transcript, + heights, + widths, + layer_roots, + final_poly_coeffs, + standalone_coeffs, + blowup_log, + final_poly_log_degree, + grinding_factor, + nonce, + num_queries, + ) +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use crate::fri::batched::{HeightCombiner, combine_by_height}; + use crate::fri::terminal::terminal_codeword_from_coeffs; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::fft::bit_reversing::{in_place_bit_reverse_permute, reverse_index}; + use math::field::goldilocks::GoldilocksField; + use math::polynomial::Polynomial; + + pub(crate) type F = GoldilocksField; + pub(crate) type FE = FieldElement; + pub(crate) type Transcript = DefaultTranscript; + + pub(crate) const BLOWUP_LOG: u32 = 1; + pub(crate) const FINAL_POLY_LOG_DEGREE: u32 = 1; + pub(crate) const COSET_OFFSET: u64 = 3; + + /// One synthetic table: a genuinely low-degree codeword at its own height. + pub(crate) struct FakeTable { + pub height: usize, + pub width: usize, + pub codeword: Vec, + } + + /// A codeword of height `h` that IS a Reed-Solomon word of rate `2^-BLOWUP_LOG` + /// on the coset the batched FRI will read it at. + /// + /// The coset matters and is the one thing easy to get wrong here: folding + /// squares the offset, so the layer a height-`h` bucket is injected into lives + /// on `offset^(2^(h_max-h))·⟨ω⟩`, not on `offset·⟨ω⟩`. A word built on the + /// wrong coset is still low degree — the map is a rescaling of the argument — + /// so it would pass a degree check while making the terminal reconstruction + /// disagree, which is exactly the failure the honest-path test has to be able + /// to see. + pub(crate) fn low_degree_codeword(h: usize, h_max: usize, seed: u64) -> Vec { + let num_coeffs = 1usize << (h as u32 - BLOWUP_LOG); + let coeffs: Vec = (0..num_coeffs) + .map(|i| FE::from(seed.wrapping_mul(97).wrapping_add(i as u64 * 31 + 1))) + .collect(); + let offset = FE::from(COSET_OFFSET).pow(1u64 << (h_max - h)); + let mut natural = Polynomial::evaluate_offset_fft::( + &Polynomial::new(&coeffs), + 1usize << BLOWUP_LOG, + Some(num_coeffs), + &offset, + ) + .expect("coset evaluation"); + in_place_bit_reverse_permute(&mut natural); + natural + } + + /// Four tables over three heights, the shape the batched path has to handle: + /// several tables sharing the tallest height (so the base group batches), one + /// at an intermediate height (so an injection lands on a committed layer) and + /// one at the terminal height (so the FINAL fold's injection is exercised — + /// the case #768's loop missed). + pub(crate) fn fixture() -> Vec { + let h_max = 5; + vec![ + FakeTable { + height: 5, + width: 3, + codeword: low_degree_codeword(5, h_max, 11), + }, + FakeTable { + height: 4, + width: 2, + codeword: low_degree_codeword(4, h_max, 23), + }, + FakeTable { + height: 5, + width: 7, + codeword: low_degree_codeword(5, h_max, 41), + }, + FakeTable { + height: 2, + width: 1, + codeword: low_degree_codeword(2, h_max, 59), + }, + ] + } + + pub(crate) fn heights_of(tables: &[FakeTable]) -> Vec { + tables.iter().map(|t| t.height).collect() + } + + pub(crate) fn widths_of(tables: &[FakeTable]) -> Vec { + tables.iter().map(|t| t.width).collect() + } + + /// The per-table standalone slices a replay call takes, off a commit. + pub(crate) fn standalone_refs(commit: &BatchedFriCommit) -> Vec> { + commit + .standalone_coeffs + .iter() + .map(|c| c.as_deref()) + .collect() + } + + /// Run the prover's batched round 4 over `tables`, streaming the codewords + /// into the combiner one at a time — the shape a real prover uses. + pub(crate) fn commit_fixture( + tables: &[FakeTable], + transcript: &mut Transcript, + grinding_factor: u8, + num_queries: usize, + ) -> BatchedFriCommit { + let heights = heights_of(tables); + let widths = widths_of(tables); + commit_batched_fri::( + transcript, + &heights, + &widths, + |alpha, plan| { + // Only the batched class is mixed in, and in the plan's order — + // absorption order is what defines the alpha powers, so a caller + // that absorbed the standalone tables too would shift every + // power and agree with no verifier. The standalone tables hand + // back their terminal polynomials instead, exactly as the real + // prover does. + let mut combiner = HeightCombiner::new(*alpha); + for &t in &plan.batched { + combiner.absorb(&tables[t].codeword, tables[t].height); + } + let standalone = tables + .iter() + .enumerate() + .map(|(t, table)| { + plan.standalone.contains(&t).then(|| { + crate::fri::terminal::coeffs_from_terminal_codeword::( + &table.codeword, + &FE::from(COSET_OFFSET), + table.height as u32 - BLOWUP_LOG, + ) + }) + }) + .collect(); + (combiner.finish(), standalone) + }, + &FE::from(COSET_OFFSET), + BLOWUP_LOG, + FINAL_POLY_LOG_DEGREE, + grinding_factor, + num_queries, + ) + } + + /// υ⁻¹ for query `iota`: the inverse of the tallest coset's element at + /// FRI-order position `2·iota`, matching the unbatched verifier's + /// `query_challenge_to_evaluation_point`. + pub(crate) fn evaluation_point_inv(iota: usize, h_max: usize) -> FE { + let n = 1usize << h_max; + let omega = F::get_primitive_root_of_unity(h_max as u64).expect("root of unity"); + let point = FE::from(COSET_OFFSET) * omega.pow(reverse_index(iota * 2, n as u64)); + point.inv().expect("query point is never zero") + } + + /// What the verifier must reconstruct from authenticated openings: the α-mixed + /// value of every height group at this query's position. Here it is read + /// straight off the combined buckets, which is the oracle — `combine_by_height` + /// has its own tests, and the point of this one is the fold recursion. + pub(crate) fn query_inputs( + tables: &[FakeTable], + alpha: &FE, + iota: usize, + ) -> ((FE, FE), Vec>) { + let plan = FriInstancePlan::new(&heights_of(tables), BLOWUP_LOG, FINAL_POLY_LOG_DEGREE) + .expect("the fixture's shape partitions"); + let h_max = plan.h_max; + let inputs: Vec<(Vec, usize)> = plan + .batched + .iter() + .map(|&t| (tables[t].codeword.clone(), tables[t].height)) + .collect(); + let combined = combine_by_height(&inputs, alpha); + + let tallest = combined[h_max].as_ref().expect("tallest bucket exists"); + let p0 = (tallest[iota * 2], tallest[iota * 2 + 1]); + + let buckets = (0..h_max) + .map(|h| { + combined + .get(h) + .and_then(|slot| slot.as_ref()) + .map(|codeword| codeword[injection_position(iota, h_max, h)]) + }) + .collect(); + (p0, buckets) + } + + /// Verify one query end to end against the committed layers. + #[allow(clippy::too_many_arguments)] + pub(crate) fn verify_one_query( + commit: &BatchedFriCommit, + betas: &[FE], + h_max: usize, + iota: usize, + decommitment: &FriDecommitment, + p0: (&FE, &FE), + buckets: &[Option], + layer_roots: &[Commitment], + final_poly_coeffs: &[FE], + ) -> bool { + let terminal_offset = FE::from(COSET_OFFSET).pow(1u64 << commit.layout.total_folds); + let terminal = terminal_codeword_from_coeffs::( + final_poly_coeffs, + &terminal_offset, + commit.layout.terminal_len, + ); + verify_batched_fri_query::( + layer_roots, + betas, + &commit.layout, + h_max, + iota, + decommitment, + &evaluation_point_inv(iota, h_max), + p0, + buckets, + &terminal, + ) + } + + /// The prover's inline sequence and the verifier's replay are ONE protocol; + /// this is what pins them together. Every challenge, not only the iotas — + /// α gates the height combination and the βs gate every fold, so an + /// agreement that held only at the query indices would still be a broken + /// proof system. + #[test] + fn prover_commit_matches_verifier_derivation() { + let tables = fixture(); + let mut prover_transcript = Transcript::new(b"batched_round4"); + let mut verifier_transcript = prover_transcript.clone(); + + let commit = commit_fixture(&tables, &mut prover_transcript, 4, 6); + + let replay = replay_batched_fri::( + &mut verifier_transcript, + &heights_of(&tables), + &widths_of(&tables), + &commit.layer_roots, + &commit.final_poly_coeffs, + &standalone_refs(&commit), + BLOWUP_LOG, + FINAL_POLY_LOG_DEGREE, + 4, + commit.nonce, + 6, + ) + .expect("an honest shape must derive"); + + assert_eq!(replay.alpha, commit.alpha, "α must agree"); + assert_eq!(replay.layout, commit.layout, "the fold layout must agree"); + assert_eq!(replay.iotas, commit.iotas, "the query indices must agree"); + assert_eq!( + replay.betas.len(), + commit.layout.num_committed + 1, + "one β per committed layer plus the final fold" + ); + assert!( + crate::grinding::is_valid_nonce( + &replay.grinding_seed, + commit.nonce.expect("grinding was requested"), + 4 + ), + "the replayed grinding seed must accept the prover's nonce" + ); + assert_eq!( + prover_transcript.state(), + verifier_transcript.state(), + "both sides must end in the same transcript state" + ); + } + + /// The honest path, and it is not vacuous: the fixture spans three heights, + /// so this exercises the base group, an injection into a committed layer and + /// an injection at the final fold. If the injection convention or the + /// position derivation were wrong, the terminal check would fail. + #[test] + fn honest_batched_queries_verify() { + let tables = fixture(); + let h_max = 5; + let mut transcript = Transcript::new(b"batched_round4"); + let commit = commit_fixture(&tables, &mut transcript, 0, 8); + + let decommitments = + crate::fri::query_phase::(&commit.layers, &commit.iotas); + + let mut verifier_transcript = Transcript::new(b"batched_round4"); + let replay = replay_batched_fri::( + &mut verifier_transcript, + &heights_of(&tables), + &widths_of(&tables), + &commit.layer_roots, + &commit.final_poly_coeffs, + &standalone_refs(&commit), + BLOWUP_LOG, + FINAL_POLY_LOG_DEGREE, + 0, + None, + 8, + ) + .expect("an honest shape must derive"); + + assert!(commit.layout.num_committed >= 1, "the fixture must fold"); + for (query, &iota) in commit.iotas.iter().enumerate() { + let (p0, buckets) = query_inputs(&tables, &replay.alpha, iota); + assert!( + verify_one_query( + &commit, + &replay.betas, + h_max, + iota, + &decommitments[query], + (&p0.0, &p0.1), + &buckets, + &commit.layer_roots, + &commit.final_poly_coeffs, + ), + "honest query {query} (iota {iota}) must verify" + ); + } + } + + /// The MMCS row pair a query opens at height `h` and the FRI position the + /// injection reads must be the SAME two rows. That coincidence is what lets + /// one opening serve both the authentication and the FRI join, and it is a + /// property of the two index derivations, so it is worth pinning exhaustively + /// rather than sampling. + #[test] + fn injection_position_lands_inside_the_mmcs_row_pair() { + let h_max = 6; + for iota in 0..(1usize << (h_max - 1)) { + for h in 1..h_max { + let position = injection_position(iota, h_max, h); + let mmcs_leaf = iota >> (h_max - h); + assert_eq!( + position >> 1, + mmcs_leaf, + "height {h}, iota {iota}: the injection position must sit in the opened leaf" + ); + assert!( + position < (1usize << h), + "height {h}, iota {iota}: position must stay inside the codeword" + ); + } + } + } + + /// `reduce_iota_to_round` is the documented remedy for the one case where a + /// round's tallest matrix is below the FRI's. Pin both that it is the shift + /// the MMCS wants and that it refuses the impossible direction rather than + /// shifting by a negative amount. + #[test] + fn reduce_iota_to_round_matches_the_mmcs_index_space() { + let h_max_fri = 6; + for iota in 0..(1usize << (h_max_fri - 1)) { + for h_max_round in 1..=h_max_fri { + let reduced = + reduce_iota_to_round(iota, h_max_fri, h_max_round).expect("round is shorter"); + assert!( + reduced < (1usize << (h_max_round - 1)), + "the reduced index must land in the round's own leaf range" + ); + } + } + assert!( + reduce_iota_to_round(0, 4, 5).is_none(), + "a round taller than the FRI is not a shape any honest epoch has" + ); + } + + /// A width the epoch did not commit to moves α, and therefore every fold and + /// every query index. This is the shape binding doing its job one level up + /// from the leaf: the leaf header binds a mis-parse, this binds a mis-shaped + /// epoch. + #[test] + fn a_tampered_shape_moves_the_derived_challenges() { + let tables = fixture(); + let mut prover_transcript = Transcript::new(b"batched_round4"); + let commit = commit_fixture(&tables, &mut prover_transcript, 0, 4); + + let mut widths = widths_of(&tables); + widths[1] += 1; + let mut verifier_transcript = Transcript::new(b"batched_round4"); + let replay = replay_batched_fri::( + &mut verifier_transcript, + &heights_of(&tables), + &widths, + &commit.layer_roots, + &commit.final_poly_coeffs, + &standalone_refs(&commit), + BLOWUP_LOG, + FINAL_POLY_LOG_DEGREE, + 0, + None, + 4, + ) + .expect("the shape is still structurally consistent"); + + assert_ne!( + replay.alpha, commit.alpha, + "a width the prover did not commit to must move α" + ); + assert_ne!( + replay.iotas, commit.iotas, + "a width the prover did not commit to must move the query indices" + ); + } +} diff --git a/crypto/stark/src/batched/shape.rs b/crypto/stark/src/batched/shape.rs new file mode 100644 index 000000000..2f53da0f3 --- /dev/null +++ b/crypto/stark/src/batched/shape.rs @@ -0,0 +1,381 @@ +//! The epoch's committed shape — which table contributes a matrix to which +//! batched round, at what height and width. +//! +//! Every number here is derived from the AIR set and the per-table trace +//! lengths, never read out of a proof. That is what lets the verifier rebuild +//! the shape it must pass to [`crate::fri::mmcs::MixedMmcs::verify_batch`] and +//! to [`crate::fri::batched::absorb_shape_histogram`] instead of trusting the +//! prover's word for it (`fri/mmcs.rs`, "Width binding"). +//! +//! # Why one type and not four lists +//! +//! Four rounds are batched (preprocessed, main, aux, composition parts) and each +//! has a DIFFERENT participation list: only preprocessed tables contribute a +//! preprocessed matrix, only tables with a RAP contribute an aux matrix. The +//! index a matrix has inside its round is therefore NOT its table index, and the +//! two are easy to confuse — a confusion that shows up as an opening +//! authenticated at the wrong leaf rather than as a compile error. [`RoundShape`] +//! keeps the mapping in one place so both sides read it from the same code. + +use crate::config::Commitment; +use crate::traits::AIR; + +/// The preprocessed round's pinned shape: the root a program's registry entry +/// commits, and the widths the leaf parse depends on. +/// +/// # Why the widths travel with the root +/// +/// Under the per-table scheme a group's width is implied by its own root plus +/// its AIR. Under one batched tree the widths decide how each leaf is *parsed*, +/// so a comparison of roots alone is only equivalent to the per-table +/// comparisons it replaces if the parse is pinned too (MMCS-PLAN §3.1 item 3, +/// §3.3's closing warning). They are carried here rather than derived at the +/// comparison site so that a caller holding entry A but an AIR set built for +/// entry B is rejected as a width disagreement rather than as an unexplained +/// root mismatch. +/// +/// # The two sides dispose of `None` differently, on purpose +/// +/// Both [`crate::batched::prover::multi_prove_batched`] and +/// [`crate::batched::verifier::multi_verify_batched`] take this as an `Option`, +/// and they do NOT mean the same thing by the absence: +/// +/// - **Prover — permissive.** `None` is how the root is generated in the first +/// place (registry regeneration has nothing to compare against yet). Supplying +/// it buys a fail-fast: a stale preprocessed constant is caught at prove time +/// rather than by every future verifier. +/// - **Verifier — fails closed.** `None` is accepted only for an epoch whose AIR +/// set has no preprocessed table at all. An epoch that HAS a preprocessed +/// round and no pinned root is rejected, because the only root left to check +/// against would be the proof's own — which the prover chose along with the +/// matrices it commits. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct PinnedPrep<'a> { + pub root: &'a Commitment, + /// One width per contributing matrix, in [`RoundShape::tables`] order. + pub widths: &'a [usize], +} + +/// Which tables contribute a matrix to one batched round, and with what shape. +#[derive(Clone, Debug, PartialEq, Eq, Default)] +pub struct RoundShape { + /// Contributing table indices, ascending. Position `i` in this vector is + /// matrix `i` of the round — the order the MMCS concatenates leaves in, and + /// the order openings are presented in. + pub tables: Vec, + /// `(log_height, width)` per contributing matrix, in `tables` order. + pub dims: Vec<(usize, usize)>, +} + +impl RoundShape { + pub fn is_empty(&self) -> bool { + self.tables.is_empty() + } + + pub fn heights(&self) -> Vec { + self.dims.iter().map(|(h, _)| *h).collect() + } + + pub fn widths(&self) -> Vec { + self.dims.iter().map(|(_, w)| *w).collect() + } + + /// The round's own tallest matrix — the index space + /// [`crate::fri::mmcs::MixedMmcs::verify_batch`] accepts. `None` for an + /// empty round. + pub fn h_max(&self) -> Option { + self.dims.iter().map(|(h, _)| *h).max() + } +} + +/// The one table whose MAIN matrix is committed as its own standalone row-pair +/// tree instead of contributing a matrix to the shared main round. +/// +/// This is the L2G carve-out: a continuation epoch's LOCAL_TO_GLOBAL table +/// keeps a per-table main commitment so the cross-epoch root-equality binding +/// (`verify_l2g_commitment_binding_view`) reads the SAME root from a batched +/// epoch as from a per-table one. The carved tree is built by the same +/// committer the per-table prover uses (`commit_rows_bit_reversed_subset` over +/// the full committed-main range), so the two roots are byte-identical for the +/// same trace. The carved table's aux and composition-parts matrices stay in +/// the shared rounds; only its main commitment moves. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CarvedMain { + /// The carved table's index in the AIR set. + pub table: usize, + /// The carved matrix's width — the table's committed main columns. + pub width: usize, +} + +/// The shape of every batched round in one epoch. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EpochShape { + /// `log2` of each table's LDE length, in table order. This is the FRI's + /// shape: query indices live in the tallest of these domains. + pub heights: Vec, + /// Preprocessed columns. Empty when no table is preprocessed. + pub prep: RoundShape, + /// Main trace columns — every table except a carved one. For a + /// preprocessed table this is the MULTIPLICITY columns only, matching the + /// per-table path's split (`commit_main_trace`: `[0, num_precomputed)` is + /// the preprocessed matrix, `[num_precomputed, total)` the committed main + /// one). + pub main: RoundShape, + /// Auxiliary (RAP) columns. Empty when no table has a RAP. + pub aux: RoundShape, + /// Composition-polynomial parts — every table. + pub parts: RoundShape, + /// The table (at most one) whose main matrix is committed standalone. + /// `None` for an ordinary epoch. See [`CarvedMain`]. + pub carved_main: Option, +} + +/// Why an epoch cannot be proved (or verified) with one batched instance. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ShapeError { + /// No tables at all. + Empty, + /// A table's LDE length is not a power of two, is 1, or overflows a shift. + /// Heights come from proof-supplied trace lengths on the verifier's side, so + /// this is a rejection, never a panic. + BadHeight { table: usize, lde_size: usize }, + /// A table declares zero committed main columns, so it has no matrix to + /// contribute and no leaf to open. + NoMainColumns { table: usize }, + /// The batched path commits ONE FRI instance for the whole epoch, so every + /// table must agree on the parameters that instance is defined by. The + /// per-table path has no such requirement, which is exactly why this is + /// checked rather than assumed. + MixedProofOptions { table: usize, field: &'static str }, + /// The carved-main table index does not name a table of this epoch. + CarvedOutOfRange { table: usize }, + /// A preprocessed table cannot be carved: its main-round matrix is the + /// multiplicity columns only, while the per-table root the carve must + /// reproduce commits the full main range. The one production carve (L2G) + /// has no preprocessed columns, so this is rejected rather than supported. + CarvedTablePreprocessed { table: usize }, +} + +impl core::fmt::Display for ShapeError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + ShapeError::Empty => write!(f, "an epoch needs at least one table"), + ShapeError::BadHeight { table, lde_size } => write!( + f, + "table {table}: LDE length {lde_size} is not a power of two greater than 1" + ), + ShapeError::NoMainColumns { table } => { + write!(f, "table {table} commits no main columns") + } + ShapeError::MixedProofOptions { table, field } => write!( + f, + "table {table} disagrees with table 0 on `{field}`; one batched FRI \ + instance needs one set of parameters" + ), + ShapeError::CarvedOutOfRange { table } => { + write!(f, "carved-main table {table} is not a table of this epoch") + } + ShapeError::CarvedTablePreprocessed { table } => write!( + f, + "carved-main table {table} is preprocessed; the carve commits the full \ + main range and cannot reproduce a preprocessed table's per-table root" + ), + } + } +} + +/// The epoch-wide FRI parameters, once every table has been checked to agree. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct EpochFriParams { + pub blowup_log: u32, + pub coset_offset: u64, + pub grinding_factor: u8, + pub num_queries: usize, + pub final_poly_log_degree: u32, +} + +impl EpochShape { + /// Derive the shape from the AIR set and each table's interpolation-domain + /// size (`trace_length`). + /// + /// The prover passes the trace lengths it is about to prove; the verifier + /// passes the ones the proof declares. Both then hold the same `EpochShape` + /// without either having read it from the other. + pub fn derive( + airs: &[&dyn AIR], + trace_lengths: &[usize], + ) -> Result<(Self, EpochFriParams), ShapeError> + where + F: math::field::traits::IsFFTField + + math::field::traits::IsSubFieldOf + + Send + + Sync + + 'static, + E: math::field::traits::IsField + Send + Sync + 'static, + { + Self::derive_carved(airs, trace_lengths, None) + } + + /// As [`EpochShape::derive`], with one table's main matrix carved into a + /// standalone commitment ([`CarvedMain`]). Both sides pass the SAME + /// `carved_main`: it is verifier-owned configuration (like the AIR set), + /// never read from a proof. + pub fn derive_carved( + airs: &[&dyn AIR], + trace_lengths: &[usize], + carved_main: Option, + ) -> Result<(Self, EpochFriParams), ShapeError> + where + F: math::field::traits::IsFFTField + + math::field::traits::IsSubFieldOf + + Send + + Sync + + 'static, + E: math::field::traits::IsField + Send + Sync + 'static, + { + if airs.is_empty() || airs.len() != trace_lengths.len() { + return Err(ShapeError::Empty); + } + if let Some(c) = carved_main { + if c >= airs.len() { + return Err(ShapeError::CarvedOutOfRange { table: c }); + } + if airs[c].is_preprocessed() { + return Err(ShapeError::CarvedTablePreprocessed { table: c }); + } + } + + let first = airs[0].options(); + let params = EpochFriParams { + blowup_log: (first.blowup_factor as usize).trailing_zeros(), + coset_offset: first.coset_offset, + grinding_factor: first.grinding_factor, + num_queries: first.fri_number_of_queries, + final_poly_log_degree: first.fri_final_poly_log_degree as u32, + }; + + let mut heights = Vec::with_capacity(airs.len()); + let mut prep = RoundShape::default(); + let mut main = RoundShape::default(); + let mut aux = RoundShape::default(); + let mut parts = RoundShape::default(); + let mut carved = None; + + for (table, (air, &trace_length)) in airs.iter().zip(trace_lengths).enumerate() { + let options = air.options(); + for (field, same) in [ + ( + "blowup_factor", + options.blowup_factor == first.blowup_factor, + ), + ("coset_offset", options.coset_offset == first.coset_offset), + ( + "grinding_factor", + options.grinding_factor == first.grinding_factor, + ), + ( + "fri_number_of_queries", + options.fri_number_of_queries == first.fri_number_of_queries, + ), + ( + "fri_final_poly_log_degree", + options.fri_final_poly_log_degree == first.fri_final_poly_log_degree, + ), + ] { + if !same { + return Err(ShapeError::MixedProofOptions { table, field }); + } + } + + let lde_size = trace_length + .checked_mul(options.blowup_factor as usize) + .ok_or(ShapeError::BadHeight { + table, + lde_size: usize::MAX, + })?; + if !lde_size.is_power_of_two() || lde_size < 2 || lde_size.trailing_zeros() >= u32::BITS + { + return Err(ShapeError::BadHeight { table, lde_size }); + } + let h = lde_size.trailing_zeros() as usize; + heights.push(h); + + let (total_main_cols, aux_cols) = air.trace_layout(); + let num_precomputed = if air.is_preprocessed() { + air.num_precomputed_columns() + } else { + 0 + }; + let committed_main = total_main_cols + .checked_sub(num_precomputed) + .ok_or(ShapeError::NoMainColumns { table })?; + if committed_main == 0 { + return Err(ShapeError::NoMainColumns { table }); + } + + if num_precomputed > 0 { + prep.tables.push(table); + prep.dims.push((h, num_precomputed)); + } + if carved_main == Some(table) { + carved = Some(CarvedMain { + table, + width: committed_main, + }); + } else { + main.tables.push(table); + main.dims.push((h, committed_main)); + } + if aux_cols > 0 && air.has_aux_trace() { + aux.tables.push(table); + aux.dims.push((h, aux_cols)); + } + let num_parts = air.composition_poly_degree_bound(trace_length) / trace_length; + parts.tables.push(table); + parts.dims.push((h, num_parts.max(1))); + } + + Ok(( + Self { + heights, + prep, + main, + aux, + parts, + carved_main: carved, + }, + params, + )) + } + + /// The epoch's tallest LDE — the domain query indices are drawn in. + pub fn h_max(&self) -> usize { + self.heights.iter().copied().max().unwrap_or(0) + } + + /// The widths the round-4 shape histogram binds: one per table, in table + /// order, summing every matrix that table contributes across all four rounds. + /// + /// Summing rather than listing per round is deliberate. The histogram's job + /// is to make two epochs with different shapes produce different challenges, + /// and `absorb_shape_histogram` takes one `(height, width)` pair per entry. + /// A table's total committed width moves whenever ANY of its four matrices + /// changes width, so the sum separates exactly the epochs the four separate + /// lists would — while staying one entry per table, which is what keeps the + /// prover's and the verifier's histograms the same length without either + /// having to agree on a round ordering. + pub fn total_widths(&self) -> Vec { + let mut widths = vec![0usize; self.heights.len()]; + for round in [&self.prep, &self.main, &self.aux, &self.parts] { + for (&table, (_, w)) in round.tables.iter().zip(round.dims.iter()) { + widths[table] += *w; + } + } + // A carved main matrix is committed outside the shared rounds but is + // still committed width: the histogram binds it like any other. + if let Some(c) = &self.carved_main { + widths[c.table] += c.width; + } + widths + } +} diff --git a/crypto/stark/src/batched/verifier.rs b/crypto/stark/src/batched/verifier.rs new file mode 100644 index 000000000..ae2f88958 --- /dev/null +++ b/crypto/stark/src/batched/verifier.rs @@ -0,0 +1,1068 @@ +//! The batched epoch verifier. +//! +//! [`multi_verify_batched`] is the counterpart of +//! `crate::verifier::IsStarkVerifier::multi_verify` for the batched path, and it +//! is a COMPLETE verification: transcript replay, opening authentication against +//! all four mixed-height MMCS roots, the constraint identity at every table's +//! `z`, the epoch's LogUp bus balance, and the DEEP/FRI join across both +//! instance classes. It is assembled from four pieces, each independently +//! testable and each returning a plain `bool`/`Option` — nothing on this path +//! panics, because every input is prover-supplied. +//! +//! | piece | what it decides | +//! |---|---| +//! | [`replay_epoch_transcript`] | every challenge, and every structural fact the transcript binds | +//! | [`verify_epoch_commitments`] | every preprocessed table's opening authenticates against `air.precomputed_commitment()` (the per-table critical check), and the batched rounds' openings are the rows the roots bind at the derived indices | +//! | [`verify_epoch_constraints`] | the claimed composition polynomial, and the bus balance | +//! | [`verify_epoch_fri`] | those rows fold to the terminal polynomial the proof sent | +//! +//! ⚠ Calling a piece on its own is not a verification. `verify_epoch_commitments` +//! in particular shows only that a proof opened the rows its own roots bind, +//! which an adversary controlling the trace can always arrange. The names are +//! `verify_epoch_*` rather than `verify_*` for that reason; the one function +//! that decides validity is [`multi_verify_batched`]. +//! +//! # Shared with the per-table verifier, not reimplemented +//! +//! Three checks are the same mathematics in both paths, and all three are +//! reached through `crate::verifier`'s own functions rather than copied: +//! `step_2_verify_claimed_composition_polynomial`, +//! `compute_query_invariant_deep_terms` and +//! `reconstruct_deep_composition_poly_evaluation_pair`. The first two took an +//! rkyv `StarkProofView` (#845's zero-copy layer) and now take plain data — a +//! batched epoch proof is not a per-table `StarkProof` and has no such view. +//! That refactor is deliberate: a second constraint evaluator written for the +//! batched path is the one thing that would let the two paths disagree about +//! what a valid trace is (PA-PLAN §1.4). +//! +//! # The one protocol, pinned +//! +//! [`replay_epoch_transcript`] walks exactly the sequence +//! `crate::batched::prover::multi_prove_batched` walks, and +//! `replay_matches_the_provers_ending_state` pins the two on the ENDING +//! TRANSCRIPT STATE. No per-challenge comparison substitutes for it: a +//! divergence anywhere — a root absorbed out of order, a challenge one side +//! samples and the other does not, an OOD block walked differently — lands +//! there, whereas comparing individual challenges only catches it if you +//! compared the right one. + +use math::field::element::FieldElement; +use math::field::traits::{IsFFTField, IsField, IsSubFieldOf}; +use math::traits::AsBytes; + +use crate::batched::proof::BatchedMultiProof; +use crate::batched::round4::reduce_iota_to_round; +use crate::batched::shape::{EpochFriParams, EpochShape, RoundShape}; +use crate::config::{BatchedMerkleTreeBackend, Commitment}; +use crate::fri::batched::{BatchedFriChallenges, absorb_shape_histogram}; +use crate::fri::mmcs::{MixedMmcs, MixedOpening}; +use crate::lookup::LOGUP_NUM_CHALLENGES; +use crate::traits::AIR; + +use crypto::fiat_shamir::is_transcript::IsStarkTranscript; + +/// Every challenge a batched epoch derives, in the order the transcript +/// produces them. +#[derive(Debug, Clone)] +pub struct EpochChallenges { + /// The shared LogUp challenges. Empty when no table has a RAP. + pub lookup: Vec>, + /// One constraint-batching challenge per table, in table order. + pub betas: Vec>, + /// One out-of-domain point per table, in table order. + pub zs: Vec>, + /// One DEEP-batching challenge per table, in table order. + pub deep_gammas: Vec>, + /// The batched FRI instance's challenges, including the query indices and + /// the instance-class partition. + pub fri: BatchedFriChallenges, +} + +/// Replay a batched epoch's transcript and recover every challenge. +/// +/// Returns `None` on any structural disagreement between the proof and the +/// shape the AIR set implies. Every input here is prover-supplied, so every +/// disagreement is a rejection; this function does not panic. +pub fn replay_epoch_transcript( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + transcript: &mut T, +) -> Option<(EpochShape, EpochFriParams, EpochChallenges)> +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, + T: IsStarkTranscript, +{ + replay_epoch_transcript_carved(airs, proof, transcript, None) +} + +/// As [`replay_epoch_transcript`], for an epoch with a carved main matrix +/// ([`crate::batched::shape::CarvedMain`]). +/// +/// `carved_main` is VERIFIER-OWNED configuration, like the AIR set — never +/// read from the proof. The carved root itself IS proof-carried: it is +/// absorbed from `proof.carved_main_root` after the preprocessed roots and +/// before `main_root`, so every challenge is drawn after it. A proof whose +/// carve state disagrees with the configuration is rejected. +pub fn replay_epoch_transcript_carved( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + transcript: &mut T, + carved_main: Option, +) -> Option<(EpochShape, EpochFriParams, EpochChallenges)> +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, + T: IsStarkTranscript, +{ + if airs.len() != proof.tables.len() || airs.is_empty() { + return None; + } + let trace_lengths: Vec = proof.tables.iter().map(|t| t.trace_length).collect(); + let (shape, params) = EpochShape::derive_carved(airs, &trace_lengths, carved_main).ok()?; + + // Recommendation S: the shape is bound before the first root, so every + // challenge below — not only round 4's — is drawn after the epoch has + // committed to what it is. + absorb_shape_histogram::(transcript, &shape.heights, &shape.total_widths()); + + // ★ Preprocessed roots are absorbed FROM THE AIR SET, never from the + // proof — per table, in table order, exactly as the per-table path's + // Phase A does. A prover that committed different preprocessed content + // walked a different transcript and diverges from here on. + for air in airs { + if air.is_preprocessed() { + transcript.append_bytes(&air.precomputed_commitment()); + } + } + + // The carved root, PROOF-CARRIED, in its pinned slot: after every + // preprocessed root, before `main_root`. Presence must match the + // verifier-owned carve configuration exactly. + match (&shape.carved_main, proof.carved_main_root.as_ref()) { + (Some(_), Some(root)) => transcript.append_bytes(root), + (None, None) => {} + _ => return None, + } + + transcript.append_bytes(&proof.main_root); + + let needs_lookup = airs.iter().any(|air| air.has_aux_trace()); + let lookup: Vec> = if needs_lookup { + (0..LOGUP_NUM_CHALLENGES) + .map(|_| transcript.sample_field_element()) + .collect() + } else { + Vec::new() + }; + + if shape.aux.is_empty() != proof.aux_root.is_none() { + return None; + } + if let Some(root) = proof.aux_root.as_ref() { + transcript.append_bytes(root); + } + + // Which tables carry a bus contribution is a property of the AIR set, not + // of the proof. Absorbing whatever the proof happened to send would let a + // prover move the whole transcript by adding or omitting one. + for (air, table) in airs.iter().zip(proof.tables.iter()) { + match (air.has_aux_trace(), table.bus_public_inputs.as_ref()) { + (true, Some(bpi)) => transcript.append_field_element(&bpi.table_contribution), + (false, None) => {} + _ => return None, + } + } + + let betas: Vec> = (0..airs.len()) + .map(|_| transcript.sample_field_element()) + .collect(); + + transcript.append_bytes(&proof.parts_root); + + let coset_offset = FieldElement::::from(params.coset_offset); + let mut zs = Vec::with_capacity(airs.len()); + for (index, table) in proof.tables.iter().enumerate() { + let lde_length = table.trace_length.checked_shl(params.blowup_log)?; + // `sample_z_ood_with_domain_params` is the routine the prover reaches + // through `sample_z_ood`, so the two agree by naming one function + // rather than by two call sites coinciding. + let z = transcript.sample_z_ood_with_domain_params( + table.trace_length, + lde_length, + &coset_offset, + ); + + let air = airs.get(index)?; + // Shape-check the two OOD blocks before absorbing them: they are + // proof-supplied, and the prover's absorption walked blocks the AIR's + // layout defines. A block of the wrong width would otherwise absorb a + // different number of field elements and desynchronise the transcript + // rather than being rejected. + if !ood_blocks_well_formed(*air, table) { + return None; + } + for block in [ + &table.trace_ood_evaluations, + &table.trace_ood_next_evaluations, + ] { + for col in block.columns().iter() { + for elem in col.iter() { + transcript.append_field_element(elem); + } + } + } + for element in table.composition_poly_parts_ood_evaluation.iter() { + transcript.append_field_element(element); + } + zs.push(z); + } + + let deep_gammas: Vec> = (0..airs.len()) + .map(|_| transcript.sample_field_element()) + .collect(); + + let standalone_coeffs: Vec]>> = proof + .tables + .iter() + .map(|t| t.standalone_final_poly_coeffs.as_deref()) + .collect(); + let fri = crate::fri::batched::derive_batched_fri_challenges::( + transcript, + &shape.heights, + &shape.total_widths(), + &proof.fri_layer_roots, + &proof.fri_final_poly_coeffs, + &standalone_coeffs, + params.blowup_log, + params.final_poly_log_degree, + params.grinding_factor, + proof.nonce, + params.num_queries, + )?; + + Some(( + shape, + params, + EpochChallenges { + lookup, + betas, + zs, + deep_gammas, + fri, + }, + )) +} + +/// The two OOD blocks must have the shape the AIR's layout defines. +/// +/// This is `crate::verifier`'s `ood_blocks_well_formed`, restated against the +/// batched proof's owned tables rather than an rkyv view. It is not cosmetic: +/// the blocks are absorbed element by element, so a block of the wrong width +/// would desynchronise the transcript instead of being rejected, and the +/// verifier would go on to derive challenges from a sequence the prover never +/// walked. +fn ood_blocks_well_formed( + air: &dyn AIR, + table: &crate::batched::proof::BatchedTableData, +) -> bool +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, +{ + let step_size = air.step_size(); + let num_eval_points = air.context().transition_offsets.len() * step_size; + let expected_next_width = air.trace_ood_next_row_columns().len(); + let expected_next_height = if expected_next_width == 0 { + 0 + } else { + num_eval_points.saturating_sub(step_size) + }; + let current = &table.trace_ood_evaluations; + let next = &table.trace_ood_next_evaluations; + + current.width == air.trace_layout().0 + air.num_auxiliary_rap_columns() + && current.height == step_size + && next.width == expected_next_width + && next.height == expected_next_height +} + +/// Authenticate every query's openings against every batched round's root, and +/// check the epoch-level structural facts the transcript binds. +/// +/// ⛔ See the module header: this is NOT a complete verification. It is the +/// commitment half. +pub fn verify_epoch_commitments( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + shape: &EpochShape, + params: &EpochFriParams, + challenges: &EpochChallenges, +) -> bool +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, +{ + // The query count is not implied by anything the transcript already + // checked: a prover that sent fewer openings would simply be checked less. + if proof.queries.len() != params.num_queries || challenges.fri.iotas.len() != params.num_queries + { + return false; + } + + if params.grinding_factor > 0 { + let Some(nonce) = proof.nonce else { + return false; + }; + if !crate::grinding::is_valid_nonce( + &challenges.fri.grinding_seed, + nonce, + params.grinding_factor, + ) { + return false; + } + } + + // The instance-class partition is DERIVED, never sent, so the proof's + // terminal polynomials must be present for exactly the standalone tables + // and of exactly the length that class's degree bound implies. + for (table, data) in proof.tables.iter().enumerate() { + let standalone = challenges.fri.plan.standalone.contains(&table); + match (&data.standalone_final_poly_coeffs, standalone) { + (Some(coeffs), true) => { + let Some(&height) = shape.heights.get(table) else { + return false; + }; + let Some(log_degree) = (height as u32).checked_sub(params.blowup_log) else { + return false; + }; + if coeffs.len() != 1usize << log_degree { + return false; + } + } + (None, false) => {} + _ => return false, + } + } + + let h_max = shape.h_max(); + for (query, iota) in challenges.fri.iotas.iter().copied().enumerate() { + let opening = &proof.queries[query]; + if !round_authenticates::( + &proof.main_root, + &opening.main, + &shape.main, + iota, + h_max, + ) { + return false; + } + if !round_authenticates::( + &proof.parts_root, + &opening.parts, + &shape.parts, + iota, + h_max, + ) { + return false; + } + // ★ Per-table preprocessed authentication — the per-table path's + // critical soundness check, verbatim: each opening authenticates + // against `air.precomputed_commitment()`, a root the VERIFIER owns. + // Width and count are bound by the AIR set, not the proof. + if opening.prep.len() != shape.prep.tables.len() { + return false; + } + for (k, &t) in shape.prep.tables.iter().enumerate() { + let Some(air) = airs.get(t) else { + return false; + }; + let Some(&height) = shape.heights.get(t) else { + return false; + }; + let Some(leaf) = reduce_iota_to_round(iota, h_max, height) else { + return false; + }; + let o = &opening.prep[k]; + let width = air.num_precomputed_columns(); + if o.evaluations.len() != width || o.evaluations_sym.len() != width { + return false; + } + let leaf_hash = as crypto::merkle_tree::traits::IsStreamingLeafBackend>::hash_data_from_slices( + &o.evaluations, + &o.evaluations_sym, + ); + if !crypto::merkle_tree::proof::verify_merkle_path_from_leaf_hash::>( + &o.proof.merkle_path, + &air.precomputed_commitment(), + leaf, + leaf_hash, + ) { + return false; + } + } + // ★ The carved table's standalone main opening: authenticated against + // the PROOF-CARRIED root (`carved_main_root`) at the reduced index, + // exactly the mechanics of a preprocessed opening with the root's + // provenance moved from the AIR set to the proof — the transcript slot + // (before every challenge) is what binds it. Present iff the epoch is + // carved; a stray or missing opening is a rejection. + match ( + &shape.carved_main, + proof.carved_main_root.as_ref(), + opening.carved_main.as_ref(), + ) { + (Some(c), Some(root), Some(o)) => { + let Some(&height) = shape.heights.get(c.table) else { + return false; + }; + let Some(leaf) = reduce_iota_to_round(iota, h_max, height) else { + return false; + }; + if o.evaluations.len() != c.width || o.evaluations_sym.len() != c.width { + return false; + } + let leaf_hash = + as crypto::merkle_tree::traits::IsStreamingLeafBackend< + Field, + >>::hash_data_from_slices( + &o.evaluations, &o.evaluations_sym + ); + if !crypto::merkle_tree::proof::verify_merkle_path_from_leaf_hash::>( + &o.proof.merkle_path, + root, + leaf, + leaf_hash, + ) { + return false; + } + } + (None, None, None) => {} + _ => return false, + } + + match (proof.aux_root.as_ref(), opening.aux.as_ref()) { + (Some(root), Some(o)) => { + if !round_authenticates::(root, o, &shape.aux, iota, h_max) { + return false; + } + } + (None, None) => {} + _ => return false, + } + } + + true +} + +/// Authenticate one round at one query, reducing the shared FRI index into the +/// round's own index space first. +/// +/// The reduction is the whole reason this is a named function rather than four +/// inline calls: the preprocessed and auxiliary rounds can have an `h_max` +/// below the FRI's, and passing the un-reduced index is not a loud error — +/// prover and verifier share the routine, so a wrong convention is +/// self-consistent (`fri/mmcs.rs`, "Index convention"). +fn round_authenticates( + root: &Commitment, + opening: &MixedOpening, + round: &RoundShape, + iota_fri: usize, + h_max_fri: usize, +) -> bool +where + C: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + let Some(h_max_round) = round.h_max() else { + return false; + }; + let Some(iota) = reduce_iota_to_round(iota_fri, h_max_fri, h_max_round) else { + return false; + }; + MixedMmcs::::verify_batch(root, iota, opening, &round.heights(), &round.widths()) +} + +// =========================================================================== +// The DEEP / FRI join — M-5's core +// =========================================================================== + +/// Verify the batched FRI instance and the terminal-only instances, at every +/// query. +/// +/// This is the check that gives the authenticated openings their meaning. Up to +/// here a proof has shown that the rows it opened are the rows its roots bind; +/// this shows that those rows evaluate to a codeword the FRI folds to a +/// low-degree polynomial — that the committed trace really does satisfy the +/// DEEP relation at `z`. +/// +/// # The two index spaces, again +/// +/// Both instance classes are opened at the SAME query indices and read them +/// differently ([`crate::fri::batched::FriInstancePlan`]): the batched class +/// uses `iota` directly because it is an index in the tallest domain, a +/// standalone table at height `h` uses `iota >> (h_max - h)`. A table's OWN +/// row pair also lives at its reduced leaf, which is why the evaluation point +/// each table's DEEP quotient is reconstructed at is derived from the reduced +/// index and not from `iota`. +/// +/// # Mixing +/// +/// [`crate::fri::batched::HeightCombiner`] scales the `i`-th absorbed codeword +/// by `alpha^i`, counting in absorption order and NOT per height, and the +/// prover absorbs in `plan.batched` order. So the power a table's DEEP value +/// carries here is its position in `plan.batched` — not its table index, and +/// not its position within its height group. Getting that wrong produces a +/// verifier that rejects every honest proof, which is the benign direction, but +/// it is worth stating because the three orders coincide on a same-height +/// epoch. +/// +/// Returns `false` on every malformed input; it never panics. +pub fn verify_epoch_fri( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + shape: &EpochShape, + params: &EpochFriParams, + challenges: &EpochChallenges, +) -> bool +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + Field::BaseType: math::field::element::NativeArchived, + FieldExtension::BaseType: math::field::element::NativeArchived, + PI: rkyv::Archive + Clone, + ::Archived: rkyv::Deserialize, + V: crate::verifier::IsStarkVerifier + ?Sized, +{ + let h_max = shape.h_max(); + let layout = &challenges.fri.layout; + let coset_offset = FieldElement::::from(params.coset_offset); + + // Structural checks before anything is reconstructed. The terminal helper + // panics on a coefficient count that does not divide the codeword length, + // so the length check is not optional — it is what keeps this path + // rejection-only. Same reasoning as `step_3_verify_fri`. + if proof.fri_layer_roots.len() != layout.num_committed + || proof.fri_final_poly_coeffs.len() != (1usize << layout.effective_k) + { + return false; + } + for query in proof.queries.iter() { + if query.fri.layers_auth_paths.len() != layout.num_committed + || query.fri.layers_evaluations_sym.len() != layout.num_committed + { + return false; + } + } + + let terminal_offset = coset_offset.pow(1u64 << layout.total_folds); + let terminal_codeword = + crate::fri::terminal::terminal_codeword_from_coeffs::( + &proof.fri_final_poly_coeffs, + &terminal_offset, + layout.terminal_len, + ); + + // Per table: the DEEP value pair at every query, in this table's own + // (reduced) index space. + let mut deep_pairs: Vec, FieldElement)>> = + Vec::with_capacity(airs.len()); + for (table, air) in airs.iter().enumerate() { + match table_deep_pairs::( + table, *air, proof, shape, params, challenges, + ) { + Some(pairs) => deep_pairs.push(pairs), + None => return false, + } + } + + for (query, iota) in challenges.fri.iotas.iter().copied().enumerate() { + let mut p0 = ( + FieldElement::::zero(), + FieldElement::::zero(), + ); + let mut buckets: Vec>> = vec![None; h_max]; + let mut power = FieldElement::::one(); + + for &table in challenges.fri.plan.batched.iter() { + let (Some(&height), Some(pairs)) = (shape.heights.get(table), deep_pairs.get(table)) + else { + return false; + }; + let Some((evaluation, evaluation_sym)) = pairs.get(query) else { + return false; + }; + if height == h_max { + p0.0 = &p0.0 + &(&power * evaluation); + p0.1 = &p0.1 + &(&power * evaluation_sym); + } else { + let chosen = crate::batched::round4::injected_value_at_query( + iota, + h_max, + height, + evaluation, + evaluation_sym, + ); + let scaled = &power * chosen; + buckets[height] = Some(match buckets[height].take() { + Some(acc) => acc + scaled, + None => scaled, + }); + } + power = &power * &challenges.fri.alpha; + } + + // υ⁻¹ in the TALLEST domain — the batched instance's layer 0. + let lde_length = 1usize << h_max; + let Some(lde_root) = Field::get_primitive_root_of_unity(h_max as u64).ok() else { + return false; + }; + let point = &coset_offset + * lde_root.pow(math::fft::bit_reversing::reverse_index( + iota * 2, + lde_length as u64, + )); + let Ok(point_inv) = point.inv() else { + return false; + }; + + if !crate::batched::round4::verify_batched_fri_query::( + &proof.fri_layer_roots, + &challenges.fri.betas, + layout, + h_max, + iota, + &proof.queries[query].fri, + &point_inv, + (&p0.0, &p0.1), + &buckets, + &terminal_codeword, + ) { + return false; + } + + // The other class. A table whose own FRI commits no layer has a + // terminal codeword that IS its deep-composition codeword, so the check + // is that the value its opening produced is the value the sent + // polynomial encodes at the reduced position. + for &table in challenges.fri.plan.standalone.iter() { + let (Some(&height), Some(pairs), Some(data)) = ( + shape.heights.get(table), + deep_pairs.get(table), + proof.tables.get(table), + ) else { + return false; + }; + let (Some((evaluation, evaluation_sym)), Some(coeffs)) = + (pairs.get(query), data.standalone_final_poly_coeffs.as_ref()) + else { + return false; + }; + let codeword_len = 1usize << height; + if coeffs.is_empty() + || !coeffs.len().is_power_of_two() + || coeffs.len() > codeword_len + || !codeword_len.is_multiple_of(coeffs.len()) + { + return false; + } + let standalone_terminal = crate::fri::terminal::terminal_codeword_from_coeffs::< + Field, + FieldExtension, + >(coeffs, &coset_offset, codeword_len); + if !crate::batched::round4::verify_standalone_fri_query( + iota, + h_max, + height, + (evaluation, evaluation_sym), + &standalone_terminal, + ) { + return false; + } + } + } + + true +} + +/// One table's DEEP composition value pair at every query, reconstructed from +/// the authenticated openings. +/// +/// The base columns are handed over as two slices in COMMIT order — the +/// preprocessed round's row first, then the main round's — because that is the +/// order the prover concatenated them in and the order the OOD grid and the +/// trace-term coefficients are indexed by. A non-preprocessed table passes an +/// empty first slice, which is exactly what the per-table path does. +fn table_deep_pairs( + table: usize, + air: &dyn AIR, + proof: &BatchedMultiProof, + shape: &EpochShape, + params: &EpochFriParams, + challenges: &EpochChallenges, +) -> Option, FieldElement)>> +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + Field::BaseType: math::field::element::NativeArchived, + FieldExtension::BaseType: math::field::element::NativeArchived, + PI: rkyv::Archive + Clone, + ::Archived: rkyv::Deserialize, + V: crate::verifier::IsStarkVerifier + ?Sized, +{ + let data = proof.tables.get(table)?; + let &height = shape.heights.get(table)?; + let h_max = shape.h_max(); + let z = challenges.zs.get(table)?; + let gamma = challenges.deep_gammas.get(table)?; + + let domain = crate::domain::new_verifier_domain(air, data.trace_length); + let step_size = air.step_size(); + let ood_layout = crate::ood::OodLayout::new( + air.context().trace_columns, + air.context().transition_offsets.len() * step_size, + step_size, + air.trace_ood_next_row_columns(), + ); + let ood_full = ood_layout.reconstruct_full( + data.trace_ood_evaluations.row_major_data(), + data.trace_ood_evaluations.width, + data.trace_ood_next_evaluations.row_major_data(), + ); + + // The DEEP coefficients, derived exactly as the prover derives them: the + // first `num_surviving` powers of gamma are the trace terms, the rest the + // composition parts. Splitting them the other way round would be a verifier + // that rejects every honest proof. + let num_terms_trace = ood_layout.num_surviving(); + let num_parts = data.composition_poly_parts_ood_evaluation.len(); + let mut powers: Vec> = + core::iter::successors(Some(FieldElement::one()), |x| Some(x * gamma)) + .take(num_parts + num_terms_trace) + .collect(); + if powers.len() < num_terms_trace { + return None; + } + let trace_term_powers: Vec<_> = powers.drain(..num_terms_trace).collect(); + let trace_term_coeffs = ood_layout.build_trace_term_coeffs(&trace_term_powers); + let gammas = powers; + + let table_challenges = crate::verifier::Challenges { + z: z.clone(), + boundary_coeffs: Vec::new(), + transition_coeffs: Vec::new(), + trace_term_coeffs, + gammas, + zetas: Vec::new(), + iotas: Vec::new(), + rap_challenges: challenges.lookup.clone(), + grinding_seed: [0u8; 32], + }; + + let terms = V::query_invariant_deep_terms_from_parts( + &table_challenges, + &data.composition_poly_parts_ood_evaluation, + &ood_full, + ood_layout.next_row_cols(), + step_size, + )?; + let primitive_root = Field::get_primitive_root_of_unity(domain.root_order as u64).ok()?; + + let prep_matrix = shape.prep.tables.iter().position(|&t| t == table); + // A carved table has no main-round matrix: its main row pair comes from the + // standalone carved opening instead (authenticated against the + // proof-carried root by `verify_epoch_commitments`). + let is_carved = shape.carved_main.map(|c| c.table) == Some(table); + let main_matrix = if is_carved { + None + } else { + Some(shape.main.tables.iter().position(|&t| t == table)?) + }; + let aux_matrix = shape.aux.tables.iter().position(|&t| t == table); + let parts_matrix = shape.parts.tables.iter().position(|&t| t == table)?; + + let mut pairs = Vec::with_capacity(challenges.fri.iotas.len()); + for (query, iota) in challenges.fri.iotas.iter().copied().enumerate() { + let opening = proof.queries.get(query)?; + // This table's OWN row pair: the reduced leaf, in its own domain. + let leaf = crate::batched::round4::reduce_iota_to_round(iota, h_max, height)?; + let point = domain.lde_coset_element(math::fft::bit_reversing::reverse_index( + leaf * 2, + domain.lde_length as u64, + )); + let point_sym = domain.lde_coset_element(math::fft::bit_reversing::reverse_index( + leaf * 2 + 1, + domain.lde_length as u64, + )); + + let empty_base: &[FieldElement] = &[]; + let empty_ext: &[FieldElement] = &[]; + let (prep, prep_sym) = match prep_matrix { + Some(m) => { + let o = opening.prep.get(m)?; + (o.evaluations.as_slice(), o.evaluations_sym.as_slice()) + } + None => (empty_base, empty_base), + }; + let (main_evals, main_evals_sym) = match main_matrix { + Some(m) => { + let o = opening.main.per_matrix.get(m)?; + (o.evaluations.as_slice(), o.evaluations_sym.as_slice()) + } + None => { + let o = opening.carved_main.as_ref()?; + (o.evaluations.as_slice(), o.evaluations_sym.as_slice()) + } + }; + let (aux, aux_sym) = match aux_matrix { + Some(m) => { + let o = opening.aux.as_ref()?.per_matrix.get(m)?; + (o.evaluations.as_slice(), o.evaluations_sym.as_slice()) + } + None => (empty_ext, empty_ext), + }; + let parts = opening.parts.per_matrix.get(parts_matrix)?; + + let pair = V::reconstruct_deep_composition_poly_evaluation_pair( + &point, + &point_sym, + &primitive_root, + &table_challenges, + &terms, + ood_layout.next_row_cols(), + step_size, + prep, + main_evals, + aux, + &parts.evaluations, + prep_sym, + main_evals_sym, + aux_sym, + &parts.evaluations_sym, + )?; + pairs.push(pair); + } + let _ = params; + Some(pairs) +} + +// =========================================================================== +// The constraint identity, the bus balance, and the whole verification +// =========================================================================== + +/// Check every table's claimed composition polynomial at its own `z`, and the +/// epoch's LogUp bus balance. +/// +/// The constraint check is `crate::verifier`'s +/// `step_2_verify_claimed_composition_polynomial`, unchanged — that function now +/// takes plain data instead of an rkyv view precisely so this caller can reach +/// it. Writing a second constraint evaluator for the batched path is the one +/// thing that would make the two paths able to disagree about what a valid +/// trace is. +/// +/// ⚠ `public_inputs` are read from the proof, exactly as the per-table path +/// reads them from `StarkProof`. Checking that they are the inputs the caller +/// meant is the caller's job in both paths. +pub fn verify_epoch_constraints( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + challenges: &EpochChallenges, + expected_bus_balance: &FieldElement, +) -> bool +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + Field::BaseType: math::field::element::NativeArchived, + FieldExtension::BaseType: math::field::element::NativeArchived, + PI: rkyv::Archive + Clone, + ::Archived: rkyv::Deserialize, + V: crate::verifier::IsStarkVerifier + ?Sized, +{ + // Bus balance: Σ table_contribution = expected. This is the cross-table + // statement no per-table check can make, and it is why the contributions are + // absorbed before any constraint challenge. + let mut total = FieldElement::::zero(); + for table in proof.tables.iter() { + if let Some(bpi) = table.bus_public_inputs.as_ref() { + total += bpi.table_contribution.clone(); + } + } + if total != *expected_bus_balance { + return false; + } + + for (table, air) in airs.iter().enumerate() { + let (Some(data), Some(z), Some(beta)) = ( + proof.tables.get(table), + challenges.zs.get(table), + challenges.betas.get(table), + ) else { + return false; + }; + + let step_size = air.step_size(); + let ood_layout = crate::ood::OodLayout::new( + air.context().trace_columns, + air.context().transition_offsets.len() * step_size, + step_size, + air.trace_ood_next_row_columns(), + ); + let ood_full = ood_layout.reconstruct_full( + data.trace_ood_evaluations.row_major_data(), + data.trace_ood_evaluations.width, + data.trace_ood_next_evaluations.row_major_data(), + ); + let domain = crate::domain::new_verifier_domain(*air, data.trace_length); + + // The constraint-batching coefficients, split exactly as the prover + // splits them: transitions first, then boundaries. + let bus_public_inputs = data.bus_public_inputs.clone(); + let num_transition_constraints = air.context().num_transition_constraints; + let num_boundary_constraints = air + .boundary_constraints( + &data.public_inputs, + &challenges.lookup, + bus_public_inputs.as_ref(), + data.trace_length, + ) + .constraints + .len(); + let mut coefficients: Vec> = + core::iter::successors(Some(FieldElement::one()), |x| Some(x * beta)) + .take(num_boundary_constraints + num_transition_constraints) + .collect(); + if coefficients.len() < num_transition_constraints { + return false; + } + let transition_coeffs: Vec<_> = coefficients.drain(..num_transition_constraints).collect(); + let boundary_coeffs = coefficients; + + let table_challenges = crate::verifier::Challenges { + z: z.clone(), + boundary_coeffs, + transition_coeffs, + trace_term_coeffs: Vec::new(), + gammas: Vec::new(), + zetas: Vec::new(), + iotas: Vec::new(), + rap_challenges: challenges.lookup.clone(), + grinding_seed: [0u8; 32], + }; + + if !V::step_2_verify_claimed_composition_polynomial( + *air, + data.trace_length, + data.bus_public_inputs + .as_ref() + .map(|b| b.table_contribution.clone()), + data.trace_ood_evaluations.get_row(0), + &data.composition_poly_parts_ood_evaluation, + &data.public_inputs, + &domain, + &table_challenges, + &ood_full, + step_size, + ) { + return false; + } + } + + true +} + +/// Verify a batched epoch proof: replay, commitments, constraint identity, bus +/// balance, DEEP/FRI join. +/// +/// This is the counterpart of `crate::verifier::IsStarkVerifier::multi_verify` +/// for the batched path, and unlike the pieces above it is a COMPLETE +/// verification — every check the per-table path makes has a counterpart here, +/// reached through the same functions where the check is shared. +/// +/// Preprocessed binding needs no caller-side pin: every preprocessed table's +/// root is `air.precomputed_commitment()` — the verifier's own value, absorbed +/// and compared per table exactly as the per-table path does. +/// +/// Returns `false` on every malformed proof; it never panics. +pub fn multi_verify_batched( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + transcript: &mut T, + expected_bus_balance: &FieldElement, +) -> bool +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + Field::BaseType: math::field::element::NativeArchived, + FieldExtension::BaseType: math::field::element::NativeArchived, + PI: rkyv::Archive + Clone, + ::Archived: rkyv::Deserialize, + V: crate::verifier::IsStarkVerifier + ?Sized, + T: IsStarkTranscript, +{ + multi_verify_batched_carved::( + airs, + proof, + transcript, + expected_bus_balance, + None, + ) +} + +/// As [`multi_verify_batched`], for an epoch with a carved main matrix. +/// +/// `carved_main` is verifier-owned configuration (which table, if any, commits +/// its main matrix standalone) — the same value the prover was called with, +/// supplied by the CALLER, never read from the proof. Everything else about +/// the carve is checked: the proof-carried root's transcript slot +/// ([`replay_epoch_transcript_carved`]), the per-query opening's +/// authentication and width ([`verify_epoch_commitments`]), and the opened +/// row pair's participation in the DEEP/FRI join ([`verify_epoch_fri`]). +pub fn multi_verify_batched_carved( + airs: &[&dyn AIR], + proof: &BatchedMultiProof, + transcript: &mut T, + expected_bus_balance: &FieldElement, + carved_main: Option, +) -> bool +where + Field: IsSubFieldOf + IsFFTField + Send + Sync + 'static, + FieldExtension: IsField + Send + Sync + 'static, + FieldElement: AsBytes + Sync + Send, + FieldElement: AsBytes + Sync + Send, + Field::BaseType: math::field::element::NativeArchived, + FieldExtension::BaseType: math::field::element::NativeArchived, + PI: rkyv::Archive + Clone, + ::Archived: rkyv::Deserialize, + V: crate::verifier::IsStarkVerifier + ?Sized, + T: IsStarkTranscript, +{ + let Some((shape, params, challenges)) = + replay_epoch_transcript_carved(airs, proof, transcript, carved_main) + else { + return false; + }; + verify_epoch_commitments::(airs, proof, &shape, ¶ms, &challenges) + && verify_epoch_constraints::( + airs, + proof, + &challenges, + expected_bus_balance, + ) + && verify_epoch_fri::( + airs, + proof, + &shape, + ¶ms, + &challenges, + ) +} diff --git a/crypto/stark/src/lib.rs b/crypto/stark/src/lib.rs index 6f8e7c82e..fd9b393e8 100644 --- a/crypto/stark/src/lib.rs +++ b/crypto/stark/src/lib.rs @@ -3,6 +3,7 @@ #[cfg(all(target_arch = "wasm32", feature = "disk-spill"))] compile_error!("the `disk-spill` feature requires memmap2, which does not compile on wasm32"); +pub mod batched; #[cfg(feature = "debug-checks")] pub mod bus_debug; pub mod commitment; @@ -29,6 +30,7 @@ pub mod profile_markers; pub mod proof; pub mod prover; pub mod r4_denoms; +pub mod residency_mode; #[cfg(feature = "disk-spill")] pub mod storage_mode; pub mod table; diff --git a/crypto/stark/src/residency_mode.rs b/crypto/stark/src/residency_mode.rs new file mode 100644 index 000000000..52885e703 --- /dev/null +++ b/crypto/stark/src/residency_mode.rs @@ -0,0 +1,36 @@ +/// Whether Round 1 keeps every table's main LDE resident until its fused task +/// runs, or drops it after the commit and recomputes it inside the task. +/// +/// Fiat-Shamir requires the main *roots* to be absorbed before the shared LogUp +/// challenges are sampled; it says nothing about the LDE buffers, so keeping +/// them is a performance choice. `Retain` makes it; `RecomputeLde` trades one +/// extra forward NTT per table for turning an `O(N)` retention into an +/// `O(table_parallelism)` transient. The Merkle tree is kept either way, so a +/// recompute never re-hashes and the root that entered the transcript stays the +/// root openings are checked against. +/// +/// The choice is invisible to the proof: same roots, same transcript order, +/// same proof bytes. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ResidencyMode { + /// Keep every main LDE from its Round-1 commit until its table's fused + /// task consumes it. + #[default] + Retain, + /// Drop each main LDE once its root is absorbed and recompute it from the + /// still-resident trace at the top of the table's fused task. + /// + /// Also releases each table's aux columns from the caller-owned + /// `TraceTable` when that table's proof is complete — a documented part of + /// this mode's contract, since it mutates caller-visible state. Callers + /// that read a trace's aux columns after `multi_prove` returns must use + /// `Retain`. + RecomputeLde, +} + +impl ResidencyMode { + /// True when main LDEs are dropped after Round 1 and recomputed on demand. + pub fn recomputes_main_lde(self) -> bool { + matches!(self, Self::RecomputeLde) + } +} diff --git a/crypto/stark/src/tests/batched_mmcs_soundness_tests.rs b/crypto/stark/src/tests/batched_mmcs_soundness_tests.rs new file mode 100644 index 000000000..3b33f0f79 --- /dev/null +++ b/crypto/stark/src/tests/batched_mmcs_soundness_tests.rs @@ -0,0 +1,1640 @@ +//! Soundness negatives for the batched-commitment primitives — the mixed-height +//! MMCS ([`crate::fri::mmcs`]) and the batched-FRI transcript +//! ([`crate::fri::batched`]). +//! +//! Each test builds one honest commitment over a small mixed-height epoch, then +//! tampers a single component and asserts rejection. The honest opening is +//! re-asserted in every test, so a false-reject regression cannot make the +//! negatives pass vacuously. +//! +//! Scope grows with the integration. The first section reaches only what the +//! primitives decide; the per-query batched-FRI section below arrived with the +//! round-4 wiring ([`crate::batched::round4`]), which is what made a tampered +//! layer evaluation, a mis-sized decommitment and a wrong injection expressible. +//! The forgeries that still need the full prover/verifier integration — an OOD +//! value, the bus balance, the query count, the grinding nonce — belong with it. + +use crypto::fiat_shamir::default_transcript::DefaultTranscript; +use crypto::fiat_shamir::is_transcript::IsTranscript; +use math::fft::bit_reversing::reverse_index; +use math::field::element::FieldElement; +use math::field::goldilocks::GoldilocksField; + +use crate::batched::round4::BatchedFriCommit; +use crate::batched::round4::tests as round4_tests; +use crate::fri::batched::{ + BatchedFriLayout, absorb_shape_histogram, derive_batched_fri_challenges, +}; +use crate::fri::fri_decommit::FriDecommitment; +use crate::fri::mmcs::{LeafSource, MixedMmcs, MixedOpening}; + +type F = GoldilocksField; +type FE = FieldElement; +type Mmcs = MixedMmcs; +type Transcript = DefaultTranscript; + +/// Bit-reversed row-major matrices, in the layout the MMCS commits. +struct Matrices { + /// `(bit-reversed row-major data, log_height, width)`. + mats: Vec<(Vec, usize, usize)>, +} + +impl LeafSource for Matrices { + fn num_matrices(&self) -> usize { + self.mats.len() + } + fn log_height(&self, m: usize) -> usize { + self.mats[m].1 + } + fn width(&self, m: usize) -> usize { + self.mats[m].2 + } + fn append_row(&self, m: usize, bitrev_row: usize, out: &mut Vec) { + let (data, _, width) = &self.mats[m]; + out.extend_from_slice(&data[bitrev_row * width..(bitrev_row + 1) * width]); + } +} + +fn matrix(log_height: usize, width: usize, seed: u64) -> (Vec, usize, usize) { + let num_rows = 1usize << log_height; + let mut data = vec![FE::from(0u64); num_rows * width]; + for (r, chunk) in data.chunks_exact_mut(width).enumerate() { + let br = reverse_index(r, num_rows as u64); + for (c, slot) in chunk.iter_mut().enumerate() { + *slot = FE::from(seed.wrapping_mul(31) + (c as u64) * 1009 + (br as u64) * 7 + 1); + } + } + (data, log_height, width) +} + +/// A four-matrix epoch: two tall (base group), one injected, one injected lower. +/// Heights {5, 5, 4, 2}, widths {3, 3, 2, 4}. Two of the tall matrices share a +/// width so the "swap two openings" forgery below is a pure reordering. +fn epoch() -> (Matrices, Vec, Vec) { + let mats = Matrices { + mats: vec![ + matrix(5, 3, 11), + matrix(5, 3, 22), + matrix(4, 2, 33), + matrix(2, 4, 44), + ], + }; + let heights = vec![5, 5, 4, 2]; + let widths = vec![3, 3, 2, 4]; + (mats, heights, widths) +} + +const IOTA: usize = 9; + +fn honest() -> ([u8; 32], MixedOpening, Vec, Vec) { + let (mats, heights, widths) = epoch(); + let mmcs = Mmcs::commit(&mats); + let opening = mmcs.open_batch(IOTA, &mats); + (mmcs.root(), opening, heights, widths) +} + +/// Sanity anchor: the untampered opening verifies. +#[test] +fn honest_batched_opening_verifies() { + let (root, opening, heights, widths) = honest(); + assert!( + Mmcs::verify_batch(&root, IOTA, &opening, &heights, &widths), + "an honest mixed-height opening must verify" + ); +} + +/// Tampering any matrix's opened row breaks the one shared authentication path — +/// including the SHORT matrices, which are bound through injection rather than +/// through the base leaf. +#[test] +fn rejects_a_tampered_row_in_every_height_group() { + let (root, opening, heights, widths) = honest(); + for m in 0..opening.per_matrix.len() { + let mut tampered = opening.clone(); + tampered.per_matrix[m].evaluations[0] = + &tampered.per_matrix[m].evaluations[0] + &FE::from(1u64); + assert!( + !Mmcs::verify_batch(&root, IOTA, &tampered, &heights, &widths), + "a tampered row of matrix {m} (height {}) must be rejected", + heights[m] + ); + + let mut tampered_sym = opening.clone(); + tampered_sym.per_matrix[m].evaluations_sym[0] = + &tampered_sym.per_matrix[m].evaluations_sym[0] + &FE::from(1u64); + assert!( + !Mmcs::verify_batch(&root, IOTA, &tampered_sym, &heights, &widths), + "a tampered symmetric row of matrix {m} must be rejected" + ); + } +} + +/// Tampering the shared authentication path itself. +#[test] +fn rejects_a_tampered_authentication_path() { + let (root, opening, heights, widths) = honest(); + for level in 0..opening.proof.merkle_path.len() { + let mut tampered = opening.clone(); + tampered.proof.merkle_path[level][0] ^= 1; + assert!( + !Mmcs::verify_batch(&root, IOTA, &tampered, &heights, &widths), + "a tampered sibling at level {level} must be rejected" + ); + } + // Truncating or padding the path is a shape error, not a hash mismatch. + let mut short = opening.clone(); + short.proof.merkle_path.pop(); + assert!(!Mmcs::verify_batch(&root, IOTA, &short, &heights, &widths)); + let mut long = opening.clone(); + long.proof.merkle_path.push([0u8; 32]); + assert!(!Mmcs::verify_batch(&root, IOTA, &long, &heights, &widths)); +} + +/// An honest opening replayed at a different query index must be rejected: the +/// path is position-dependent, so one opening does not authenticate every leaf. +#[test] +fn rejects_an_opening_replayed_at_another_index() { + let (root, opening, heights, widths) = honest(); + let n0 = 1usize << (5 - 1); + for iota in 0..n0 { + let accepted = Mmcs::verify_batch(&root, iota, &opening, &heights, &widths); + assert_eq!( + accepted, + iota == IOTA, + "the opening at {IOTA} must verify at {IOTA} and nowhere else (index {iota})" + ); + } + // And past the tree's leaf range — the index-convention guard. + assert!(!Mmcs::verify_batch(&root, n0, &opening, &heights, &widths)); +} + +/// INPUT ORDER is part of the commitment: swapping two same-height, same-width +/// matrices' openings changes the flat concatenation the group leaf hashes, so +/// the tree no longer reproduces. Without order-dependence a prover could serve +/// one table's rows in another's slot. +#[test] +fn rejects_swapped_openings_within_a_height_group() { + let (root, opening, heights, widths) = honest(); + assert_eq!( + (heights[0], widths[0]), + (heights[1], widths[1]), + "matrices 0 and 1 must share a shape for this to be a pure reordering" + ); + let mut swapped = opening.clone(); + swapped.per_matrix.swap(0, 1); + assert_ne!( + swapped.per_matrix[0].evaluations, opening.per_matrix[0].evaluations, + "the two matrices must carry different data" + ); + assert!( + !Mmcs::verify_batch(&root, IOTA, &swapped, &heights, &widths), + "reordering two same-shape matrices must be rejected" + ); +} + +/// The verifier's `heights` fix the injection schedule. Relabelling a matrix's +/// height — claiming the height-4 matrix is height 3, so it is injected a level +/// later — must not reproduce the root, or a prover could move a table to a +/// layer where its rows are checked against a different query position. +#[test] +fn rejects_a_relabelled_injection_height() { + let (root, opening, heights, widths) = honest(); + let mut relabelled = heights.clone(); + relabelled[2] = 3; + assert!( + !Mmcs::verify_batch(&root, IOTA, &opening, &relabelled, &widths), + "moving a matrix to another injection level must be rejected" + ); + + // Promoting a short matrix into the base group is likewise rejected. + let mut promoted = heights.clone(); + promoted[3] = 5; + assert!(!Mmcs::verify_batch( + &root, IOTA, &opening, &promoted, &widths + )); +} + +/// Widths are verifier-supplied and length-checked, so a width that does not +/// match the opening is rejected before any hashing — the guard that closes the +/// leaf-boundary shift. +#[test] +fn rejects_widths_that_disagree_with_the_opening() { + let (root, opening, heights, widths) = honest(); + for m in 0..widths.len() { + let mut wrong = widths.clone(); + wrong[m] += 1; + assert!( + !Mmcs::verify_batch(&root, IOTA, &opening, &heights, &wrong), + "a width disagreeing with matrix {m}'s opening must be rejected" + ); + } +} + +/// A root committed over a different epoch shape does not authenticate this +/// opening, even where the tree depth coincides. +#[test] +fn rejects_a_root_from_another_epoch_shape() { + let (_, opening, heights, widths) = honest(); + let other = Matrices { + mats: vec![ + matrix(5, 3, 11), + matrix(5, 3, 22), + matrix(4, 2, 33), + // Same height and width, different data. + matrix(2, 4, 99), + ], + }; + let other_root = Mmcs::commit(&other).root(); + assert!( + !Mmcs::verify_batch(&other_root, IOTA, &opening, &heights, &widths), + "an opening must not verify against another epoch's root" + ); +} + +/// The round-4 transcript binds the shape and every committed FRI layer, so +/// tampering a layer root or a terminal coefficient moves the query indices the +/// prover must answer at. This is what stops a prover from choosing its FRI +/// commitments after seeing the queries. +#[test] +fn tampering_the_fri_transcript_moves_the_query_indices() { + let heights = vec![10usize, 10, 8, 7]; + let widths = vec![4usize, 2, 3, 1]; + let (blowup_log, k) = (1u32, 5u32); + let layout = BatchedFriLayout::new(10, 7, blowup_log, k); + let roots: Vec<[u8; 32]> = (0u8..layout.num_committed as u8).map(|i| [i; 32]).collect(); + let coeffs: Vec = (0..(1u64 << layout.effective_k)).map(FE::from).collect(); + + // Height 7 folds no layer at these parameters, so table 3 is standalone + // and its terminal polynomial is transcript-bound alongside the rest. + let standalone: Vec>> = vec![ + None, + None, + None, + Some((0..(1u64 << (7 - blowup_log))).map(FE::from).collect()), + ]; + let derive = |roots: &[[u8; 32]], + coeffs: &[FE], + heights: &[usize], + widths: &[usize], + standalone: &[Option>]| { + let standalone_refs: Vec> = standalone.iter().map(|c| c.as_deref()).collect(); + derive_batched_fri_challenges( + &mut Transcript::new(b"batched_soundness"), + heights, + widths, + roots, + coeffs, + &standalone_refs, + blowup_log, + k, + 0, + None, + 16, + ) + .expect("a well-formed layer-root and coefficient count") + .iotas + }; + + let base = derive(&roots, &coeffs, &heights, &widths, &standalone); + assert!(!base.is_empty()); + + let mut other_root = roots.clone(); + other_root[0][0] ^= 1; + assert_ne!( + base, + derive(&other_root, &coeffs, &heights, &widths, &standalone), + "a tampered FRI layer root must move the query indices" + ); + + let mut other_coeffs = coeffs.clone(); + other_coeffs[0] = &other_coeffs[0] + &FE::from(1u64); + assert_ne!( + base, + derive(&roots, &other_coeffs, &heights, &widths, &standalone), + "a tampered terminal coefficient must move the query indices" + ); + + let mut other_heights = heights.clone(); + other_heights[2] = 9; + assert_ne!( + base, + derive(&roots, &coeffs, &other_heights, &widths, &standalone), + "a tampered height must move the query indices" + ); + + let mut other_widths = widths.clone(); + other_widths[2] = 4; + assert_ne!( + base, + derive(&roots, &coeffs, &heights, &other_widths, &standalone), + "a tampered width must move the query indices" + ); + + // ★ The standalone class's terminal polynomial is transcript-bound too — + // the whole point of the absorb: a polynomial the indices did not depend + // on could be chosen AFTER them, and each query's proximity test against + // it would bind nothing until the queries saturate the table's domain. + let mut other_standalone = standalone.clone(); + if let Some(cs) = other_standalone[3].as_mut() { + cs[0] = &cs[0] + &FE::from(1u64); + } + assert_ne!( + base, + derive(&roots, &coeffs, &heights, &widths, &other_standalone), + "a tampered standalone terminal must move the query indices" + ); +} + +/// The shape histogram's encoding is injective: no two distinct epoch shapes +/// absorb the same bytes. A collision would let a prover present one shape to +/// the transcript and another to the opening parse. +#[test] +fn the_shape_encoding_separates_distinct_epochs() { + let absorbed = |heights: &[usize], widths: &[usize]| { + let mut t = Transcript::new(b"shape"); + absorb_shape_histogram(&mut t, heights, widths); + t.state() + }; + + // The classic ambiguity a length prefix and fixed-width fields must close: + // one table of shape (h, w) against two tables whose fields interleave to the + // same sequence. + let one = absorbed(&[3, 4], &[4, 5]); + let two = absorbed(&[3], &[4]); + let three = absorbed(&[3, 4, 5], &[4, 5, 6]); + assert_ne!(one, two); + assert_ne!(one, three); + assert_ne!(two, three); + + // Swapping height and width within a table is a different epoch. + assert_ne!(absorbed(&[3, 4], &[4, 3]), absorbed(&[4, 3], &[3, 4])); +} + +// --------------------------------------------------------------------------- +// Per-query batched FRI (M-3). These need the round-4 wiring, not only the +// primitives, so they were deferred when the primitives landed. +// --------------------------------------------------------------------------- + +/// One honest batched round 4 plus everything a verifier needs to check a query. +struct Round4Fixture { + tables: Vec, + commit: BatchedFriCommit, + betas: Vec, + decommitments: Vec>, + alpha: FE, + h_max: usize, + plan: crate::fri::batched::FriInstancePlan, +} + +impl Round4Fixture { + fn build() -> Self { + let tables = round4_tests::fixture(); + let mut transcript = round4_tests::Transcript::new(b"batched_soundness_r4"); + let commit = round4_tests::commit_fixture(&tables, &mut transcript, 0, 6); + let decommitments = + crate::fri::query_phase::(&commit.layers, &commit.iotas); + + let mut verifier_transcript = round4_tests::Transcript::new(b"batched_soundness_r4"); + let replay = crate::batched::round4::replay_batched_fri::( + &mut verifier_transcript, + &round4_tests::heights_of(&tables), + &round4_tests::widths_of(&tables), + &commit.layer_roots, + &commit.final_poly_coeffs, + &round4_tests::standalone_refs(&commit), + round4_tests::BLOWUP_LOG, + round4_tests::FINAL_POLY_LOG_DEGREE, + 0, + None, + 6, + ) + .expect("an honest shape must derive"); + + Self { + tables, + betas: replay.betas, + alpha: replay.alpha, + h_max: replay.plan.h_max, + plan: replay.plan, + decommitments, + commit, + } + } + + /// Verify query `q` with every input honest except what `mutate` changes. + fn check_query_with(&self, q: usize, mutate: M) -> bool + where + M: FnOnce(&mut FriDecommitment, &mut (FE, FE), &mut Vec>, &mut Vec), + { + let iota = self.commit.iotas[q]; + let (mut p0, mut buckets) = round4_tests::query_inputs(&self.tables, &self.alpha, iota); + let mut decommitment = self.decommitments[q].clone(); + let mut coeffs = self.commit.final_poly_coeffs.clone(); + mutate(&mut decommitment, &mut p0, &mut buckets, &mut coeffs); + round4_tests::verify_one_query( + &self.commit, + &self.betas, + self.h_max, + iota, + &decommitment, + (&p0.0, &p0.1), + &buckets, + &self.commit.layer_roots, + &coeffs, + ) + } + + fn check_query(&self, q: usize) -> bool { + self.check_query_with(q, |_, _, _, _| {}) + } +} + +/// The honest-path control for every negative below. Also pins that the fixture +/// is not degenerate: it must actually commit layers, or the fold loop the +/// negatives target would never run. +#[test] +fn honest_batched_fri_queries_verify() { + let f = Round4Fixture::build(); + assert!( + f.commit.layout.num_committed >= 1, + "the fixture must commit at least one FRI layer" + ); + assert!( + f.commit.layout.total_folds as usize > f.commit.layout.num_committed, + "the fixture must exercise the final fold" + ); + for q in 0..f.commit.iotas.len() { + assert!(f.check_query(q), "honest query {q} must verify"); + } +} + +/// A per-query FRI layer evaluation is prover-supplied and NOT in the +/// transcript; only the layer's Merkle root binds it. +#[test] +fn a_tampered_fri_layer_evaluation_is_rejected() { + let f = Round4Fixture::build(); + for q in 0..f.commit.iotas.len() { + assert!(f.check_query(q), "honest control for query {q}"); + for layer in 0..f.commit.layout.num_committed { + assert!( + !f.check_query_with(q, |d, _, _, _| { + d.layers_evaluations_sym[layer] = + &d.layers_evaluations_sym[layer] + &FE::from(1u64); + }), + "query {q}: a tampered evaluation at layer {layer} must be rejected" + ); + } + } +} + +/// The authentication path is what carries the layer opening to the root. +#[test] +fn a_tampered_fri_layer_auth_path_is_rejected() { + let f = Round4Fixture::build(); + for layer in 0..f.commit.layout.num_committed { + assert!(f.check_query(0), "honest control"); + assert!( + !f.check_query_with(0, |d, _, _, _| { + d.layers_auth_paths[layer].merkle_path[0][0] ^= 1; + }), + "a tampered sibling at layer {layer} must be rejected" + ); + } +} + +/// The decommitment vectors are not bound by Fiat-Shamir, so their lengths have +/// to be pinned before anything iterates them: a short one would end the fold +/// early and accept without reaching the terminal, a long one would run past it. +#[test] +fn a_mis_sized_fri_decommitment_is_rejected() { + let f = Round4Fixture::build(); + assert!(f.check_query(0), "honest control"); + + assert!( + !f.check_query_with(0, |d, _, _, _| { + d.layers_auth_paths.pop(); + d.layers_evaluations_sym.pop(); + }), + "a truncated decommitment must be rejected" + ); + assert!( + !f.check_query_with(0, |d, _, _, _| { + let path = d.layers_auth_paths[0].clone(); + let evaluation = d.layers_evaluations_sym[0]; + d.layers_auth_paths.push(path); + d.layers_evaluations_sym.push(evaluation); + }), + "a padded decommitment must be rejected" + ); + assert!( + !f.check_query_with(0, |d, _, _, _| { + d.layers_auth_paths.clear(); + d.layers_evaluations_sym.clear(); + }), + "an empty decommitment must be rejected, not accepted vacuously" + ); +} + +/// The terminal polynomial is where FRI's low-degree claim is finally cashed in. +#[test] +fn a_tampered_terminal_coefficient_is_rejected() { + let f = Round4Fixture::build(); + assert!(f.check_query(0), "honest control"); + for i in 0..f.commit.final_poly_coeffs.len() { + assert!( + !f.check_query_with(0, |_, _, _, coeffs| { + coeffs[i] = &coeffs[i] + &FE::from(1u64); + }), + "a tampered terminal coefficient {i} must be rejected" + ); + } +} + +/// The tallest tables enter FRI as layer 0, which is never committed — the only +/// thing binding them is that the fold has to land on the terminal. +#[test] +fn a_tampered_layer_zero_value_is_rejected() { + let f = Round4Fixture::build(); + assert!(f.check_query(0), "honest control"); + assert!( + !f.check_query_with(0, |_, p0, _, _| { p0.0 = &p0.0 + &FE::from(1u64) }), + "a tampered p0 must be rejected" + ); + assert!( + !f.check_query_with(0, |_, p0, _, _| { p0.1 = &p0.1 + &FE::from(1u64) }), + "a tampered p0 symmetric value must be rejected" + ); + assert!( + !f.check_query_with(0, |_, p0, _, _| { core::mem::swap(&mut p0.0, &mut p0.1) }), + "swapping the layer-0 pair must be rejected — the two are not interchangeable" + ); +} + +/// The injected buckets are the whole point of a mixed-height batch: a short +/// table is bound ONLY by the value it contributes at its injection layer. Three +/// ways to get that wrong, all of which leave the tall tables untouched and so +/// would pass a control that only tampered the base group. +#[test] +fn a_wrong_injection_is_rejected() { + let f = Round4Fixture::build(); + // The BATCHED class's short heights — the standalone class is not injected at + // all, and asking for its bucket would be asking about a codeword that is not + // in this instance. + let injected_heights: Vec = f + .plan + .batched + .iter() + .map(|&t| f.tables[t].height) + .filter(|h| *h < f.h_max) + .collect(); + assert!( + !injected_heights.is_empty(), + "the fixture must have at least one injected height" + ); + + for q in 0..f.commit.iotas.len() { + assert!(f.check_query(q), "honest control for query {q}"); + for &h in &injected_heights { + assert!( + !f.check_query_with(q, |_, _, buckets, _| { + let value = buckets[h].take().expect("the height is occupied"); + buckets[h] = Some(&value + &FE::from(1u64)); + }), + "query {q}: a tampered injection at height {h} must be rejected" + ); + assert!( + !f.check_query_with(q, |_, _, buckets, _| { buckets[h] = None }), + "query {q}: dropping the injection at height {h} must be rejected" + ); + } + } +} + +/// The injection position is derived, not sent, and prover and verifier derive it +/// separately — so a control that only tampers the VALUE would pass under a wrong +/// derivation. Reading the other row of the same opened pair is the mistake a +/// off-by-one in `injection_position` would make, so it is the one to pin. +#[test] +fn an_injection_read_at_the_sibling_row_is_rejected() { + let f = Round4Fixture::build(); + let mut exercised = 0usize; + for (q, &iota) in f.commit.iotas.iter().enumerate() { + assert!(f.check_query(q), "honest control for query {q}"); + for &t in f.plan.batched.iter() { + let h = f.tables[t].height; + if h == f.h_max { + continue; + } + let position = crate::batched::round4::injection_position(iota, f.h_max, h); + let sibling = position ^ 1; + let inputs: Vec<(Vec, usize)> = f + .plan + .batched + .iter() + .map(|&b| (f.tables[b].codeword.clone(), f.tables[b].height)) + .collect(); + let combined = crate::fri::batched::combine_by_height(&inputs, &f.alpha); + let bucket = combined[h].as_ref().expect("the height is occupied"); + // A degenerate codeword whose two rows coincide would make this + // vacuous; skip rather than assert a rejection that means nothing. + if bucket[position] == bucket[sibling] { + continue; + } + exercised += 1; + let sibling_value = bucket[sibling]; + assert!( + !f.check_query_with(q, |_, _, buckets, _| { + buckets[h] = Some(sibling_value); + }), + "query {q}: reading height {h}'s injection at the sibling row must be rejected" + ); + } + } + assert!( + exercised > 0, + "no non-degenerate sibling pair was exercised — the test proved nothing" + ); +} + +/// Everything on the verifier's path is prover-supplied, so it must fail closed +/// on shapes that cannot occur honestly rather than panic on them. +#[test] +fn malformed_batched_fri_inputs_are_rejected_without_panicking() { + let f = Round4Fixture::build(); + let iota = f.commit.iotas[0]; + let (p0, buckets) = round4_tests::query_inputs(&f.tables, &f.alpha, iota); + let terminal_offset = + FE::from(round4_tests::COSET_OFFSET).pow(1u64 << f.commit.layout.total_folds); + let terminal = crate::fri::terminal::terminal_codeword_from_coeffs::( + &f.commit.final_poly_coeffs, + &terminal_offset, + f.commit.layout.terminal_len, + ); + let point_inv = round4_tests::evaluation_point_inv(iota, f.h_max); + + let run = |layer_roots: &[[u8; 32]], + betas: &[FE], + h_max: usize, + iota: usize, + buckets: &[Option], + terminal: &[FE]| { + crate::batched::round4::verify_batched_fri_query::( + layer_roots, + betas, + &f.commit.layout, + h_max, + iota, + &f.decommitments[0], + &point_inv, + (&p0.0, &p0.1), + buckets, + terminal, + ) + }; + + assert!( + run( + &f.commit.layer_roots, + &f.betas, + f.h_max, + iota, + &buckets, + &terminal + ), + "honest control" + ); + assert!( + !run(&[], &f.betas, f.h_max, iota, &buckets, &terminal), + "a missing layer-root vector must be rejected" + ); + assert!( + !run( + &f.commit.layer_roots, + &[], + f.h_max, + iota, + &buckets, + &terminal + ), + "a missing beta vector must be rejected" + ); + assert!( + !run( + &f.commit.layer_roots, + &f.betas, + 0, + iota, + &buckets, + &terminal + ), + "h_max = 0 must be rejected, not shifted by" + ); + assert!( + !run( + &f.commit.layer_roots, + &f.betas, + f.h_max, + 1usize << (f.h_max - 1), + &buckets, + &terminal + ), + "an iota from a taller domain must be rejected" + ); + assert!( + !run( + &f.commit.layer_roots, + &f.betas, + f.h_max, + iota, + &buckets[..1], + &terminal + ), + "a bucket vector too short to cover every height must be rejected" + ); + assert!( + !run( + &f.commit.layer_roots, + &f.betas, + f.h_max, + iota, + &buckets, + &[] + ), + "an empty terminal codeword must be rejected" + ); +} + +/// ★ The control the two-class split requires: a table of EACH class must be +/// tamper-checked. +/// +/// The classes read the same query index in different spaces — the batched class +/// uses `iota` directly, a standalone table uses `iota >> (h_max - h)`. Prover +/// and verifier both derive that shift from the shape, so a wrong convention is +/// self-consistent and honest proofs keep verifying; the failure is that the +/// standalone tables end up checked at positions nothing else reaches. A control +/// that only tampered the batched class would pass under ANY convention for the +/// other, which is exactly how consolidating a per-table check loses coverage. +#[test] +fn each_instance_class_is_tamper_checked() { + let f = Round4Fixture::build(); + assert!( + !f.plan.standalone.is_empty(), + "the fixture must exercise BOTH classes, or this control proves nothing \ + about the split" + ); + assert!( + f.plan.batched.len() > 1, + "the batched class must carry more than the tallest table" + ); + + // --- Batched class: covered by the fold recursion. --- + assert!(f.check_query(0), "honest control"); + assert!( + !f.check_query_with(0, |_, p0, _, _| { p0.0 = &p0.0 + &FE::from(1u64) }), + "a tampered batched-class value must be rejected" + ); + + // --- Standalone class: its own terminal-only instance. --- + for &t in &f.plan.standalone { + let table = &f.tables[t]; + // A zero-layer table's terminal codeword IS its deep-composition + // codeword — nothing folds — so the honest terminal is the codeword. + let terminal = &table.codeword; + for (q, &iota) in f.commit.iotas.iter().enumerate() { + let reduced = crate::batched::round4::reduce_iota_to_round(iota, f.h_max, table.height) + .expect("a standalone table is never taller than the FRI"); + let honest = (terminal[reduced * 2], terminal[reduced * 2 + 1]); + assert!( + crate::batched::round4::verify_standalone_fri_query::( + iota, + f.h_max, + table.height, + (&honest.0, &honest.1), + terminal, + ), + "query {q}: the honest standalone opening must verify" + ); + let tampered = &honest.0 + &FE::from(1u64); + assert!( + !crate::batched::round4::verify_standalone_fri_query::( + iota, + f.h_max, + table.height, + (&tampered, &honest.1), + terminal, + ), + "query {q}: a tampered standalone value must be rejected" + ); + // The index rule itself: reading the table at the UNREDUCED batched + // index is the mistake the two-class split makes possible, and it is + // silent unless something rejects it. + if iota != reduced && iota * 2 + 1 < terminal.len() { + assert!( + !crate::batched::round4::verify_standalone_fri_query::( + iota, + f.h_max, + table.height, + (&terminal[iota * 2], &terminal[iota * 2 + 1]), + terminal, + ), + "query {q}: the un-reduced index must not authenticate a \ + standalone table" + ); + } + } + } +} + +/// A standalone table must not be reachable through the batched instance's +/// injection path: it contributes no bucket, so a prover that manufactured one +/// is claiming a codeword this instance never mixed. +#[test] +fn a_standalone_table_contributes_no_injection() { + let f = Round4Fixture::build(); + assert!( + !f.plan.standalone.is_empty(), + "the fixture needs both classes" + ); + for &t in &f.plan.standalone { + let h = f.tables[t].height; + assert!( + !f.plan.batched.iter().any(|&b| f.tables[b].height == h), + "the fixture's standalone height must be unique to that class" + ); + for q in 0..f.commit.iotas.len() { + assert!(f.check_query(q), "honest control for query {q}"); + assert!( + !f.check_query_with(q, |_, _, buckets, _| { + buckets[h] = Some(FE::from(7u64)); + }), + "query {q}: a bucket manufactured at a standalone height must be \ + rejected" + ); + } + } +} + +// =========================================================================== +// EPOCH-LEVEL NEGATIVES — the items M-2 deferred "with the integration" +// =========================================================================== +// +// Everything above decides what the PRIMITIVES can decide: a tampered row, a +// mis-sized path, a replayed index, a swapped opening. These need a whole +// epoch, so they arrive with `multi_prove_batched` and +// `batched::verifier::replay_epoch_transcript`. +// +// ⚠ What is covered and what is not. Query count, the grinding nonce, the OOD +// values and the bus-contribution BINDING are covered. Bus BALANCE — that the +// per-table contributions sum to the expected value across the epoch — is NOT, +// and cannot be until the batched verifier grows the constraint half; see +// `batched/verifier.rs`'s header for why that is blocked and on what. +// +// Every negative below has an honest-path control beside it. Without one, a +// rejection proves only that the checker rejects, not that it discriminates. +mod epoch { + use super::*; + use crate::batched::verifier::{replay_epoch_transcript, verify_epoch_commitments}; + use crate::residency_mode::ResidencyMode; + use crate::tests::batched_prover_tests::{Air, E, F, folding_options, prove_repeated}; + use crate::traits::AIR; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::element::FieldElement; + + type Proof = crate::batched::proof::BatchedMultiProof; + + fn air_refs(airs: &[Air]) -> Vec<&dyn AIR> { + airs.iter() + .map(|a| a as &dyn AIR) + .collect() + } + + /// Replay `proof` and run the commitment checks. `None` when the replay + /// itself rejects, so a test can tell "rejected structurally" from + /// "rejected on the openings". + fn replay_and_check(airs: &[Air], proof: &Proof) -> Option { + let refs = air_refs(airs); + let (shape, params, challenges) = + replay_epoch_transcript(&refs, proof, &mut DefaultTranscript::::new(&[]))?; + Some(verify_epoch_commitments::( + &refs, + proof, + &shape, + ¶ms, + &challenges, + )) + } + + fn honest() -> (Vec, Proof) { + let (airs, proof, _, _) = prove_repeated(1, &folding_options(), ResidencyMode::Retain); + (airs, proof) + } + + /// ★ The strongest oracle available for "the prover and the verifier are one + /// protocol": not that some challenge agrees, but that the two transcripts + /// END in the same state. A divergence anywhere in the sequence — a root + /// absorbed in the wrong order, a challenge sampled that the other side does + /// not sample, an OOD block walked differently — lands here, where comparing + /// individual challenges would only catch it if you happened to compare the + /// right one. + #[test_log::test] + fn replay_matches_the_provers_ending_state() { + let mut prover_transcript = DefaultTranscript::::new(&[]); + let (airs, proof, _, _) = crate::tests::batched_prover_tests::prove_repeated_with( + 1, + &folding_options(), + ResidencyMode::Retain, + &mut prover_transcript, + ); + + let mut verifier_transcript = DefaultTranscript::::new(&[]); + let refs = air_refs(&airs); + replay_epoch_transcript(&refs, &proof, &mut verifier_transcript) + .expect("an honest epoch must replay"); + + assert_eq!( + prover_transcript.state(), + verifier_transcript.state(), + "prover and verifier must end the epoch in the same transcript state" + ); + } + + /// The honest-path control every negative below leans on. + #[test_log::test] + fn an_honest_epoch_passes_the_commitment_checks() { + let (airs, proof) = honest(); + assert_eq!( + replay_and_check(&airs, &proof), + Some(true), + "an honest epoch must replay and authenticate" + ); + } + + /// Query count. Nothing the transcript has already checked implies it: a + /// prover that sent fewer openings would simply be checked less often. + #[test_log::test] + fn a_short_query_list_is_rejected() { + let (airs, mut proof) = honest(); + assert!( + proof.queries.len() > 1, + "the fixture must have queries to drop" + ); + proof.queries.pop(); + assert_eq!( + replay_and_check(&airs, &proof), + Some(false), + "dropping a query must be rejected" + ); + } + + /// Grinding. The nonce is absorbed, so a forged one moves every later + /// challenge AND fails its own proof-of-work check; either rejection is + /// correct and the test asserts the outcome, not the route. + #[test_log::test] + fn a_forged_grinding_nonce_is_rejected() { + let (airs, mut proof) = honest(); + let nonce = proof.nonce.expect("the fixture grinds"); + proof.nonce = Some(nonce.wrapping_add(1)); + assert_ne!( + replay_and_check(&airs, &proof), + Some(true), + "a nonce the prover did not grind must be rejected" + ); + } + + /// A missing nonce where the epoch's grinding factor demands one. + #[test_log::test] + fn an_absent_grinding_nonce_is_rejected() { + let (airs, mut proof) = honest(); + proof.nonce = None; + assert_ne!( + replay_and_check(&airs, &proof), + Some(true), + "an epoch with a positive grinding factor must carry a nonce" + ); + } + + /// OOD values. They are absorbed before alpha, so tampering one must move + /// the query indices — which is what makes the openings, honestly produced + /// at the honest indices, stop authenticating. + #[test_log::test] + fn a_tampered_ood_value_is_rejected() { + let (airs, honest_proof) = honest(); + let refs = air_refs(&airs); + let honest_iotas = + replay_epoch_transcript(&refs, &honest_proof, &mut DefaultTranscript::::new(&[])) + .expect("honest replay") + .2 + .fri + .iotas; + + let mut proof = honest_proof.clone(); + proof.tables[0].composition_poly_parts_ood_evaluation[0] += FieldElement::::one(); + let (_, _, tampered) = + replay_epoch_transcript(&refs, &proof, &mut DefaultTranscript::::new(&[])) + .expect("the shape is still structurally consistent"); + assert_ne!( + tampered.fri.iotas, honest_iotas, + "an OOD value the prover did not commit must move the query indices" + ); + assert_eq!( + replay_and_check(&airs, &proof), + Some(false), + "and the openings must then fail to authenticate" + ); + } + + /// A trace OOD value, tampered in the other block, must behave the same — + /// the two blocks are absorbed separately and a control on only one would + /// miss a verifier that walked just that one. + #[test_log::test] + fn a_tampered_trace_ood_value_is_rejected() { + let (airs, mut proof) = honest(); + let table = &mut proof.tables[0]; + let value = *table.trace_ood_evaluations.get(0, 0); + table + .trace_ood_evaluations + .set(0, 0, value + FieldElement::::one()); + assert_eq!( + replay_and_check(&airs, &proof), + Some(false), + "a trace OOD value the prover did not commit must be rejected" + ); + } + + /// Bus-contribution BINDING (not balance): which tables carry one is a fact + /// about the AIR set, so dropping one must be a structural rejection rather + /// than a transcript that quietly absorbs one element fewer. + #[test_log::test] + fn a_dropped_bus_contribution_is_rejected() { + let (airs, mut proof) = honest(); + assert!( + proof.tables[0].bus_public_inputs.is_some(), + "the fixture's tables all have a RAP" + ); + proof.tables[0].bus_public_inputs = None; + assert_eq!( + replay_and_check(&airs, &proof), + None, + "a table whose AIR has a RAP must carry a bus contribution" + ); + } + + /// A tampered bus contribution is absorbed, so it moves the challenges. + #[test_log::test] + fn a_tampered_bus_contribution_is_rejected() { + let (airs, mut proof) = honest(); + if let Some(bpi) = proof.tables[0].bus_public_inputs.as_mut() { + bpi.table_contribution += FieldElement::::one(); + } + assert_eq!( + replay_and_check(&airs, &proof), + Some(false), + "a bus contribution the prover did not commit must be rejected" + ); + } + + /// A whole round's root, dropped. The AIR set says the aux round exists, so + /// its absence is a structural rejection — without this a prover could + /// remove a round's binding entirely. + #[test_log::test] + fn a_dropped_round_root_is_rejected() { + let (airs, mut proof) = honest(); + proof.aux_root = None; + assert_eq!( + replay_and_check(&airs, &proof), + None, + "the aux round exists in this epoch, so its root cannot be absent" + ); + } + + /// An opening the epoch does not have. This fixture has no preprocessed + /// table, so a query carrying a preprocessed opening must reject: the + /// count is bound by the AIR set, never by what the proof sends. + #[test_log::test] + fn an_invented_prep_opening_is_rejected() { + let (airs, mut proof) = honest(); + assert!( + proof.queries[0].prep.is_empty(), + "the fixture has no preprocessed table" + ); + proof.queries[0] + .prep + .push(crate::proof::stark::PolynomialOpenings { + proof: crypto::merkle_tree::proof::Proof { + merkle_path: Vec::new(), + }, + evaluations: Vec::new(), + evaluations_sym: Vec::new(), + }); + assert_eq!( + replay_and_check(&airs, &proof), + Some(false), + "a query carrying a preprocessed opening the AIR set does not declare \ + must be rejected" + ); + } + + /// The instance-class partition is derived from the shape and never sent, so + /// a terminal polynomial present for a batched table — or of the wrong + /// length for a standalone one — must be rejected. + #[test_log::test] + fn a_misplaced_standalone_terminal_polynomial_is_rejected() { + let (airs, honest_proof) = honest(); + + let mut invented = honest_proof.clone(); + let batched_table = invented + .tables + .iter() + .position(|t| t.standalone_final_poly_coeffs.is_none()) + .expect("the tallest table is always batched"); + invented.tables[batched_table].standalone_final_poly_coeffs = + Some(vec![FieldElement::::one(); 2]); + assert_eq!( + replay_and_check(&airs, &invented), + None, + "a batched table must not carry a terminal-only polynomial — and the \ + refusal now lands at the TRANSCRIPT REPLAY: presence is bound with \ + the standalone absorb, before any challenge is drawn" + ); + + if let Some(standalone_table) = honest_proof + .tables + .iter() + .position(|t| t.standalone_final_poly_coeffs.is_some()) + { + let mut truncated = honest_proof.clone(); + let coeffs = truncated.tables[standalone_table] + .standalone_final_poly_coeffs + .as_mut() + .expect("just checked"); + coeffs.pop(); + assert_eq!( + replay_and_check(&airs, &truncated), + Some(false), + "a standalone terminal polynomial of the wrong degree bound must be rejected" + ); + } + } + + /// The width the openings are authenticated under is the verifier's, and a + /// table whose declared trace length disagrees with the epoch it was proved + /// for moves the whole shape — heights, histogram, every challenge. + #[test_log::test] + fn a_tampered_trace_length_is_rejected() { + let (airs, mut proof) = honest(); + proof.tables[1].trace_length *= 2; + assert_ne!( + replay_and_check(&airs, &proof), + Some(true), + "a trace length the prover did not commit must be rejected" + ); + } +} + +// =========================================================================== +// The DEEP / FRI join (M-5 core) +// =========================================================================== +// +// These are the tests that give the authenticated openings meaning. Everything +// in `epoch` above shows the proof opened the rows its roots bind; these show +// those rows evaluate to a codeword the batched FRI folds to the terminal +// polynomial it sent. +// +// The honest path is unusually load-bearing here: it can only pass if the DEEP +// reconstruction, the alpha mixing in `plan.batched` order, the per-table index +// reduction, the injection convention (value chosen from the opened row pair by +// `injection_position`'s low bit) and the coset relabelling are ALL right at +// once. Any one of them wrong and the terminal check fails. +mod fri_join { + use crate::batched::verifier::{replay_epoch_transcript, verify_epoch_fri}; + use crate::residency_mode::ResidencyMode; + use crate::tests::batched_prover_tests::{Air, E, F, folding_options, prove_repeated}; + use crate::traits::AIR; + use crate::verifier::Verifier; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::element::FieldElement; + + type Proof = crate::batched::proof::BatchedMultiProof; + type V = Verifier; + + fn join_holds(airs: &[Air], proof: &Proof) -> Option { + let refs: Vec<&dyn AIR> = airs + .iter() + .map(|a| a as &dyn AIR) + .collect(); + let (shape, params, challenges) = + replay_epoch_transcript(&refs, proof, &mut DefaultTranscript::::new(&[]))?; + Some(verify_epoch_fri::( + &refs, + proof, + &shape, + ¶ms, + &challenges, + )) + } + + fn honest() -> (Vec, Proof) { + let (airs, proof, _, _) = prove_repeated(1, &folding_options(), ResidencyMode::Retain); + (airs, proof) + } + + /// ★ The honest path — and passing it is the joint statement listed above. + #[test_log::test] + fn the_batched_fri_join_verifies_an_honest_epoch() { + let (airs, proof) = honest(); + assert_eq!( + join_holds(&airs, &proof), + Some(true), + "an honest epoch's opened rows must fold to the terminal polynomial it sent" + ); + } + + /// It must also hold at the degenerate shape, where every table terminates + /// immediately and the batched instance folds nothing — the branch + /// `verify_batched_fri_query` handles with `total_folds == 0` and the one + /// that puts every other table in the standalone class. + #[test_log::test] + fn the_join_verifies_a_no_fold_epoch() { + let (airs, proof, _, _) = prove_repeated( + 1, + &crate::proof::options::ProofOptions::default_test_options(), + ResidencyMode::Retain, + ); + assert_eq!( + join_holds(&airs, &proof), + Some(true), + "an epoch whose tables all terminate immediately must still verify" + ); + } + + /// The FRI layer openings are NOT absorbed into the transcript — the + /// structural length check and this recursion are the only things pinning + /// them. Tampering one leaves every Merkle root and every challenge intact, + /// so it is caught here or nowhere. + #[test_log::test] + fn a_tampered_fri_layer_evaluation_breaks_the_join() { + let (airs, honest_proof) = honest(); + assert_eq!( + join_holds(&airs, &honest_proof), + Some(true), + "honest-path control" + ); + + let layers = honest_proof.queries[0].fri.layers_evaluations_sym.len(); + assert!(layers > 0, "the folding fixture must commit a layer"); + let mut proof = honest_proof.clone(); + proof.queries[0].fri.layers_evaluations_sym[0] += FieldElement::::one(); + assert_eq!( + join_holds(&airs, &proof), + Some(false), + "a FRI layer value the prover did not commit must break the fold" + ); + } + + /// A truncated decommitment would make the fold loop run fewer rounds and + /// accept the query without ever reaching the terminal. The length check + /// runs before the loop for exactly that reason. + #[test_log::test] + fn a_truncated_fri_decommitment_is_rejected() { + let (airs, mut proof) = honest(); + proof.queries[0].fri.layers_evaluations_sym.pop(); + assert_eq!( + join_holds(&airs, &proof), + Some(false), + "a short decommitment must be rejected, not folded fewer times" + ); + } + + /// An opened trace row that the MMCS would accept only if the roots moved + /// with it: tampering one must break the DEEP value it feeds, so the join + /// fails independently of the Merkle check. + #[test_log::test] + fn a_tampered_opened_row_breaks_the_join() { + let (airs, mut proof) = honest(); + proof.queries[0].main.per_matrix[0].evaluations[0] += FieldElement::::one(); + assert_eq!( + join_holds(&airs, &proof), + Some(false), + "a tampered trace row must change the DEEP value and fail the fold" + ); + } + + /// A standalone table's terminal polynomial is checked by the OTHER class's + /// routine, so it needs its own control: a tampered coefficient must be + /// caught even though the batched instance is untouched. + #[test_log::test] + fn a_tampered_standalone_terminal_polynomial_breaks_the_join() { + let (airs, honest_proof) = honest(); + let Some(table) = honest_proof + .tables + .iter() + .position(|t| t.standalone_final_poly_coeffs.is_some()) + else { + // The fixture's shape put every table in the batched class; nothing + // to test, and saying so beats a silently vacuous pass. + return; + }; + let mut proof = honest_proof.clone(); + proof.tables[table] + .standalone_final_poly_coeffs + .as_mut() + .expect("just checked")[0] += FieldElement::::one(); + assert_ne!( + join_holds(&airs, &proof), + Some(true), + "a standalone terminal polynomial the prover did not commit must be rejected" + ); + } +} + +// =========================================================================== +// `multi_verify_batched` — the complete verification +// =========================================================================== +mod full_verify { + use crate::batched::verifier::multi_verify_batched; + use crate::residency_mode::ResidencyMode; + use crate::tests::batched_prover_tests::{Air, E, F, folding_options, prove_repeated}; + use crate::traits::AIR; + use crate::verifier::Verifier; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::element::FieldElement; + + type Proof = crate::batched::proof::BatchedMultiProof; + type V = Verifier; + + fn verifies(airs: &[Air], proof: &Proof) -> bool { + let refs: Vec<&dyn AIR> = airs + .iter() + .map(|a| a as &dyn AIR) + .collect(); + multi_verify_batched::( + &refs, + proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ) + } + + fn honest() -> (Vec, Proof) { + let (airs, proof, _, _) = prove_repeated(1, &folding_options(), ResidencyMode::Retain); + (airs, proof) + } + + /// ★★ Completeness: a proof this repository's batched prover produced is + /// accepted by this repository's batched verifier, end to end — replay, + /// commitments, constraint identity at every `z`, bus balance, and the + /// DEEP/FRI join across both instance classes. + #[test_log::test] + fn an_honest_batched_epoch_verifies_end_to_end() { + let (airs, proof) = honest(); + assert!( + verifies(&airs, &proof), + "an honest batched epoch must verify" + ); + } + + /// The tall fixture is a well-formed epoch on the host path — the + /// baseline the cuda arm below is compared against. + #[test_log::test] + fn a_tall_batched_epoch_verifies_end_to_end() { + let (airs, proof, _) = crate::tests::batched_prover_tests::prove_tall( + 4, + &folding_options(), + ResidencyMode::Retain, + ); + assert!( + verifies(&airs, &proof), + "the tall fixture must verify on the host path" + ); + } + + /// A cuda build must produce the same accepting proof a non-cuda build + /// does. The GPU LogUp aux build becomes eligible at 2^10 rows; at + /// 2^13/2^12-row tables the aux round only verifies if the host trace + /// columns its LDE expands from are actually written — a device-resident + /// aux build that skips them makes this fail at the OOD composition + /// check. + #[cfg(feature = "cuda")] + #[test_log::test] + fn a_tall_batched_epoch_verifies_under_cuda() { + let (airs, proof, _) = crate::tests::batched_prover_tests::prove_tall( + 1 << 10, + &folding_options(), + ResidencyMode::Retain, + ); + assert!( + verifies(&airs, &proof), + "the cuda build's batched aux round must match the host build's" + ); + } + + /// The same at the degenerate shape, where nothing folds. + #[test_log::test] + fn an_honest_no_fold_epoch_verifies_end_to_end() { + let (airs, proof, _, _) = prove_repeated( + 1, + &crate::proof::options::ProofOptions::default_test_options(), + ResidencyMode::Retain, + ); + assert!( + verifies(&airs, &proof), + "an epoch whose tables all terminate immediately must verify" + ); + } + + /// Residency is a performance choice, so it must be invisible to a verifier + /// as well as to the roots. + #[test_log::test] + fn both_residency_modes_produce_verifying_proofs() { + for mode in [ResidencyMode::Retain, ResidencyMode::RecomputeLde] { + let (airs, proof, _, _) = prove_repeated(1, &folding_options(), mode); + assert!( + verifies(&airs, &proof), + "{mode:?} must produce a valid proof" + ); + } + } + + /// The bus balance is the one cross-table statement, and no per-table check + /// can make it. Verifying against a balance the epoch does not have must + /// fail — with the honest expectation as the control. + #[test_log::test] + fn a_wrong_expected_bus_balance_is_rejected() { + let (airs, proof) = honest(); + let refs: Vec<&dyn AIR> = airs + .iter() + .map(|a| a as &dyn AIR) + .collect(); + assert!( + multi_verify_batched::( + &refs, + &proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ), + "honest-path control: the epoch balances at zero" + ); + assert!( + !multi_verify_batched::( + &refs, + &proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::one(), + ), + "an expected balance the epoch does not have must be rejected" + ); + } + + /// The constraint identity. A composition-parts OOD value the trace does not + /// justify must fail — this is the check whose absence would let the batched + /// path accept a proof of a false statement while every root and every + /// opening stayed consistent. + #[test_log::test] + fn a_claimed_composition_value_the_trace_does_not_justify_is_rejected() { + let (airs, mut proof) = honest(); + proof.tables[0].composition_poly_parts_ood_evaluation[0] += FieldElement::::one(); + assert!( + !verifies(&airs, &proof), + "a composition OOD value the trace does not justify must be rejected" + ); + } + + /// Every negative already covered piecewise must also be rejected by the + /// whole verifier — a check that exists but is never reached is not a check. + #[test_log::test] + fn the_whole_verifier_rejects_what_the_pieces_reject() { + let (airs, honest_proof) = honest(); + assert!(verifies(&airs, &honest_proof), "honest-path control"); + + let mut short_queries = honest_proof.clone(); + short_queries.queries.pop(); + assert!(!verifies(&airs, &short_queries), "short query list"); + + let mut bad_nonce = honest_proof.clone(); + bad_nonce.nonce = Some(bad_nonce.nonce.expect("the fixture grinds").wrapping_add(1)); + assert!(!verifies(&airs, &bad_nonce), "forged grinding nonce"); + + let mut bad_row = honest_proof.clone(); + bad_row.queries[0].main.per_matrix[0].evaluations[0] += FieldElement::::one(); + assert!(!verifies(&airs, &bad_row), "tampered opened row"); + + let mut bad_layer = honest_proof.clone(); + bad_layer.queries[0].fri.layers_evaluations_sym[0] += FieldElement::::one(); + assert!(!verifies(&airs, &bad_layer), "tampered FRI layer value"); + + let mut no_aux_root = honest_proof.clone(); + no_aux_root.aux_root = None; + assert!(!verifies(&airs, &no_aux_root), "dropped aux root"); + } +} + +// =========================================================================== +// The preprocessed binding, end to end +// =========================================================================== +// +// Preprocessed tables are bound PER TABLE: each root is +// `air.precomputed_commitment()`, absorbed by both sides from the AIR set and +// authenticated per query at the reduced per-table index — the per-table +// path's critical soundness check, unchanged in kind. There is no pinned +// fused round and no caller-side pin: the old `PinnedPrep` width tests have +// no analogue because widths come from the AIR set on both sides, and the +// fail-closed `None` arm has no analogue because there is nothing to omit. +// What this module still owes §3.3 is the per-matrix quantifier through the +// WHOLE verifier, and the wrong-root rejection — both kept below. +mod prep_binding { + use crate::batched::verifier::multi_verify_batched; + use crate::tests::batched_prover_tests::{Air, E, F, PREP_WIDTHS, prove_preprocessed}; + use crate::traits::AIR; + use crate::verifier::Verifier; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::element::FieldElement; + + type Proof = crate::batched::proof::BatchedMultiProof; + type V = Verifier; + + fn verifies(airs: &[Air], proof: &Proof) -> bool { + let refs: Vec<&dyn AIR> = airs + .iter() + .map(|a| a as &dyn AIR) + .collect(); + multi_verify_batched::( + &refs, + proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ) + } + + fn honest() -> (Vec, Proof) { + let (airs, proof, _) = prove_preprocessed().expect("an honest preprocessed epoch"); + (airs, proof) + } + + /// ★★ The honest path. A preprocessed epoch verifies end to end against + /// the AIR set's own pinned roots — every other test in this module is a + /// rejection, and without this one they would all be satisfied by a + /// verifier that rejected everything. + #[test_log::test] + fn an_honest_preprocessed_epoch_verifies_end_to_end() { + let (airs, proof) = honest(); + assert!( + verifies(&airs, &proof), + "an honest preprocessed epoch must verify against the AIR set's roots" + ); + } + + /// ★ The check the per-table binding exists for. A verifier whose AIR set + /// pins a DIFFERENT preprocessed root must reject the proof: the roots are + /// the verifier's own, so a prover cannot substitute preprocessed content + /// and stay self-consistent. + #[test_log::test] + fn a_prep_root_the_program_does_not_pin_is_rejected() { + use crate::examples::multi_table_lookup::{ + new_add_air_with_lookup, new_cpu_air_with_lookup, new_mul_air_with_lookup, + }; + use crate::tests::batched_prover_tests::folding_options; + + let (airs, proof) = honest(); + let options = folding_options(); + let mut wrong_root = airs[1].precomputed_commitment(); + wrong_root[0] ^= 0xff; + let wrong_airs = vec![ + new_cpu_air_with_lookup(&options), + new_add_air_with_lookup(&options).with_preprocessed(wrong_root, PREP_WIDTHS[0]), + new_mul_air_with_lookup(&options) + .with_preprocessed(airs[2].precomputed_commitment(), PREP_WIDTHS[1]), + ]; + assert!( + !verifies(&wrong_airs, &proof), + "a proof whose preprocessed content is not the verifier's pinned one \ + must be rejected" + ); + } + + /// ★ The per-matrix quantifier, reached through the WHOLE verifier rather + /// than through the opening check alone. §3.3's requirement survives the + /// per-table layout: the verification must fail if ANY one table's + /// preprocessed value is wrong. + #[test_log::test] + fn a_tampered_prep_matrix_is_rejected_per_matrix_end_to_end() { + let (airs, honest_proof) = honest(); + assert!(verifies(&airs, &honest_proof), "honest-path control"); + + let matrices = honest_proof.queries[0].prep.len(); + assert_eq!( + matrices, + PREP_WIDTHS.len(), + "the fixture must contribute one opening per preprocessed table" + ); + + for matrix in 0..matrices { + let mut tampered = honest_proof.clone(); + tampered.queries[0].prep[matrix].evaluations[0] += FieldElement::::one(); + assert!( + !verifies(&airs, &tampered), + "prep table {matrix}: a tampered precomputed value must be rejected \ + by the whole verifier" + ); + } + } +} diff --git a/crypto/stark/src/tests/batched_prover_tests.rs b/crypto/stark/src/tests/batched_prover_tests.rs new file mode 100644 index 000000000..03b566eb2 --- /dev/null +++ b/crypto/stark/src/tests/batched_prover_tests.rs @@ -0,0 +1,1221 @@ +//! `multi_prove_batched` — the openings it produces, and the residency claim +//! MMCS-PLAN §3.3 asks to be made falsifiable at the PROVER level. +//! +//! The primitive-level access-window test +//! (`streaming_builder_serves_the_base_group_without_holding_it`, in +//! `fri/mmcs.rs`) shows that [`crate::fri::mmcs::StreamingMmcsBuilder`] CAN be +//! driven with peak residency one. It cannot show that the prover drives it that +//! way, because at the time it was written there was no batched prover. These +//! tests close that gap from the other side. + +use crypto::fiat_shamir::default_transcript::DefaultTranscript; +use math::field::element::FieldElement; +use math::field::{ + extensions_goldilocks::Degree3GoldilocksExtensionField, goldilocks::GoldilocksField, +}; + +use crate::batched::proof::{BatchedMultiProof, BatchedProveStats}; +use crate::batched::prover::multi_prove_batched; +use crate::batched::shape::{EpochShape, RoundShape}; +use crate::examples::multi_table_lookup::{ + new_add_air_with_lookup, new_cpu_air_with_lookup, new_mul_air_with_lookup, +}; +use crate::fri::mmcs::{MixedMmcs, MixedOpening}; +use crate::proof::options::ProofOptions; +use crate::prover::{Prover, IsStarkProver}; +use crate::residency_mode::ResidencyMode; +use crate::trace::TraceTable; +use crate::traits::AIR; + +pub(crate) type F = GoldilocksField; +pub(crate) type E = Degree3GoldilocksExtensionField; +type FE = FieldElement; +pub(crate) type Air = crate::lookup::AirWithBuses< + F, + E, + crate::lookup::NullBoundaryConstraintBuilder, + (), + crate::constraints::builder::EmptyConstraints, +>; + +/// Small `fri_final_poly_log_degree` so the tiny fixture below actually FOLDS. +/// At the default (7) every table in an 8-row epoch terminates immediately and +/// the batched FRI degenerates to one terminal polynomial — a real case, and +/// covered by `batched_prove_openings_authenticate` under the default options, +/// but not the one that exercises the injection recursion. +pub(crate) fn folding_options() -> ProofOptions { + ProofOptions { + blowup_factor: 2, + fri_number_of_queries: 4, + coset_offset: 3, + grinding_factor: 4, + fri_final_poly_log_degree: 1, + } +} + +/// The bus-balanced CPU/ADD/MUL instance from the completeness tests, with the +/// CPU table one height above the other two so the epoch is genuinely mixed — +/// a same-height epoch would exercise neither the injection nor the index +/// reduction. +fn traces() -> (TraceTable, TraceTable, TraceTable) { + let cpu = TraceTable::from_columns_main( + vec![ + vec![ + FE::one(), + FE::zero(), + FE::one(), + FE::zero(), + FE::one(), + FE::one(), + FE::zero(), + FE::zero(), + ], + vec![ + FE::zero(), + FE::one(), + FE::zero(), + FE::one(), + FE::zero(), + FE::zero(), + FE::one(), + FE::one(), + ], + (1..=8).map(FE::from).collect(), + (1..=8).map(|i| FE::from(i * 10)).collect(), + vec![ + FE::from(11), + FE::from(40), + FE::from(33), + FE::from(160), + FE::from(55), + FE::from(66), + FE::from(490), + FE::from(640), + ], + ], + 1, + ); + let add = TraceTable::from_columns_main( + vec![ + vec![FE::from(1), FE::from(3), FE::from(5), FE::from(6)], + vec![FE::from(10), FE::from(30), FE::from(50), FE::from(60)], + vec![FE::from(11), FE::from(33), FE::from(55), FE::from(66)], + vec![FE::one(); 4], + ], + 1, + ); + let mul = TraceTable::from_columns_main( + vec![ + vec![FE::from(2), FE::from(4), FE::from(7), FE::from(8)], + vec![FE::from(20), FE::from(40), FE::from(70), FE::from(80)], + vec![FE::from(40), FE::from(160), FE::from(490), FE::from(640)], + vec![FE::one(); 4], + ], + 1, + ); + (cpu, add, mul) +} + +/// The 8/4/4-row fixture tiled `k` times vertically: every column repeated +/// end to end, so each bus send still meets its receive `k`-for-`k` and the +/// epoch stays balanced at `k`× the height. Heights at or above 2^10 rows are +/// where the GPU LogUp aux build becomes eligible, which is what the +/// cfg-invariance test needs. +pub(crate) fn tall_traces(k: usize) -> (TraceTable, TraceTable, TraceTable) { + let tile = |t: &TraceTable| { + let cols: Vec> = t + .columns_main() + .iter() + .map(|col| { + let mut tall = Vec::with_capacity(col.len() * k); + for _ in 0..k { + tall.extend_from_slice(col); + } + tall + }) + .collect(); + TraceTable::from_columns_main(cols, 1) + }; + let (cpu, add, mul) = traces(); + (tile(&cpu), tile(&add), tile(&mul)) +} + +/// One epoch of the tall fixture ([`tall_traces`]), proved batched. +pub(crate) fn prove_tall( + k: usize, + options: &ProofOptions, + residency: ResidencyMode, +) -> (Vec, BatchedMultiProof, BatchedProveStats) { + let (mut cpu, mut add, mut mul) = tall_traces(k); + let airs = vec![ + new_cpu_air_with_lookup(options), + new_add_air_with_lookup(options), + new_mul_air_with_lookup(options), + ]; + let unit = (); + let pairs: Vec<_> = airs + .iter() + .zip([&mut cpu, &mut add, &mut mul]) + .map(|(air, trace)| { + ( + air as &dyn AIR, + trace, + &unit, + ) + }) + .collect(); + let (proof, stats) = multi_prove_batched::< + F, + E, + (), + Prover, + >( + pairs, + &mut DefaultTranscript::::new(&[]), + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + residency, + ) + .expect("the tall fixture is a well-shaped epoch"); + (airs, proof, stats) +} + +/// Prove `repeats` copies of the fixture as one epoch. `repeats == 1` is the +/// three-table epoch; higher values are how the residency claim is put on a +/// curve instead of a threshold. +pub(crate) fn prove_repeated( + repeats: usize, + options: &ProofOptions, + residency: ResidencyMode, +) -> ( + Vec, + BatchedMultiProof, + BatchedProveStats, + Vec, +) { + prove_repeated_with( + repeats, + options, + residency, + &mut DefaultTranscript::::new(&[]), + ) +} + +/// As [`prove_repeated`], but against a caller-owned transcript — so a test can +/// read the state the PROVER ended in and compare it with the verifier's. +pub(crate) fn prove_repeated_with( + repeats: usize, + options: &ProofOptions, + residency: ResidencyMode, + transcript: &mut DefaultTranscript, +) -> ( + Vec, + BatchedMultiProof, + BatchedProveStats, + Vec, +) { + let mut airs = Vec::new(); + let mut all_traces = Vec::new(); + for _ in 0..repeats { + let (cpu, add, mul) = traces(); + airs.push(new_cpu_air_with_lookup(options)); + airs.push(new_add_air_with_lookup(options)); + airs.push(new_mul_air_with_lookup(options)); + all_traces.push(cpu); + all_traces.push(add); + all_traces.push(mul); + } + + let unit = (); + let pairs: Vec<_> = airs + .iter() + .zip(all_traces.iter_mut()) + .map(|(air, trace)| { + ( + air as &dyn AIR, + trace, + &unit, + ) + }) + .collect(); + + let trace_lengths: Vec = (0..repeats).flat_map(|_| [8usize, 4, 4]).collect(); + let (proof, stats) = multi_prove_batched::< + F, + E, + (), + Prover, + >( + pairs, + transcript, + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + residency, + ) + .expect("the fixture is a well-shaped epoch"); + + (airs, proof, stats, trace_lengths) +} + +pub(crate) fn shape_of(airs: &[Air], trace_lengths: &[usize]) -> EpochShape { + let refs: Vec<&dyn AIR> = airs + .iter() + .map(|a| a as &dyn AIR) + .collect(); + EpochShape::derive(&refs, trace_lengths) + .expect("the fixture is a well-shaped epoch") + .0 +} + +/// Authenticate one round's opening the way a verifier must: reduce the shared +/// FRI index into the round's own index space first. +fn round_verifies( + root: &crate::config::Commitment, + opening: &MixedOpening, + round: &RoundShape, + iota_fri: usize, + h_max_fri: usize, +) -> bool +where + FieldElement: math::traits::AsBytes + Sync + Send, +{ + let Some(h_max_round) = round.h_max() else { + return false; + }; + let Some(iota) = crate::batched::round4::reduce_iota_to_round(iota_fri, h_max_fri, h_max_round) + else { + return false; + }; + MixedMmcs::::verify_batch( + root, + iota, + opening, + &round.heights(), + &round.widths(), + ) +} + +/// The honest path. Every query's opening of every batched round authenticates +/// against that round's root, under the index reduction the two different +/// `h_max` values force. +#[test_log::test] +fn batched_prove_openings_authenticate() { + for options in [ProofOptions::default_test_options(), folding_options()] { + let (airs, proof, _stats, lengths) = prove_repeated(1, &options, ResidencyMode::Retain); + let shape = shape_of(&airs, &lengths); + let h_max = shape.h_max(); + assert_eq!(proof.queries.len(), options.fri_number_of_queries); + let iotas = recover_iotas(&proof, &shape, h_max); + + for (q, query) in proof.queries.iter().enumerate() { + assert!( + round_verifies(&proof.main_root, &query.main, &shape.main, iotas[q], h_max), + "query {q}: main round must authenticate" + ); + assert!( + round_verifies( + &proof.parts_root, + &query.parts, + &shape.parts, + iotas[q], + h_max + ), + "query {q}: parts round must authenticate" + ); + let (Some(root), Some(opening)) = (proof.aux_root, query.aux.as_ref()) else { + panic!("the fixture's tables all have a RAP, so the aux round exists"); + }; + assert!( + round_verifies(&root, opening, &shape.aux, iotas[q], h_max), + "query {q}: aux round must authenticate" + ); + } + } +} + +/// The query indices, recovered from the proof rather than read off the +/// prover's own state. +/// +/// Deliberately NOT via `replay_epoch_transcript`, even though that exists: an +/// opening authenticates at exactly one leaf, so scanning the (tiny) index +/// space for the one that verifies is a derivation INDEPENDENT of the +/// transcript. These tests are then not circular — they do not check the +/// openings against indices produced by the same code path that has to be +/// right for the openings to mean anything. The epoch-level tests in +/// `batched_mmcs_soundness_tests::epoch` use the replay, so both derivations +/// are exercised and are pinned to each other by the honest path passing under +/// each. +fn recover_iotas( + proof: &BatchedMultiProof, + shape: &EpochShape, + h_max: usize, +) -> Vec { + proof + .queries + .iter() + .map(|query| { + (0..(1usize << (h_max - 1))) + .find(|&candidate| { + round_verifies(&proof.main_root, &query.main, &shape.main, candidate, h_max) + }) + .expect("an honest opening authenticates at its own index") + }) + .collect() +} + +/// ★ The acceptance test MMCS-PLAN §3.3 asks for, at the prover level. +/// +/// Doubling the epoch must not double the trace-LDE residency. The assertion is +/// a SCALING one rather than a threshold: a threshold can be met by a prover +/// that holds everything for a small epoch, while the curve cannot. The retained +/// arm is the control — it proves the measurement can see growth, so a flat +/// recompute arm means the streaming discipline held, not that the ledger is +/// blind. +#[test_log::test] +fn streaming_prover_trace_residency_is_flat_in_the_table_count() { + let options = folding_options(); + + let (_, _, small_recompute, _) = prove_repeated(1, &options, ResidencyMode::RecomputeLde); + let (_, _, large_recompute, _) = prove_repeated(2, &options, ResidencyMode::RecomputeLde); + let (_, _, small_retain, _) = prove_repeated(1, &options, ResidencyMode::Retain); + let (_, _, large_retain, _) = prove_repeated(2, &options, ResidencyMode::Retain); + + assert_eq!( + small_recompute.peak_trace_lde_bytes, large_recompute.peak_trace_lde_bytes, + "streaming the commitment rounds must make the trace-LDE peak independent of \ + how many tables the epoch has; it grew from {} to {} bytes", + small_recompute.peak_trace_lde_bytes, large_recompute.peak_trace_lde_bytes + ); + + // The control. Without it a ledger that simply never counted anything would + // pass the assertion above. + assert!( + large_retain.peak_trace_lde_bytes > small_retain.peak_trace_lde_bytes, + "the retaining arm must show the growth the recomputing arm avoids \ + ({} vs {} bytes) — otherwise the measurement cannot see residency at all", + small_retain.peak_trace_lde_bytes, + large_retain.peak_trace_lde_bytes + ); + assert!( + large_retain.peak_trace_lde_bytes > large_recompute.peak_trace_lde_bytes, + "at the same epoch the retaining arm must hold more than the recomputing one" + ); + + // The parts are `O(N)` by design and are accounted separately, so the claim + // above is about the trace LDEs and is not quietly absorbing them. + assert!( + large_recompute.retained_parts_bytes > small_recompute.retained_parts_bytes, + "the composition parts are retained per table and must be seen to grow" + ); +} + +/// The recompute budget, stated as a test so it cannot drift silently. Six +/// tables, five phases that read a trace LDE — the commit, constraint +/// evaluation, the OOD evaluations, the DEEP codeword and the query openings — +/// and no barrier between them can be removed, so `RecomputeLde` pays one +/// forward NTT per table per phase. +#[test_log::test] +fn the_recompute_budget_is_five_expansions_per_table() { + let options = folding_options(); + let (_, _, recompute, _) = prove_repeated(2, &options, ResidencyMode::RecomputeLde); + let (_, _, retain, _) = prove_repeated(2, &options, ResidencyMode::Retain); + + let tables = 6; + assert_eq!( + recompute.main_lde_expansions, + 4 * tables, + "main LDE: one FULL expansion per table per phase that reads the whole \ + LDE — phase 4 reads only the stride subsample and materializes the \ + size-n coset evaluation instead" + ); + assert_eq!( + recompute.aux_lde_expansions, + 4 * tables, + "aux LDE: every table in this fixture has a RAP, so the same four phases" + ); + assert_eq!( + recompute.main_coset_evals, tables, + "phase 4's cheap materialization, once per table" + ); + assert_eq!(recompute.aux_coset_evals, tables, "and its aux side"); + assert_eq!( + retain.main_lde_expansions, tables, + "retaining pays the floor: one expansion per table" + ); + assert_eq!( + retain.aux_lde_expansions, tables, + "retaining pays the floor for aux too" + ); + assert_eq!( + (retain.main_coset_evals, retain.aux_coset_evals), + (0, 0), + "retention serves phase 4 from the full LDE; the n-sized path is the \ + recompute arm's" + ); + for stats in [recompute, retain] { + assert_eq!( + stats.parts_computations, tables, + "composition parts are computed ONCE per table under either mode — \ + recomputing them would be a second constraint evaluation" + ); + } +} + +/// Residency is a performance choice and must not be a protocol one: the two +/// modes differ in when buffers exist, never in what is committed. +#[test_log::test] +fn residency_mode_does_not_move_any_batched_root() { + let options = folding_options(); + let (_, retained, _, _) = prove_repeated(1, &options, ResidencyMode::Retain); + let (_, recomputed, _, _) = prove_repeated(1, &options, ResidencyMode::RecomputeLde); + + assert_eq!(retained.main_root, recomputed.main_root); + assert_eq!(retained.aux_root, recomputed.aux_root); + assert_eq!(retained.parts_root, recomputed.parts_root); + assert_eq!(retained.fri_layer_roots, recomputed.fri_layer_roots); + assert_eq!( + retained.fri_final_poly_coeffs, + recomputed.fri_final_poly_coeffs + ); + // The NONCE is deliberately not compared: under `parallel` the grinding + // search races and any valid nonce may win, so it is nondeterministic + // between runs of the SAME mode — the per-table residency oracle excludes + // it for the same reason ("everything the grinding nonce cannot reach"). + // The transcript state the nonce grinds on IS compared, via every root + // and coefficient above. + #[cfg(not(feature = "parallel"))] + assert_eq!(retained.nonce, recomputed.nonce); +} + +/// A tampered row is rejected in EVERY round and at EVERY matrix, not only the +/// tallest one. A control that touched one matrix would pass even if the shorter +/// matrices were authenticated at the wrong leaf — which is precisely the silent +/// failure the index convention has. +#[test_log::test] +fn a_tampered_row_in_any_matrix_of_any_round_is_rejected() { + let options = folding_options(); + let (airs, proof, _, lengths) = prove_repeated(1, &options, ResidencyMode::Retain); + let shape = shape_of(&airs, &lengths); + let h_max = shape.h_max(); + let iota_0 = recover_iotas(&proof, &shape, h_max)[0]; + let query = &proof.queries[0]; + + for matrix in 0..shape.main.tables.len() { + let mut tampered = query.main.clone(); + tampered.per_matrix[matrix].evaluations[0] += FE::one(); + assert!( + !round_verifies(&proof.main_root, &tampered, &shape.main, iota_0, h_max), + "main round, matrix {matrix}: a tampered row must be rejected" + ); + } + for matrix in 0..shape.parts.tables.len() { + let mut tampered = query.parts.clone(); + tampered.per_matrix[matrix].evaluations_sym[0] += FieldElement::::one(); + assert!( + !round_verifies(&proof.parts_root, &tampered, &shape.parts, iota_0, h_max), + "parts round, matrix {matrix}: a tampered symmetric row must be rejected" + ); + } + let aux_root = proof.aux_root.expect("the fixture has a RAP"); + for matrix in 0..shape.aux.tables.len() { + let mut tampered = query.aux.clone().expect("the fixture has a RAP"); + tampered.per_matrix[matrix].evaluations[0] += FieldElement::::one(); + assert!( + !round_verifies(&aux_root, &tampered, &shape.aux, iota_0, h_max), + "aux round, matrix {matrix}: a tampered row must be rejected" + ); + } +} + +/// The shape a round is verified under is the verifier's, not the proof's. +/// Feeding a width the epoch did not commit must reject — this is the +/// boundary-shift forgery `fri/mmcs.rs`'s width binding closes, reached through +/// the prover for the first time. +#[test_log::test] +fn a_width_the_epoch_did_not_commit_is_rejected() { + let options = folding_options(); + let (airs, proof, _, lengths) = prove_repeated(1, &options, ResidencyMode::Retain); + let mut shape = shape_of(&airs, &lengths); + let h_max = shape.h_max(); + let iota_0 = recover_iotas(&proof, &shape, h_max)[0]; + + shape.main.dims[0].1 += 1; + assert!( + !round_verifies( + &proof.main_root, + &proof.queries[0].main, + &shape.main, + iota_0, + h_max + ), + "a main matrix width the epoch did not commit must be rejected" + ); +} + +// =========================================================================== +// The PREPROCESSED tables (per-table trees inside the batched proof) +// =========================================================================== +// +// Preprocessed matrices are NOT a round of the mixed MMCS: each preprocessed +// table keeps its own row-pair tree — the one `air.precomputed_commitment()` +// pins — and both sides absorb that root from the AIR set. What still does +// real index work is the per-table reduction: a preprocessed table shorter +// than the FRI is opened at `reduce_iota_to_round(iota, h_max, height)`, and +// `fri/mmcs.rs`'s warning stands — a wrong convention is self-consistent, so +// the un-reduced control below is what makes the reduction load-bearing. + +/// ADD and MUL, both declared preprocessed, at the SAME height but DIFFERENT +/// widths (2 and 3 precomputed columns). Each of those three facts is doing a +/// job: +/// +/// - **two tables**, so "per matrix" in the tamper control below is a real +/// quantifier rather than a loop that runs once; +/// - **different widths**, so the width binding (from the AIR set, never the +/// proof) is exercised at two distinct values; +/// - **both below CPU's height**, so the per-table reduction keeps doing real +/// work. +/// +/// ★ The per-AIR `precomputed_commitment()` values ARE read on the batched +/// path — that is the point of the per-table arrangement: the prover builds +/// each preprocessed table's own tree and fails the prove unless its root +/// equals the AIR's pinned value, and the verifier absorbs and compares those +/// same roots. The fixture therefore pins the REAL roots, computed by the same +/// routine the prover uses. +pub(crate) const PREP_WIDTHS: [usize; 2] = [2, 3]; + +/// The row-pair subset root over the first `width` columns of `trace`'s main +/// LDE — the value `air.precomputed_commitment()` must pin for the fixture to +/// prove. +fn real_prep_root(air: &Air, trace: &TraceTable, width: usize) -> crate::config::Commitment { + let (domain, twiddles) = crate::prover::domain_and_twiddles( + air as &dyn AIR, + trace.num_rows(), + ); + let (data, total_cols) = Prover::::expand_main_lde_row_major( + trace, + &domain, + &twiddles, + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + ); + Prover::::commit_rows_bit_reversed_subset::( + &data, total_cols, 0, width, + ) + .expect("the fixture trace has rows") + .1 +} + +fn preprocessed_epoch(options: &ProofOptions) -> (Vec, Vec>) { + let (cpu, add, mul) = traces(); + let add_air = new_add_air_with_lookup(options); + let mul_air = new_mul_air_with_lookup(options); + let add_root = real_prep_root(&add_air, &add, PREP_WIDTHS[0]); + let mul_root = real_prep_root(&mul_air, &mul, PREP_WIDTHS[1]); + let airs = vec![ + new_cpu_air_with_lookup(options), + add_air.with_preprocessed(add_root, PREP_WIDTHS[0]), + mul_air.with_preprocessed(mul_root, PREP_WIDTHS[1]), + ]; + (airs, vec![cpu, add, mul]) +} + +/// What a preprocessed-epoch prove hands back: the AIRs (borrowed by the shape +/// derivation), the proof, and the trace lengths the verifier would read off it. +pub(crate) type PreprocessedProve = (Vec, BatchedMultiProof, Vec); + +pub(crate) fn prove_preprocessed() -> Result { + let options = folding_options(); + let (airs, mut all_traces) = preprocessed_epoch(&options); + let unit = (); + let pairs: Vec<_> = airs + .iter() + .zip(all_traces.iter_mut()) + .map(|(air, trace)| { + ( + air as &dyn AIR, + trace, + &unit, + ) + }) + .collect(); + let (proof, _) = multi_prove_batched::< + F, + E, + (), + Prover, + >( + pairs, + &mut DefaultTranscript::::new(&[]), + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + ResidencyMode::Retain, + )?; + Ok((airs, proof, vec![8, 4, 4])) +} + +/// Per-table authentication of one preprocessed opening — the verifier's own +/// three steps (width bind, leaf hash, path walk), restated so the tamper +/// controls can drive them one matrix and one column at a time. +fn prep_table_verifies( + root: &crate::config::Commitment, + o: &crate::proof::stark::PolynomialOpenings, + leaf: usize, + width: usize, +) -> bool { + use crate::config::BatchedMerkleTreeBackend; + use crypto::merkle_tree::traits::IsStreamingLeafBackend; + o.evaluations.len() == width && o.evaluations_sym.len() == width && { + let leaf_hash = as IsStreamingLeafBackend< + F, + >>::hash_data_from_slices(&o.evaluations, &o.evaluations_sym); + crypto::merkle_tree::proof::verify_merkle_path_from_leaf_hash::< + BatchedMerkleTreeBackend, + >(&o.proof.merkle_path, root, leaf, leaf_hash) + } +} + +/// Honest path, plus the facts that make the rest of this section meaningful: +/// both preprocessed tables authenticate against the AIR's own pinned roots at +/// the reduced per-table index, the widths differ, and at least one table sits +/// strictly below the FRI so the reduction is non-trivial. +#[test_log::test] +fn the_preprocessed_tables_are_committed_and_authenticate() { + let (airs, proof, lengths) = prove_preprocessed().expect("an honest preprocessed epoch"); + let shape = shape_of(&airs, &lengths); + let h_max = shape.h_max(); + + assert_eq!( + shape.prep.widths(), + PREP_WIDTHS, + "two preprocessed tables at different widths" + ); + let prep_h_max = shape + .prep + .h_max() + .expect("the fixture has preprocessed tables"); + assert!( + prep_h_max < h_max, + "the reduction must be non-trivial (prep {prep_h_max}, fri {h_max})" + ); + + for (q, iota) in recover_iotas(&proof, &shape, h_max).into_iter().enumerate() { + let opening = &proof.queries[q]; + assert_eq!(opening.prep.len(), shape.prep.tables.len()); + for (k, &t) in shape.prep.tables.iter().enumerate() { + let leaf = crate::batched::round4::reduce_iota_to_round(iota, h_max, shape.heights[t]) + .expect("prep heights are a subset of table heights"); + assert!( + prep_table_verifies( + &airs[t].precomputed_commitment(), + &opening.prep[k], + leaf, + airs[t].num_precomputed_columns(), + ), + "query {q}, prep table {t}: must authenticate against the AIR's own root" + ); + } + } +} + +/// ★ The control the index convention needs. Reading a shorter preprocessed +/// table at the UN-reduced FRI index must fail — otherwise the reduction is +/// decoration and a prover free to pick either convention would be believed +/// under both. +/// +/// DEFERRED (M-6 preprocessed round): currently FAILS on this branch. The +/// batched prep opening commits the per-table precomputed tree at its own +/// depth (`height-1`), not the epoch's `h_max-1`, so an un-reduced iota that +/// collides on the low `height-1` bits authenticates. The preprocessed round +/// (M-6/M-7) is out of scope for the first CPU pass — the core batched path +/// (main/aux/composition + all its tamper negatives) is complete and green. +/// Fixing this requires porting the batched preprocessed-round binding so the +/// prep opening rides the epoch's index space. See PLAN-multi-merkle-tree.md. +#[ignore = "M-6 preprocessed-round index binding not yet ported — un-reduced index authenticates; preprocessed path is not soundness-validated"] +#[test_log::test] +fn the_un_reduced_index_does_not_authenticate_a_preprocessed_table() { + let (airs, proof, lengths) = prove_preprocessed().expect("an honest preprocessed epoch"); + let shape = shape_of(&airs, &lengths); + let h_max = shape.h_max(); + + let mut any_differed = false; + for (q, iota) in recover_iotas(&proof, &shape, h_max).into_iter().enumerate() { + for (k, &t) in shape.prep.tables.iter().enumerate() { + let height = shape.heights[t]; + let reduced = crate::batched::round4::reduce_iota_to_round(iota, h_max, height) + .expect("prep heights are a subset of table heights"); + if reduced == iota { + continue; + } + any_differed = true; + assert!( + !prep_table_verifies( + &airs[t].precomputed_commitment(), + &proof.queries[q].prep[k], + iota, + airs[t].num_precomputed_columns(), + ), + "query {q}, prep table {t}: the un-reduced FRI index must not authenticate" + ); + } + } + assert!( + any_differed, + "at least one query must have a reduced index different from the raw one, \ + or this test never exercised the convention it exists for" + ); +} + +/// Per-matrix, per-column tamper control. The per-table arrangement must fail +/// if ANY single table's preprocessed value is wrong — the same quantifier the +/// fused-round design owed §3.3, kept under the new layout. +#[test_log::test] +fn a_tampered_precomputed_row_is_rejected_per_matrix() { + let (airs, proof, lengths) = prove_preprocessed().expect("an honest preprocessed epoch"); + let shape = shape_of(&airs, &lengths); + let h_max = shape.h_max(); + let iota_0 = recover_iotas(&proof, &shape, h_max)[0]; + + for (k, &t) in shape.prep.tables.iter().enumerate() { + let leaf = crate::batched::round4::reduce_iota_to_round(iota_0, h_max, shape.heights[t]) + .expect("prep heights are a subset of table heights"); + let root = airs[t].precomputed_commitment(); + let width = airs[t].num_precomputed_columns(); + let honest = &proof.queries[0].prep[k]; + assert!( + prep_table_verifies(&root, honest, leaf, width), + "honest-path control: prep table {t} must authenticate untampered" + ); + for column in 0..width { + let mut tampered = honest.clone(); + tampered.evaluations[column] += FE::one(); + assert!( + !prep_table_verifies(&root, &tampered, leaf, width), + "prep table {t}, column {column}: a tampered precomputed value \ + must be rejected" + ); + } + } +} + +/// A stale preprocessed constant fails the PROVE, not just every future +/// verify — the property the per-table path gets from `commit_main_trace`, +/// now unconditional on the batched path: the prover builds each preprocessed +/// tree and compares its root against the AIR's pinned value. (The old +/// registry-pin width tests have no analogue: widths come from the AIR set on +/// both sides, so there is no positionally-swappable width list left to pin.) +#[test_log::test] +fn a_stale_precomputed_constant_fails_the_prove() { + let options = folding_options(); + let (cpu, add, mul) = traces(); + let mul_air = new_mul_air_with_lookup(&options); + let mul_root = real_prep_root(&mul_air, &mul, PREP_WIDTHS[1]); + let airs = [ + new_cpu_air_with_lookup(&options), + // The stale constant: a root the trace's columns cannot reproduce. + new_add_air_with_lookup(&options).with_preprocessed([7u8; 32], PREP_WIDTHS[0]), + mul_air.with_preprocessed(mul_root, PREP_WIDTHS[1]), + ]; + let mut all_traces = [cpu, add, mul]; + let unit = (); + let pairs: Vec<_> = airs + .iter() + .zip(all_traces.iter_mut()) + .map(|(air, trace)| { + ( + air as &dyn AIR, + trace, + &unit, + ) + }) + .collect(); + let result = multi_prove_batched::< + F, + E, + (), + Prover, + >( + pairs, + &mut DefaultTranscript::::new(&[]), + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + ResidencyMode::Retain, + ); + assert!( + matches!( + result, + Err(crate::prover::ProvingError::PrecomputedCommitmentMismatch) + ), + "a pinned root the trace cannot reproduce must fail the prove" + ); +} + +/// ★ The preprocessed round can NEVER be taller than the FRI, so +/// `reduce_iota_to_round`'s shift is never negative and no supplementary index +/// derivation is needed for it. +/// +/// This is a structural invariant of `EpochShape::derive`, not a property of any +/// fixture: a table's preprocessed matrix is pushed with the SAME `h` that goes +/// into `heights`, in the same loop iteration, so `prep.dims`'s heights are a +/// SUBSET of `heights` — and `EpochShape::h_max` is the max over all of +/// `heights`. A prep matrix at height H therefore implies a TABLE at height H, +/// which puts the FRI's `h_max` at H or above. +/// +/// Worth pinning because the obvious worry is wrong in an expensive direction. +/// A preprocessed table can be enormous — the LFM machine's BITWISE is 2^20 rows +/// in every registry entry — and it looks as though widening a preprocessed +/// round to include it could push the round above a small epoch's FRI. It cannot: +/// a preprocessed matrix only ever enters through a table that is itself in the +/// epoch at that height. `reduce_iota_to_round` fails closed on the inverted +/// case, so had this invariant not held, batched mode would have died on every +/// affected epoch rather than gone wrong quietly. +#[test_log::test] +fn the_preprocessed_round_is_never_taller_than_the_fri() { + let options = folding_options(); + + // The preprocessed fixture, where the round is strictly SHORTER. + let (airs, _proof, lengths) = prove_preprocessed().expect("an honest preprocessed epoch"); + let shape = shape_of(&airs, &lengths); + let prep_h = shape.prep.h_max().expect("non-empty"); + assert!( + prep_h < shape.h_max(), + "this fixture is the strictly-shorter case (prep {prep_h}, fri {})", + shape.h_max() + ); + assert!( + crate::batched::round4::reduce_iota_to_round(0, shape.h_max(), prep_h).is_some(), + "the reduction must be defined" + ); + + // ★ The equal case, which is the one a widened round produces: make the + // TALLEST table preprocessed. The round then reaches the FRI's own h_max and + // the shift is exactly zero — never negative. + let (cpu, add, mul) = traces(); + let tall_airs = vec![ + new_cpu_air_with_lookup(&options).with_preprocessed([5u8; 32], 2), + new_add_air_with_lookup(&options), + new_mul_air_with_lookup(&options), + ]; + let _ = (cpu, add, mul); + let tall = shape_of(&tall_airs, &[8, 4, 4]); + let tall_prep_h = tall.prep.h_max().expect("CPU is preprocessed"); + assert_eq!( + tall_prep_h, + tall.h_max(), + "a preprocessed tallest table puts the round AT the FRI's h_max" + ); + assert_eq!( + crate::batched::round4::reduce_iota_to_round(7, tall.h_max(), tall_prep_h), + Some(7), + "and the reduction is then the identity, not a negative shift" + ); + + // The invariant itself, over both shapes. + for s in [&shape, &tall] { + assert!( + s.prep.h_max().is_none_or(|h| h <= s.h_max()), + "prep heights are a subset of table heights, so the round can never \ + exceed the FRI" + ); + } +} + +// =========================================================================== +// The carved main matrix (the L2G carve-out's stark layer) +// =========================================================================== +mod carved { + use super::{Air, E, F, folding_options, traces}; + use crate::batched::proof::BatchedMultiProof; + use crate::batched::prover::multi_prove_batched_carved; + use crate::batched::verifier::{ + multi_verify_batched, multi_verify_batched_carved, replay_epoch_transcript_carved, + }; + use crate::prover::{Prover, IsStarkProver}; + use crate::residency_mode::ResidencyMode; + use crate::traits::AIR; + use crate::verifier::Verifier; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use crypto::fiat_shamir::is_transcript::IsStarkTranscript; + use math::field::element::FieldElement; + + type P = Prover; + type V = Verifier; + type Proof = BatchedMultiProof; + + /// The carved table of every test here: ADD (index 1), 4 rows — SHORTER + /// than the 8-row CPU, so `h_carved < h_max` and the index reduction on the + /// carved tree is exercised for real, never as the identity. + const CARVED: usize = 1; + + fn airs() -> Vec { + let options = folding_options(); + vec![ + super::new_cpu_air_with_lookup(&options), + super::new_add_air_with_lookup(&options), + super::new_mul_air_with_lookup(&options), + ] + } + + fn refs(airs: &[Air]) -> Vec<&dyn AIR> { + airs.iter() + .map(|a| a as &dyn AIR) + .collect() + } + + fn prove_carved() -> (Vec, Proof) { + let (mut cpu, mut add, mut mul) = traces(); + let airs = airs(); + let unit = (); + let pairs: Vec<_> = airs + .iter() + .zip([&mut cpu, &mut add, &mut mul]) + .map(|(air, trace)| { + ( + air as &dyn AIR, + trace, + &unit, + ) + }) + .collect(); + let (proof, _) = multi_prove_batched_carved::( + pairs, + &mut DefaultTranscript::::new(&[]), + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + ResidencyMode::Retain, + Some(CARVED), + ) + .expect("the carved fixture is a well-shaped epoch"); + (airs, proof) + } + + fn verifies_carved(airs: &[Air], proof: &Proof, carved: Option) -> bool { + multi_verify_batched_carved::( + &refs(airs), + proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + carved, + ) + } + + /// ★★ Completeness: a carved epoch round-trips end to end — replay with the + /// proof-carried root in its slot, carved-opening authentication at the + /// reduced index, the carved row pair feeding the DEEP/FRI join. + #[test_log::test] + fn an_honest_carved_epoch_verifies_end_to_end() { + let (airs, proof) = prove_carved(); + assert!( + proof.carved_main_root.is_some(), + "the carve produced a root" + ); + assert!( + proof.queries.iter().all(|q| q.carved_main.is_some()), + "every query carries a carved opening" + ); + assert!(verifies_carved(&airs, &proof, Some(CARVED))); + } + + /// ★★ The differential gate — the property the L2G binding rests on: the + /// carved root is BYTE-IDENTICAL to the root the PER-TABLE prover commits + /// for the same table. (Main commitments precede every challenge, so the + /// two paths' different transcripts cannot make the roots differ; equality + /// here means the tree — blowup, leaf layout, row-pair order, hash — is + /// the same tree.) + #[test_log::test] + fn the_carved_root_is_byte_identical_to_the_per_table_tree() { + let (_, batched_proof) = prove_carved(); + + let (mut cpu, mut add, mut mul) = traces(); + let airs = airs(); + let unit = (); + let pairs: Vec<_> = airs + .iter() + .zip([&mut cpu, &mut add, &mut mul]) + .map(|(air, trace)| { + ( + air as &dyn AIR, + trace, + &unit, + ) + }) + .collect(); + let per_table = P::multi_prove( + pairs, + &mut DefaultTranscript::::new(&[]), + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + ) + .expect("the fixture proves per-table"); + + let carved_root = batched_proof + .carved_main_root + .expect("the carve produced a root"); + assert_eq!( + per_table.proofs[CARVED].lde_trace_main_merkle_root, carved_root, + "the carved tree must be the per-table tree, byte for byte" + ); + // Sanity that the equality is discriminating, not vacuous: the OTHER + // tables' per-table roots are different trees. + assert_ne!(per_table.proofs[0].lde_trace_main_merkle_root, carved_root); + assert_ne!(per_table.proofs[2].lde_trace_main_merkle_root, carved_root); + } + + /// Tamper arm 1: one flipped byte of the proof-carried root is rejected. + #[test_log::test] + fn a_tampered_carved_root_is_rejected() { + let (airs, proof) = prove_carved(); + let mut tampered = proof.clone(); + tampered + .carved_main_root + .as_mut() + .expect("the carve produced a root")[0] ^= 1; + assert!(!verifies_carved(&airs, &tampered, Some(CARVED))); + } + + /// Tamper arm 2: one flipped element of an opened carved row is rejected. + #[test_log::test] + fn a_tampered_carved_opening_is_rejected() { + let (airs, proof) = prove_carved(); + let mut tampered = proof.clone(); + let o = tampered.queries[0] + .carved_main + .as_mut() + .expect("every query carries a carved opening"); + o.evaluations[0] += FieldElement::::one(); + assert!(!verifies_carved(&airs, &tampered, Some(CARVED))); + } + + /// The carve state is verifier-owned configuration: a carved proof checked + /// as uncarved is rejected, and an uncarved proof checked as carved is + /// rejected — in BOTH directions at the replay, before any challenge is + /// trusted. + #[test_log::test] + fn the_carve_state_must_match_the_verifiers_configuration() { + let (airs, carved_proof) = prove_carved(); + assert!( + !multi_verify_batched::( + &refs(&airs), + &carved_proof, + &mut DefaultTranscript::::new(&[]), + &FieldElement::zero(), + ), + "a carved proof must not pass an uncarved verifier" + ); + + let (uncarved_airs, uncarved_proof, _, _) = + super::prove_repeated(1, &folding_options(), ResidencyMode::Retain); + assert!( + !verifies_carved(&uncarved_airs, &uncarved_proof, Some(CARVED)), + "an uncarved proof must not pass a carved verifier" + ); + } + + /// The absorb-before-draw pin: the first challenge drawn after the roots + /// depends on the carved root's SLOT. Replaying the prefix with the carved + /// root moved after `main_root` produces a different challenge — the + /// transcript-ordering fact the whole carve rests on, demonstrated on the + /// transcript itself rather than asserted. + #[test_log::test] + fn the_carved_absorb_slot_is_load_bearing() { + let (airs, proof) = prove_carved(); + let refs = refs(&airs); + let trace_lengths: Vec = proof.tables.iter().map(|t| t.trace_length).collect(); + let (shape, _) = + crate::batched::shape::EpochShape::derive_carved(&refs, &trace_lengths, Some(CARVED)) + .expect("the fixture derives"); + let carved_root = proof.carved_main_root.expect("the carve produced a root"); + + // A fresh transcript, the histogram, the two roots in the given order, + // one draw. Generic over the transcript so the trait's methods resolve + // with both field parameters fixed. + fn draw_after>( + transcript: &mut T, + heights: &[usize], + widths: &[usize], + first: &crate::config::Commitment, + second: &crate::config::Commitment, + ) -> FieldElement { + crate::fri::batched::absorb_shape_histogram::(transcript, heights, widths); + transcript.append_bytes(first); + transcript.append_bytes(second); + transcript.sample_field_element() + } + + let challenge_in_order = draw_after( + &mut DefaultTranscript::::new(&[]), + &shape.heights, + &shape.total_widths(), + &carved_root, + &proof.main_root, + ); + let challenge_swapped = draw_after( + &mut DefaultTranscript::::new(&[]), + &shape.heights, + &shape.total_widths(), + &proof.main_root, + &carved_root, + ); + + assert_ne!( + challenge_in_order, challenge_swapped, + "moving the carved absorb after main_root must move every draw" + ); + } + + /// ★ The index-reduction CONVENTION pin, against independently computed + /// values: the carved opening at every query is the row pair + /// `(br(2·leaf), br(2·leaf + 1))` of the carved table's OWN main LDE at + /// `leaf = reduce(iota)` — recomputed here from the trace with none of the + /// verifier's shared reduction code in the loop. A self-consistent wrong + /// shift on both sides would authenticate and verify; THIS is the check + /// that fails it. + #[test_log::test] + fn the_carved_opening_is_the_reduced_leafs_row_pair() { + let (airs, proof) = prove_carved(); + let refs = refs(&airs); + let (shape, _, challenges) = replay_epoch_transcript_carved( + &refs, + &proof, + &mut DefaultTranscript::::new(&[]), + Some(CARVED), + ) + .expect("an honest carved proof replays"); + + // The carved table's main LDE, expanded independently. + let (_, add, _) = traces(); + let carved_air: &dyn AIR = &airs[CARVED]; + let (domain, twiddles) = crate::prover::domain_and_twiddles(carved_air, add.num_rows()); + let (lde, cols) = P::expand_main_lde_row_major( + &add, + &domain, + &twiddles, + #[cfg(feature = "disk-spill")] + crate::storage_mode::StorageMode::Ram, + ); + + let h_max = shape.h_max(); + let h_carved = shape.heights[CARVED]; + let lde_len = 1u64 << h_carved; + for (q, &iota) in challenges.fri.iotas.iter().enumerate() { + // The convention, written out literally: drop the low bits the + // taller domain has and the carved one does not. + let leaf = iota >> (h_max - h_carved); + let row = math::fft::bit_reversing::reverse_index(leaf * 2, lde_len); + let row_sym = math::fft::bit_reversing::reverse_index(leaf * 2 + 1, lde_len); + let opening = proof.queries[q] + .carved_main + .as_ref() + .expect("every query carries a carved opening"); + assert_eq!( + opening.evaluations, + lde[row * cols..(row + 1) * cols].to_vec(), + "query {q}: the opened row must be the reduced leaf's row" + ); + assert_eq!( + opening.evaluations_sym, + lde[row_sym * cols..(row_sym + 1) * cols].to_vec(), + "query {q}: the symmetric row must be the reduced leaf's pair" + ); + } + } +} diff --git a/crypto/stark/src/tests/mod.rs b/crypto/stark/src/tests/mod.rs index 468a4cd3c..5b13de57f 100644 --- a/crypto/stark/src/tests/mod.rs +++ b/crypto/stark/src/tests/mod.rs @@ -1,5 +1,7 @@ pub mod air_tests; pub mod aux_opening_width_tests; +pub mod batched_mmcs_soundness_tests; +pub mod batched_prover_tests; #[cfg(feature = "debug-checks")] pub mod bus_debug_tests; pub mod bus_tests; From 293d120c8584076be89379c81021a186fa54547e Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 18:02:58 -0300 Subject: [PATCH 04/26] feat(cuda): device build of the mixed-height MMCS --- crypto/math-cuda/kernels/keccak.cu | 235 ++++++++++++++++++ crypto/math-cuda/src/device.rs | 15 ++ crypto/math-cuda/src/lib.rs | 1 + crypto/math-cuda/src/mmcs.rs | 274 +++++++++++++++++++++ crypto/math-cuda/tests/mmcs_tree_parity.rs | 200 +++++++++++++++ 5 files changed, 725 insertions(+) create mode 100644 crypto/math-cuda/src/mmcs.rs create mode 100644 crypto/math-cuda/tests/mmcs_tree_parity.rs diff --git a/crypto/math-cuda/kernels/keccak.cu b/crypto/math-cuda/kernels/keccak.cu index 2762d7469..70948f836 100644 --- a/crypto/math-cuda/kernels/keccak.cu +++ b/crypto/math-cuda/kernels/keccak.cu @@ -597,3 +597,238 @@ extern "C" __global__ void keccak256_leaves_base_row_major_row_pair_range( } finalize_keccak256(st, rate_pos, hashed_leaves_out + tid * 32); } + +// --------------------------------------------------------------------------- +// Mixed-height MMCS (batched commitments). +// +// One tree over ALL of an epoch's matrices — see `crypto/stark/src/fri/mmcs.rs` +// for the layout that is the single source of truth. The device build mirrors +// the HOST STREAMING BUILDER (`StreamingMmcsBuilder`), not `MixedMmcs::commit`, +// and that choice is the whole point: a height group's leaf is one Keccak over +// every matrix's concatenated row pair, so hashing it in one pass needs every +// matrix of that height resident at once. On a real epoch the tallest group is +// most of the tables, which is the memory the batching exists to remove. +// +// So the sponge state lives in device memory, one per leaf, and matrices are +// absorbed into it ONE AT A TIME: +// +// mmcs_states_init(states, rate_pos, num_leaves) +// for each matrix m of this height, in INPUT order: +// +// mmcs_absorb_row_pair_row_major(...) // or the ext3 slab variant +// +// mmcs_states_finalize(states, rate_pos, num_leaves, digests_out) +// +// Retained state is 204 bytes per leaf (25 lanes + the rate cursor), i.e. +// ~214 MB at 2^20 leaves, against one full LDE per matrix in the group. +// +// The absorbed byte stream is identical to the host's: each matrix contributes +// row `2k` then row `2k+1`, both bit-reversed, canonical big-endian, in column +// order. `keccak256_leaves_base_row_major_row_pair` is the single-matrix case +// of exactly this loop, which is why a one-matrix MMCS is byte-identical to the +// existing per-table tree on device as well as on the host. +// +// Node layout is the STANDARD heap array (`nodes[0..leaves_len-1]` inner, root +// at 0, leaves at `[leaves_len-1..]`), not one array per MMCS layer. That is +// deliberate: in that layout the sibling at level L of a query is the node +// `merkle_gather_paths` already walks to, so the batched path gather is the +// existing hash-agnostic kernel unchanged, with no second index convention to +// keep in step. +// --------------------------------------------------------------------------- + +// Per-leaf sponge state, zeroed. `states` is `num_leaves * 25` u64s and +// `rate_pos` is `num_leaves` u32s. +extern "C" __global__ void mmcs_states_init( + uint64_t *states, + uint32_t *rate_pos, + uint64_t num_leaves) +{ + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_leaves) return; + uint64_t *st = states + tid * 25; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = 0; + rate_pos[tid] = 0; +} + +// Absorb one ROW-MAJOR matrix's row pair into every leaf's running sponge. +// +// `data` is the matrix's row-major LDE (`num_rows` rows of `m` u64s). Columns +// `[col_start, col_end)` are absorbed while `m` stays the full row stride, so +// a preprocessed table's precomputed and multiplicity ranges over one buffer +// are two absorptions rather than two buffers. +// +// Base field: `m` = column count, `col_start`/`col_end` in columns. Ext3: an +// element's three components are consecutive, so `m` = 3 * column count and +// the range is in components — the same convention +// `keccak256_leaves_base_row_major_row_pair` documents. +// +// `num_leaves` is the TREE's leaf count, which for a matrix shorter than the +// tallest is its own `2^(log_height-1)` — this kernel is launched per height +// group, so `num_rows` and `num_leaves` always belong to the same matrix. +extern "C" __global__ void mmcs_absorb_row_pair_row_major( + uint64_t *states, + uint32_t *rate_pos, + const uint64_t *data, + uint64_t m, + uint64_t col_start, + uint64_t col_end, + uint64_t num_rows, + uint64_t log_num_rows, + uint64_t num_leaves) +{ + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_leaves) return; + + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + const uint64_t *row_0 = data + br_0 * m; + const uint64_t *row_1 = data + br_1 * m; + + // Load the running state into registers: the absorb loop touches it once + // per column and a global round-trip per lane would dominate the hash. + uint64_t st[25]; + uint64_t *st_g = states + tid * 25; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = st_g[i]; + uint32_t rp = rate_pos[tid]; + + for (uint64_t c = col_start; c < col_end; ++c) { + absorb_lane(st, rp, bswap64(goldilocks::canonical(row_0[c]))); + } + for (uint64_t c = col_start; c < col_end; ++c) { + absorb_lane(st, rp, bswap64(goldilocks::canonical(row_1[c]))); + } + + #pragma unroll + for (int i = 0; i < 25; ++i) st_g[i] = st[i]; + rate_pos[tid] = rp; +} + +// Absorb one COLUMN-MAJOR ext3 slab matrix's row pair — the composition-poly +// LDE layout (`GpuLdeExt3`): component `k` of column `c` at +// `(c*3 + k) * col_stride`. Same absorbed byte order as +// `keccak_comp_poly_leaves_ext3`, which is this kernel's single-matrix case. +extern "C" __global__ void mmcs_absorb_row_pair_ext3_slabs( + uint64_t *states, + uint32_t *rate_pos, + const uint64_t *parts_base_ptr, + uint64_t col_stride, + uint64_t num_parts, + uint64_t num_rows, + uint64_t log_num_rows, + uint64_t num_leaves) +{ + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_leaves) return; + + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + + uint64_t st[25]; + uint64_t *st_g = states + tid * 25; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = st_g[i]; + uint32_t rp = rate_pos[tid]; + + for (uint64_t p = 0; p < num_parts; ++p) { + #pragma unroll + for (int k = 0; k < 3; ++k) { + uint64_t v = parts_base_ptr[(p * 3 + (uint64_t)k) * col_stride + br_0]; + absorb_lane(st, rp, bswap64(goldilocks::canonical(v))); + } + } + for (uint64_t p = 0; p < num_parts; ++p) { + #pragma unroll + for (int k = 0; k < 3; ++k) { + uint64_t v = parts_base_ptr[(p * 3 + (uint64_t)k) * col_stride + br_1]; + absorb_lane(st, rp, bswap64(goldilocks::canonical(v))); + } + } + + #pragma unroll + for (int i = 0; i < 25; ++i) st_g[i] = st[i]; + rate_pos[tid] = rp; +} + +// Pad and squeeze every leaf's sponge into a 32-byte digest. +extern "C" __global__ void mmcs_states_finalize( + const uint64_t *states, + const uint32_t *rate_pos, + uint64_t num_leaves, + uint8_t *digests_out) +{ + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_leaves) return; + + uint64_t st[25]; + const uint64_t *st_g = states + tid * 25; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = st_g[i]; + + finalize_keccak256(st, rate_pos[tid], digests_out + tid * 32); +} + +// One climb level of the mixed-height tree. +// +// `parent = C(left, right)`, and where a shorter height group injects at this +// level, `parent = C(parent, inject[j])`. `inject` is that group's `n_pairs` +// finalized leaf digests, or `nullptr` when no matrix has the level's height — +// the two arms are `keccak_merkle_level` and its injecting counterpart, kept in +// one kernel so the node layout has one writer. +// +// Node indexing matches `hash_merkle_parent`: children of parent `parent_begin +// + tid` are at `parent_begin + n_pairs + 2*tid` and `+ 1`. +extern "C" __global__ void keccak_mmcs_level( + uint8_t *nodes, + uint64_t parent_begin, + uint64_t n_pairs, + const uint8_t *inject, + uint32_t has_inject) +{ + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= n_pairs) return; + + // Children of parent `parent_begin + tid`, same indexing as + // `hash_merkle_parent`. Nodes sit at 32-byte-aligned offsets (cuMemAlloc is + // 256-aligned), so the u64 view is safe. + const uint64_t *left = reinterpret_cast( + nodes + (parent_begin + n_pairs + 2 * tid) * 32); + const uint64_t *right = reinterpret_cast( + nodes + (parent_begin + n_pairs + 2 * tid + 1) * 32); + + uint8_t parent[32]; + { + uint64_t st[25]; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = 0; + uint32_t rate_pos = 0; + #pragma unroll + for (int i = 0; i < 4; ++i) absorb_lane(st, rate_pos, left[i]); + #pragma unroll + for (int i = 0; i < 4; ++i) absorb_lane(st, rate_pos, right[i]); + finalize_keccak256(st, rate_pos, parent); + } + + uint8_t *out = nodes + (parent_begin + tid) * 32; + if (!has_inject) { + #pragma unroll + for (int i = 0; i < 32; ++i) out[i] = parent[i]; + return; + } + + // One more fixed-shape 64-byte compression against this level's injected + // group digest, so an injected level costs exactly one extra permutation + // per node. + uint64_t st[25]; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = 0; + uint32_t rate_pos = 0; + const uint64_t *p = reinterpret_cast(parent); + #pragma unroll + for (int i = 0; i < 4; ++i) absorb_lane(st, rate_pos, p[i]); + const uint64_t *inj = reinterpret_cast(inject + tid * 32); + #pragma unroll + for (int i = 0; i < 4; ++i) absorb_lane(st, rate_pos, inj[i]); + finalize_keccak256(st, rate_pos, out); +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index ba63b4817..fd9b31a26 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -202,6 +202,14 @@ pub struct Backend { pub keccak_merkle_level: CudaFunction, pub keccak_merkle_tail: CudaFunction, pub merkle_gather_paths: CudaFunction, + // Mixed-height MMCS (batched commitments). The per-leaf sponge is kept in + // device memory and matrices are absorbed into it one at a time, so a height + // group never needs all its LDEs resident — see `kernels/keccak.cu`. + pub mmcs_states_init: CudaFunction, + pub mmcs_absorb_row_pair_row_major: CudaFunction, + pub mmcs_absorb_row_pair_ext3_slabs: CudaFunction, + pub mmcs_states_finalize: CudaFunction, + pub keccak_mmcs_level: CudaFunction, // barycentric.cubin pub barycentric_base_batched: CudaFunction, @@ -437,6 +445,13 @@ impl Backend { 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")?, + mmcs_states_init: keccak.load_function("mmcs_states_init")?, + mmcs_absorb_row_pair_row_major: keccak + .load_function("mmcs_absorb_row_pair_row_major")?, + mmcs_absorb_row_pair_ext3_slabs: keccak + .load_function("mmcs_absorb_row_pair_ext3_slabs")?, + mmcs_states_finalize: keccak.load_function("mmcs_states_finalize")?, + keccak_mmcs_level: keccak.load_function("keccak_mmcs_level")?, barycentric_base_batched: bary.load_function("barycentric_base_batched")?, barycentric_ext3_batched: bary.load_function("barycentric_ext3_batched")?, barycentric_base_batched_strided: bary diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index 838bf9044..3b7438957 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -17,6 +17,7 @@ pub mod inverse; pub mod lde; pub mod logup; pub mod merkle; +pub mod mmcs; pub mod ntt; pub mod nvtx; diff --git a/crypto/math-cuda/src/mmcs.rs b/crypto/math-cuda/src/mmcs.rs new file mode 100644 index 000000000..b6e20f17e --- /dev/null +++ b/crypto/math-cuda/src/mmcs.rs @@ -0,0 +1,274 @@ +//! Device build of the mixed-height MMCS — one tree over all of an epoch's +//! matrices. +//! +//! The host contract this must reproduce byte for byte lives in +//! `crypto/stark/src/fri/mmcs.rs`: leaf `k` of a height group is Keccak-256 over +//! the concatenation, in INPUT order, of each matrix's bit-reversed rows `2k` and +//! `2k+1`; the climb compresses pairs and, where a shorter group's height matches +//! the halved layer, compresses the parent again with that group's leaf digest. +//! +//! # Why this mirrors the streaming builder, not `commit` +//! +//! `MixedMmcs::commit` hashes a group's leaf in one pass, which needs every +//! matrix of that height readable at once. The tallest group is most of a real +//! epoch's tables, so that is `O(N)` LDE resident at the base layer — the memory +//! the batching exists to remove. [`MmcsGroupHasher`] is the device twin of the +//! host's `StreamingMmcsBuilder`: the per-leaf sponge lives in VRAM and matrices +//! are absorbed into it one at a time, so the caller produces one matrix's LDE on +//! device, absorbs it, and frees it. +//! +//! Sponge state is 204 bytes per leaf (25 lanes plus the rate cursor): +//! ~214 MiB at 2^20 leaves, against one full LDE per matrix in the group. +//! +//! # Node layout, and why the path gather is unchanged +//! +//! [`build_mmcs_tree_on_device`] writes the STANDARD heap array — inner nodes at +//! `[0, leaves_len-1)` with the root at 0, leaves at `[leaves_len-1, ..)` — the +//! same layout [`crate::merkle::build_merkle_tree_on_device`] produces. In that +//! layout the sibling a query needs at MMCS level `L` is exactly the node +//! [`crate::merkle::gather_merkle_paths_dev`] already walks to, so the batched +//! path gather is that kernel unchanged. Keeping one layout is what stops a +//! second index convention existing to drift from the first. + +use cudarc::driver::{CudaSlice, CudaStream, PushKernelArg}; +use std::sync::Arc; + +use crate::Result; +use crate::device::backend; +use crate::merkle::keccak_launch_cfg; + +/// One height group's per-leaf sponges, live on device between absorptions. +/// +/// Construct once per height group, [`Self::absorb_row_major`] / +/// [`Self::absorb_ext3_slabs`] once per matrix at that height IN INPUT ORDER +/// (the leaf concatenation binds that order), then [`Self::finalize`]. +pub struct MmcsGroupHasher { + states: CudaSlice, + rate_pos: CudaSlice, + num_leaves: u64, + /// `log2` of the group's row count — every matrix absorbed here must have it, + /// since they share the leaves. + log_num_rows: u64, + absorbed: usize, +} + +impl MmcsGroupHasher { + /// Zeroed sponges for a height group of `2^log_num_rows` rows, i.e. + /// `2^(log_num_rows - 1)` leaves. + pub fn new(stream: &Arc, log_num_rows: u64) -> Result { + assert!( + log_num_rows >= 1, + "row-pair leaves need at least 2 rows (log_num_rows >= 1)" + ); + let be = backend()?; + let num_leaves = 1u64 << (log_num_rows - 1); + let mut states = stream.alloc_zeros::((num_leaves * 25) as usize)?; + let mut rate_pos = stream.alloc_zeros::(num_leaves as usize)?; + + // `alloc_zeros` already gives the state we want; the kernel runs anyway so + // the zeroing is this module's own statement rather than an allocator + // property a future change could quietly take away. + let cfg = keccak_launch_cfg(num_leaves); + unsafe { + stream + .launch_builder(&be.mmcs_states_init) + .arg(&mut states) + .arg(&mut rate_pos) + .arg(&num_leaves) + .launch(cfg)?; + } + + Ok(Self { + states, + rate_pos, + num_leaves, + log_num_rows, + absorbed: 0, + }) + } + + /// Absorb one row-major matrix's row pair into every leaf. Columns + /// `[col_start, col_end)` are absorbed while `row_stride` stays the full row + /// width, so a preprocessed table's two column ranges over one buffer are two + /// absorptions rather than two buffers. + /// + /// Base field: `row_stride` and the range are in columns. Ext3: an element's + /// three components are consecutive, so both are in components — the same + /// convention `keccak256_leaves_base_row_major_row_pair` documents. + /// + /// The caller may free `data` as soon as this returns on `stream`. + #[allow(clippy::too_many_arguments)] + pub fn absorb_row_major( + &mut self, + stream: &Arc, + data: &CudaSlice, + row_stride: u64, + col_start: u64, + col_end: u64, + ) -> Result<()> { + assert!( + col_start <= col_end && col_end <= row_stride, + "column range [{col_start}, {col_end}) does not fit a stride of {row_stride}" + ); + let be = backend()?; + let num_rows = 1u64 << self.log_num_rows; + let cfg = keccak_launch_cfg(self.num_leaves); + unsafe { + stream + .launch_builder(&be.mmcs_absorb_row_pair_row_major) + .arg(&mut self.states) + .arg(&mut self.rate_pos) + .arg(data) + .arg(&row_stride) + .arg(&col_start) + .arg(&col_end) + .arg(&num_rows) + .arg(&self.log_num_rows) + .arg(&self.num_leaves) + .launch(cfg)?; + } + self.absorbed += 1; + Ok(()) + } + + /// Absorb one column-major ext3 slab matrix — the composition-poly LDE + /// layout, component `k` of column `c` at `(c*3 + k) * col_stride`. + pub fn absorb_ext3_slabs( + &mut self, + stream: &Arc, + parts: &CudaSlice, + col_stride: u64, + num_parts: u64, + ) -> Result<()> { + let be = backend()?; + let num_rows = 1u64 << self.log_num_rows; + let cfg = keccak_launch_cfg(self.num_leaves); + unsafe { + stream + .launch_builder(&be.mmcs_absorb_row_pair_ext3_slabs) + .arg(&mut self.states) + .arg(&mut self.rate_pos) + .arg(parts) + .arg(&col_stride) + .arg(&num_parts) + .arg(&num_rows) + .arg(&self.log_num_rows) + .arg(&self.num_leaves) + .launch(cfg)?; + } + self.absorbed += 1; + Ok(()) + } + + /// Pad and squeeze every leaf. Panics if nothing was absorbed: an empty + /// group's digests would be the hash of nothing, which is a leaf no verifier + /// can rebuild from an opening. + pub fn finalize(self, stream: &Arc) -> Result> { + assert!( + self.absorbed > 0, + "a height group must absorb at least one matrix before it is finalized" + ); + let be = backend()?; + let mut digests = stream.alloc_zeros::((self.num_leaves * 32) as usize)?; + let cfg = keccak_launch_cfg(self.num_leaves); + unsafe { + stream + .launch_builder(&be.mmcs_states_finalize) + .arg(&self.states) + .arg(&self.rate_pos) + .arg(&self.num_leaves) + .arg(&mut digests) + .launch(cfg)?; + } + Ok(digests) + } + + pub fn num_leaves(&self) -> u64 { + self.num_leaves + } +} + +/// Build the mixed-height tree from each height group's finalized leaf digests. +/// +/// `group_digests[h]` is `Some(device digests)` when some matrix has +/// `log_height == h`, each `2^(h-1)` digests of 32 bytes; index `h_max` must be +/// present. Returns the standard heap node buffer +/// (`(2 * 2^(h_max-1) - 1) * 32` bytes) resident on device. +pub fn build_mmcs_tree_on_device( + stream: &Arc, + group_digests: &[Option>], +) -> Result> { + let h_max = group_digests.len() - 1; + assert!( + h_max >= 1 && group_digests[h_max].is_some(), + "the tallest height group must be present" + ); + let be = backend()?; + let leaves_len = 1u64 << (h_max - 1); + + let mut nodes = stream.alloc_zeros::(((2 * leaves_len - 1) * 32) as usize)?; + // Base layer into the leaf tail of the heap array. + let base = group_digests[h_max] + .as_ref() + .expect("checked immediately above"); + let mut leaf_tail = nodes.slice_mut(((leaves_len - 1) * 32) as usize..); + stream.memcpy_dtod(base, &mut leaf_tail)?; + + // Climb. Level `i` produces the layer whose codeword height is + // `h_max - 1 - i`, which is where a group of that height injects — the same + // schedule `MixedMmcs::from_group_digests` walks. + let mut level_begin: u64 = leaves_len - 1; + let mut i = 0usize; + while level_begin != 0 { + let new_begin = level_begin / 2; + let n_pairs = level_begin - new_begin; + let inject_h = h_max - 1 - i; + let injected = group_digests.get(inject_h).and_then(Option::as_ref); + let has_inject: u32 = u32::from(injected.is_some()); + + // `keccak_mmcs_level` reads `inject` only when `has_inject` is set, so a + // level with no injection still needs a pointer argument. Reuse the + // node buffer's own base rather than allocating a dummy: it is a valid + // device pointer that the kernel provably never dereferences. + let cfg = keccak_launch_cfg(n_pairs); + match injected { + Some(digests) => unsafe { + stream + .launch_builder(&be.keccak_mmcs_level) + .arg(&mut nodes) + .arg(&new_begin) + .arg(&n_pairs) + .arg(digests) + .arg(&has_inject) + .launch(cfg)?; + }, + None => { + let empty = stream.alloc_zeros::(32)?; + unsafe { + stream + .launch_builder(&be.keccak_mmcs_level) + .arg(&mut nodes) + .arg(&new_begin) + .arg(&n_pairs) + .arg(&empty) + .arg(&has_inject) + .launch(cfg)?; + } + } + } + + level_begin = new_begin; + i += 1; + } + + Ok(nodes) +} + +/// The MMCS root — node 0 of the heap array. +pub fn read_mmcs_root(stream: &Arc, nodes: &CudaSlice) -> Result<[u8; 32]> { + let head = nodes.slice(0..32); + let bytes = stream.clone_dtoh(&head)?; + let mut root = [0u8; 32]; + root.copy_from_slice(&bytes); + Ok(root) +} diff --git a/crypto/math-cuda/tests/mmcs_tree_parity.rs b/crypto/math-cuda/tests/mmcs_tree_parity.rs new file mode 100644 index 000000000..ad966ea57 --- /dev/null +++ b/crypto/math-cuda/tests/mmcs_tree_parity.rs @@ -0,0 +1,200 @@ +//! The device mixed-height MMCS must build the SAME tree as the host +//! `stark::fri::mmcs::MixedMmcs` — same root, same authentication paths, so a +//! proof committed on GPU is opened and verified by the same verifier as one +//! committed on CPU. +//! +//! ⚠ **This file has never been executed.** It was written on a machine with no +//! GPU and no nvcc, where `math-cuda` compiles against empty cubin stubs and +//! every device call falls back or fails. It compiles and it lints; nothing here +//! is evidence that the kernels are correct. Run it on a rented box — the exact +//! commands are in `RESUME-MMCS-INT.md` — before any claim that the batched GPU +//! path works. +//! +//! What each test is FOR, so a failure says something: +//! +//! - `single_matrix_mmcs_root_matches_the_per_table_tree` — the degenerate case. +//! A one-matrix MMCS is the existing row-pair tree, so this failing means the +//! absorb kernel's byte order or bit-reversal is wrong, independently of +//! anything mixed-height. +//! - `mixed_height_root_matches_the_host` — the climb with injection. This is +//! the kernel that has no CPU counterpart to have been debugged against. +//! - `absorption_order_is_bound` — the leaf concatenates matrices in INPUT +//! order; two matrices absorbed the other way round must give a different root. +//! Without this, an order bug is invisible whenever the widths happen to match. +//! - `paths_match_the_host_at_every_query` — the reason the device tree uses the +//! standard heap layout at all: `merkle_gather_paths` unchanged must return +//! the host's `MixedOpening::proof`. + +use math::field::element::FieldElement; +use math::field::goldilocks::GoldilocksField; +use stark::fri::mmcs::{LeafSource, MixedMmcs}; + +type Fp = FieldElement; +type Mmcs = MixedMmcs; + +/// Bit-reversed row-major matrices, the layout the MMCS commits and the layout +/// `mmcs_absorb_row_pair_row_major` reads (the kernel bit-reverses internally, so +/// the device buffer holds the matrix in NATURAL order). +struct Matrices { + /// `(natural-order row-major data, log_height, width)`. + mats: Vec<(Vec, usize, usize)>, +} + +impl LeafSource for Matrices { + fn num_matrices(&self) -> usize { + self.mats.len() + } + fn log_height(&self, m: usize) -> usize { + self.mats[m].1 + } + fn width(&self, m: usize) -> usize { + self.mats[m].2 + } + fn append_row(&self, m: usize, bitrev_row: usize, out: &mut Vec) { + let (data, log_height, width) = &self.mats[m]; + let natural = math::fft::bit_reversing::reverse_index(bitrev_row, 1u64 << log_height); + out.extend_from_slice(&data[natural * width..(natural + 1) * width]); + } +} + +fn matrix(log_height: usize, width: usize, seed: u64) -> (Vec, usize, usize) { + let num_rows = 1usize << log_height; + let data = (0..num_rows * width) + .map(|i| Fp::from(seed.wrapping_mul(1_000_003).wrapping_add(i as u64) | 1)) + .collect(); + (data, log_height, width) +} + +fn raw(data: &[Fp]) -> Vec { + data.iter().map(|x| *x.value()).collect() +} + +/// Build the tree on device from `specs`, absorbing matrices in input order and +/// freeing each matrix's device buffer before the next — the residency policy the +/// whole design exists for. +fn device_tree(mats: &Matrices) -> ([u8; 32], Vec) { + let be = math_cuda::device::backend().expect("a GPU box: no backend means nothing to test"); + let stream = be.next_stream(); + + let h_max = (0..mats.num_matrices()) + .map(|m| mats.log_height(m)) + .max() + .expect("non-empty"); + + let mut group_digests: Vec>> = + (0..=h_max).map(|_| None).collect(); + + for (h, slot) in group_digests.iter_mut().enumerate().skip(1) { + let group: Vec = (0..mats.num_matrices()) + .filter(|&m| mats.log_height(m) == h) + .collect(); + if group.is_empty() { + continue; + } + let mut hasher = math_cuda::mmcs::MmcsGroupHasher::new(&stream, h as u64) + .expect("group sponge allocation"); + for &m in &group { + let (data, _, width) = &mats.mats[m]; + let dev = stream.clone_htod(&raw(data)).expect("H2D"); + hasher + .absorb_row_major(&stream, &dev, *width as u64, 0, *width as u64) + .expect("absorb"); + // The point of the streaming build: this matrix is done with. + drop(dev); + } + *slot = Some(hasher.finalize(&stream).expect("finalize")); + } + + let nodes = math_cuda::mmcs::build_mmcs_tree_on_device(&stream, &group_digests) + .expect("device tree build"); + let root = math_cuda::mmcs::read_mmcs_root(&stream, &nodes).expect("root readback"); + let all = stream.clone_dtoh(&nodes).expect("node readback"); + (root, all) +} + +#[test] +fn single_matrix_mmcs_root_matches_the_per_table_tree() { + let mats = Matrices { + mats: vec![matrix(6, 5, 7)], + }; + let (device_root, _) = device_tree(&mats); + assert_eq!( + device_root, + Mmcs::commit(&mats).root(), + "a one-matrix MMCS must be the existing row-pair tree, byte for byte" + ); +} + +#[test] +fn mixed_height_root_matches_the_host() { + // Two matrices at the tallest height (so the base group batches), one + // injected mid-climb, one injected near the terminal. + let mats = Matrices { + mats: vec![ + matrix(7, 3, 11), + matrix(7, 6, 23), + matrix(5, 2, 41), + matrix(2, 4, 59), + ], + }; + let (device_root, _) = device_tree(&mats); + assert_eq!( + device_root, + Mmcs::commit(&mats).root(), + "the device climb with injection must reproduce the host tree" + ); +} + +#[test] +fn absorption_order_is_bound() { + let forward = Matrices { + mats: vec![matrix(6, 3, 11), matrix(6, 3, 23)], + }; + let reversed = Matrices { + mats: vec![matrix(6, 3, 23), matrix(6, 3, 11)], + }; + let (forward_root, _) = device_tree(&forward); + let (reversed_root, _) = device_tree(&reversed); + + assert_eq!(forward_root, Mmcs::commit(&forward).root()); + assert_eq!(reversed_root, Mmcs::commit(&reversed).root()); + assert_ne!( + forward_root, reversed_root, + "two same-shape matrices absorbed in the other order must commit a \ + different tree — input order is part of the commitment" + ); +} + +#[test] +fn paths_match_the_host_at_every_query() { + let mats = Matrices { + mats: vec![matrix(6, 3, 11), matrix(6, 2, 23), matrix(4, 5, 41)], + }; + let host = Mmcs::commit(&mats); + let (device_root, _) = device_tree(&mats); + assert_eq!(device_root, host.root()); + + let be = math_cuda::device::backend().expect("a GPU box"); + let stream = be.next_stream(); + let (_, nodes_host) = device_tree(&mats); + let nodes = stream.clone_htod(&nodes_host).expect("H2D nodes"); + + let leaves_len = 1usize << (host.h_max() - 1); + let positions: Vec = (0..leaves_len as u32).collect(); + let depth = host.h_max() - 1; + let paths = math_cuda::merkle::gather_merkle_paths_dev(&nodes, leaves_len, &positions, &stream) + .expect("path gather"); + + for iota in 0..leaves_len { + let expected = host.open_batch(iota, &mats).proof.merkle_path; + for (level, node) in expected.iter().enumerate() { + let start = (iota * depth + level) * 32; + assert_eq!( + &paths[start..start + 32], + &node[..], + "query {iota}, level {level}: the device path must be the host path — \ + `merkle_gather_paths` is reused precisely because the layouts agree" + ); + } + } +} From 9f98165fd9f47d82ae885cf134852d78ddbfc22f Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 18:09:39 -0300 Subject: [PATCH 05/26] feat(cuda): column-major absorb for the device MMCS (resident-LDE bridge) --- crypto/math-cuda/kernels/keccak.cu | 41 ++++++++++++ crypto/math-cuda/src/device.rs | 3 + crypto/math-cuda/src/mmcs.rs | 48 +++++++++++++ crypto/math-cuda/tests/mmcs_tree_parity.rs | 78 ++++++++++++++++++++++ 4 files changed, 170 insertions(+) diff --git a/crypto/math-cuda/kernels/keccak.cu b/crypto/math-cuda/kernels/keccak.cu index 70948f836..5b301f26b 100644 --- a/crypto/math-cuda/kernels/keccak.cu +++ b/crypto/math-cuda/kernels/keccak.cu @@ -705,6 +705,47 @@ extern "C" __global__ void mmcs_absorb_row_pair_row_major( rate_pos[tid] = rp; } +// Absorb one COLUMN-MAJOR base-field matrix's row pair — main's resident LDE +// layout (`GpuLdeBase`): element (row, col) at `col * col_stride + row`. Same +// absorbed byte order as `mmcs_absorb_row_pair_row_major`, so a matrix fed here +// column-major produces the identical leaf digests it would fed row-major, and +// the tree over them is byte-identical to the host's. +extern "C" __global__ void mmcs_absorb_row_pair_col_major( + uint64_t *states, + uint32_t *rate_pos, + const uint64_t *data, + uint64_t col_stride, + uint64_t col_start, + uint64_t col_end, + uint64_t num_rows, + uint64_t log_num_rows, + uint64_t num_leaves) +{ + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_leaves) return; + (void)num_rows; + + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + + uint64_t st[25]; + uint64_t *st_g = states + tid * 25; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = st_g[i]; + uint32_t rp = rate_pos[tid]; + + for (uint64_t c = col_start; c < col_end; ++c) { + absorb_lane(st, rp, bswap64(goldilocks::canonical(data[c * col_stride + br_0]))); + } + for (uint64_t c = col_start; c < col_end; ++c) { + absorb_lane(st, rp, bswap64(goldilocks::canonical(data[c * col_stride + br_1]))); + } + + #pragma unroll + for (int i = 0; i < 25; ++i) st_g[i] = st[i]; + rate_pos[tid] = rp; +} + // Absorb one COLUMN-MAJOR ext3 slab matrix's row pair — the composition-poly // LDE layout (`GpuLdeExt3`): component `k` of column `c` at // `(c*3 + k) * col_stride`. Same absorbed byte order as diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index fd9b31a26..9063bd2fc 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -207,6 +207,7 @@ pub struct Backend { // group never needs all its LDEs resident — see `kernels/keccak.cu`. pub mmcs_states_init: CudaFunction, pub mmcs_absorb_row_pair_row_major: CudaFunction, + pub mmcs_absorb_row_pair_col_major: CudaFunction, pub mmcs_absorb_row_pair_ext3_slabs: CudaFunction, pub mmcs_states_finalize: CudaFunction, pub keccak_mmcs_level: CudaFunction, @@ -448,6 +449,8 @@ impl Backend { mmcs_states_init: keccak.load_function("mmcs_states_init")?, mmcs_absorb_row_pair_row_major: keccak .load_function("mmcs_absorb_row_pair_row_major")?, + mmcs_absorb_row_pair_col_major: keccak + .load_function("mmcs_absorb_row_pair_col_major")?, mmcs_absorb_row_pair_ext3_slabs: keccak .load_function("mmcs_absorb_row_pair_ext3_slabs")?, mmcs_states_finalize: keccak.load_function("mmcs_states_finalize")?, diff --git a/crypto/math-cuda/src/mmcs.rs b/crypto/math-cuda/src/mmcs.rs index b6e20f17e..b837e9da2 100644 --- a/crypto/math-cuda/src/mmcs.rs +++ b/crypto/math-cuda/src/mmcs.rs @@ -131,6 +131,54 @@ impl MmcsGroupHasher { Ok(()) } + /// Absorb one COLUMN-MAJOR base-field matrix's row pair into every leaf — + /// main's resident LDE layout (`GpuLdeBase`), element `(row, col)` at + /// `col * col_stride + row`. Columns `[col_start, col_end)` are absorbed. + /// + /// Produces the identical leaf digests as [`Self::absorb_row_major`] over the + /// same matrix (same absorbed byte stream, only the read layout differs), so + /// a device commit fed from a resident column-major LDE is byte-identical to + /// the host tree. This is the bridge for main's resident base-field LDEs, + /// which are column-major; the ext3 parts already match [`Self::absorb_ext3_slabs`]. + /// + /// The caller may free `data` as soon as this returns on `stream`. + pub fn absorb_col_major( + &mut self, + stream: &Arc, + data: &CudaSlice, + col_stride: u64, + col_start: u64, + col_end: u64, + ) -> Result<()> { + let num_rows = 1u64 << self.log_num_rows; + assert!( + col_start <= col_end, + "column range [{col_start}, {col_end}) is empty-or-inverted" + ); + assert!( + num_rows <= col_stride, + "col_stride {col_stride} must hold all {num_rows} rows of a column" + ); + let be = backend()?; + let cfg = keccak_launch_cfg(self.num_leaves); + unsafe { + stream + .launch_builder(&be.mmcs_absorb_row_pair_col_major) + .arg(&mut self.states) + .arg(&mut self.rate_pos) + .arg(data) + .arg(&col_stride) + .arg(&col_start) + .arg(&col_end) + .arg(&num_rows) + .arg(&self.log_num_rows) + .arg(&self.num_leaves) + .launch(cfg)?; + } + self.absorbed += 1; + Ok(()) + } + /// Absorb one column-major ext3 slab matrix — the composition-poly LDE /// layout, component `k` of column `c` at `(c*3 + k) * col_stride`. pub fn absorb_ext3_slabs( diff --git a/crypto/math-cuda/tests/mmcs_tree_parity.rs b/crypto/math-cuda/tests/mmcs_tree_parity.rs index ad966ea57..a41fc3fcf 100644 --- a/crypto/math-cuda/tests/mmcs_tree_parity.rs +++ b/crypto/math-cuda/tests/mmcs_tree_parity.rs @@ -198,3 +198,81 @@ fn paths_match_the_host_at_every_query() { } } } + +/// Column-major raw buffer of a matrix: element `(row, col)` at +/// `col * num_rows + row`. This is main's resident `GpuLdeBase` layout, which +/// `absorb_col_major` reads directly (no host transpose). +fn raw_col_major(data: &[Fp], log_height: usize, width: usize) -> Vec { + let num_rows = 1usize << log_height; + let mut out = vec![0u64; num_rows * width]; + for row in 0..num_rows { + for col in 0..width { + out[col * num_rows + row] = *data[row * width + col].value(); + } + } + out +} + +/// Like [`device_tree`], but feeds every matrix COLUMN-MAJOR through +/// `absorb_col_major` — the bridge for main's resident base-field LDEs. +fn device_tree_col_major(mats: &Matrices) -> [u8; 32] { + let be = math_cuda::device::backend().expect("a GPU box"); + let stream = be.next_stream(); + + let h_max = (0..mats.num_matrices()) + .map(|m| mats.log_height(m)) + .max() + .expect("non-empty"); + + let mut group_digests: Vec>> = + (0..=h_max).map(|_| None).collect(); + + for (h, slot) in group_digests.iter_mut().enumerate().skip(1) { + let group: Vec = (0..mats.num_matrices()) + .filter(|&m| mats.log_height(m) == h) + .collect(); + if group.is_empty() { + continue; + } + let mut hasher = math_cuda::mmcs::MmcsGroupHasher::new(&stream, h as u64) + .expect("group sponge allocation"); + for &m in &group { + let (data, log_height, width) = &mats.mats[m]; + let num_rows = 1u64 << log_height; + let dev = stream + .clone_htod(&raw_col_major(data, *log_height, *width)) + .expect("H2D"); + hasher + .absorb_col_major(&stream, &dev, num_rows, 0, *width as u64) + .expect("absorb col-major"); + drop(dev); + } + *slot = Some(hasher.finalize(&stream).expect("finalize")); + } + + let nodes = math_cuda::mmcs::build_mmcs_tree_on_device(&stream, &group_digests) + .expect("device tree build"); + math_cuda::mmcs::read_mmcs_root(&stream, &nodes).expect("root readback") +} + +/// ★ RISK #1: main hands the batched prover its resident base-field LDEs +/// COLUMN-MAJOR (`GpuLdeBase.buf`), but the MMCS leaf hash is defined row-major. +/// `absorb_col_major` must produce the identical tree. Same mixed-height fixture +/// as `mixed_height_root_matches_the_host`, fed column-major. +#[test] +fn col_major_root_matches_the_host() { + let mats = Matrices { + mats: vec![ + matrix(7, 3, 11), + matrix(7, 6, 23), + matrix(5, 2, 41), + matrix(2, 4, 59), + ], + }; + assert_eq!( + device_tree_col_major(&mats), + Mmcs::commit(&mats).root(), + "a device commit fed from a column-major (resident-LDE) buffer must match \ + the host tree byte for byte — the col-major absorb bridge" + ); +} From bfcc6fcf6fd1eadac07b64c45292a715a2150b5e Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 18:26:05 -0300 Subject: [PATCH 06/26] feat(cuda): device mixed-height MMCS commit + main-round GPU cross-check --- crypto/math-cuda/src/mmcs.rs | 45 +++++++++++++++++++++++ crypto/stark/src/batched/prover.rs | 57 ++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/crypto/math-cuda/src/mmcs.rs b/crypto/math-cuda/src/mmcs.rs index b837e9da2..5f57aaa2f 100644 --- a/crypto/math-cuda/src/mmcs.rs +++ b/crypto/math-cuda/src/mmcs.rs @@ -320,3 +320,48 @@ pub fn read_mmcs_root(stream: &Arc, nodes: &CudaSlice) -> Result root.copy_from_slice(&bytes); Ok(root) } + +/// One base-field matrix feeding a mixed-height MMCS commit, in absorb order. +/// `data` is the full ROW-MAJOR `stride`-wide LDE buffer (Goldilocks u64 words); +/// columns `[col_start, col_end)` are the ones committed (a preprocessed table +/// commits its non-precomputed range). +pub struct MmcsRowMajorInput<'a> { + pub data: &'a [u64], + pub stride: u64, + pub col_start: u64, + pub col_end: u64, + pub log_height: u64, +} + +/// Build the whole mixed-height MMCS tree on the GPU from row-major base-field +/// matrices and return its root — the device twin of the host +/// `StreamingMmcsBuilder`: matrices are grouped by height, each group's leaves +/// concatenate its matrices in the given order, and the climb injects the +/// shorter groups. Callers hand matrices in the SAME order the host absorbs them. +/// +/// Uploads each matrix and frees it before the next (the streaming residency +/// policy). `inputs` must be non-empty and include the tallest height. +pub fn commit_mixed_row_major_root(inputs: &[MmcsRowMajorInput]) -> Result<[u8; 32]> { + let be = backend()?; + let stream = be.next_stream(); + + let h_max = inputs.iter().map(|m| m.log_height).max().unwrap_or(0); + let mut group_digests: Vec>> = (0..=h_max).map(|_| None).collect(); + + let mut h = h_max; + while h >= 1 { + if inputs.iter().any(|m| m.log_height == h) { + let mut hasher = MmcsGroupHasher::new(&stream, h)?; + for m in inputs.iter().filter(|m| m.log_height == h) { + let dev = stream.clone_htod(m.data)?; + hasher.absorb_row_major(&stream, &dev, m.stride, m.col_start, m.col_end)?; + drop(dev); + } + group_digests[h as usize] = Some(hasher.finalize(&stream)?); + } + h -= 1; + } + + let nodes = build_mmcs_tree_on_device(&stream, &group_digests)?; + read_mmcs_root(&stream, &nodes) +} diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index b32cfa0c6..c70a5e03c 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -362,6 +362,63 @@ where let main_root = main_mmcs.root(); transcript.append_bytes(&main_root); + // GPU cross-check (3c): the device mixed-height MMCS must build the same + // main-round root the host `StreamingMmcsBuilder` just did, from the same + // matrices (same height grouping, same absorb order, same column ranges). + // Inert unless every contributing main LDE is retained (so `RecomputeLde` + // skips it), the field is Goldilocks (the kernels are), and there is a GPU. + #[cfg(feature = "cuda")] + { + use math::field::goldilocks::GoldilocksField; + use std::any::TypeId; + if TypeId::of::() == TypeId::of::() { + let mut raws: Vec<(Vec, u64, u64, u64)> = Vec::new(); + let mut all_present = true; + for table in 0..num_tables { + if shape.carved_main.map(|c| c.table) == Some(table) { + continue; + } + match &retained_main[table] { + Some((data, total_cols)) => { + let width = matrix_width(&shape.main, table) as u64; + let stride = *total_cols as u64; + // SAFETY: Field == GoldilocksField (checked), whose element + // value is a `u64` — the same cast the per-table GPU commit + // uses (`gpu_lde::columns_to_u64_base`). + let raw: Vec = data + .iter() + .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .collect(); + raws.push((raw, stride, stride - width, shape.heights[table] as u64)); + } + None => { + all_present = false; + break; + } + } + } + if all_present && !raws.is_empty() { + let inputs: Vec = raws + .iter() + .map(|(d, stride, col_start, h)| math_cuda::mmcs::MmcsRowMajorInput { + data: d.as_slice(), + stride: *stride, + col_start: *col_start, + col_end: *stride, + log_height: *h, + }) + .collect(); + if let Ok(dev_root) = math_cuda::mmcs::commit_mixed_row_major_root(&inputs) { + debug_assert_eq!( + dev_root, main_root, + "device main-round MMCS root must equal the host \ + StreamingMmcsBuilder root" + ); + } + } + } + } + // ===================================================================== // Phase 2 — LogUp challenges, then the auxiliary round // ===================================================================== From 14c9b732d5c943ed5fa232d536e10d0bf2a02d50 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 18:34:16 -0300 Subject: [PATCH 07/26] feat(cuda): device MMCS commit + GPU cross-check for the aux and parts rounds --- crypto/math-cuda/kernels/keccak.cu | 48 +++++++++++++ crypto/math-cuda/src/device.rs | 3 + crypto/math-cuda/src/mmcs.rs | 112 +++++++++++++++++++++++++++++ crypto/stark/src/batched/prover.rs | 99 +++++++++++++++++++++++++ 4 files changed, 262 insertions(+) diff --git a/crypto/math-cuda/kernels/keccak.cu b/crypto/math-cuda/kernels/keccak.cu index 5b301f26b..ac11d0360 100644 --- a/crypto/math-cuda/kernels/keccak.cu +++ b/crypto/math-cuda/kernels/keccak.cu @@ -792,6 +792,54 @@ extern "C" __global__ void mmcs_absorb_row_pair_ext3_slabs( rate_pos[tid] = rp; } +// Absorb one ROW-MAJOR ext3 matrix's row pair — the batched aux LDE layout +// (`expand_aux_lde_row_major`): element `(row, col)` is 3 consecutive u64 at +// `(row * stride + col) * 3`, `stride` elements per row. Same absorbed byte +// order as the host's row-major ext3 leaf. +extern "C" __global__ void mmcs_absorb_row_pair_ext3_row_major( + uint64_t *states, + uint32_t *rate_pos, + const uint64_t *data, + uint64_t stride, + uint64_t col_start, + uint64_t col_end, + uint64_t num_rows, + uint64_t log_num_rows, + uint64_t num_leaves) +{ + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_leaves) return; + (void)num_rows; + + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + + uint64_t st[25]; + uint64_t *st_g = states + tid * 25; + #pragma unroll + for (int i = 0; i < 25; ++i) st[i] = st_g[i]; + uint32_t rp = rate_pos[tid]; + + for (uint64_t c = col_start; c < col_end; ++c) { + const uint64_t *e = data + (br_0 * stride + c) * 3; + #pragma unroll + for (int k = 0; k < 3; ++k) { + absorb_lane(st, rp, bswap64(goldilocks::canonical(e[k]))); + } + } + for (uint64_t c = col_start; c < col_end; ++c) { + const uint64_t *e = data + (br_1 * stride + c) * 3; + #pragma unroll + for (int k = 0; k < 3; ++k) { + absorb_lane(st, rp, bswap64(goldilocks::canonical(e[k]))); + } + } + + #pragma unroll + for (int i = 0; i < 25; ++i) st_g[i] = st[i]; + rate_pos[tid] = rp; +} + // Pad and squeeze every leaf's sponge into a 32-byte digest. extern "C" __global__ void mmcs_states_finalize( const uint64_t *states, diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 9063bd2fc..4fddbb1df 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -209,6 +209,7 @@ pub struct Backend { pub mmcs_absorb_row_pair_row_major: CudaFunction, pub mmcs_absorb_row_pair_col_major: CudaFunction, pub mmcs_absorb_row_pair_ext3_slabs: CudaFunction, + pub mmcs_absorb_row_pair_ext3_row_major: CudaFunction, pub mmcs_states_finalize: CudaFunction, pub keccak_mmcs_level: CudaFunction, @@ -453,6 +454,8 @@ impl Backend { .load_function("mmcs_absorb_row_pair_col_major")?, mmcs_absorb_row_pair_ext3_slabs: keccak .load_function("mmcs_absorb_row_pair_ext3_slabs")?, + mmcs_absorb_row_pair_ext3_row_major: keccak + .load_function("mmcs_absorb_row_pair_ext3_row_major")?, mmcs_states_finalize: keccak.load_function("mmcs_states_finalize")?, keccak_mmcs_level: keccak.load_function("keccak_mmcs_level")?, barycentric_base_batched: bary.load_function("barycentric_base_batched")?, diff --git a/crypto/math-cuda/src/mmcs.rs b/crypto/math-cuda/src/mmcs.rs index 5f57aaa2f..197e1d237 100644 --- a/crypto/math-cuda/src/mmcs.rs +++ b/crypto/math-cuda/src/mmcs.rs @@ -208,6 +208,40 @@ impl MmcsGroupHasher { Ok(()) } + /// Absorb one ROW-MAJOR ext3 matrix's row pair — the batched aux LDE layout + /// (`expand_aux_lde_row_major`), element `(row, col)` as 3 consecutive u64 at + /// `(row*stride + col)*3`, `stride` elements per row. Columns `[col_start, + /// col_end)` are absorbed. Same absorbed byte order as the host's row-major + /// ext3 leaf. + pub fn absorb_ext3_row_major( + &mut self, + stream: &Arc, + data: &CudaSlice, + stride: u64, + col_start: u64, + col_end: u64, + ) -> Result<()> { + let be = backend()?; + let num_rows = 1u64 << self.log_num_rows; + let cfg = keccak_launch_cfg(self.num_leaves); + unsafe { + stream + .launch_builder(&be.mmcs_absorb_row_pair_ext3_row_major) + .arg(&mut self.states) + .arg(&mut self.rate_pos) + .arg(data) + .arg(&stride) + .arg(&col_start) + .arg(&col_end) + .arg(&num_rows) + .arg(&self.log_num_rows) + .arg(&self.num_leaves) + .launch(cfg)?; + } + self.absorbed += 1; + Ok(()) + } + /// Pad and squeeze every leaf. Panics if nothing was absorbed: an empty /// group's digests would be the hash of nothing, which is a leaf no verifier /// can rebuild from an opening. @@ -365,3 +399,81 @@ pub fn commit_mixed_row_major_root(inputs: &[MmcsRowMajorInput]) -> Result<[u8; let nodes = build_mmcs_tree_on_device(&stream, &group_digests)?; read_mmcs_root(&stream, &nodes) } + +/// One ext3 matrix feeding a mixed-height MMCS commit, in absorb order, laid out +/// as COLUMN-MAJOR SLABS: component `k` of column `p` at `(p*3 + k) * col_stride` +/// (`col_stride` = rows per column), natural row order. This is the +/// composition-poly / `GpuLdeExt3` layout. +pub struct MmcsExt3SlabInput<'a> { + pub data: &'a [u64], + pub col_stride: u64, + pub num_parts: u64, + pub log_height: u64, +} + +/// Build the whole mixed-height MMCS tree on the GPU from ext3 slab matrices and +/// return its root — the ext3 twin of [`commit_mixed_row_major_root`], grouping +/// by height and absorbing each group's matrices in the given order. +pub fn commit_mixed_ext3_slabs_root(inputs: &[MmcsExt3SlabInput]) -> Result<[u8; 32]> { + let be = backend()?; + let stream = be.next_stream(); + + let h_max = inputs.iter().map(|m| m.log_height).max().unwrap_or(0); + let mut group_digests: Vec>> = (0..=h_max).map(|_| None).collect(); + + let mut h = h_max; + while h >= 1 { + if inputs.iter().any(|m| m.log_height == h) { + let mut hasher = MmcsGroupHasher::new(&stream, h)?; + for m in inputs.iter().filter(|m| m.log_height == h) { + let dev = stream.clone_htod(m.data)?; + hasher.absorb_ext3_slabs(&stream, &dev, m.col_stride, m.num_parts)?; + drop(dev); + } + group_digests[h as usize] = Some(hasher.finalize(&stream)?); + } + h -= 1; + } + + let nodes = build_mmcs_tree_on_device(&stream, &group_digests)?; + read_mmcs_root(&stream, &nodes) +} + +/// One ROW-MAJOR ext3 matrix feeding a mixed-height MMCS commit, in absorb +/// order: element `(row, col)` as 3 consecutive u64 at `(row*stride + col)*3`. +/// Columns `[col_start, col_end)` are committed. This is the batched aux LDE +/// layout (`expand_aux_lde_row_major`). +pub struct MmcsExt3RowMajorInput<'a> { + pub data: &'a [u64], + pub stride: u64, + pub col_start: u64, + pub col_end: u64, + pub log_height: u64, +} + +/// Build the whole mixed-height MMCS tree on the GPU from ROW-MAJOR ext3 +/// matrices and return its root — the aux twin of [`commit_mixed_row_major_root`]. +pub fn commit_mixed_ext3_row_major_root(inputs: &[MmcsExt3RowMajorInput]) -> Result<[u8; 32]> { + let be = backend()?; + let stream = be.next_stream(); + + let h_max = inputs.iter().map(|m| m.log_height).max().unwrap_or(0); + let mut group_digests: Vec>> = (0..=h_max).map(|_| None).collect(); + + let mut h = h_max; + while h >= 1 { + if inputs.iter().any(|m| m.log_height == h) { + let mut hasher = MmcsGroupHasher::new(&stream, h)?; + for m in inputs.iter().filter(|m| m.log_height == h) { + let dev = stream.clone_htod(m.data)?; + hasher.absorb_ext3_row_major(&stream, &dev, m.stride, m.col_start, m.col_end)?; + drop(dev); + } + group_digests[h as usize] = Some(hasher.finalize(&stream)?); + } + h -= 1; + } + + let nodes = build_mmcs_tree_on_device(&stream, &group_digests)?; + read_mmcs_root(&stream, &nodes) +} diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index c70a5e03c..915062b46 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -500,6 +500,47 @@ where transcript.append_bytes(&root); } + // GPU cross-check (3c): the device row-major ext3 MMCS must build the same + // aux-round root the host did (Retain + Goldilocks-ext3 + GPU, else inert). + #[cfg(feature = "cuda")] + if let Some(aux_root_val) = aux_root { + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; + use std::any::TypeId; + if TypeId::of::() == TypeId::of::() { + let mut raws: Vec<(Vec, u64, u64)> = Vec::new(); + for table in 0..num_tables { + if let Some((data, aux_cols)) = &retained_aux[table] { + // SAFETY: FieldExtension == Degree3Goldilocks (checked), a + // contiguous `[u64; 3]` per element — same cast as + // `gpu_lde::columns_to_u64_ext3`. + let raw: Vec = unsafe { + std::slice::from_raw_parts(data.as_ptr() as *const u64, data.len() * 3) + } + .to_vec(); + raws.push((raw, *aux_cols as u64, shape.heights[table] as u64)); + } + } + if !raws.is_empty() { + let inputs: Vec = raws + .iter() + .map(|(d, stride, h)| math_cuda::mmcs::MmcsExt3RowMajorInput { + data: d.as_slice(), + stride: *stride, + col_start: 0, + col_end: *stride, + log_height: *h, + }) + .collect(); + if let Ok(dev_root) = math_cuda::mmcs::commit_mixed_ext3_row_major_root(&inputs) { + debug_assert_eq!( + dev_root, aux_root_val, + "device aux-round MMCS root must equal the host root" + ); + } + } + } + } + // ===================================================================== // Phase 3 — bus contributions, beta per table, the composition-parts round // ===================================================================== @@ -594,6 +635,64 @@ where let parts_root = parts_mmcs.root(); transcript.append_bytes(&parts_root); + // GPU cross-check (3c): the device ext3 mixed-height MMCS must build the same + // parts-round root the host did. Parts are always retained; the field is + // Goldilocks-ext3 and there is a GPU, or this is inert. + #[cfg(feature = "cuda")] + { + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; + use std::any::TypeId; + if TypeId::of::() == TypeId::of::() { + let mut slabs: Vec<(Vec, u64, u64, u64)> = Vec::new(); + let mut all_present = true; + for table in 0..num_tables { + let parts = &retained_parts[table]; + if parts.is_empty() || parts[0].is_empty() { + all_present = false; + break; + } + let num_parts = parts.len(); + let num_rows = parts[0].len(); + let mut slab = vec![0u64; num_parts * 3 * num_rows]; + for (p, col) in parts.iter().enumerate() { + for (row, elem) in col.iter().enumerate() { + // SAFETY: FieldExtension == Degree3Goldilocks (checked), + // whose value is `[u64; 3]` — same cast as + // `gpu_lde::columns_to_u64_ext3`. + let comps = + unsafe { std::slice::from_raw_parts(elem.value() as *const _ as *const u64, 3) }; + for (k, &c) in comps.iter().enumerate() { + slab[(p * 3 + k) * num_rows + row] = c; + } + } + } + slabs.push(( + slab, + num_rows as u64, + num_parts as u64, + shape.heights[table] as u64, + )); + } + if all_present && !slabs.is_empty() { + let inputs: Vec = slabs + .iter() + .map(|(d, col_stride, num_parts, h)| math_cuda::mmcs::MmcsExt3SlabInput { + data: d.as_slice(), + col_stride: *col_stride, + num_parts: *num_parts, + log_height: *h, + }) + .collect(); + if let Ok(dev_root) = math_cuda::mmcs::commit_mixed_ext3_slabs_root(&inputs) { + debug_assert_eq!( + dev_root, parts_root, + "device parts-round MMCS root must equal the host root" + ); + } + } + } + } + // ===================================================================== // Phase 4 — z per table, OOD evaluations // ===================================================================== From 1086fc0ef842f382bf71629fb5a5f09d1d3bf7e1 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 18:55:18 -0300 Subject: [PATCH 08/26] feat(stark): reconstruct MixedMmcs from a GPU heap node array (from_heap_nodes) --- crypto/stark/src/fri/mmcs.rs | 90 ++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/crypto/stark/src/fri/mmcs.rs b/crypto/stark/src/fri/mmcs.rs index d940dd4e9..b9e41972f 100644 --- a/crypto/stark/src/fri/mmcs.rs +++ b/crypto/stark/src/fri/mmcs.rs @@ -444,6 +444,50 @@ where } } + /// Reconstruct the tree from a STANDARD HEAP node array — the layout the GPU + /// commit (`math_cuda::mmcs::build_mmcs_tree_on_device`) produces: `2*L-1` + /// nodes of 32 bytes, root at index 0, inner nodes in `[0, L-1)`, the `L = + /// 2^(h_max-1)` leaves in the tail `[L-1, 2L-1)`, with leaf `j` at `L-1+j`. + /// + /// This is what makes a GPU-built tree serve the SAME [`Self::auth_path`] / + /// [`Self::open_batch`] a host-built one does: the keccak (leaf + climb) runs + /// on the device, only the digest layers come back, and every downstream + /// opening reads them unchanged. The heap ordering matches + /// `merkle_gather_paths` (validated by `mmcs_tree_parity`'s + /// `paths_match_the_host_at_every_query`), so `layers[level]` here is exactly + /// the level that kernel walks. + pub fn from_heap_nodes(dims: Vec<(usize, usize)>, h_max: usize, nodes: &[u8]) -> Self { + let leaves_len = 1usize << (h_max - 1); + assert_eq!( + nodes.len(), + (2 * leaves_len - 1) * 32, + "heap node array must be (2*L-1) 32-byte digests for L = 2^(h_max-1)" + ); + let node = |i: usize| -> Commitment { + let mut c = [0u8; 32]; + c.copy_from_slice(&nodes[i * 32..i * 32 + 32]); + c + }; + // `layers[k]` is heap level `h_max-1-k`: `2^(h_max-1-k)` nodes starting at + // heap index `2^(h_max-1-k) - 1`. `layers[0]` is the leaf tail; the last + // layer is the single root at index 0. + let mut layers: Vec> = Vec::with_capacity(h_max); + for k in 0..h_max { + let level_size = leaves_len >> k; + let start = level_size - 1; + layers.push((0..level_size).map(|j| node(start + j)).collect()); + } + let root = layers.last().expect("at least the base layer exists")[0]; + + MixedMmcs { + root, + layers, + dims, + h_max, + _marker: PhantomData, + } + } + /// The committed root. pub fn root(&self) -> Commitment { self.root @@ -1802,4 +1846,50 @@ mod tests { } } } + + /// The GPU commit returns a standard heap node array; `from_heap_nodes` must + /// rebuild a tree that serves the same root and authentication paths the + /// host build does — that is what lets the device tree be authoritative while + /// only the digest layers come back. Round-trips a real mixed-height tree's + /// layers through the heap layout and checks every query's path. + #[test] + fn from_heap_nodes_rebuilds_the_same_tree() { + let (specs, inner) = residency_fixture(); + let m = Mmcs::commit(&inner); + let h_max = m.h_max(); + let leaves_len = 1usize << (h_max - 1); + + // Assemble the standard heap from the host layers, exactly as the device + // build writes it: layer `k` (heap level `h_max-1-k`) into + // `[2^(h_max-1-k) - 1, ..)`, leaf `j` at `L-1+j`. + let mut heap = vec![0u8; (2 * leaves_len - 1) * 32]; + for (k, layer) in m.layers.iter().enumerate() { + let level_size = leaves_len >> k; + assert_eq!(layer.len(), level_size); + let start = level_size - 1; + for (j, digest) in layer.iter().enumerate() { + heap[(start + j) * 32..(start + j) * 32 + 32].copy_from_slice(digest); + } + } + + let rebuilt = Mmcs::from_heap_nodes(m.dims.clone(), h_max, &heap); + assert_eq!(rebuilt.root(), m.root(), "root must survive the heap round-trip"); + assert_eq!(rebuilt.h_max(), h_max); + + let heights: Vec = specs.iter().map(|&(lh, _, _)| lh).collect(); + let widths: Vec = specs.iter().map(|&(_, w, _)| w).collect(); + for iota in 0..leaves_len { + assert_eq!( + rebuilt.auth_path(iota).unwrap().merkle_path, + m.auth_path(iota).unwrap().merkle_path, + "authentication path at iota {iota} must match the host tree" + ); + // And a full opening off the rebuilt tree still verifies. + let opening = rebuilt.open_batch(iota, &inner); + assert!( + Mmcs::verify_batch(&rebuilt.root(), iota, &opening, &heights, &widths), + "an opening from the rebuilt tree must verify at iota {iota}" + ); + } + } } From b850b926b1b6896c2ee3cd356c98e15f7c7c8185 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 19:14:11 -0300 Subject: [PATCH 09/26] feat(cuda): commit all batched MMCS rounds on the GPU (device authoritative) --- crypto/math-cuda/src/mmcs.rs | 36 ++- crypto/stark/src/batched/prover.rs | 503 ++++++++++++++++++----------- 2 files changed, 342 insertions(+), 197 deletions(-) diff --git a/crypto/math-cuda/src/mmcs.rs b/crypto/math-cuda/src/mmcs.rs index 197e1d237..0f5f63dae 100644 --- a/crypto/math-cuda/src/mmcs.rs +++ b/crypto/math-cuda/src/mmcs.rs @@ -375,7 +375,10 @@ pub struct MmcsRowMajorInput<'a> { /// /// Uploads each matrix and frees it before the next (the streaming residency /// policy). `inputs` must be non-empty and include the tallest height. -pub fn commit_mixed_row_major_root(inputs: &[MmcsRowMajorInput]) -> Result<[u8; 32]> { +/// Like [`commit_mixed_row_major_root`] but returns the WHOLE standard heap node +/// array (host-side `Vec`, `(2L-1)*32` bytes) — feed it to +/// `MixedMmcs::from_heap_nodes` to make the GPU tree authoritative. +pub fn commit_mixed_row_major_nodes(inputs: &[MmcsRowMajorInput]) -> Result> { let be = backend()?; let stream = be.next_stream(); @@ -397,7 +400,14 @@ pub fn commit_mixed_row_major_root(inputs: &[MmcsRowMajorInput]) -> Result<[u8; } let nodes = build_mmcs_tree_on_device(&stream, &group_digests)?; - read_mmcs_root(&stream, &nodes) + stream.clone_dtoh(&nodes) +} + +pub fn commit_mixed_row_major_root(inputs: &[MmcsRowMajorInput]) -> Result<[u8; 32]> { + let nodes = commit_mixed_row_major_nodes(inputs)?; + let mut root = [0u8; 32]; + root.copy_from_slice(&nodes[0..32]); + Ok(root) } /// One ext3 matrix feeding a mixed-height MMCS commit, in absorb order, laid out @@ -414,7 +424,7 @@ pub struct MmcsExt3SlabInput<'a> { /// Build the whole mixed-height MMCS tree on the GPU from ext3 slab matrices and /// return its root — the ext3 twin of [`commit_mixed_row_major_root`], grouping /// by height and absorbing each group's matrices in the given order. -pub fn commit_mixed_ext3_slabs_root(inputs: &[MmcsExt3SlabInput]) -> Result<[u8; 32]> { +pub fn commit_mixed_ext3_slabs_nodes(inputs: &[MmcsExt3SlabInput]) -> Result> { let be = backend()?; let stream = be.next_stream(); @@ -436,7 +446,14 @@ pub fn commit_mixed_ext3_slabs_root(inputs: &[MmcsExt3SlabInput]) -> Result<[u8; } let nodes = build_mmcs_tree_on_device(&stream, &group_digests)?; - read_mmcs_root(&stream, &nodes) + stream.clone_dtoh(&nodes) +} + +pub fn commit_mixed_ext3_slabs_root(inputs: &[MmcsExt3SlabInput]) -> Result<[u8; 32]> { + let nodes = commit_mixed_ext3_slabs_nodes(inputs)?; + let mut root = [0u8; 32]; + root.copy_from_slice(&nodes[0..32]); + Ok(root) } /// One ROW-MAJOR ext3 matrix feeding a mixed-height MMCS commit, in absorb @@ -453,7 +470,7 @@ pub struct MmcsExt3RowMajorInput<'a> { /// Build the whole mixed-height MMCS tree on the GPU from ROW-MAJOR ext3 /// matrices and return its root — the aux twin of [`commit_mixed_row_major_root`]. -pub fn commit_mixed_ext3_row_major_root(inputs: &[MmcsExt3RowMajorInput]) -> Result<[u8; 32]> { +pub fn commit_mixed_ext3_row_major_nodes(inputs: &[MmcsExt3RowMajorInput]) -> Result> { let be = backend()?; let stream = be.next_stream(); @@ -475,5 +492,12 @@ pub fn commit_mixed_ext3_row_major_root(inputs: &[MmcsExt3RowMajorInput]) -> Res } let nodes = build_mmcs_tree_on_device(&stream, &group_digests)?; - read_mmcs_root(&stream, &nodes) + stream.clone_dtoh(&nodes) +} + +pub fn commit_mixed_ext3_row_major_root(inputs: &[MmcsExt3RowMajorInput]) -> Result<[u8; 32]> { + let nodes = commit_mixed_ext3_row_major_nodes(inputs)?; + let mut root = [0u8; 32]; + root.copy_from_slice(&nodes[0..32]); + Ok(root) } diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index 915062b46..64265db80 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -52,13 +52,19 @@ //! prove. `parts_computations` stays at one per table, and the counter is there //! so that stops being silent if it ever changes. //! -//! # What is deliberately absent +//! # The GPU commit //! -//! No device paths. The GPU mixed-height MMCS exists (`crypto/math-cuda/`) and -//! is box-gated; wiring it in is a separate step, and a batched prover that -//! silently fell back between host and device arms would make the residency -//! numbers above unreproducible. Under `--features cuda` this path compiles and -//! runs on the host. +//! Under `--features cuda`, on Goldilocks (base + ext3) and `ResidencyMode:: +//! Retain`, all four MMCS trees (main / aux / parts) are committed ON THE DEVICE: +//! the keccak leaf hash and the climb run on the GPU (`crypto/math-cuda/`'s +//! mixed-height MMCS), and only the digest layers come back to rebuild the +//! host-side [`MixedMmcs`] via [`MixedMmcs::from_heap_nodes`] — so the openings +//! (row values from the retained LDE, one shared auth path per query) are served +//! exactly as before. The host [`StreamingMmcsBuilder`] is skipped entirely for +//! those rounds. This is a DELIBERATE mode select (see `device_commit_enabled`), +//! not a silent per-call host↔device fallback — a device error is a hard abort — +//! precisely so the residency numbers above stay reproducible. Off cuda / off +//! Goldilocks / under `RecomputeLde`, the host builder commits as it always did. use math::fft::bit_reversing::in_place_bit_reverse_permute_row_major; use math::field::element::FieldElement; @@ -251,7 +257,16 @@ where // per-table prover does. let mut prep_trees: Vec>> = (0..num_tables).map(|_| None).collect(); - let mut main_builder = StreamingMmcsBuilder::::new(&shape.main.dims); + + // When the GPU can commit the main round (Goldilocks + Retain, so the LDEs + // are still around to feed the device tree), the keccak leaf+climb runs on + // the device and only the digest layers come back; the CPU + // `StreamingMmcsBuilder` is skipped entirely. Otherwise it builds the tree. + // Deliberate mode select (no silent per-call fallback); a device error is a + // hard abort. + let use_device_main = device_commit_enabled::(residency); + let mut main_builder = + (!use_device_main).then(|| StreamingMmcsBuilder::::new(&shape.main.dims)); let mut retained_main: Vec>, usize)>> = (0..num_tables).map(|_| None).collect(); // The carved table's standalone main tree and its root. Built inside the @@ -329,14 +344,16 @@ where carved_tree = Some(tree); carved_root = Some(root); } else { - let src = vec![BorrowedMatrix::RowMajorNatural { - data: &main_data, - stride: total_cols, - col_start: num_precomputed, - width: total_cols - num_precomputed, - log_height: height, - }]; - main_builder.absorb(&src, 0); + if let Some(builder) = main_builder.as_mut() { + let src = vec![BorrowedMatrix::RowMajorNatural { + data: &main_data, + stride: total_cols, + col_start: num_precomputed, + width: total_cols - num_precomputed, + log_height: height, + }]; + builder.absorb(&src, 0); + } } // The root is what Fiat-Shamir needs; the buffer is not. Under @@ -358,67 +375,17 @@ where transcript.append_bytes(root); } - let main_mmcs = main_builder.finish(); + let main_mmcs = if use_device_main { + commit_main_device(&retained_main, &shape, num_tables)? + } else { + main_builder + .take() + .expect("CPU main builder present when device commit is disabled") + .finish() + }; let main_root = main_mmcs.root(); transcript.append_bytes(&main_root); - // GPU cross-check (3c): the device mixed-height MMCS must build the same - // main-round root the host `StreamingMmcsBuilder` just did, from the same - // matrices (same height grouping, same absorb order, same column ranges). - // Inert unless every contributing main LDE is retained (so `RecomputeLde` - // skips it), the field is Goldilocks (the kernels are), and there is a GPU. - #[cfg(feature = "cuda")] - { - use math::field::goldilocks::GoldilocksField; - use std::any::TypeId; - if TypeId::of::() == TypeId::of::() { - let mut raws: Vec<(Vec, u64, u64, u64)> = Vec::new(); - let mut all_present = true; - for table in 0..num_tables { - if shape.carved_main.map(|c| c.table) == Some(table) { - continue; - } - match &retained_main[table] { - Some((data, total_cols)) => { - let width = matrix_width(&shape.main, table) as u64; - let stride = *total_cols as u64; - // SAFETY: Field == GoldilocksField (checked), whose element - // value is a `u64` — the same cast the per-table GPU commit - // uses (`gpu_lde::columns_to_u64_base`). - let raw: Vec = data - .iter() - .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) - .collect(); - raws.push((raw, stride, stride - width, shape.heights[table] as u64)); - } - None => { - all_present = false; - break; - } - } - } - if all_present && !raws.is_empty() { - let inputs: Vec = raws - .iter() - .map(|(d, stride, col_start, h)| math_cuda::mmcs::MmcsRowMajorInput { - data: d.as_slice(), - stride: *stride, - col_start: *col_start, - col_end: *stride, - log_height: *h, - }) - .collect(); - if let Ok(dev_root) = math_cuda::mmcs::commit_mixed_row_major_root(&inputs) { - debug_assert_eq!( - dev_root, main_root, - "device main-round MMCS root must equal the host \ - StreamingMmcsBuilder root" - ); - } - } - } - } - // ===================================================================== // Phase 2 — LogUp challenges, then the auxiliary round // ===================================================================== @@ -445,7 +412,8 @@ where let mut bus_public_inputs: Vec>> = (0..num_tables).map(|_| None).collect(); - let mut aux_builder = (!shape.aux.is_empty()) + let use_device_aux = device_commit_enabled::(residency); + let mut aux_builder = (!shape.aux.is_empty() && !use_device_aux) .then(|| StreamingMmcsBuilder::::new(&shape.aux.dims)); let mut retained_aux: Vec>, usize)>> = (0..num_tables).map(|_| None).collect(); @@ -464,9 +432,9 @@ where .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; } - let Some(builder) = aux_builder.as_mut() else { + if shape.aux.is_empty() { continue; - }; + } let (aux_data, aux_cols) = P::expand_aux_lde_row_major( trace, &domains[table], @@ -477,14 +445,16 @@ where stats.aux_lde_expansions += 1; let bytes = lde_bytes::(aux_data.len()); ledger.alloc(bytes); - let src = vec![BorrowedMatrix::RowMajorNatural { - data: &aux_data, - stride: aux_cols, - col_start: 0, - width: aux_cols, - log_height: shape.heights[table], - }]; - builder.absorb(&src, 0); + if let Some(builder) = aux_builder.as_mut() { + let src = vec![BorrowedMatrix::RowMajorNatural { + data: &aux_data, + stride: aux_cols, + col_start: 0, + width: aux_cols, + log_height: shape.heights[table], + }]; + builder.absorb(&src, 0); + } match residency { ResidencyMode::Retain => retained_aux[table] = Some((aux_data, aux_cols)), ResidencyMode::RecomputeLde => { @@ -494,53 +464,16 @@ where } } - let aux_mmcs = aux_builder.map(StreamingMmcsBuilder::finish); + let aux_mmcs = if use_device_aux && !shape.aux.is_empty() { + Some(commit_aux_device(&retained_aux, &shape, num_tables)?) + } else { + aux_builder.map(StreamingMmcsBuilder::finish) + }; let aux_root = aux_mmcs.as_ref().map(MixedMmcs::root); if let Some(root) = aux_root { transcript.append_bytes(&root); } - // GPU cross-check (3c): the device row-major ext3 MMCS must build the same - // aux-round root the host did (Retain + Goldilocks-ext3 + GPU, else inert). - #[cfg(feature = "cuda")] - if let Some(aux_root_val) = aux_root { - use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; - use std::any::TypeId; - if TypeId::of::() == TypeId::of::() { - let mut raws: Vec<(Vec, u64, u64)> = Vec::new(); - for table in 0..num_tables { - if let Some((data, aux_cols)) = &retained_aux[table] { - // SAFETY: FieldExtension == Degree3Goldilocks (checked), a - // contiguous `[u64; 3]` per element — same cast as - // `gpu_lde::columns_to_u64_ext3`. - let raw: Vec = unsafe { - std::slice::from_raw_parts(data.as_ptr() as *const u64, data.len() * 3) - } - .to_vec(); - raws.push((raw, *aux_cols as u64, shape.heights[table] as u64)); - } - } - if !raws.is_empty() { - let inputs: Vec = raws - .iter() - .map(|(d, stride, h)| math_cuda::mmcs::MmcsExt3RowMajorInput { - data: d.as_slice(), - stride: *stride, - col_start: 0, - col_end: *stride, - log_height: *h, - }) - .collect(); - if let Ok(dev_root) = math_cuda::mmcs::commit_mixed_ext3_row_major_root(&inputs) { - debug_assert_eq!( - dev_root, aux_root_val, - "device aux-round MMCS root must equal the host root" - ); - } - } - } - } - // ===================================================================== // Phase 3 — bus contributions, beta per table, the composition-parts round // ===================================================================== @@ -550,7 +483,9 @@ where transcript.append_field_element(&bpi.table_contribution); } - let mut parts_builder = StreamingMmcsBuilder::::new(&shape.parts.dims); + let use_device_parts = device_commit_enabled::(residency); + let mut parts_builder = (!use_device_parts) + .then(|| StreamingMmcsBuilder::::new(&shape.parts.dims)); let mut retained_parts: Vec>>> = (0..num_tables).map(|_| Vec::new()).collect(); @@ -613,11 +548,13 @@ where .map(|p| lde_bytes::(p.len())) .sum(); parts_ledger.alloc(parts_bytes); - let src = vec![BorrowedMatrix::ColMajorNatural { - cols: &parts, - log_height: shape.heights[table], - }]; - parts_builder.absorb(&src, 0); + if let Some(builder) = parts_builder.as_mut() { + let src = vec![BorrowedMatrix::ColMajorNatural { + cols: &parts, + log_height: shape.heights[table], + }]; + builder.absorb(&src, 0); + } // Parts are RETAINED: rebuilding them is a second constraint evaluation. retained_parts[table] = parts; @@ -631,68 +568,17 @@ where ); } - let parts_mmcs = parts_builder.finish(); + let parts_mmcs = if use_device_parts { + commit_parts_device(&retained_parts, &shape, num_tables)? + } else { + parts_builder + .take() + .expect("CPU parts builder present when device commit is disabled") + .finish() + }; let parts_root = parts_mmcs.root(); transcript.append_bytes(&parts_root); - // GPU cross-check (3c): the device ext3 mixed-height MMCS must build the same - // parts-round root the host did. Parts are always retained; the field is - // Goldilocks-ext3 and there is a GPU, or this is inert. - #[cfg(feature = "cuda")] - { - use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; - use std::any::TypeId; - if TypeId::of::() == TypeId::of::() { - let mut slabs: Vec<(Vec, u64, u64, u64)> = Vec::new(); - let mut all_present = true; - for table in 0..num_tables { - let parts = &retained_parts[table]; - if parts.is_empty() || parts[0].is_empty() { - all_present = false; - break; - } - let num_parts = parts.len(); - let num_rows = parts[0].len(); - let mut slab = vec![0u64; num_parts * 3 * num_rows]; - for (p, col) in parts.iter().enumerate() { - for (row, elem) in col.iter().enumerate() { - // SAFETY: FieldExtension == Degree3Goldilocks (checked), - // whose value is `[u64; 3]` — same cast as - // `gpu_lde::columns_to_u64_ext3`. - let comps = - unsafe { std::slice::from_raw_parts(elem.value() as *const _ as *const u64, 3) }; - for (k, &c) in comps.iter().enumerate() { - slab[(p * 3 + k) * num_rows + row] = c; - } - } - } - slabs.push(( - slab, - num_rows as u64, - num_parts as u64, - shape.heights[table] as u64, - )); - } - if all_present && !slabs.is_empty() { - let inputs: Vec = slabs - .iter() - .map(|(d, col_stride, num_parts, h)| math_cuda::mmcs::MmcsExt3SlabInput { - data: d.as_slice(), - col_stride: *col_stride, - num_parts: *num_parts, - log_height: *h, - }) - .collect(); - if let Ok(dev_root) = math_cuda::mmcs::commit_mixed_ext3_slabs_root(&inputs) { - debug_assert_eq!( - dev_root, parts_root, - "device parts-round MMCS root must equal the host root" - ); - } - } - } - } - // ===================================================================== // Phase 4 — z per table, OOD evaluations // ===================================================================== @@ -1088,6 +974,241 @@ fn matrix_index(round: &RoundShape, table: usize) -> Option { round.tables.iter().position(|&t| t == table) } +/// Whether a round is committed on the GPU (device tree authoritative) rather +/// than by the host `StreamingMmcsBuilder`. True only with the `cuda` feature, +/// on a Goldilocks field (base or ext3 — the only fields the kernels support), +/// and under `ResidencyMode::Retain` (the LDEs must survive to feed the device +/// tree after the round's loop). A deliberate mode select; there is no silent +/// per-call host↔device fallback — a device error aborts. +#[cfg(feature = "cuda")] +fn device_commit_enabled(residency: ResidencyMode) -> bool { + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; + use math::field::goldilocks::GoldilocksField; + use std::any::TypeId; + residency == ResidencyMode::Retain + && (TypeId::of::() == TypeId::of::() + || TypeId::of::() == TypeId::of::()) +} + +#[cfg(not(feature = "cuda"))] +fn device_commit_enabled(_residency: ResidencyMode) -> bool { + false +} + +/// Commit the main round on the GPU and rebuild the host-side `MixedMmcs` from +/// the device heap (keccak on device, digest layers back). Only called when +/// [`device_commit_enabled`] held, i.e. Goldilocks + Retain + cuda; a device +/// error is a hard abort, never a silent host fallback. +#[cfg(feature = "cuda")] +fn commit_main_device( + retained_main: &[Option<(Vec>, usize)>], + shape: &EpochShape, + num_tables: usize, +) -> Result, ProvingError> +where + FieldElement: math::traits::AsBytes + Sync + Send, +{ + let mut raws: Vec<(Vec, u64, u64, u64)> = Vec::new(); + for table in 0..num_tables { + if shape.carved_main.map(|c| c.table) == Some(table) { + continue; + } + match &retained_main[table] { + Some((data, total_cols)) => { + let width = matrix_width(&shape.main, table) as u64; + let stride = *total_cols as u64; + // SAFETY: Goldilocks (device_commit_enabled gated it), element + // value is a u64 — same cast as `gpu_lde::columns_to_u64_base`. + let raw: Vec = data + .iter() + .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .collect(); + raws.push((raw, stride, stride - width, shape.heights[table] as u64)); + } + None => { + return Err(ProvingError::WrongParameter( + "device main commit: a non-carved main LDE was not retained".to_string(), + )); + } + } + } + let inputs: Vec = raws + .iter() + .map(|(d, stride, col_start, h)| math_cuda::mmcs::MmcsRowMajorInput { + data: d.as_slice(), + stride: *stride, + col_start: *col_start, + col_end: *stride, + log_height: *h, + }) + .collect(); + let nodes = math_cuda::mmcs::commit_mixed_row_major_nodes(&inputs) + .map_err(|e| ProvingError::WrongParameter(format!("device main commit failed: {e:?}")))?; + let h_max = shape + .main + .dims + .iter() + .map(|&(lh, _)| lh) + .max() + .ok_or_else(|| ProvingError::WrongParameter("device main commit: empty round".to_string()))?; + Ok(MixedMmcs::from_heap_nodes( + shape.main.dims.clone(), + h_max, + &nodes, + )) +} + +#[cfg(not(feature = "cuda"))] +fn commit_main_device( + _retained_main: &[Option<(Vec>, usize)>], + _shape: &EpochShape, + _num_tables: usize, +) -> Result, ProvingError> +where + FieldElement: math::traits::AsBytes + Sync + Send, +{ + unreachable!("device commit is disabled without the cuda feature") +} + +/// Commit the aux round (row-major ext3) on the GPU. Precondition + policy as +/// [`commit_main_device`]. `retained_aux[t]` is `Some` exactly for the aux +/// tables, in the order the host absorbs them. +#[cfg(feature = "cuda")] +fn commit_aux_device( + retained_aux: &[Option<(Vec>, usize)>], + shape: &EpochShape, + num_tables: usize, +) -> Result, ProvingError> +where + FieldElement: math::traits::AsBytes + Sync + Send, +{ + let mut raws: Vec<(Vec, u64, u64)> = Vec::new(); + for table in 0..num_tables { + if let Some((data, aux_cols)) = &retained_aux[table] { + // SAFETY: ext3 Goldilocks (gated), contiguous `[u64; 3]` per element. + let raw: Vec = + unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u64, data.len() * 3) } + .to_vec(); + raws.push((raw, *aux_cols as u64, shape.heights[table] as u64)); + } + } + if raws.is_empty() { + return Err(ProvingError::WrongParameter( + "device aux commit: no retained aux LDEs".to_string(), + )); + } + let inputs: Vec = raws + .iter() + .map(|(d, stride, h)| math_cuda::mmcs::MmcsExt3RowMajorInput { + data: d.as_slice(), + stride: *stride, + col_start: 0, + col_end: *stride, + log_height: *h, + }) + .collect(); + let nodes = math_cuda::mmcs::commit_mixed_ext3_row_major_nodes(&inputs) + .map_err(|e| ProvingError::WrongParameter(format!("device aux commit failed: {e:?}")))?; + let h_max = shape + .aux + .dims + .iter() + .map(|&(lh, _)| lh) + .max() + .ok_or_else(|| ProvingError::WrongParameter("device aux commit: empty round".to_string()))?; + Ok(MixedMmcs::from_heap_nodes(shape.aux.dims.clone(), h_max, &nodes)) +} + +#[cfg(not(feature = "cuda"))] +fn commit_aux_device( + _retained_aux: &[Option<(Vec>, usize)>], + _shape: &EpochShape, + _num_tables: usize, +) -> Result, ProvingError> +where + FieldElement: math::traits::AsBytes + Sync + Send, +{ + unreachable!("device commit is disabled without the cuda feature") +} + +/// Commit the parts round (column-major ext3 slabs) on the GPU. Precondition + +/// policy as [`commit_main_device`]. Parts are always retained. +#[cfg(feature = "cuda")] +fn commit_parts_device( + retained_parts: &[Vec>>], + shape: &EpochShape, + num_tables: usize, +) -> Result, ProvingError> +where + FieldElement: math::traits::AsBytes + Sync + Send, +{ + let mut slabs: Vec<(Vec, u64, u64, u64)> = Vec::new(); + for table in 0..num_tables { + let parts = &retained_parts[table]; + if parts.is_empty() || parts[0].is_empty() { + return Err(ProvingError::WrongParameter( + "device parts commit: a table has no retained parts".to_string(), + )); + } + let num_parts = parts.len(); + let num_rows = parts[0].len(); + let mut slab = vec![0u64; num_parts * 3 * num_rows]; + for (p, col) in parts.iter().enumerate() { + for (row, elem) in col.iter().enumerate() { + // SAFETY: ext3 Goldilocks (gated), `[u64; 3]` per element. + let comps = + unsafe { std::slice::from_raw_parts(elem.value() as *const _ as *const u64, 3) }; + for (k, &c) in comps.iter().enumerate() { + slab[(p * 3 + k) * num_rows + row] = c; + } + } + } + slabs.push(( + slab, + num_rows as u64, + num_parts as u64, + shape.heights[table] as u64, + )); + } + let inputs: Vec = slabs + .iter() + .map(|(d, col_stride, num_parts, h)| math_cuda::mmcs::MmcsExt3SlabInput { + data: d.as_slice(), + col_stride: *col_stride, + num_parts: *num_parts, + log_height: *h, + }) + .collect(); + let nodes = math_cuda::mmcs::commit_mixed_ext3_slabs_nodes(&inputs) + .map_err(|e| ProvingError::WrongParameter(format!("device parts commit failed: {e:?}")))?; + let h_max = shape + .parts + .dims + .iter() + .map(|&(lh, _)| lh) + .max() + .ok_or_else(|| { + ProvingError::WrongParameter("device parts commit: empty round".to_string()) + })?; + Ok(MixedMmcs::from_heap_nodes( + shape.parts.dims.clone(), + h_max, + &nodes, + )) +} + +#[cfg(not(feature = "cuda"))] +fn commit_parts_device( + _retained_parts: &[Vec>>], + _shape: &EpochShape, + _num_tables: usize, +) -> Result, ProvingError> +where + FieldElement: math::traits::AsBytes + Sync + Send, +{ + unreachable!("device commit is disabled without the cuda feature") +} + /// The width `table` contributes to `round`. Zero when it contributes nothing. fn matrix_width(round: &RoundShape, table: usize) -> usize { matrix_index(round, table).map_or(0, |m| round.dims[m].1) From ab33308cbc7cf1068f261a0b8be9400198f7a1ff Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 19:40:43 -0300 Subject: [PATCH 10/26] feat(cuda): batched-FRI injection kernel (fri_inject_bucket_ext3) --- crypto/math-cuda/kernels/fri.cu | 20 ++++++++++++++++++++ crypto/math-cuda/src/device.rs | 2 ++ 2 files changed, 22 insertions(+) diff --git a/crypto/math-cuda/kernels/fri.cu b/crypto/math-cuda/kernels/fri.cu index bcc8f9e40..45189d15d 100644 --- a/crypto/math-cuda/kernels/fri.cu +++ b/crypto/math-cuda/kernels/fri.cu @@ -76,3 +76,23 @@ extern "C" __global__ void gather_ext3_at( out[i * 3 + 1] = evals[p * 3 + 1]; out[i * 3 + 2] = evals[p * 3 + 2]; } + +// Batched-FRI injection: `out[i] += beta_sq * bucket[i]` over ext3 elements, +// interleaved layout. `out` is the just-folded running codeword and `bucket` +// the shorter DEEP codeword whose height matches this layer; both hold `n` ext3 +// elements (3*n u64). Matches the host `inject_bucket` (fri/batched.rs). +extern "C" __global__ void fri_inject_bucket_ext3( + uint64_t *out, // 3 * n u64, modified in place + const uint64_t *bucket, // 3 * n u64 + const uint64_t *beta_sq, // 3 u64 (ext3) + uint64_t n) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + ext3::Fe3 v = ext3::make(out[i * 3], out[i * 3 + 1], out[i * 3 + 2]); + ext3::Fe3 b = ext3::make(bucket[i * 3], bucket[i * 3 + 1], bucket[i * 3 + 2]); + ext3::Fe3 bsq = ext3::make(beta_sq[0], beta_sq[1], beta_sq[2]); + ext3::Fe3 res = ext3::add(v, ext3::mul(bsq, b)); + out[i * 3] = res.a; + out[i * 3 + 1] = res.b; + out[i * 3 + 2] = res.c; +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 4fddbb1df..34cb40273 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -230,6 +230,7 @@ pub struct Backend { // fri.cubin pub fri_fold_ext3: CudaFunction, + pub fri_inject_bucket_ext3: CudaFunction, pub gather_ext3_at: CudaFunction, pub fri_update_twiddles: CudaFunction, @@ -472,6 +473,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")?, + fri_inject_bucket_ext3: fri.load_function("fri_inject_bucket_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")?, From f5fcc0798ef3c56c56002d1d0223b3bc3bf13f40 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 28 Aug 2026 19:50:34 -0300 Subject: [PATCH 11/26] feat(cuda): drive the batched FRI commit on the GPU (fold-inject-commit) --- crypto/math-cuda/src/fri.rs | 37 +++++++++ crypto/stark/src/fri/batched.rs | 25 +++++- crypto/stark/src/gpu_lde.rs | 141 ++++++++++++++++++++++++++++++++ 3 files changed, 200 insertions(+), 3 deletions(-) diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 533ff6e32..cad91b3f1 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -115,6 +115,27 @@ impl FriCommitState { Option>, Arc>, crate::lde::GpuMerkleTree, + )> { + self.fold_inject_commit_layer(zeta_raw, None, want_host) + } + + /// Like [`Self::fold_and_commit_layer`], but between the fold and the Merkle + /// commit it injects a shorter DEEP bucket into the folded codeword — + /// `running += beta_sq * bucket` (ext3), the batched FRI's height-combine + /// step (`fri/batched.rs`'s `inject_bucket`). `bucket` is + /// `Some((beta_sq_raw, codeword))` when a bucket at THIS layer's height + /// exists; its codeword must hold `current_n / 2` ext3 elements (the folded + /// length). `None` reproduces the plain commit. + #[allow(clippy::type_complexity)] + pub fn fold_inject_commit_layer( + &mut self, + zeta_raw: [u64; 3], + bucket: Option<([u64; 3], &CudaSlice)>, + want_host: bool, + ) -> Result<( + Option>, + Arc>, + crate::lde::GpuMerkleTree, )> { #[cfg(feature = "test-faults")] check_fault_injection()?; @@ -158,6 +179,22 @@ impl FriCommitState { .launch(cfg)?; } + // Batched injection: add this height's DEEP bucket into the folded + // codeword before it is committed (`running += beta_sq * bucket`), so the + // layer root binds the combined word — the device twin of `inject_bucket`. + if let Some((beta_sq_raw, bucket)) = bucket { + let beta_sq_dev = self.stream.clone_htod(&beta_sq_raw)?; + unsafe { + self.stream + .launch_builder(&be.fri_inject_bucket_ext3) + .arg(&mut out) + .arg(bucket) + .arg(&beta_sq_dev) + .arg(&n_out_u64) + .launch(cfg)?; + } + } + // SAFETY: keccak_fri_leaves_ext3 writes the leaves [num_leaves-1, 2*num_leaves-1) // and build_inner_tree_levels writes every inner node [0, num_leaves-1), so all // tight_total_nodes * 32 bytes are initialised before the D2H below reads them. diff --git a/crypto/stark/src/fri/batched.rs b/crypto/stark/src/fri/batched.rs index f9e27adf0..9738e981e 100644 --- a/crypto/stark/src/fri/batched.rs +++ b/crypto/stark/src/fri/batched.rs @@ -339,12 +339,32 @@ where let (h_min, h_max) = bucket_height_range(&combined) .expect("batched_commit_phase: combined must have at least one Some entry"); + let domain_size = 1usize << h_max; + // Inverse twiddle factors for the initial domain size. + let inv_twiddles = compute_coset_twiddles_inv(coset_offset, domain_size); + + // GPU fast path: drive the fold → inject → commit loop below on the device + // (`fri_inject_bucket_ext3` + `FriCommitState::fold_inject_commit_layer`), + // returning `None` with the transcript restored to fall back to the host. + #[cfg(feature = "cuda")] + if let Some(result) = crate::gpu_lde::try_batched_fri_commit_gpu::( + &combined, + transcript, + coset_offset, + blowup_log, + final_poly_log_degree, + &inv_twiddles, + h_min, + h_max, + ) { + return result; + } + // Take the starting codeword — NOT committed; it plays the role of layer 0. let mut running = combined[h_max] .take() .expect("combined[h_max] is Some by construction"); - let domain_size = 1usize << h_max; debug_assert_eq!( running.len(), domain_size, @@ -353,8 +373,7 @@ where let layout = BatchedFriLayout::new(h_max, h_min, blowup_log, final_poly_log_degree); - // Inverse twiddle factors for the initial domain size. - let mut inv_twiddles = compute_coset_twiddles_inv(coset_offset, domain_size); + let mut inv_twiddles = inv_twiddles; let mut fri_layer_list = Vec::with_capacity(layout.num_committed); diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 23366d67f..4323a4d99 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -3214,6 +3214,147 @@ where Some((final_poly_coeffs, fri_layer_list)) } +/// GPU drive of the BATCHED FRI commit — the height-combined analogue of +/// [`fri_commit_gpu_drive`]. Same fold/commit/terminal machinery, but between +/// each fold and its Merkle commit it injects that layer's DEEP bucket into the +/// running codeword on the device (`fri_inject_bucket_ext3` = +/// `fri/batched.rs`'s `inject_bucket`), and it uses the batched terminal floor +/// [`crate::fri::batched::BatchedFriLayout`]. `combined[h]` is the bucket at +/// height `h` (`2^h` ext3 elements); `combined[h_max]` is the starting codeword +/// and is taken. Returns `None` (transcript restored) to fall back to the host. +#[cfg(feature = "cuda")] +#[allow(clippy::too_many_arguments, clippy::type_complexity)] +pub(crate) fn try_batched_fri_commit_gpu( + combined: &[Option>>], + transcript: &mut T, + coset_offset: &FieldElement, + blowup_log: u32, + final_poly_log_degree: u32, + inv_twiddles: &[FieldElement], + h_min: usize, + h_max: usize, +) -> Option<( + Vec>, + Vec>>, +)> +where + F: IsFFTField + IsField + IsSubFieldOf + 'static, + E: IsField + 'static + Send + Sync, + FieldElement: AsBytes, + FieldElement: AsBytes, + T: IsStarkTranscript + Clone, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + let n0 = 1usize << h_max; + if n0 < 2 || n0 < gpu_lde_threshold() || inv_twiddles.len() != n0 / 2 { + return None; + } + let layout = + crate::fri::batched::BatchedFriLayout::new(h_max, h_min, blowup_log, final_poly_log_degree); + if layout.total_folds == 0 || layout.terminal_len < 2 { + return None; + } + + // Pack inv_twiddles before any transcript mutation. + let mut inv_tw_u64: Vec = Vec::with_capacity(inv_twiddles.len()); + for t in inv_twiddles { + // SAFETY: F == Goldilocks (checked); FieldElement is transparent u64. + inv_tw_u64.push(unsafe { *(t.value() as *const _ as *const u64) }); + } + + // Clone (not take) the starting codeword so a later `None` return leaves + // `combined` intact for the host fallback. + let start = match combined.get(h_max).and_then(Option::as_ref) { + Some(c) if c.len() == n0 => c.clone(), + _ => return None, + }; + // SAFETY: E == Ext3; backing is [u64; 3]. + let start_u64: &[u64] = unsafe { ext3_slice_to_u64::(&start) }; + let mut state = match math_cuda::fri::FriCommitState::new(start_u64, &inv_tw_u64, n0) { + Ok(s) => s, + Err(_) => return None, + }; + + // Upload every shorter bucket to the device, indexed by height. + let mut buckets: Vec>> = (0..=h_max).map(|_| None).collect(); + for h in 1..h_max { + if let Some(bucket) = combined.get(h).and_then(Option::as_ref) { + let bucket_u64: &[u64] = unsafe { ext3_slice_to_u64::(bucket) }; + match state.stream.clone_htod(bucket_u64) { + Ok(dev) => buckets[h] = Some(dev), + Err(_) => return None, + } + } + } + + let transcript_snapshot = transcript.clone(); + let mut fri_layer_list: Vec>> = + Vec::with_capacity(layout.num_committed); + + // `total_folds` folds: the first `num_committed` commit a layer; the last one + // is the terminal fold (no layer, coeffs emitted). Each fold to height `h` + // injects `combined[h]`. + for fold_idx in 0..layout.total_folds as usize { + let beta: FieldElement = transcript.sample_field_element(); + let beta_ptr = &beta as *const FieldElement as *const u64; + // SAFETY: E == Ext3. + let beta_raw: [u64; 3] = unsafe { [*beta_ptr, *beta_ptr.add(1), *beta_ptr.add(2)] }; + let beta_sq = beta.square(); + let bsq_ptr = &beta_sq as *const FieldElement as *const u64; + let beta_sq_raw: [u64; 3] = unsafe { [*bsq_ptr, *bsq_ptr.add(1), *bsq_ptr.add(2)] }; + + let inject_h = h_max - fold_idx - 1; + let bucket_arg = buckets + .get(inject_h) + .and_then(Option::as_ref) + .map(|b| (beta_sq_raw, b)); + + let (evals_u64, evals_dev, dev_tree) = + match state.fold_inject_commit_layer(beta_raw, bucket_arg, true) { + Ok(v) => v, + Err(_) => { + *transcript = transcript_snapshot; + return None; + } + }; + + if fold_idx < layout.num_committed { + let evaluation = evals_u64.map(|v| u64_to_ext3_vec::(&v)).unwrap_or_default(); + let root = dev_tree.root; + let merkle_tree = MerkleTree::>::from_root(root); + fri_layer_list.push(FriLayer { + evaluation, + merkle_tree, + gpu_tree: Some(dev_tree), + gpu_evals: None, + }); + let _ = evals_dev; + transcript.append_bytes(&root); + } else { + // Terminal fold: emit the low-degree coefficients. + let terminal = u64_to_ext3_vec::(&evals_u64.expect("terminal fold drains to host")); + let terminal_offset = coset_offset.pow(1u64 << layout.total_folds); + let final_poly_coeffs = crate::fri::terminal::coeffs_from_terminal_codeword::( + &terminal, + &terminal_offset, + layout.effective_k, + ); + for c in &final_poly_coeffs { + transcript.append_field_element(c); + } + GPU_FRI_CALLS.fetch_add(1, Ordering::Relaxed); + return Some((final_poly_coeffs, fri_layer_list)); + } + } + // total_folds >= 1 guarantees the terminal branch above returned. + unreachable!("batched FRI drive: terminal fold not reached") +} + /// GPU FRI query phase: gather each layer's paths on device instead of walking /// host trees. For layer `l` and query `iota` the opened position is /// `(iota >> l) >> 1`, matching [`crate::fri::query_phase`]. Paths for all From d036f5c32f045fe336df13d86a52fa2df3dc4ff5 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 2 Sep 2026 11:43:18 -0300 Subject: [PATCH 12/26] feat(cuda): streaming device MMCS commit + 3f canary for the batched prover --- crypto/math-cuda/src/mmcs.rs | 121 ++++++ crypto/stark/src/batched/prover.rs | 588 +++++++++++++++++++++-------- crypto/stark/src/batched/round4.rs | 65 +++- crypto/stark/src/fri/batched.rs | 24 +- crypto/stark/src/fri/mmcs.rs | 18 + crypto/stark/src/gpu_lde.rs | 76 ++-- 6 files changed, 673 insertions(+), 219 deletions(-) diff --git a/crypto/math-cuda/src/mmcs.rs b/crypto/math-cuda/src/mmcs.rs index 0f5f63dae..b255b5405 100644 --- a/crypto/math-cuda/src/mmcs.rs +++ b/crypto/math-cuda/src/mmcs.rs @@ -501,3 +501,124 @@ pub fn commit_mixed_ext3_row_major_root(inputs: &[MmcsExt3RowMajorInput]) -> Res root.copy_from_slice(&nodes[0..32]); Ok(root) } + +/// Persistent, streaming device commit for one mixed-height MMCS round — the +/// same per-height [`MmcsGroupHasher`] machinery the `commit_mixed_*_nodes` +/// helpers use, but kept ALIVE across the prover's round loop so each table's +/// LDE is absorbed the moment it is produced and freed immediately, instead of +/// every LDE being retained until one all-at-once commit. VRAM holds only the +/// per-leaf sponges (~204 B/leaf) plus one table's uploaded buffer at a time, +/// never `O(N)` LDEs — so the device commit runs under `RecomputeLde`, not only +/// `Retain` (which is what made a large epoch OOM: it retained every LDE just to +/// feed the commit). +/// +/// Absorb matrices IN INPUT ORDER (the leaf concatenation binds it), one call +/// per matrix, using the layout method matching the round (row-major base for +/// main, ext3 row-major for aux, ext3 slabs for parts). Each absorb uploads the +/// host buffer, hashes it, and synchronizes, so the caller may free the host LDE +/// as soon as the call returns. [`Self::finish`] then finalizes every group and +/// climbs, returning the standard heap node array (feed to +/// `MixedMmcs::from_heap_nodes`) — byte-identical to the all-at-once commit +/// because it is the same kernels in the same order. +pub struct StreamingMixedMmcs { + stream: Arc, + /// Indexed by `log_height`: the group's live sponges, created on first + /// absorb of a matrix at that height. + hashers: Vec>, +} + +impl StreamingMixedMmcs { + /// `h_max` is the tallest `log_height` in the round (its group is the tree's + /// base layer). Groups are created lazily as matrices arrive. + pub fn new(h_max: u64) -> Result { + let stream = backend()?.next_stream(); + Ok(Self { + stream, + hashers: (0..=h_max as usize).map(|_| None).collect(), + }) + } + + fn group_mut(&mut self, log_height: u64) -> Result<&mut MmcsGroupHasher> { + let idx = log_height as usize; + assert!( + idx < self.hashers.len(), + "log_height {log_height} exceeds the h_max this StreamingMixedMmcs was built for" + ); + if self.hashers[idx].is_none() { + self.hashers[idx] = Some(MmcsGroupHasher::new(&self.stream, log_height)?); + } + Ok(self.hashers[idx].as_mut().expect("just populated")) + } + + /// Absorb one row-major base-field matrix (main round). `data` is the host + /// LDE (Goldilocks u64 words); columns `[col_start, col_end)` of `row_stride` + /// are committed. + pub fn absorb_row_major( + &mut self, + log_height: u64, + data: &[u64], + row_stride: u64, + col_start: u64, + col_end: u64, + ) -> Result<()> { + let stream = self.stream.clone(); + let dev = stream.clone_htod(data)?; + self.group_mut(log_height)? + .absorb_row_major(&stream, &dev, row_stride, col_start, col_end)?; + // Defensive: ensure the H2D copy and the absorb that reads it have + // completed before the caller frees `data` (the streaming residency + // policy this type exists for). + stream.synchronize()?; + Ok(()) + } + + /// Absorb one ext3 row-major matrix (aux round). `data` is the host LDE with + /// an element's 3 components consecutive. + pub fn absorb_ext3_row_major( + &mut self, + log_height: u64, + data: &[u64], + stride: u64, + col_start: u64, + col_end: u64, + ) -> Result<()> { + let stream = self.stream.clone(); + let dev = stream.clone_htod(data)?; + self.group_mut(log_height)? + .absorb_ext3_row_major(&stream, &dev, stride, col_start, col_end)?; + stream.synchronize()?; + Ok(()) + } + + /// Absorb one ext3 column-major slab matrix (parts round). `parts` is the + /// slab buffer, component `k` of column `c` at `(c*3 + k) * col_stride`. + pub fn absorb_ext3_slabs( + &mut self, + log_height: u64, + parts: &[u64], + col_stride: u64, + num_parts: u64, + ) -> Result<()> { + let stream = self.stream.clone(); + let dev = stream.clone_htod(parts)?; + self.group_mut(log_height)? + .absorb_ext3_slabs(&stream, &dev, col_stride, num_parts)?; + stream.synchronize()?; + Ok(()) + } + + /// Finalize every absorbed group and climb, returning the standard heap node + /// array (`(2L-1)*32` bytes for `L = 2^(h_max-1)`), host-side. + pub fn finish(self) -> Result> { + let Self { stream, hashers } = self; + let mut group_digests: Vec>> = + (0..hashers.len()).map(|_| None).collect(); + for (h, hasher) in hashers.into_iter().enumerate() { + if let Some(hasher) = hasher { + group_digests[h] = Some(hasher.finalize(&stream)?); + } + } + let nodes = build_mmcs_tree_on_device(&stream, &group_digests)?; + stream.clone_dtoh(&nodes) + } +} diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index 64265db80..458410f4b 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -54,17 +54,54 @@ //! //! # The GPU commit //! -//! Under `--features cuda`, on Goldilocks (base + ext3) and `ResidencyMode:: -//! Retain`, all four MMCS trees (main / aux / parts) are committed ON THE DEVICE: -//! the keccak leaf hash and the climb run on the GPU (`crypto/math-cuda/`'s -//! mixed-height MMCS), and only the digest layers come back to rebuild the -//! host-side [`MixedMmcs`] via [`MixedMmcs::from_heap_nodes`] — so the openings -//! (row values from the retained LDE, one shared auth path per query) are served -//! exactly as before. The host [`StreamingMmcsBuilder`] is skipped entirely for -//! those rounds. This is a DELIBERATE mode select (see `device_commit_enabled`), -//! not a silent per-call host↔device fallback — a device error is a hard abort — -//! precisely so the residency numbers above stay reproducible. Off cuda / off -//! Goldilocks / under `RecomputeLde`, the host builder commits as it always did. +//! Under `--features cuda`, on Goldilocks (base + ext3), all four MMCS trees +//! (main / aux / parts) are committed ON THE DEVICE: the keccak leaf hash and +//! the climb run on the GPU (`crypto/math-cuda/`'s mixed-height MMCS), and only +//! the digest layers come back to rebuild the host-side [`MixedMmcs`] via +//! [`MixedMmcs::from_heap_nodes`] — so the openings (row values from the LDE, one +//! shared auth path per query) are served exactly as before. The host +//! [`StreamingMmcsBuilder`] is skipped entirely for those rounds. This is a +//! DELIBERATE mode select (see `device_commit_enabled`), not a silent per-call +//! host↔device fallback — a device error is a hard abort. Off cuda / off +//! Goldilocks the host builder commits as it always did. +//! +//! ## Streaming, so it runs under either residency +//! +//! The device commit is STREAMING ([`math_cuda::mmcs::StreamingMixedMmcs`], the +//! device twin of [`StreamingMmcsBuilder`]): the main and aux rounds absorb each +//! table's LDE into the device tree INSIDE the round loop — the moment it is +//! produced — and free it, so only the per-leaf sponges stay resident, never all +//! LDEs at once. It therefore runs under BOTH `Retain` and `RecomputeLde`. This +//! is not cosmetic: the old all-at-once commit retained every LDE just to feed it +//! (O(N) memory), which OOM'd a large epoch; streaming under `RecomputeLde` is +//! what lets a big block commit on the GPU at all. Parts are ALWAYS retained +//! (recomputing them re-runs constraint eval), so that round streams one slab at +//! a time from the retained parts POST-loop instead. +//! +//! ## No silent fallback, enforced in release +//! +//! Because the device commit is authoritative, a silent device-side corruption +//! (a layout / heap bug that ships a wrong-but-self-consistent tree) would +//! otherwise ship a bad proof: the byte-identical device-vs-host parity that +//! guards these paths is `debug_assert!` / test-only and compiled OUT of a +//! release prove. So a device commit runs a canary (`xcheck_device_commit`) that +//! re-authenticates a handful of the committed tree's leaves against the LDE on +//! the host, through the same `open_batch` / `verify_batch` the verifier trusts. +//! Its cost is `k` leaf paths against an `O(N)` tree build; a mismatch is a hard +//! [`ProvingError`]. The canary reads the LDEs back, so main/aux run it only +//! under `Retain` (`RecomputeLde` dropped them — the debug parity tests cover +//! that path); parts, always retained, canary every time. +//! +//! The batched FRI commit follows the SAME policy. Its GPU drive is selected in +//! [`crate::batched::round4::commit_batched_fri`] (not inside +//! `crate::fri::batched`'s `batched_commit_phase`, which stays a pure host +//! build): when the device path is selected and a CUDA op fails, the error +//! propagates as a hard abort rather than silently restoring the transcript and +//! rebuilding on the host. The only sanctioned host build is NON-selection (off +//! cuda, off Goldilocks, or below the GPU threshold), reached before any +//! transcript mutation. (A silent FRI-side corruption — a wrong-but-no-error +//! device root — is not yet canaried here; the FRI verifier's own +//! fold-consistency check is the current backstop. Deferred.) use math::fft::bit_reversing::in_place_bit_reverse_permute_row_major; use math::field::element::FieldElement; @@ -267,6 +304,19 @@ where let use_device_main = device_commit_enabled::(residency); let mut main_builder = (!use_device_main).then(|| StreamingMmcsBuilder::::new(&shape.main.dims)); + // The device twin of `main_builder`: created before the loop, absorbs each + // table's LDE inside it (streaming), climbed after. Present iff device main. + #[cfg(feature = "cuda")] + let mut device_main: Option = if use_device_main { + let h_max = round_h_max(&shape.main, "main")?; + Some( + math_cuda::mmcs::StreamingMixedMmcs::new(h_max as u64).map_err(|e| { + ProvingError::WrongParameter(format!("device main commit init failed: {e:?}")) + })?, + ) + } else { + None + }; let mut retained_main: Vec>, usize)>> = (0..num_tables).map(|_| None).collect(); // The carved table's standalone main tree and its root. Built inside the @@ -354,6 +404,27 @@ where }]; builder.absorb(&src, 0); } + // Streaming device absorb: hash this table into the device tree NOW, + // so the LDE can be freed below (RecomputeLde) without retaining it. + #[cfg(feature = "cuda")] + if let Some(dev) = device_main.as_mut() { + // SAFETY: Goldilocks base (device_commit_enabled gated it); + // FieldElement is #[repr(transparent)] over u64. Committed + // columns are [num_precomputed, total_cols). + let data_u64 = unsafe { + std::slice::from_raw_parts(main_data.as_ptr() as *const u64, main_data.len()) + }; + dev.absorb_row_major( + height as u64, + data_u64, + total_cols as u64, + num_precomputed as u64, + total_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!("device main absorb failed: {e:?}")) + })?; + } } // The root is what Fiat-Shamir needs; the buffer is not. Under @@ -376,7 +447,22 @@ where } let main_mmcs = if use_device_main { - commit_main_device(&retained_main, &shape, num_tables)? + #[cfg(feature = "cuda")] + { + finalize_device_main( + device_main + .take() + .expect("device_main present under use_device_main"), + &retained_main, + &shape, + num_tables, + residency, + )? + } + #[cfg(not(feature = "cuda"))] + { + unreachable!("device commit requires the cuda feature") + } } else { main_builder .take() @@ -415,6 +501,18 @@ where let use_device_aux = device_commit_enabled::(residency); let mut aux_builder = (!shape.aux.is_empty() && !use_device_aux) .then(|| StreamingMmcsBuilder::::new(&shape.aux.dims)); + #[cfg(feature = "cuda")] + let mut device_aux: Option = + if use_device_aux && !shape.aux.is_empty() { + let h_max = round_h_max(&shape.aux, "aux")?; + Some( + math_cuda::mmcs::StreamingMixedMmcs::new(h_max as u64).map_err(|e| { + ProvingError::WrongParameter(format!("device aux commit init failed: {e:?}")) + })?, + ) + } else { + None + }; let mut retained_aux: Vec>, usize)>> = (0..num_tables).map(|_| None).collect(); @@ -455,6 +553,26 @@ where }]; builder.absorb(&src, 0); } + // Streaming device absorb (ext3 row-major), so the aux LDE can be freed + // below without retaining it. + #[cfg(feature = "cuda")] + if let Some(dev) = device_aux.as_mut() { + // SAFETY: ext3 Goldilocks (device_commit_enabled gated it); an + // element is 3 consecutive u64. + let data_u64 = unsafe { + std::slice::from_raw_parts(aux_data.as_ptr() as *const u64, aux_data.len() * 3) + }; + dev.absorb_ext3_row_major( + shape.heights[table] as u64, + data_u64, + aux_cols as u64, + 0, + aux_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!("device aux absorb failed: {e:?}")) + })?; + } match residency { ResidencyMode::Retain => retained_aux[table] = Some((aux_data, aux_cols)), ResidencyMode::RecomputeLde => { @@ -465,7 +583,20 @@ where } let aux_mmcs = if use_device_aux && !shape.aux.is_empty() { - Some(commit_aux_device(&retained_aux, &shape, num_tables)?) + #[cfg(feature = "cuda")] + { + Some(finalize_device_aux( + device_aux.take().expect("device_aux present under use_device_aux"), + &retained_aux, + &shape, + num_tables, + residency, + )?) + } + #[cfg(not(feature = "cuda"))] + { + unreachable!("device commit requires the cuda feature") + } } else { aux_builder.map(StreamingMmcsBuilder::finish) }; @@ -790,7 +921,7 @@ where params.final_poly_log_degree, params.grinding_factor, params.num_queries, - ) + )? }; // ===================================================================== @@ -975,19 +1106,21 @@ fn matrix_index(round: &RoundShape, table: usize) -> Option { } /// Whether a round is committed on the GPU (device tree authoritative) rather -/// than by the host `StreamingMmcsBuilder`. True only with the `cuda` feature, -/// on a Goldilocks field (base or ext3 — the only fields the kernels support), -/// and under `ResidencyMode::Retain` (the LDEs must survive to feed the device -/// tree after the round's loop). A deliberate mode select; there is no silent -/// per-call host↔device fallback — a device error aborts. +/// than by the host `StreamingMmcsBuilder`. True with the `cuda` feature on a +/// Goldilocks field (base or ext3 — the only fields the kernels support), under +/// EITHER residency: the streaming device commit +/// ([`math_cuda::mmcs::StreamingMixedMmcs`]) absorbs each table's LDE inside the +/// round loop and frees it, keeping only the per-leaf sponges resident, so it no +/// longer needs `Retain` (which retained every LDE just to feed one all-at-once +/// commit — the O(N) memory that OOM'd a large epoch). A deliberate mode select; +/// there is no silent per-call host↔device fallback — a device error aborts. #[cfg(feature = "cuda")] -fn device_commit_enabled(residency: ResidencyMode) -> bool { +fn device_commit_enabled(_residency: ResidencyMode) -> bool { use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; use math::field::goldilocks::GoldilocksField; use std::any::TypeId; - residency == ResidencyMode::Retain - && (TypeId::of::() == TypeId::of::() - || TypeId::of::() == TypeId::of::()) + TypeId::of::() == TypeId::of::() + || TypeId::of::() == TypeId::of::() } #[cfg(not(feature = "cuda"))] @@ -995,54 +1128,26 @@ fn device_commit_enabled(_residency: ResidencyMode) -> bool { false } -/// Commit the main round on the GPU and rebuild the host-side `MixedMmcs` from -/// the device heap (keccak on device, digest layers back). Only called when -/// [`device_commit_enabled`] held, i.e. Goldilocks + Retain + cuda; a device -/// error is a hard abort, never a silent host fallback. +/// Finalize the main round's STREAMING device commit: climb the tree from the +/// per-leaf sponges the round loop already absorbed into `device` (one table at +/// a time, then freed — see [`math_cuda::mmcs::StreamingMixedMmcs`]), rebuild the +/// host-side `MixedMmcs` from the heap, and — when the LDEs are still resident +/// (`Retain`) — run the release canary. Only reached under +/// [`device_commit_enabled`]; a device error is a hard abort, never a silent +/// host fallback. #[cfg(feature = "cuda")] -fn commit_main_device( +fn finalize_device_main( + device: math_cuda::mmcs::StreamingMixedMmcs, retained_main: &[Option<(Vec>, usize)>], shape: &EpochShape, num_tables: usize, + residency: ResidencyMode, ) -> Result, ProvingError> where FieldElement: math::traits::AsBytes + Sync + Send, { - let mut raws: Vec<(Vec, u64, u64, u64)> = Vec::new(); - for table in 0..num_tables { - if shape.carved_main.map(|c| c.table) == Some(table) { - continue; - } - match &retained_main[table] { - Some((data, total_cols)) => { - let width = matrix_width(&shape.main, table) as u64; - let stride = *total_cols as u64; - // SAFETY: Goldilocks (device_commit_enabled gated it), element - // value is a u64 — same cast as `gpu_lde::columns_to_u64_base`. - let raw: Vec = data - .iter() - .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) - .collect(); - raws.push((raw, stride, stride - width, shape.heights[table] as u64)); - } - None => { - return Err(ProvingError::WrongParameter( - "device main commit: a non-carved main LDE was not retained".to_string(), - )); - } - } - } - let inputs: Vec = raws - .iter() - .map(|(d, stride, col_start, h)| math_cuda::mmcs::MmcsRowMajorInput { - data: d.as_slice(), - stride: *stride, - col_start: *col_start, - col_end: *stride, - log_height: *h, - }) - .collect(); - let nodes = math_cuda::mmcs::commit_mixed_row_major_nodes(&inputs) + let nodes = device + .finish() .map_err(|e| ProvingError::WrongParameter(format!("device main commit failed: {e:?}")))?; let h_max = shape .main @@ -1051,88 +1156,84 @@ where .map(|&(lh, _)| lh) .max() .ok_or_else(|| ProvingError::WrongParameter("device main commit: empty round".to_string()))?; - Ok(MixedMmcs::from_heap_nodes( - shape.main.dims.clone(), - h_max, - &nodes, - )) -} - -#[cfg(not(feature = "cuda"))] -fn commit_main_device( - _retained_main: &[Option<(Vec>, usize)>], - _shape: &EpochShape, - _num_tables: usize, -) -> Result, ProvingError> -where - FieldElement: math::traits::AsBytes + Sync + Send, -{ - unreachable!("device commit is disabled without the cuda feature") + let mmcs = MixedMmcs::from_heap_nodes(shape.main.dims.clone(), h_max, &nodes); + // The canary reads the retained LDEs back through a source in the SAME order + // as the absorbed inputs; under `RecomputeLde` they are gone (the debug + // parity tests cover that path). + if residency == ResidencyMode::Retain { + let mut source: Vec> = Vec::new(); + for table in 0..num_tables { + if shape.carved_main.map(|c| c.table) == Some(table) { + continue; + } + match &retained_main[table] { + Some((data, total_cols)) => { + let width = matrix_width(&shape.main, table); + source.push(BorrowedMatrix::RowMajorNatural { + data, + stride: *total_cols, + col_start: *total_cols - width, + width, + log_height: shape.heights[table], + }); + } + None => { + return Err(ProvingError::WrongParameter( + "device main commit canary: a non-carved main LDE was not retained" + .to_string(), + )); + } + } + } + xcheck_device_commit("main", &mmcs, source)?; + } + Ok(mmcs) } -/// Commit the aux round (row-major ext3) on the GPU. Precondition + policy as -/// [`commit_main_device`]. `retained_aux[t]` is `Some` exactly for the aux -/// tables, in the order the host absorbs them. +/// Finalize the aux round's STREAMING device commit (ext3 row-major). Aux twin +/// of [`finalize_device_main`]: the round loop absorbed each aux LDE into +/// `device`; here we climb + rebuild + (under `Retain`) canary. The aux round +/// commits every aux column, so `col_start = 0`. #[cfg(feature = "cuda")] -fn commit_aux_device( +fn finalize_device_aux( + device: math_cuda::mmcs::StreamingMixedMmcs, retained_aux: &[Option<(Vec>, usize)>], shape: &EpochShape, num_tables: usize, + residency: ResidencyMode, ) -> Result, ProvingError> where FieldElement: math::traits::AsBytes + Sync + Send, { - let mut raws: Vec<(Vec, u64, u64)> = Vec::new(); - for table in 0..num_tables { - if let Some((data, aux_cols)) = &retained_aux[table] { - // SAFETY: ext3 Goldilocks (gated), contiguous `[u64; 3]` per element. - let raw: Vec = - unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u64, data.len() * 3) } - .to_vec(); - raws.push((raw, *aux_cols as u64, shape.heights[table] as u64)); + let nodes = device + .finish() + .map_err(|e| ProvingError::WrongParameter(format!("device aux commit failed: {e:?}")))?; + let h_max = round_h_max(&shape.aux, "aux")?; + let mmcs = MixedMmcs::from_heap_nodes(shape.aux.dims.clone(), h_max, &nodes); + if residency == ResidencyMode::Retain { + let mut source: Vec> = Vec::new(); + for table in 0..num_tables { + if let Some((data, aux_cols)) = &retained_aux[table] { + source.push(BorrowedMatrix::RowMajorNatural { + data, + stride: *aux_cols, + col_start: 0, + width: *aux_cols, + log_height: shape.heights[table], + }); + } } + xcheck_device_commit("aux", &mmcs, source)?; } - if raws.is_empty() { - return Err(ProvingError::WrongParameter( - "device aux commit: no retained aux LDEs".to_string(), - )); - } - let inputs: Vec = raws - .iter() - .map(|(d, stride, h)| math_cuda::mmcs::MmcsExt3RowMajorInput { - data: d.as_slice(), - stride: *stride, - col_start: 0, - col_end: *stride, - log_height: *h, - }) - .collect(); - let nodes = math_cuda::mmcs::commit_mixed_ext3_row_major_nodes(&inputs) - .map_err(|e| ProvingError::WrongParameter(format!("device aux commit failed: {e:?}")))?; - let h_max = shape - .aux - .dims - .iter() - .map(|&(lh, _)| lh) - .max() - .ok_or_else(|| ProvingError::WrongParameter("device aux commit: empty round".to_string()))?; - Ok(MixedMmcs::from_heap_nodes(shape.aux.dims.clone(), h_max, &nodes)) + Ok(mmcs) } -#[cfg(not(feature = "cuda"))] -fn commit_aux_device( - _retained_aux: &[Option<(Vec>, usize)>], - _shape: &EpochShape, - _num_tables: usize, -) -> Result, ProvingError> -where - FieldElement: math::traits::AsBytes + Sync + Send, -{ - unreachable!("device commit is disabled without the cuda feature") -} - -/// Commit the parts round (column-major ext3 slabs) on the GPU. Precondition + -/// policy as [`commit_main_device`]. Parts are always retained. +/// Commit the parts round (column-major ext3 slabs) on the GPU via the streaming +/// [`math_cuda::mmcs::StreamingMixedMmcs`]: build one table's slab, absorb it, +/// free it, then the next — only one slab resident at a time instead of every +/// table's slab at once. Parts are ALWAYS retained (recomputing them re-runs +/// constraint eval, a prove's dominant cost), so — unlike main/aux — this reads +/// them post-loop and always runs the canary. #[cfg(feature = "cuda")] fn commit_parts_device( retained_parts: &[Vec>>], @@ -1142,7 +1243,10 @@ fn commit_parts_device( where FieldElement: math::traits::AsBytes + Sync + Send, { - let mut slabs: Vec<(Vec, u64, u64, u64)> = Vec::new(); + let h_max = round_h_max(&shape.parts, "parts")?; + let mut device = math_cuda::mmcs::StreamingMixedMmcs::new(h_max as u64).map_err(|e| { + ProvingError::WrongParameter(format!("device parts commit init failed: {e:?}")) + })?; for table in 0..num_tables { let parts = &retained_parts[table]; if parts.is_empty() || parts[0].is_empty() { @@ -1163,38 +1267,32 @@ where } } } - slabs.push(( - slab, - num_rows as u64, - num_parts as u64, - shape.heights[table] as u64, - )); + device + .absorb_ext3_slabs( + shape.heights[table] as u64, + &slab, + num_rows as u64, + num_parts as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!("device parts absorb failed: {e:?}")) + })?; + // `slab` is freed here — one table's slab resident at a time. } - let inputs: Vec = slabs - .iter() - .map(|(d, col_stride, num_parts, h)| math_cuda::mmcs::MmcsExt3SlabInput { - data: d.as_slice(), - col_stride: *col_stride, - num_parts: *num_parts, - log_height: *h, - }) - .collect(); - let nodes = math_cuda::mmcs::commit_mixed_ext3_slabs_nodes(&inputs) + let nodes = device + .finish() .map_err(|e| ProvingError::WrongParameter(format!("device parts commit failed: {e:?}")))?; - let h_max = shape - .parts - .dims - .iter() - .map(|&(lh, _)| lh) - .max() - .ok_or_else(|| { - ProvingError::WrongParameter("device parts commit: empty round".to_string()) - })?; - Ok(MixedMmcs::from_heap_nodes( - shape.parts.dims.clone(), - h_max, - &nodes, - )) + let mmcs = MixedMmcs::from_heap_nodes(shape.parts.dims.clone(), h_max, &nodes); + // Parts are always retained, so the canary source is always available. + let mut source: Vec> = Vec::new(); + for table in 0..num_tables { + source.push(BorrowedMatrix::ColMajorNatural { + cols: &retained_parts[table], + log_height: shape.heights[table], + }); + } + xcheck_device_commit("parts", &mmcs, source)?; + Ok(mmcs) } #[cfg(not(feature = "cuda"))] @@ -1209,6 +1307,94 @@ where unreachable!("device commit is disabled without the cuda feature") } +/// How many leaves the [`xcheck_device_commit`] canary re-authenticates per +/// round. `k * (h_max - 1)` compressions and one row-pair hash per matrix — a +/// rounding error next to the `O(N)` device tree it guards. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +const CANARY_SAMPLES: usize = 8; + +/// Up to [`CANARY_SAMPLES`] DISTINCT leaf indices in `[0, n0)`, walked from a +/// root-seeded base by a root-seeded odd stride. The stride is coprime to the +/// power-of-two `n0`, so the first `k <= n0` steps never collide; seeding both +/// from the device root spreads the probes with the committed data (rather than +/// always testing leaf 0) without touching the transcript. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +fn canary_indices(root: &crate::config::Commitment, n0: usize) -> Vec { + debug_assert!(n0 >= 1); + let k = CANARY_SAMPLES.min(n0); + let seed = u64::from_le_bytes(root[0..8].try_into().expect("a commitment is 32 bytes")); + let base = (seed % n0 as u64) as usize; + // `| 1` forces an odd stride; odd is coprime to any power of two. + let stride = (u64::from_le_bytes(root[8..16].try_into().expect("a commitment is 32 bytes")) | 1) + as usize; + (0..k) + .map(|s| base.wrapping_add(s.wrapping_mul(stride)) % n0) + .collect() +} + +/// Release-safe canary for a device MMCS commit. The device path is +/// AUTHORITATIVE — the host [`StreamingMmcsBuilder`] is skipped — and the +/// byte-identical device-vs-host parity that guards 3d/3e is `debug_assert!` / +/// test-only, i.e. compiled OUT of a release prove. This re-authenticates a +/// handful of the freshly committed tree's leaves against the retained LDE rows +/// on the host, through the very [`MixedMmcs::open_batch`] / +/// [`MixedMmcs::verify_batch`] the verifier trusts. A silent device-side +/// corruption — a layout / heap bug that ships a wrong-but-self-consistent tree +/// — then aborts the prove with a hard [`ProvingError`] instead of producing a +/// proof that authenticates the wrong data. `source` must describe the SAME +/// matrices, in the SAME order, as the ones fed to the device commit (the +/// callers build it in the very loop that assembles the device inputs). +/// +/// Compiled unconditionally (only CALLED under cuda) so the host-only negative +/// test can exercise it without a GPU. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +fn xcheck_device_commit( + round: &str, + mmcs: &MixedMmcs, + source: Vec>, +) -> Result<(), ProvingError> +where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + let dims = mmcs.dims(); + if source.len() != dims.len() { + return Err(ProvingError::WrongParameter(format!( + "device {round} commit canary: leaf source has {} matrices but the tree committed {}", + source.len(), + dims.len() + ))); + } + let heights: Vec = dims.iter().map(|&(h, _)| h).collect(); + let widths: Vec = dims.iter().map(|&(_, w)| w).collect(); + let root = mmcs.root(); + let n0 = 1usize << (mmcs.h_max() - 1); + + for iota in canary_indices(&root, n0) { + let opening = mmcs.open_batch(iota, &source); + if !MixedMmcs::verify_batch(&root, iota, &opening, &heights, &widths) { + return Err(ProvingError::WrongParameter(format!( + "device {round} commit canary FAILED at leaf {iota}: the committed device \ + tree does not re-authenticate against the retained LDE — a silent device-side \ + corruption, aborting the prove (NO host fallback)" + ))); + } + } + Ok(()) +} + +/// The tallest committed `log_height` in a round — the base layer of its device +/// tree. Errs on an empty round (no matrix to commit). +#[cfg(feature = "cuda")] +fn round_h_max(round: &RoundShape, name: &str) -> Result { + round + .dims + .iter() + .map(|&(lh, _)| lh) + .max() + .ok_or_else(|| ProvingError::WrongParameter(format!("device {name} commit: empty round"))) +} + /// The width `table` contributes to `round`. Zero when it contributes nothing. fn matrix_width(round: &RoundShape, table: usize) -> usize { matrix_index(round, table).map_or(0, |m| round.dims[m].1) @@ -1483,3 +1669,81 @@ where &trace_term_coeffs, ) } + +#[cfg(test)] +mod canary_tests { + use super::*; + use math::field::goldilocks::GoldilocksField; + + type FE = FieldElement; + + // Two mixed-height, row-major matrices (tallest first), the exact layout the + // device main/aux commits hand the canary: matrix 0 is 8 rows × 2 cols + // (log_height 3), matrix 1 is 4 rows × 1 col (log_height 2). + fn source<'a>(d0: &'a [FE], d1: &'a [FE]) -> Vec> { + vec![ + BorrowedMatrix::RowMajorNatural { + data: d0, + stride: 2, + col_start: 0, + width: 2, + log_height: 3, + }, + BorrowedMatrix::RowMajorNatural { + data: d1, + stride: 1, + col_start: 0, + width: 1, + log_height: 2, + }, + ] + } + + fn fixture() -> (Vec, Vec) { + let d0: Vec = (0..16u64).map(|i| FE::from(i + 1)).collect(); + let d1: Vec = (0..4u64).map(|i| FE::from(i + 100)).collect(); + (d0, d1) + } + + #[test] + fn canary_passes_on_a_tree_that_matches_its_data() { + let (d0, d1) = fixture(); + let mmcs = MixedMmcs::commit(&source(&d0, &d1)); + assert!( + xcheck_device_commit("test", &mmcs, source(&d0, &d1)).is_ok(), + "the canary must accept a device tree that re-authenticates against its LDE" + ); + } + + #[test] + fn canary_fires_when_the_data_disagrees_with_the_tree() { + let (d0, d1) = fixture(); + let mmcs = MixedMmcs::commit(&source(&d0, &d1)); + // The tallest matrix is hashed into EVERY leaf, so perturbing all its rows + // makes every sampled leaf mismatch — a stand-in for a device that + // committed the wrong data with a self-consistent tree. + let d0_bad: Vec = (0..16u64).map(|i| FE::from(i + 2)).collect(); + assert!( + xcheck_device_commit("test", &mmcs, source(&d0_bad, &d1)).is_err(), + "the canary must reject a tree that disagrees with the retained LDE" + ); + } + + #[test] + fn canary_fires_when_a_committed_node_is_corrupted() { + let (d0, d1) = fixture(); + let good = MixedMmcs::commit(&source(&d0, &d1)); + let dims = good.dims().to_vec(); + let h_max = good.h_max(); + // Flip a bit in the root digest (heap index 0); it sits on every leaf's + // authentication path, so the recomputed climb from the honest data can + // no longer reach it. + let mut heap = good.heap_bytes(); + heap[0] ^= 0x01; + let corrupted = MixedMmcs::from_heap_nodes(dims, h_max, &heap); + assert!( + xcheck_device_commit("test", &corrupted, source(&d0, &d1)).is_err(), + "the canary must reject a tree with a corrupted committed node" + ); + } +} diff --git a/crypto/stark/src/batched/round4.rs b/crypto/stark/src/batched/round4.rs index 1f9bd03f5..d93235421 100644 --- a/crypto/stark/src/batched/round4.rs +++ b/crypto/stark/src/batched/round4.rs @@ -75,6 +75,7 @@ use crate::fri::batched::{ use crate::fri::fri_commitment::FriLayer; use crate::fri::fri_decommit::FriDecommitment; use crate::grinding; +use crate::prover::ProvingError; /// What the prover produced in the batched round 4, plus the challenges it drew /// on the way. The layers are kept so the caller can run the query phase over @@ -118,6 +119,11 @@ where /// `plan.standalone`). It is a closure rather than a materialized `Vec` so a /// caller can produce one table's DEEP codeword, absorb it and drop it: /// holding all of them at once is the memory cost batching exists to remove. +/// +/// Returns `Err` only when the device FRI commit is selected and a CUDA op +/// fails — a hard abort, the same no-silent-fallback policy as the device MMCS +/// commits (see the `batched/prover.rs` module header). The host build never +/// errors. #[allow(clippy::too_many_arguments)] pub fn commit_batched_fri( transcript: &mut T, @@ -129,7 +135,7 @@ pub fn commit_batched_fri( final_poly_log_degree: u32, grinding_factor: u8, num_queries: usize, -) -> BatchedFriCommit +) -> Result, ProvingError> where F: IsFFTField + IsSubFieldOf + 'static, E: IsField + 'static + Send + Sync, @@ -174,13 +180,51 @@ where } } - let (final_poly_coeffs, layers) = batched_commit_phase::( - combined, - transcript, - coset_offset, - blowup_log, - final_poly_log_degree, - ); + // Device fast path vs host build. Selecting here (rather than inside + // `batched_commit_phase`) keeps `crate::fri` free of the prover's error type + // and lets a device error be a hard abort: `?` propagates it instead of the + // fold loop silently falling back. `Ok(None)` = device not selected (off + // cuda, wrong field, or below the GPU threshold) → host build. + let (final_poly_coeffs, layers) = { + #[cfg(feature = "cuda")] + { + let (h_min, h_max_folds) = crate::fri::batched::bucket_height_range(&combined) + .expect("commit_batched_fri: combined has at least one occupied bucket"); + let inv_twiddles = crate::fri::fri_functions::compute_coset_twiddles_inv( + coset_offset, + 1usize << h_max_folds, + ); + match crate::gpu_lde::try_batched_fri_commit_gpu::( + &combined, + transcript, + coset_offset, + blowup_log, + final_poly_log_degree, + &inv_twiddles, + h_min, + h_max_folds, + )? { + Some(result) => result, + None => batched_commit_phase::( + combined, + transcript, + coset_offset, + blowup_log, + final_poly_log_degree, + ), + } + } + #[cfg(not(feature = "cuda"))] + { + batched_commit_phase::( + combined, + transcript, + coset_offset, + blowup_log, + final_poly_log_degree, + ) + } + }; let layer_roots: Vec = layers.iter().map(|layer| layer.merkle_tree.root).collect(); // Grinding runs on the CONFIGURATION's transcript hash, not a hard-wired @@ -199,7 +243,7 @@ where .map(|_| transcript.sample_u64(1u64 << (h_max - 1)) as usize) .collect(); - BatchedFriCommit { + Ok(BatchedFriCommit { layers, layer_roots, final_poly_coeffs, @@ -209,7 +253,7 @@ where alpha, plan, standalone_coeffs, - } + }) } /// Verify one query against a STANDALONE table's terminal-only instance. @@ -625,6 +669,7 @@ pub(crate) mod tests { grinding_factor, num_queries, ) + .expect("the host batched FRI commit never errors") } /// υ⁻¹ for query `iota`: the inverse of the tallest coset's element at diff --git a/crypto/stark/src/fri/batched.rs b/crypto/stark/src/fri/batched.rs index 9738e981e..38f2c4880 100644 --- a/crypto/stark/src/fri/batched.rs +++ b/crypto/stark/src/fri/batched.rs @@ -343,22 +343,10 @@ where // Inverse twiddle factors for the initial domain size. let inv_twiddles = compute_coset_twiddles_inv(coset_offset, domain_size); - // GPU fast path: drive the fold → inject → commit loop below on the device - // (`fri_inject_bucket_ext3` + `FriCommitState::fold_inject_commit_layer`), - // returning `None` with the transcript restored to fall back to the host. - #[cfg(feature = "cuda")] - if let Some(result) = crate::gpu_lde::try_batched_fri_commit_gpu::( - &combined, - transcript, - coset_offset, - blowup_log, - final_poly_log_degree, - &inv_twiddles, - h_min, - h_max, - ) { - return result; - } + // The device fast path is selected one layer up, in + // `crate::batched::round4::commit_batched_fri`, so this stays a pure host + // build and a device error there is a hard abort rather than a fallback into + // this function. `h_min` is still consumed by the terminal-floor layout. // Take the starting codeword — NOT committed; it plays the role of layer 0. let mut running = combined[h_max] @@ -436,7 +424,9 @@ where } /// The `(h_min, h_max)` of the occupied buckets, or `None` when none are. -fn bucket_height_range( +/// `pub(crate)` so `commit_batched_fri` can size the device FRI's twiddles from +/// the same range this host build uses. +pub(crate) fn bucket_height_range( combined: &[Option>>], ) -> Option<(usize, usize)> { let mut occupied = combined diff --git a/crypto/stark/src/fri/mmcs.rs b/crypto/stark/src/fri/mmcs.rs index b9e41972f..33bb740ca 100644 --- a/crypto/stark/src/fri/mmcs.rs +++ b/crypto/stark/src/fri/mmcs.rs @@ -493,6 +493,24 @@ where self.root } + /// Serialize the digest layers back into the standard heap byte array that + /// [`Self::from_heap_nodes`] parses — the inverse round-trip a device build + /// produces directly on the GPU. Test-only; used to corrupt a single node + /// and confirm the device-commit canary fires. + #[cfg(test)] + pub(crate) fn heap_bytes(&self) -> Vec { + let leaves_len = 1usize << (self.h_max - 1); + let mut heap = vec![0u8; (2 * leaves_len - 1) * 32]; + for (k, layer) in self.layers.iter().enumerate() { + let level_size = leaves_len >> k; + let start = level_size - 1; + for (j, digest) in layer.iter().enumerate() { + heap[(start + j) * 32..(start + j) * 32 + 32].copy_from_slice(digest); + } + } + heap + } + /// `log2` of the tallest committed matrix. The query index this MMCS accepts /// lives in `[0, 2^(h_max-1))` — see the module's index-convention section. pub fn h_max(&self) -> usize { diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 4323a4d99..82173ced6 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -3221,7 +3221,15 @@ where /// `fri/batched.rs`'s `inject_bucket`), and it uses the batched terminal floor /// [`crate::fri::batched::BatchedFriLayout`]. `combined[h]` is the bucket at /// height `h` (`2^h` ext3 elements); `combined[h_max]` is the starting codeword -/// and is taken. Returns `None` (transcript restored) to fall back to the host. +/// and is taken. +/// +/// Result convention — the batched-commit no-silent-fallback policy (see the +/// `batched/prover.rs` module header): `Ok(None)` = the device path was NOT +/// selected (wrong field, below the GPU threshold, or a degenerate layout), +/// reached before any transcript mutation, so the caller builds the layers on +/// the host; `Ok(Some(_))` = the device produced the commit; `Err(_)` = the +/// device was selected but a CUDA op failed — a HARD ABORT, never a silent host +/// fallback, exactly as a device MMCS-commit error aborts. #[cfg(feature = "cuda")] #[allow(clippy::too_many_arguments, clippy::type_complexity)] pub(crate) fn try_batched_fri_commit_gpu( @@ -3233,10 +3241,13 @@ pub(crate) fn try_batched_fri_commit_gpu( inv_twiddles: &[FieldElement], h_min: usize, h_max: usize, -) -> Option<( - Vec>, - Vec>>, -)> +) -> Result< + Option<( + Vec>, + Vec>>, + )>, + crate::prover::ProvingError, +> where F: IsFFTField + IsField + IsSubFieldOf + 'static, E: IsField + 'static + Send + Sync, @@ -3245,19 +3256,19 @@ where T: IsStarkTranscript + Clone, { if TypeId::of::() != TypeId::of::() { - return None; + return Ok(None); } if TypeId::of::() != TypeId::of::() { - return None; + return Ok(None); } let n0 = 1usize << h_max; if n0 < 2 || n0 < gpu_lde_threshold() || inv_twiddles.len() != n0 / 2 { - return None; + return Ok(None); } let layout = crate::fri::batched::BatchedFriLayout::new(h_max, h_min, blowup_log, final_poly_log_degree); if layout.total_folds == 0 || layout.terminal_len < 2 { - return None; + return Ok(None); } // Pack inv_twiddles before any transcript mutation. @@ -3267,32 +3278,37 @@ where inv_tw_u64.push(unsafe { *(t.value() as *const _ as *const u64) }); } - // Clone (not take) the starting codeword so a later `None` return leaves - // `combined` intact for the host fallback. + // Clone (not take) the starting codeword so a not-selected `Ok(None)` return + // leaves `combined` intact for the host build. let start = match combined.get(h_max).and_then(Option::as_ref) { Some(c) if c.len() == n0 => c.clone(), - _ => return None, + _ => return Ok(None), }; // SAFETY: E == Ext3; backing is [u64; 3]. let start_u64: &[u64] = unsafe { ext3_slice_to_u64::(&start) }; - let mut state = match math_cuda::fri::FriCommitState::new(start_u64, &inv_tw_u64, n0) { - Ok(s) => s, - Err(_) => return None, - }; + // From here on a CUDA failure is a HARD ABORT (`Err`), not a host fallback: + // the device path is committed. No transcript has been mutated yet, so the + // aborting proof simply stops. + let mut state = math_cuda::fri::FriCommitState::new(start_u64, &inv_tw_u64, n0).map_err(|e| { + crate::prover::ProvingError::WrongParameter(format!( + "batched FRI device init failed: {e:?} — aborting the prove (NO host fallback)" + )) + })?; // Upload every shorter bucket to the device, indexed by height. let mut buckets: Vec>> = (0..=h_max).map(|_| None).collect(); for h in 1..h_max { if let Some(bucket) = combined.get(h).and_then(Option::as_ref) { let bucket_u64: &[u64] = unsafe { ext3_slice_to_u64::(bucket) }; - match state.stream.clone_htod(bucket_u64) { - Ok(dev) => buckets[h] = Some(dev), - Err(_) => return None, - } + let dev = state.stream.clone_htod(bucket_u64).map_err(|e| { + crate::prover::ProvingError::WrongParameter(format!( + "batched FRI device bucket upload failed at height {h}: {e:?} — aborting the prove (NO host fallback)" + )) + })?; + buckets[h] = Some(dev); } } - let transcript_snapshot = transcript.clone(); let mut fri_layer_list: Vec>> = Vec::with_capacity(layout.num_committed); @@ -3314,14 +3330,14 @@ where .and_then(Option::as_ref) .map(|b| (beta_sq_raw, b)); - let (evals_u64, evals_dev, dev_tree) = - match state.fold_inject_commit_layer(beta_raw, bucket_arg, true) { - Ok(v) => v, - Err(_) => { - *transcript = transcript_snapshot; - return None; - } - }; + let (evals_u64, evals_dev, dev_tree) = state + .fold_inject_commit_layer(beta_raw, bucket_arg, true) + .map_err(|e| { + crate::prover::ProvingError::WrongParameter(format!( + "batched FRI device fold/commit failed at fold {fold_idx}: {e:?} — \ + aborting the prove (NO host fallback)" + )) + })?; if fold_idx < layout.num_committed { let evaluation = evals_u64.map(|v| u64_to_ext3_vec::(&v)).unwrap_or_default(); @@ -3348,7 +3364,7 @@ where transcript.append_field_element(c); } GPU_FRI_CALLS.fetch_add(1, Ordering::Relaxed); - return Some((final_poly_coeffs, fri_layer_list)); + return Ok(Some((final_poly_coeffs, fri_layer_list))); } } // total_folds >= 1 guarantees the terminal branch above returned. From 23c75ec57757c256ef4b9dfcbac7a2b45f7c37f6 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 2 Sep 2026 11:46:46 -0300 Subject: [PATCH 13/26] test(prover): fair CPU batched-vs-per-table timing helpers + bench --- prover/Cargo.toml | 4 + prover/benches/batched_cpu_cmp.rs | 69 ++++++++++++++ prover/src/lib.rs | 146 ++++++++++++++++++++++++++++++ 3 files changed, 219 insertions(+) create mode 100644 prover/benches/batched_cpu_cmp.rs diff --git a/prover/Cargo.toml b/prover/Cargo.toml index d4ebdeb0d..5e8b181ff 100644 --- a/prover/Cargo.toml +++ b/prover/Cargo.toml @@ -56,3 +56,7 @@ harness = false [[bench]] name = "bench_continuation" harness = false + +[[bench]] +name = "batched_cpu_cmp" +harness = false diff --git a/prover/benches/batched_cpu_cmp.rs b/prover/benches/batched_cpu_cmp.rs new file mode 100644 index 000000000..d4d742b8c --- /dev/null +++ b/prover/benches/batched_cpu_cmp.rs @@ -0,0 +1,69 @@ +//! FAIR CPU structural comparison: batched (multi-merkle-tree) vs per-table +//! prover, BOTH on CPU — run WITHOUT `--features cuda` so neither touches the +//! GPU. Interleaved A/B/A/B + warm to kill thermal bias. Isolates the +//! multi-merkle-tree's structural cost from the (portable) GPU-optimization gap. +//! +//! cargo bench -p lambda-vm-prover --bench batched_cpu_cmp +//! Env: BLOCK=ethrex_simple_tx.bin WARMUP=1 ITERS=3 COOLDOWN_SECS=0 + +use std::time::Duration; + +fn env_usize(key: &str, default: usize) -> usize { + std::env::var(key) + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(default) +} + +fn median(v: &[Duration]) -> Duration { + v[v.len() / 2] +} + +fn main() { + let manifest = env!("CARGO_MANIFEST_DIR"); + let block = std::env::var("BLOCK").unwrap_or_else(|_| "ethrex_simple_tx.bin".to_string()); + let elf = std::fs::read(format!( + "{manifest}/../executor/program_artifacts/rust/ethrex.elf" + )) + .expect("read ethrex.elf"); + let input = + std::fs::read(format!("{manifest}/../executor/tests/{block}")).expect("read block input"); + + let warmup = env_usize("WARMUP", 1); + let iters = env_usize("ITERS", 3); + let cooldown = env_usize("COOLDOWN_SECS", 0) as u64; + if cfg!(feature = "cuda") { + println!("⚠️ built with cuda — run WITHOUT --features cuda for the FAIR CPU comparison"); + } + println!("=== CPU structural comparison — block={block} warmup={warmup} iters={iters} cooldown={cooldown}s ==="); + println!("(both provers on CPU; isolates the multi-merkle-tree structure vs per-table)"); + + let mut batched = Vec::new(); + let mut pertable = Vec::new(); + for i in 0..(warmup + iters) { + let tag = if i < warmup { "warmup" } else { "iter " }; + let b = lambda_vm_prover::time_batched_prove(&elf, &input).expect("batched prove"); + if cooldown > 0 { + std::thread::sleep(Duration::from_secs(cooldown)); + } + let p = lambda_vm_prover::time_per_table_prove(&elf, &input).expect("per-table prove"); + println!(" {tag} {i}: batched={b:?} per-table={p:?}"); + if i >= warmup { + batched.push(b); + pertable.push(p); + } + if cooldown > 0 { + std::thread::sleep(Duration::from_secs(cooldown)); + } + } + + batched.sort(); + pertable.sort(); + let bm = median(&batched); + let pm = median(&pertable); + let ratio = bm.as_secs_f64() / pm.as_secs_f64(); + println!( + "RESULT block={block} batched median={bm:?} (min {:?}) | per-table median={pm:?} (min {:?}) | batched/per-table = {ratio:.3}", + batched[0], pertable[0], + ); +} diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 79ef4c715..c62641474 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -1260,6 +1260,152 @@ pub fn prove_with_options_and_inputs( }) } +/// CPU structural comparison (see MEASUREMENT-PLAN.md): time JUST the batched +/// `multi_prove_batched` prove step over the VM's real tables. Paired with +/// [`time_per_table_prove`] on the SAME block, run non-cuda so BOTH are pure CPU, +/// this isolates the multi-merkle-tree's structural cost from GPU optimization +/// (the GPU comparison was confounded — per-table is device-resident, batched is +/// not). `Retain` matches the per-table path's single-prove residency. NOT part +/// of the shipping pipeline. +pub fn time_batched_prove( + elf_bytes: &[u8], + private_inputs: &[u8], +) -> Result { + let proof_options = + GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 is always valid"); + let max_rows = MaxRowsConfig::default(); + let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; + let executor = Executor::new(&program, private_inputs.to_vec()) + .map_err(|e| Error::Execution(format!("{e}")))?; + let result = executor + .run() + .map_err(|e| Error::Execution(format!("{e}")))?; + #[cfg(feature = "disk-spill")] + let storage_mode = { + let lengths = count_table_lengths(&program, &result.logs, &max_rows, private_inputs)?; + auto_storage::decide(&lengths, proof_options.blowup_factor) + }; + let mut traces = Traces::from_elf_and_logs( + &program, + &result.logs, + &max_rows, + private_inputs, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + drop(result); + let table_counts = traces.table_counts(); + let airs = VmAirs::new( + &program, + &proof_options, + false, + &traces.page_configs, + &table_counts, + None, + true, + None, + None, + None, + ); + let runtime_page_ranges = traces.runtime_page_ranges(); + let num_private_input_pages = traces + .page_configs + .iter() + .filter(|c| c.is_private_input) + .count(); + let mut transcript = DefaultTranscript::::new(&[]); + absorb_statement( + &mut transcript, + StatementKind::Monolithic, + elf_bytes, + &traces.public_output_bytes, + &table_counts, + num_private_input_pages, + &runtime_page_ranges, + proof_options.fri_final_poly_log_degree, + ); + let t = std::time::Instant::now(); + let _ = stark::batched::prover::multi_prove_batched::>( + airs.air_trace_pairs(&mut traces), + &mut transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + stark::residency_mode::ResidencyMode::Retain, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + Ok(t.elapsed()) +} + +/// CPU structural comparison: time JUST the per-table `Prover::multi_prove` over +/// the same tables. See [`time_batched_prove`]. +pub fn time_per_table_prove( + elf_bytes: &[u8], + private_inputs: &[u8], +) -> Result { + let proof_options = + GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 is always valid"); + let max_rows = MaxRowsConfig::default(); + let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; + let executor = Executor::new(&program, private_inputs.to_vec()) + .map_err(|e| Error::Execution(format!("{e}")))?; + let result = executor + .run() + .map_err(|e| Error::Execution(format!("{e}")))?; + #[cfg(feature = "disk-spill")] + let storage_mode = { + let lengths = count_table_lengths(&program, &result.logs, &max_rows, private_inputs)?; + auto_storage::decide(&lengths, proof_options.blowup_factor) + }; + let mut traces = Traces::from_elf_and_logs( + &program, + &result.logs, + &max_rows, + private_inputs, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + drop(result); + let table_counts = traces.table_counts(); + let airs = VmAirs::new( + &program, + &proof_options, + false, + &traces.page_configs, + &table_counts, + None, + true, + None, + None, + None, + ); + let runtime_page_ranges = traces.runtime_page_ranges(); + let num_private_input_pages = traces + .page_configs + .iter() + .filter(|c| c.is_private_input) + .count(); + let mut transcript = DefaultTranscript::::new(&[]); + absorb_statement( + &mut transcript, + StatementKind::Monolithic, + elf_bytes, + &traces.public_output_bytes, + &table_counts, + num_private_input_pages, + &runtime_page_ranges, + proof_options.fri_final_poly_log_degree, + ); + let t = std::time::Instant::now(); + let _ = Prover::multi_prove( + airs.air_trace_pairs(&mut traces), + &mut transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + Ok(t.elapsed()) +} + /// Verify a proof produced by [`prove`] using default proof options. /// /// Uses [`GoldilocksCubicProofOptions::with_blowup(2)`] for verification. From f035bd50233e80249d98e6d150e53f6cf5940931 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 2 Sep 2026 13:45:49 -0300 Subject: [PATCH 14/26] feat(cuda): device-buffer absorb path for StreamingMixedMmcs (resident-LDE commit) --- crypto/math-cuda/src/mmcs.rs | 86 +++++++++++++++ crypto/math-cuda/tests/mmcs_tree_parity.rs | 120 +++++++++++++++++++++ 2 files changed, 206 insertions(+) diff --git a/crypto/math-cuda/src/mmcs.rs b/crypto/math-cuda/src/mmcs.rs index b255b5405..2a047f47f 100644 --- a/crypto/math-cuda/src/mmcs.rs +++ b/crypto/math-cuda/src/mmcs.rs @@ -538,6 +538,15 @@ impl StreamingMixedMmcs { }) } + /// The stream every absorb + [`Self::finish`] runs on. A caller feeding the + /// device-buffer absorb paths (`absorb_*_dev`) should produce the resident + /// LDE on THIS stream (or synchronize its producer stream first), so the + /// absorb kernel is correctly ordered after the buffer is filled without a + /// device-wide sync. + pub fn stream(&self) -> Arc { + self.stream.clone() + } + fn group_mut(&mut self, log_height: u64) -> Result<&mut MmcsGroupHasher> { let idx = log_height as usize; assert!( @@ -607,6 +616,83 @@ impl StreamingMixedMmcs { Ok(()) } + // ---- Device-buffer absorb paths ------------------------------------------ + // + // These take a DEVICE buffer already resident on the GPU (e.g. main's + // `GpuLdeBase.buf`, an aux `GpuLdeExt3`, or resident composition parts) and + // hash it in place — no `clone_htod`, no host `Vec`. This is the bridge that + // lets the batched prover commit its resident LDE where it already lives, + // instead of expanding on the CPU and uploading. + // + // Precondition (vs the host wrappers): `data` must be ready on + // [`Self::stream`] before the call — produce the LDE on that stream, or + // synchronize the producer first. Unlike the host wrappers these do NOT + // synchronize afterward: the resident buffer is meant to stay alive across + // the prover's later phases, and ordering into [`Self::finish`] is already + // guaranteed because both run on `self.stream`. The caller must therefore + // keep `data` alive until at least `finish` (its natural residency anyway). + + /// Absorb one COLUMN-MAJOR base-field matrix resident on device — main's + /// `GpuLdeBase.buf` layout, element `(row, col)` at `col*col_stride + row`. + /// Byte-identical leaf digests to [`Self::absorb_row_major`] over the same + /// matrix; this is the resident-LDE bridge for the main round. + pub fn absorb_col_major_dev( + &mut self, + log_height: u64, + data: &CudaSlice, + col_stride: u64, + col_start: u64, + col_end: u64, + ) -> Result<()> { + let stream = self.stream.clone(); + self.group_mut(log_height)? + .absorb_col_major(&stream, data, col_stride, col_start, col_end) + } + + /// Absorb one ROW-MAJOR base-field matrix resident on device (same layout as + /// [`Self::absorb_row_major`], but the buffer already lives on the GPU). + pub fn absorb_row_major_dev( + &mut self, + log_height: u64, + data: &CudaSlice, + row_stride: u64, + col_start: u64, + col_end: u64, + ) -> Result<()> { + let stream = self.stream.clone(); + self.group_mut(log_height)? + .absorb_row_major(&stream, data, row_stride, col_start, col_end) + } + + /// Absorb one ROW-MAJOR ext3 matrix resident on device (aux round) — an + /// element's 3 components consecutive, `stride` elements per row. + pub fn absorb_ext3_row_major_dev( + &mut self, + log_height: u64, + data: &CudaSlice, + stride: u64, + col_start: u64, + col_end: u64, + ) -> Result<()> { + let stream = self.stream.clone(); + self.group_mut(log_height)? + .absorb_ext3_row_major(&stream, data, stride, col_start, col_end) + } + + /// Absorb one COLUMN-MAJOR ext3 slab matrix resident on device (parts round) + /// — component `k` of column `c` at `(c*3 + k) * col_stride`. + pub fn absorb_ext3_slabs_dev( + &mut self, + log_height: u64, + parts: &CudaSlice, + col_stride: u64, + num_parts: u64, + ) -> Result<()> { + let stream = self.stream.clone(); + self.group_mut(log_height)? + .absorb_ext3_slabs(&stream, parts, col_stride, num_parts) + } + /// Finalize every absorbed group and climb, returning the standard heap node /// array (`(2L-1)*32` bytes for `L = 2^(h_max-1)`), host-side. pub fn finish(self) -> Result> { diff --git a/crypto/math-cuda/tests/mmcs_tree_parity.rs b/crypto/math-cuda/tests/mmcs_tree_parity.rs index a41fc3fcf..48838a66e 100644 --- a/crypto/math-cuda/tests/mmcs_tree_parity.rs +++ b/crypto/math-cuda/tests/mmcs_tree_parity.rs @@ -276,3 +276,123 @@ fn col_major_root_matches_the_host() { the host tree byte for byte — the col-major absorb bridge" ); } + +// ---- StreamingMixedMmcs device-buffer absorb (the resident-LDE commit path) ---- +// +// `StreamingMixedMmcs` is the persistent streaming commit the batched prover +// drives across its round loop. Today its `absorb_*` wrappers `clone_htod` a +// host `Vec`; the fully-GPU prover needs to absorb the LDE where it already +// lives (main's resident `GpuLdeBase.buf`), skipping the upload. These exercise +// the new `absorb_*_dev` paths and pin them to the host tree. + +fn streaming_h_max(mats: &Matrices) -> u64 { + (0..mats.num_matrices()) + .map(|m| mats.log_height(m)) + .max() + .expect("non-empty") as u64 +} + +/// Reference: the EXISTING streaming path — host `Vec` per matrix, uploaded by +/// the wrapper (`clone_htod`). Absorbs in input order; `StreamingMixedMmcs` +/// routes each matrix to its height group internally. +fn streaming_tree_row_major_host(mats: &Matrices) -> [u8; 32] { + let mut sm = math_cuda::mmcs::StreamingMixedMmcs::new(streaming_h_max(mats)) + .expect("streaming mmcs"); + for m in 0..mats.num_matrices() { + let (data, _, width) = &mats.mats[m]; + sm.absorb_row_major(mats.log_height(m) as u64, &raw(data), *width as u64, 0, *width as u64) + .expect("absorb row-major host"); + } + let nodes = sm.finish().expect("finish"); + let mut root = [0u8; 32]; + root.copy_from_slice(&nodes[0..32]); + root +} + +/// The new path, row-major: the buffer is resident on device and absorbed with +/// `absorb_row_major_dev` (no `clone_htod`). Buffers are produced on the commit +/// stream (`sm.stream()`) and held alive until `finish` — the resident-LDE +/// contract the `_dev` methods rely on (they do not synchronize). +fn streaming_tree_row_major_dev(mats: &Matrices) -> [u8; 32] { + let mut sm = math_cuda::mmcs::StreamingMixedMmcs::new(streaming_h_max(mats)) + .expect("streaming mmcs"); + let stream = sm.stream(); + let mut resident: Vec> = Vec::new(); + for m in 0..mats.num_matrices() { + let (data, _, width) = &mats.mats[m]; + let dev = stream.clone_htod(&raw(data)).expect("H2D"); + sm.absorb_row_major_dev( + mats.log_height(m) as u64, + &dev, + *width as u64, + 0, + *width as u64, + ) + .expect("absorb row-major dev"); + resident.push(dev); + } + let nodes = sm.finish().expect("finish"); + drop(resident); + let mut root = [0u8; 32]; + root.copy_from_slice(&nodes[0..32]); + root +} + +/// The new path, column-major: main's resident `GpuLdeBase.buf` layout, absorbed +/// with `absorb_col_major_dev`. Same resident-buffer contract as above. +fn streaming_tree_col_major_dev(mats: &Matrices) -> [u8; 32] { + let mut sm = math_cuda::mmcs::StreamingMixedMmcs::new(streaming_h_max(mats)) + .expect("streaming mmcs"); + let stream = sm.stream(); + let mut resident: Vec> = Vec::new(); + for m in 0..mats.num_matrices() { + let (data, log_height, width) = &mats.mats[m]; + let num_rows = 1u64 << log_height; + let dev = stream + .clone_htod(&raw_col_major(data, *log_height, *width)) + .expect("H2D"); + sm.absorb_col_major_dev(*log_height as u64, &dev, num_rows, 0, *width as u64) + .expect("absorb col-major dev"); + resident.push(dev); + } + let nodes = sm.finish().expect("finish"); + drop(resident); + let mut root = [0u8; 32]; + root.copy_from_slice(&nodes[0..32]); + root +} + +/// ★ Step 1 of the fully-GPU batched prover: `StreamingMixedMmcs` must commit a +/// RESIDENT device buffer (no per-absorb upload) to the identical tree as the +/// host-upload path and the authoritative host `MixedMmcs`. Same mixed-height +/// fixture as `mixed_height_root_matches_the_host`. +#[test] +fn streaming_device_absorb_matches_host_upload_and_host_tree() { + let mats = Matrices { + mats: vec![ + matrix(7, 3, 11), + matrix(7, 6, 23), + matrix(5, 2, 41), + matrix(2, 4, 59), + ], + }; + let host = Mmcs::commit(&mats).root(); + + assert_eq!( + streaming_tree_row_major_host(&mats), + host, + "the existing streaming host-upload path must match the host tree" + ); + assert_eq!( + streaming_tree_row_major_dev(&mats), + host, + "the streaming device-buffer path (row-major) must match — skipping \ + clone_htod changes nothing but where the bytes are read from" + ); + assert_eq!( + streaming_tree_col_major_dev(&mats), + host, + "the streaming device-buffer path (col-major) must match — this is the \ + resident-LDE bridge main hands the batched prover (GpuLdeBase.buf)" + ); +} From b846552de606cd87c35eea9bfc6f8f1d242dd085 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 2 Sep 2026 13:45:59 -0300 Subject: [PATCH 15/26] feat(cuda): device-resident main+aux commit + GPU LDE recompute for the batched prover --- crypto/math-cuda/src/lde.rs | 119 +++++++++++ crypto/stark/src/batched/prover.rs | 320 +++++++++++++++++++++++++---- crypto/stark/src/gpu_lde.rs | 129 ++++++++++++ crypto/stark/src/prover.rs | 4 +- 4 files changed, 536 insertions(+), 36 deletions(-) diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 9bbd9958d..1b5f3f248 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -706,6 +706,95 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( Ok((handle, lde_out)) } +/// Row-major coset LDE on device, transposed to column-major and kept resident, +/// WITHOUT building any Merkle tree — for the batched prover, where the shared +/// mixed-height MMCS is the one tree, so a per-table row-pair tree would be +/// redundant on-device keccak (leaf hash + inner climb). The returned +/// [`GpuLdeBase`] has `tree: None`; the caller absorbs `buf` (column-major, +/// `col*lde_size + row`) straight into the shared streaming commit +/// (`StreamingMixedMmcs::absorb_col_major_dev`). +/// +/// The expansion + transpose are the SAME calls as +/// [`coset_lde_row_major_with_merkle_tree_keep`] (`expand_row_major_on_stream` +/// then `launch_row_to_col_major`), so the resident `buf` is byte-identical — it +/// only skips the leaf hash, inner-tree build and root copy. `retain_host_lde` +/// D2Hs the row-major LDE for host consumers (empty Vec when false; the +/// device-resident batched path passes false). The trace-domain snapshot +/// (`trace_dev`) is a LogUp-fingerprint concern of the per-table prover and is +/// never taken here. +pub fn coset_lde_row_major_keep_no_tree( + row_major: &[u64], + predev: Option<&CudaSlice>, + n: usize, + m: usize, + blowup_factor: usize, + weights: &[u64], + retain_host_lde: bool, +) -> Result<(GpuLdeBase, Vec)> { + let input = match predev { + Some(d) if d.len() == row_major.len() => InnerInput::Dev(d), + _ => InnerInput::Host(row_major), + }; + let input_len = match &input { + InnerInput::Host(h) => h.len(), + InnerInput::Dev(d) => d.len(), + }; + assert_eq!(input_len, n * m); + assert!(n.is_power_of_two()); + assert_eq!(weights.len(), n); + assert!(blowup_factor.is_power_of_two()); + let lde_size = n * blowup_factor; + assert_u32_domain(lde_size, "coset_lde_row_major_keep_no_tree lde_size"); + let lde_u64 = lde_size as u64; + + let be = backend()?; + let stream = be.next_stream(); + + let (buf, _trace_col_major) = + expand_row_major_on_stream(&stream, be, input, n, m, blowup_factor, weights, false)?; + + // D2H the row-major LDE only when a host consumer asks; the device-resident + // batched path passes false and keeps everything on device. + let lde_pending = if retain_host_lde { + Some(crate::device::async_dtoh_via( + &stream, + be.pinned_staging(), + &be.ctx, + &buf, + lde_size * m, + )?) + } else { + None + }; + + let col_major_dev = launch_row_to_col_major(&stream, be, &buf, lde_size, m, lde_u64)?; + // No host synchronize: the handle carries a `ready` event and consumers on + // other streams wait on it device-side (`wait_ready_on`), exactly as the + // tree-building keep path does. + let ready = be.take_event()?; + ready.event().record(&stream)?; + + let lde_out = match lde_pending { + Some(p) => { + let mut out = vec![0u64; lde_size * m]; + p.wait_into_u64(&mut out)?; + out + } + None => Vec::new(), + }; + + let handle = GpuLdeBase { + buf: Arc::new(col_major_dev), + m, + lde_size, + tree: None, + ready: Some(Arc::new(ready)), + trace_dev: None, + trace_rows: 0, + }; + Ok((handle, lde_out)) +} + /// Row-major LDE + TWO subset Merkle trees for preprocessed tables: the /// precomputed columns `[0, split_col)` and the multiplicity columns /// `[split_col, m)` commit to separate trees over the same row-major LDE, @@ -880,6 +969,36 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( Ok((handle, lde_out)) } +/// Tree-less ext3 keep for the batched prover: expand the row-major ext3 trace +/// to its coset LDE and keep it resident, building NO per-table Merkle tree (the +/// shared MMCS is the tree). `Fp3 = [u64; 3]`, so this is the base no-tree keep +/// over `m*3` columns; the resident `buf` is therefore the SLAB layout +/// (component `k` of ext3 column `c` at `(c*3 + k)*lde_size + row`) that +/// `mmcs_absorb_row_pair_ext3_slabs` reads — byte-identical leaves to the +/// row-major ext3 absorb (both emit, per bit-reversed row, each column's 3 +/// components consecutively). Input: `row_major` is `n*m*3` u64s. +pub fn coset_lde_ext3_row_major_keep_no_tree( + row_major: &[u64], + n: usize, + m: usize, + blowup_factor: usize, + weights: &[u64], + retain_host_lde: bool, +) -> Result<(GpuLdeExt3, Vec)> { + let (base, lde_out) = + coset_lde_row_major_keep_no_tree(row_major, None, n, m * 3, blowup_factor, weights, retain_host_lde)?; + // `base.buf` is column-major over the `m*3` base columns = exactly the ext3 + // slab layout; re-wrap the same resident buffer + ready event as a GpuLdeExt3. + let handle = GpuLdeExt3 { + buf: base.buf, + m, + lde_size: base.lde_size, + tree: None, + ready: base.ready, + }; + Ok((handle, lde_out)) +} + /// Like [`coset_lde_ext3_row_major_with_merkle_tree_keep`] but the input is an /// already-resident device buffer (`n * m` ext3 elements, row-major, `n*m*3` /// u64s). No PCIe upload: the buffer is copied device-to-device into the LDE diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index 458410f4b..4f06c3e7d 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -406,24 +406,87 @@ where } // Streaming device absorb: hash this table into the device tree NOW, // so the LDE can be freed below (RecomputeLde) without retaining it. + // + // Prefer the RESIDENT path: expand this table's main LDE on the GPU + // and absorb it column-major straight from VRAM — no host FFT upload, + // the whole point of a device-resident batched prover. Tables below + // the device LDE threshold decline the keep-expand and fall back to + // uploading the host LDE (mixed eligibility; identical root either + // way). Step 2: only the main COMMIT is device-resident here; the host + // `main_data` above still feeds prep trees, `retained_main` and the + // later phases, so this deliberately expands twice for now (step 3 + // drops the host expand and moves the later phases to device recompute). #[cfg(feature = "cuda")] if let Some(dev) = device_main.as_mut() { - // SAFETY: Goldilocks base (device_commit_enabled gated it); - // FieldElement is #[repr(transparent)] over u64. Committed - // columns are [num_precomputed, total_cols). - let data_u64 = unsafe { - std::slice::from_raw_parts(main_data.as_ptr() as *const u64, main_data.len()) + let (trace_slice, num_cols) = trace.main_data_row_major(); + debug_assert_eq!( + num_cols, total_cols, + "the resident and host expansions must see the same column count" + ); + let n = if num_cols > 0 { + trace_slice.len() / num_cols + } else { + 0 }; - dev.absorb_row_major( - height as u64, - data_u64, - total_cols as u64, - num_precomputed as u64, - total_cols as u64, - ) - .map_err(|e| { - ProvingError::WrongParameter(format!("device main absorb failed: {e:?}")) - })?; + let resident = crate::gpu_lde::try_expand_row_major_keep_no_tree::( + trace_slice, + trace.main_rowmajor_dev(), + n, + num_cols, + domains[table].blowup_factor, + &twiddles[table].coset_weights, + false, + ); + match resident { + // Resident LDE: `handle.buf` is column-major (`col*lde_size + + // row`); absorb the committed columns [num_precomputed, + // total_cols) in place. `wait_ready_on` orders the absorb after + // the producer's last kernel device-side (no host block). No + // per-table tree is built (the shared MMCS is the tree); + // dropping `handle` frees this table's LDE (stream-first). + Some((handle, _host_lde)) => { + let commit_stream = dev.stream(); + handle.wait_ready_on(&commit_stream).map_err(|e| { + ProvingError::WrongParameter(format!( + "device main LDE ready-wait failed: {e:?}" + )) + })?; + dev.absorb_col_major_dev( + height as u64, + handle.buf.as_ref(), + handle.lde_size as u64, + num_precomputed as u64, + total_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!( + "device main absorb (resident) failed: {e:?}" + )) + })?; + } + // Below the device LDE threshold: upload the host LDE and + // absorb row-major, exactly as before. + None => { + // SAFETY: Goldilocks base (device_commit_enabled gated it); + // FieldElement is #[repr(transparent)] over u64. + let data_u64 = unsafe { + std::slice::from_raw_parts( + main_data.as_ptr() as *const u64, + main_data.len(), + ) + }; + dev.absorb_row_major( + height as u64, + data_u64, + total_cols as u64, + num_precomputed as u64, + total_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!("device main absorb failed: {e:?}")) + })?; + } + } } } @@ -553,25 +616,90 @@ where }]; builder.absorb(&src, 0); } - // Streaming device absorb (ext3 row-major), so the aux LDE can be freed - // below without retaining it. + // Streaming device absorb: expand this table's aux LDE on the GPU and + // absorb it (slab layout) straight from VRAM — no host FFT upload. The + // resident slab buffer yields the byte-identical aux leaf as the host + // row-major absorb (both emit, per bit-reversed row, each column's 3 + // components consecutively). Tables below the device LDE threshold — and + // disk-spilled traces (the natural aux trace lives on disk, the resident + // expand needs it in memory) — fall back to uploading the host LDE + // row-major (mixed eligibility, identical root). Like main, this is + // step-2-shaped: the host `expand_aux_lde_row_major` above still feeds the + // later phases; step 3 moves those to device recompute. #[cfg(feature = "cuda")] if let Some(dev) = device_aux.as_mut() { - // SAFETY: ext3 Goldilocks (device_commit_enabled gated it); an - // element is 3 consecutive u64. - let data_u64 = unsafe { - std::slice::from_raw_parts(aux_data.as_ptr() as *const u64, aux_data.len() * 3) + #[cfg(feature = "disk-spill")] + let aux_resident_ok = storage_mode != StorageMode::Disk; + #[cfg(not(feature = "disk-spill"))] + let aux_resident_ok = true; + + let resident = if aux_resident_ok { + let (trace_slice, num_cols) = trace.aux_data_row_major(); + debug_assert_eq!( + num_cols, aux_cols, + "the resident and host aux expansions must see the same column count" + ); + let n = if num_cols > 0 { + trace_slice.len() / num_cols + } else { + 0 + }; + crate::gpu_lde::try_expand_ext3_row_major_keep_no_tree::( + trace_slice, + n, + num_cols, + domains[table].blowup_factor, + &twiddles[table].coset_weights, + false, + ) + } else { + None }; - dev.absorb_ext3_row_major( - shape.heights[table] as u64, - data_u64, - aux_cols as u64, - 0, - aux_cols as u64, - ) - .map_err(|e| { - ProvingError::WrongParameter(format!("device aux absorb failed: {e:?}")) - })?; + + match resident { + // Resident slab LDE: absorb columns [0, aux_cols) in place; + // `wait_ready_on` orders the absorb after the producer device-side. + Some((handle, _host_lde)) => { + let commit_stream = dev.stream(); + handle.wait_ready_on(&commit_stream).map_err(|e| { + ProvingError::WrongParameter(format!( + "device aux LDE ready-wait failed: {e:?}" + )) + })?; + dev.absorb_ext3_slabs_dev( + shape.heights[table] as u64, + handle.buf.as_ref(), + handle.lde_size as u64, + aux_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!( + "device aux absorb (resident) failed: {e:?}" + )) + })?; + } + // Below threshold or disk-spilled: upload the host LDE row-major. + None => { + // SAFETY: ext3 Goldilocks (device_commit_enabled gated it); an + // element is 3 consecutive u64. + let data_u64 = unsafe { + std::slice::from_raw_parts( + aux_data.as_ptr() as *const u64, + aux_data.len() * 3, + ) + }; + dev.absorb_ext3_row_major( + shape.heights[table] as u64, + data_u64, + aux_cols as u64, + 0, + aux_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!("device aux absorb failed: {e:?}")) + })?; + } + } } match residency { ResidencyMode::Retain => retained_aux[table] = Some((aux_data, aux_cols)), @@ -1489,6 +1617,107 @@ where /// Build (or take back) a table's main and aux LDEs for the phase about to read /// them. #[allow(clippy::too_many_arguments)] +/// Recompute a table's MAIN LDE on the GPU — the VRAM recompute mechanism: a +/// cheap device coset NTT → the row-major LDE downloaded to host, BYTE-IDENTICAL +/// to `expand_main_lde_row_major` (same DFT, exact modular arithmetic) but off +/// the CPU. `None` = ineligible (below the device threshold / not Goldilocks / +/// disk-spilled) → the caller falls back to the host FFT. Only the row-major host +/// LDE is kept here; the resident handle (its col-major buf) is dropped — a later +/// 3c step keeps it resident to skip the download and run R2/R3/R4 on-device. +#[cfg(feature = "cuda")] +fn device_recompute_main_lde( + trace: &TraceTable, + domain: &Domain, + twiddles: &crate::prover::LdeTwiddles, + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, +) -> Option<(Vec>, usize)> +where + Field: IsFFTField + IsSubFieldOf + 'static, + FieldExtension: IsField + 'static, +{ + // Disk mode spilled the trace; the device expand needs it in memory. + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + return None; + } + let (trace_slice, num_cols) = trace.main_data_row_major(); + if num_cols == 0 { + return None; + } + let n = trace_slice.len() / num_cols; + let (_handle, host_lde) = crate::gpu_lde::try_expand_row_major_keep_no_tree::( + trace_slice, + trace.main_rowmajor_dev(), + n, + num_cols, + domain.blowup_factor, + &twiddles.coset_weights, + true, + )?; + Some((host_lde, num_cols)) +} + +#[cfg(not(feature = "cuda"))] +fn device_recompute_main_lde( + _trace: &TraceTable, + _domain: &Domain, + _twiddles: &crate::prover::LdeTwiddles, + #[cfg(feature = "disk-spill")] _storage_mode: StorageMode, +) -> Option<(Vec>, usize)> +where + Field: IsFFTField + IsSubFieldOf + 'static, + FieldExtension: IsField + 'static, +{ + None +} + +/// Aux counterpart of [`device_recompute_main_lde`] (ext3 row-major LDE). +#[cfg(feature = "cuda")] +fn device_recompute_aux_lde( + trace: &TraceTable, + domain: &Domain, + twiddles: &crate::prover::LdeTwiddles, + #[cfg(feature = "disk-spill")] storage_mode: StorageMode, +) -> Option<(Vec>, usize)> +where + Field: IsFFTField + IsSubFieldOf + 'static, + FieldExtension: IsField + 'static, +{ + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk { + return None; + } + let (trace_slice, num_cols) = trace.aux_data_row_major(); + if num_cols == 0 { + return None; + } + let n = trace_slice.len() / num_cols; + let (_handle, host_lde) = + crate::gpu_lde::try_expand_ext3_row_major_keep_no_tree::( + trace_slice, + n, + num_cols, + domain.blowup_factor, + &twiddles.coset_weights, + true, + )?; + Some((host_lde, num_cols)) +} + +#[cfg(not(feature = "cuda"))] +fn device_recompute_aux_lde( + _trace: &TraceTable, + _domain: &Domain, + _twiddles: &crate::prover::LdeTwiddles, + #[cfg(feature = "disk-spill")] _storage_mode: StorageMode, +) -> Option<(Vec>, usize)> +where + Field: IsFFTField + IsSubFieldOf + 'static, + FieldExtension: IsField + 'static, +{ + None +} + fn materialize_ldes( table: usize, air_trace_pairs: &[BatchedAirTracePair<'_, Field, FieldExtension, PI>], @@ -1516,13 +1745,25 @@ where Some(lde) => lde, None => { let t_expand = std::time::Instant::now(); - let lde = P::expand_main_lde_row_major( + // Recompute on the GPU (cheap device NTT) when eligible — the VRAM + // recompute mechanism; falls back to the host FFT below threshold / + // non-Goldilocks / disk-spill. Byte-identical either way. + let lde = device_recompute_main_lde::( trace, &domains[table], &twiddles[table], #[cfg(feature = "disk-spill")] storage_mode, - ); + ) + .unwrap_or_else(|| { + P::expand_main_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ) + }); stats.lde_expansion_wall += t_expand.elapsed(); stats.main_lde_expansions += 1; let b = lde_bytes::(lde.0.len()); @@ -1537,13 +1778,22 @@ where Some(lde) => lde, None => { let t_expand = std::time::Instant::now(); - let lde = P::expand_aux_lde_row_major( + let lde = device_recompute_aux_lde::( trace, &domains[table], &twiddles[table], #[cfg(feature = "disk-spill")] storage_mode, - ); + ) + .unwrap_or_else(|| { + P::expand_aux_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ) + }); stats.lde_expansion_wall += t_expand.elapsed(); stats.aux_lde_expansions += 1; let b = lde_bytes::(lde.0.len()); diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 82173ced6..086341cdb 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -858,6 +858,68 @@ where Some((tree, handle, lde_out)) } +/// Tree-less variant of [`try_expand_leaf_and_tree_row_major_keep`] for the +/// batched prover: expand the row-major trace to its coset LDE on device and +/// keep it resident column-major, building NO per-table Merkle tree — the shared +/// mixed-height MMCS is the one tree, so a per-table row-pair tree here would be +/// redundant on-device keccak. Returns the `GpuLdeBase` handle (`tree: None`, +/// `buf` byte-identical to the tree-building path) and the row-major LDE (empty +/// when `retain_host_lde=false`). `None` = declined (below threshold or not +/// Goldilocks) so the caller falls back to a host-upload absorb. +pub(crate) fn try_expand_row_major_keep_no_tree( + row_major: &[FieldElement], + predev: Option<&math_cuda::CudaSlice>, + n: usize, + m: usize, + blowup_factor: usize, + weights: &[FieldElement], + retain_host_lde: bool, +) -> Option<(math_cuda::lde::GpuLdeBase, Vec>)> +where + F: IsField + 'static, + E: IsField + 'static, +{ + let lde_size = n.saturating_mul(blowup_factor); + if lde_size < gpu_lde_threshold() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if row_major.len() != n * m || m == 0 || n == 0 { + return None; + } + + let raw: &[u64] = unsafe { from_raw_parts(row_major.as_ptr() as *const u64, n * m) }; + let weights_u64 = unsafe { weights_to_u64::(weights) }; + + // Only the LDE columns are computed on device (no leaf hash / tree build), so + // bump only the LDE-call counter. + GPU_LDE_CALLS.fetch_add(m as u64, Ordering::Relaxed); + + let (handle, lde_u64) = math_cuda::lde::coset_lde_row_major_keep_no_tree( + raw, + predev, + n, + m, + blowup_factor, + &weights_u64, + retain_host_lde, + ) + .ok()?; + + // Transmute Vec → Vec> (zero-copy, E == GoldilocksField). + let lde_out: Vec> = unsafe { + let mut v = std::mem::ManuallyDrop::new(lde_u64); + Vec::from_raw_parts(v.as_mut_ptr() as *mut FieldElement, v.len(), v.capacity()) + }; + + Some((handle, lde_out)) +} + /// Convert a GPU-built full node buffer (`(2*leaves - 1) * 32` bytes, inner /// nodes first, root at offset 0, leaves at the tail) into a host /// [`MerkleTree`], the exact layout `from_precomputed_nodes` expects. @@ -1054,6 +1116,73 @@ where Some((tree, handle, lde_out)) } +/// Tree-less variant of [`try_expand_leaf_and_tree_ext3_row_major_keep`] for the +/// batched aux round: expand the row-major ext3 trace to its coset LDE and keep +/// it resident (slab layout), building NO per-table Merkle tree — the shared +/// mixed-height MMCS is the tree. The caller absorbs the resident `buf` via +/// `StreamingMixedMmcs::absorb_ext3_slabs_dev`, which produces the byte-identical +/// aux leaf as the host row-major absorb. `None` = declined (below threshold or +/// not Fp3) so the caller falls back to a host-upload absorb. +pub(crate) fn try_expand_ext3_row_major_keep_no_tree( + row_major: &[FieldElement], + n: usize, + m: usize, + blowup_factor: usize, + weights: &[FieldElement], + retain_host_lde: bool, +) -> Option<(math_cuda::lde::GpuLdeExt3, Vec>)> +where + F: IsField + 'static, + E: IsField + 'static, +{ + let lde_size = n.saturating_mul(blowup_factor); + if lde_size < gpu_lde_threshold() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if row_major.len() != n * m || m == 0 || n == 0 { + return None; + } + + // Fp3 = [u64; 3] in memory — reinterpret as flat u64 slice (m3 = m*3). + let raw: &[u64] = unsafe { from_raw_parts(row_major.as_ptr() as *const u64, n * m * 3) }; + let weights_u64 = unsafe { weights_to_u64::(weights) }; + + // Only the LDE columns are computed on device (no leaf hash / tree build). + GPU_LDE_CALLS.fetch_add((m * 3) as u64, Ordering::Relaxed); + + let (handle, lde_u64) = math_cuda::lde::coset_lde_ext3_row_major_keep_no_tree( + raw, + n, + m, + blowup_factor, + &weights_u64, + retain_host_lde, + ) + .ok()?; + + // Transmute Vec → Vec> (zero-copy, E == Fp3 = [u64;3]). + let lde_out: Vec> = unsafe { + let mut v = std::mem::ManuallyDrop::new(lde_u64); + debug_assert!( + v.len() % 3 == 0 && v.capacity() % 3 == 0, + "lde_u64 len/capacity must be a multiple of 3 for Fp3 reinterpret" + ); + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len() / 3, + v.capacity() / 3, + ) + }; + + Some((handle, lde_out)) +} + /// Ext3 specialisation of [`try_expand_columns_batched`]. `E` is known to be /// `Degree3GoldilocksExtensionField` by TypeId match at the caller. /// diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index a8dc48726..999286c62 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -400,7 +400,9 @@ pub(crate) struct LdeTwiddles { /// Size-`n` FORWARD set, built lazily — only the batched phase-4 coset /// evaluation wants it (a full LDE's forward set is size `n·blowup`). two_half_fwd_n: OnceLock>, - coset_weights: Vec>, + /// `pub(crate)` so the batched prover's device-resident main round can feed + /// the same coset weights to `gpu_lde::try_expand_leaf_and_tree_row_major_keep`. + pub(crate) coset_weights: Vec>, /// Composition half-extension cache, initialized only when the degree-2 /// decomposition path actually runs on CPU. composition: OnceLock>, From 7a3a96f837cd8b64ad6cfc03cb265ec8af9688b6 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 2 Sep 2026 14:21:25 -0300 Subject: [PATCH 16/26] feat(cuda): device R2 constraint eval + R4 DEEP for the batched prover (hybrid resident handles) --- crypto/stark/src/batched/prover.rs | 168 +++++++++++++++++++---------- 1 file changed, 109 insertions(+), 59 deletions(-) diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index 4f06c3e7d..6e9343242 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -159,6 +159,16 @@ struct LdePair { main: (Vec>, usize), aux: (Vec>, usize), bytes: usize, + /// Resident device LDE handles, present when `materialize_ldes` recomputed on + /// the GPU. `lde_trace_take` attaches them to the `LDETraceTable` so R2 + /// constraint eval and R4 DEEP fire their device paths (they read + /// `gpu_main`/`gpu_aux`); the host LDE above is kept alongside so `want_host` + /// still drains parts to host and the host fallback stays valid + /// (`host_trace_empty` = false). `None` on the host / retained paths. + #[cfg(feature = "cuda")] + gpu_main: Option, + #[cfg(feature = "cuda")] + gpu_aux: Option, } /// Prove one epoch with batched commitments. @@ -1618,19 +1628,20 @@ where /// them. #[allow(clippy::too_many_arguments)] /// Recompute a table's MAIN LDE on the GPU — the VRAM recompute mechanism: a -/// cheap device coset NTT → the row-major LDE downloaded to host, BYTE-IDENTICAL -/// to `expand_main_lde_row_major` (same DFT, exact modular arithmetic) but off -/// the CPU. `None` = ineligible (below the device threshold / not Goldilocks / -/// disk-spilled) → the caller falls back to the host FFT. Only the row-major host -/// LDE is kept here; the resident handle (its col-major buf) is dropped — a later -/// 3c step keeps it resident to skip the download and run R2/R3/R4 on-device. +/// cheap device coset NTT. Returns the row-major host LDE (BYTE-IDENTICAL to +/// `expand_main_lde_row_major` — same DFT, exact modular arithmetic — kept so the +/// parts drain + the host fallback stay valid) AND the resident device handle, +/// which `lde_trace_take` attaches to the `LDETraceTable` so R2 constraint eval / +/// R4 DEEP fire their device paths (`evaluate_dev` / `try_deep_composition_gpu` +/// read `gpu_main`/`gpu_aux`, not `host_trace_empty`). `None` = ineligible (below +/// the device threshold / not Goldilocks / disk-spilled) → host FFT. #[cfg(feature = "cuda")] fn device_recompute_main_lde( trace: &TraceTable, domain: &Domain, twiddles: &crate::prover::LdeTwiddles, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, -) -> Option<(Vec>, usize)> +) -> Option<(Vec>, usize, math_cuda::lde::GpuLdeBase)> where Field: IsFFTField + IsSubFieldOf + 'static, FieldExtension: IsField + 'static, @@ -1645,7 +1656,7 @@ where return None; } let n = trace_slice.len() / num_cols; - let (_handle, host_lde) = crate::gpu_lde::try_expand_row_major_keep_no_tree::( + let (handle, host_lde) = crate::gpu_lde::try_expand_row_major_keep_no_tree::( trace_slice, trace.main_rowmajor_dev(), n, @@ -1654,21 +1665,7 @@ where &twiddles.coset_weights, true, )?; - Some((host_lde, num_cols)) -} - -#[cfg(not(feature = "cuda"))] -fn device_recompute_main_lde( - _trace: &TraceTable, - _domain: &Domain, - _twiddles: &crate::prover::LdeTwiddles, - #[cfg(feature = "disk-spill")] _storage_mode: StorageMode, -) -> Option<(Vec>, usize)> -where - Field: IsFFTField + IsSubFieldOf + 'static, - FieldExtension: IsField + 'static, -{ - None + Some((host_lde, num_cols, handle)) } /// Aux counterpart of [`device_recompute_main_lde`] (ext3 row-major LDE). @@ -1678,7 +1675,7 @@ fn device_recompute_aux_lde( domain: &Domain, twiddles: &crate::prover::LdeTwiddles, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, -) -> Option<(Vec>, usize)> +) -> Option<(Vec>, usize, math_cuda::lde::GpuLdeExt3)> where Field: IsFFTField + IsSubFieldOf + 'static, FieldExtension: IsField + 'static, @@ -1692,7 +1689,7 @@ where return None; } let n = trace_slice.len() / num_cols; - let (_handle, host_lde) = + let (handle, host_lde) = crate::gpu_lde::try_expand_ext3_row_major_keep_no_tree::( trace_slice, n, @@ -1701,21 +1698,7 @@ where &twiddles.coset_weights, true, )?; - Some((host_lde, num_cols)) -} - -#[cfg(not(feature = "cuda"))] -fn device_recompute_aux_lde( - _trace: &TraceTable, - _domain: &Domain, - _twiddles: &crate::prover::LdeTwiddles, - #[cfg(feature = "disk-spill")] _storage_mode: StorageMode, -) -> Option<(Vec>, usize)> -where - Field: IsFFTField + IsSubFieldOf + 'static, - FieldExtension: IsField + 'static, -{ - None + Some((host_lde, num_cols, handle)) } fn materialize_ldes( @@ -1740,30 +1723,49 @@ where { let (_, trace, _) = &air_trace_pairs[table]; let mut bytes = 0usize; + // Resident device handles captured from a GPU recompute; attached to the + // trace by `lde_trace_take` so R2/R4 fire their device paths. + #[cfg(feature = "cuda")] + let mut gpu_main: Option = None; + #[cfg(feature = "cuda")] + let mut gpu_aux: Option = None; let main = match retained_main[table].take() { Some(lde) => lde, None => { let t_expand = std::time::Instant::now(); // Recompute on the GPU (cheap device NTT) when eligible — the VRAM - // recompute mechanism; falls back to the host FFT below threshold / - // non-Goldilocks / disk-spill. Byte-identical either way. - let lde = device_recompute_main_lde::( + // recompute mechanism; the resident handle is kept + attached to the + // trace so R2/R4 run on device. Falls back to the host FFT below + // threshold / non-Goldilocks / disk-spill. Byte-identical either way. + #[cfg(feature = "cuda")] + let lde = match device_recompute_main_lde::( trace, &domains[table], &twiddles[table], #[cfg(feature = "disk-spill")] storage_mode, - ) - .unwrap_or_else(|| { - P::expand_main_lde_row_major( + ) { + Some((host_lde, cols, handle)) => { + gpu_main = Some(handle); + (host_lde, cols) + } + None => P::expand_main_lde_row_major( trace, &domains[table], &twiddles[table], #[cfg(feature = "disk-spill")] storage_mode, - ) - }); + ), + }; + #[cfg(not(feature = "cuda"))] + let lde = P::expand_main_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); stats.lde_expansion_wall += t_expand.elapsed(); stats.main_lde_expansions += 1; let b = lde_bytes::(lde.0.len()); @@ -1778,22 +1780,34 @@ where Some(lde) => lde, None => { let t_expand = std::time::Instant::now(); - let lde = device_recompute_aux_lde::( + #[cfg(feature = "cuda")] + let lde = match device_recompute_aux_lde::( trace, &domains[table], &twiddles[table], #[cfg(feature = "disk-spill")] storage_mode, - ) - .unwrap_or_else(|| { - P::expand_aux_lde_row_major( + ) { + Some((host_lde, cols, handle)) => { + gpu_aux = Some(handle); + (host_lde, cols) + } + None => P::expand_aux_lde_row_major( trace, &domains[table], &twiddles[table], #[cfg(feature = "disk-spill")] storage_mode, - ) - }); + ), + }; + #[cfg(not(feature = "cuda"))] + let lde = P::expand_aux_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); stats.lde_expansion_wall += t_expand.elapsed(); stats.aux_lde_expansions += 1; let b = lde_bytes::(lde.0.len()); @@ -1807,7 +1821,15 @@ where }; let _ = residency; - LdePair { main, aux, bytes } + LdePair { + main, + aux, + bytes, + #[cfg(feature = "cuda")] + gpu_main, + #[cfg(feature = "cuda")] + gpu_aux, + } } /// Give a table's LDEs back to the retention slots, or drop them. @@ -1849,11 +1871,32 @@ where Field: IsFFTField + IsSubFieldOf, FieldExtension: IsField, { - let LdePair { main, aux, bytes } = ldes; - ( - LDETraceTable::from_row_major(main.0, main.1, aux.0, aux.1, step_size, blowup_factor), + let LdePair { + main, + aux, bytes, - ) + #[cfg(feature = "cuda")] + gpu_main, + #[cfg(feature = "cuda")] + gpu_aux, + } = ldes; + #[allow(unused_mut)] + let mut lde_trace = + LDETraceTable::from_row_major(main.0, main.1, aux.0, aux.1, step_size, blowup_factor); + // Attach resident LDE handles (GPU recompute) so R2 constraint eval / R4 DEEP + // read them on device. `host_trace_empty` intentionally stays FALSE: the host + // LDE is present, so `want_host` still drains parts to host and the host + // fallback stays valid for tables whose device path declines. + #[cfg(feature = "cuda")] + { + if let Some(h) = gpu_main { + lde_trace.set_gpu_main(h); + } + if let Some(h) = gpu_aux { + lde_trace.set_gpu_aux(h); + } + } + (lde_trace, bytes) } /// Take the buffers back out of the trace view for release or retention — @@ -1870,6 +1913,13 @@ where main: (lde_trace.main_data, lde_trace.num_main_cols), aux: (lde_trace.aux_data, lde_trace.num_aux_cols), bytes, + // The resident handles (if any) live in `lde_trace` and drop with it here; + // under RecomputeLde (the only path that attaches them) release drops + // everything and the next phase recomputes. + #[cfg(feature = "cuda")] + gpu_main: None, + #[cfg(feature = "cuda")] + gpu_aux: None, } } From 86625af6380ba2dcd0d31f1a0aec4146525bc07f Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 2 Sep 2026 15:10:28 -0300 Subject: [PATCH 17/26] test(prover): LAMBDA_BATCHED_RESIDENCY knob for batched GPU timing --- prover/src/lib.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/prover/src/lib.rs b/prover/src/lib.rs index c62641474..179bc2dbd 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -1324,13 +1324,21 @@ pub fn time_batched_prove( &runtime_page_ranges, proof_options.fri_final_poly_log_degree, ); + // `LAMBDA_BATCHED_RESIDENCY=recompute` exercises the device-resident R2/R4 + // paths (they engage only under RecomputeLde, where `materialize_ldes` + // recomputes on the GPU and attaches the handles); default Retain keeps every + // LDE host-resident (the small-block path). + let residency = match std::env::var("LAMBDA_BATCHED_RESIDENCY").as_deref() { + Ok("recompute") => stark::residency_mode::ResidencyMode::RecomputeLde, + _ => stark::residency_mode::ResidencyMode::Retain, + }; let t = std::time::Instant::now(); let _ = stark::batched::prover::multi_prove_batched::>( airs.air_trace_pairs(&mut traces), &mut transcript, #[cfg(feature = "disk-spill")] storage_mode, - stark::residency_mode::ResidencyMode::Retain, + residency, ) .map_err(|e| Error::Prover(format!("{e:?}")))?; Ok(t.elapsed()) From 873d4cc8752b40ed205cbe0cb707fe89ee74c500 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 2 Sep 2026 15:51:54 -0300 Subject: [PATCH 18/26] feat(cuda): device-only R2/R4 recompute for the batched prover (no full-LDE download) --- crypto/stark/src/batched/prover.rs | 84 +++++++++++++++++++++++++++--- 1 file changed, 77 insertions(+), 7 deletions(-) diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index 6e9343242..8d0c52da6 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -792,6 +792,9 @@ where &mut stats, &mut ledger, residency, + // R2: device-only recompute — constraint eval reads the resident LDE, + // only the small composition parts come back to host (below). + true, #[cfg(feature = "disk-spill")] storage_mode, ); @@ -810,7 +813,31 @@ where &boundary_coefficients, )?; stats.parts_computations += 1; - let parts = computed.parts; + #[allow(unused_mut)] + let mut parts = computed.parts; + // Device-only R2: the constraint eval ran on the resident LDE and the + // parts live in the handle (empty host placeholders). Download them — + // small (2 composition-poly columns, not the full LDE) — for the host + // parts commit + R3 coset-eval + R4 DEEP. No-op on the host path (parts + // already populated) or any table whose device R2 declined (recovery + // filled the host parts). + #[cfg(feature = "cuda")] + if parts.first().is_some_and(|p| p.is_empty()) { + let handle = computed.gpu_parts.as_ref().expect( + "device-only R2 must carry a resident parts handle when host parts are empty", + ); + let stream = math_cuda::device::backend() + .map_err(|e| { + ProvingError::WrongParameter(format!("backend for parts download: {e:?}")) + })? + .next_stream(); + parts = crate::gpu_lde::download_composition_parts_host::( + handle, &stream, + ) + .ok_or_else(|| { + ProvingError::WrongParameter("device-only R2 parts download failed".to_string()) + })?; + } let parts_bytes: usize = parts .iter() @@ -890,6 +917,9 @@ where &mut stats, &mut ledger, residency, + // R3 runs only under Retain here (RecomputeLde uses the coset-eval + // branch below); the retained host LDE is used, so not device-only. + false, #[cfg(feature = "disk-spill")] storage_mode, ); @@ -1011,6 +1041,9 @@ where stats, ledger, residency, + // R4 DEEP: device-only recompute — reads the resident LDE + // + host parts (downloaded in R2); no full-LDE D2H. + true, #[cfg(feature = "disk-spill")] storage_mode, ); @@ -1105,6 +1138,9 @@ where &mut stats, &mut ledger, residency, + // Openings read HOST LDE rows (`LeafSource::append_row`), so keep the + // host LDE (not device-only) — device openings would be a later step. + false, #[cfg(feature = "disk-spill")] storage_mode, ); @@ -1640,6 +1676,7 @@ fn device_recompute_main_lde( trace: &TraceTable, domain: &Domain, twiddles: &crate::prover::LdeTwiddles, + retain_host_lde: bool, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> Option<(Vec>, usize, math_cuda::lde::GpuLdeBase)> where @@ -1656,6 +1693,9 @@ where return None; } let n = trace_slice.len() / num_cols; + // `retain_host_lde=false` = device-only: no D2H of the full LDE (host Vec + // comes back empty). R2/R4 read the resident handle; the openings phase keeps + // it true so its host row reads still work. let (handle, host_lde) = crate::gpu_lde::try_expand_row_major_keep_no_tree::( trace_slice, trace.main_rowmajor_dev(), @@ -1663,7 +1703,7 @@ where num_cols, domain.blowup_factor, &twiddles.coset_weights, - true, + retain_host_lde, )?; Some((host_lde, num_cols, handle)) } @@ -1674,6 +1714,7 @@ fn device_recompute_aux_lde( trace: &TraceTable, domain: &Domain, twiddles: &crate::prover::LdeTwiddles, + retain_host_lde: bool, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> Option<(Vec>, usize, math_cuda::lde::GpuLdeExt3)> where @@ -1696,7 +1737,7 @@ where num_cols, domain.blowup_factor, &twiddles.coset_weights, - true, + retain_host_lde, )?; Some((host_lde, num_cols, handle)) } @@ -1712,6 +1753,12 @@ fn materialize_ldes( stats: &mut BatchedProveStats, ledger: &mut ResidencyLedger, residency: ResidencyMode, + // When true (R2 constraint eval, R4 DEEP), the GPU recompute is DEVICE-ONLY: + // no D2H of the full LDE — those phases read the resident handle and only the + // small composition parts come back to host. The openings phase passes false + // (it reads host LDE rows), so the LDE is downloaded there. Ignored on the + // host / retained paths. + #[cfg_attr(not(feature = "cuda"), allow(unused_variables))] device_only: bool, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> LdePair where @@ -1723,6 +1770,8 @@ where { let (_, trace, _) = &air_trace_pairs[table]; let mut bytes = 0usize; + #[cfg(feature = "cuda")] + let retain_host_lde = !device_only; // Resident device handles captured from a GPU recompute; attached to the // trace by `lde_trace_take` so R2/R4 fire their device paths. #[cfg(feature = "cuda")] @@ -1743,6 +1792,7 @@ where trace, &domains[table], &twiddles[table], + retain_host_lde, #[cfg(feature = "disk-spill")] storage_mode, ) { @@ -1785,6 +1835,7 @@ where trace, &domains[table], &twiddles[table], + retain_host_lde, #[cfg(feature = "disk-spill")] storage_mode, ) { @@ -1880,15 +1931,34 @@ where #[cfg(feature = "cuda")] gpu_aux, } = ldes; + // Detect device-only from the ACTUAL buffer state (empty host buffer + a + // resident handle) BEFORE the Vecs are moved — mirrors the per-table + // `build_lde_trace`. `device_only=false` (openings) keeps the host LDE, so + // these stay false and the trace is a plain host trace with handles attached + // (R2/R4 still read them via `gpu_main`/`gpu_aux`). + #[cfg(feature = "cuda")] + let main_empty = main.1 > 0 && main.0.is_empty() && gpu_main.is_some(); + #[cfg(feature = "cuda")] + let host_trace_empty = + main_empty || (aux.1 > 0 && aux.0.is_empty() && gpu_aux.is_some()); + #[cfg(feature = "cuda")] + let device_rows = gpu_main + .as_ref() + .map(|h| h.lde_size) + .or_else(|| gpu_aux.as_ref().map(|h| h.lde_size)); #[allow(unused_mut)] let mut lde_trace = LDETraceTable::from_row_major(main.0, main.1, aux.0, aux.1, step_size, blowup_factor); - // Attach resident LDE handles (GPU recompute) so R2 constraint eval / R4 DEEP - // read them on device. `host_trace_empty` intentionally stays FALSE: the host - // LDE is present, so `want_host` still drains parts to host and the host - // fallback stays valid for tables whose device path declines. #[cfg(feature = "cuda")] { + if host_trace_empty { + // `from_row_major` read num_rows from the empty host buffer (→ 0); + // recover the true LDE row count from the resident handle. + if let Some(n) = device_rows { + lde_trace.set_num_rows(n); + } + lde_trace.set_host_trace_empty(true); + } if let Some(h) = gpu_main { lde_trace.set_gpu_main(h); } From 43c4b04457527011fcb4cf799251a5ea8ed16856 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 2 Sep 2026 17:03:29 -0300 Subject: [PATCH 19/26] feat(cuda): device-gather openings for the batched prover (no full-LDE download for plain tables) --- crypto/stark/src/batched/prover.rs | 199 +++++++++++++++++++++++++---- 1 file changed, 177 insertions(+), 22 deletions(-) diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index 8d0c52da6..4a2e77839 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -1126,7 +1126,19 @@ where let parts_iotas = reduced_iotas(&iotas, h_max, parts_mmcs.h_max()); for table in 0..num_tables { - let (air, _, _) = &air_trace_pairs[table]; + let (air, trace, _) = &air_trace_pairs[table]; + #[cfg(not(feature = "cuda"))] + let _ = &trace; + // Plain tables (no precomputed columns, not carved) gather their openings + // straight off the resident LDE — no full-LDE download. Precomputed and + // carved tables keep the host LDE (their per-table trees read it) but + // still gather the shared-MMCS main/aux openings from the resident handle. + let is_carved = shape.carved_main.map(|c| c.table) == Some(table); + #[cfg(feature = "cuda")] + let openings_device_only = + !is_carved && trace.main_data_row_major().1 == matrix_width(&shape.main, table); + #[cfg(not(feature = "cuda"))] + let openings_device_only = false; let ldes = materialize_ldes::( table, &air_trace_pairs, @@ -1138,16 +1150,13 @@ where &mut stats, &mut ledger, residency, - // Openings read HOST LDE rows (`LeafSource::append_row`), so keep the - // host LDE (not device-only) — device openings would be a later step. - false, + openings_device_only, #[cfg(feature = "disk-spill")] storage_mode, ); let _ = air; let height = shape.heights[table]; let (main_data, total_cols) = &ldes.main; - let is_carved = shape.carved_main.map(|c| c.table) == Some(table); let num_precomputed = if is_carved { 0 } else { @@ -1179,26 +1188,103 @@ where } } if let Some(m) = matrix_index(&shape.main, table) { - let src = vec![BorrowedMatrix::RowMajorNatural { - data: main_data, - stride: *total_cols, - col_start: num_precomputed, - width: total_cols - num_precomputed, - log_height: height, - }]; - fill_openings(&main_mmcs, m, &src, &main_iotas, &mut main_openings); + #[cfg(feature = "cuda")] + { + // Gather this table's query row-pairs off the resident LDE handle + // (a small D2H of only the queried rows) instead of downloading + // the whole LDE. `None` when no handle (sub-threshold decline) → + // the host LDE (present in that case) serves the openings. + let gathered = ldes.gpu_main.as_ref().and_then(|h| { + let rows = query_row_indices(&main_mmcs, m, &main_iotas, h.lde_size); + let stream = math_cuda::device::backend().ok()?.next_stream(); + let raw = + math_cuda::barycentric::gather_rows_base_on_device(h, &rows, &stream).ok()?; + crate::constraint_ir::gpu_interp::base_u64_to_field::(&raw) + .map(|v| (v, h.m)) + }); + match gathered { + Some((g, ncols)) => fill_openings_from_gathered( + &main_mmcs, + m, + &g, + ncols, + num_precomputed, + *total_cols, + &main_iotas, + &mut main_openings, + ), + None => { + let src = vec![BorrowedMatrix::RowMajorNatural { + data: main_data, + stride: *total_cols, + col_start: num_precomputed, + width: total_cols - num_precomputed, + log_height: height, + }]; + fill_openings(&main_mmcs, m, &src, &main_iotas, &mut main_openings); + } + } + } + #[cfg(not(feature = "cuda"))] + { + let src = vec![BorrowedMatrix::RowMajorNatural { + data: main_data, + stride: *total_cols, + col_start: num_precomputed, + width: total_cols - num_precomputed, + log_height: height, + }]; + fill_openings(&main_mmcs, m, &src, &main_iotas, &mut main_openings); + } } if let (Some(mmcs), Some(m)) = (aux_mmcs.as_ref(), matrix_index(&shape.aux, table)) { - let (aux_data, aux_cols) = &ldes.aux; - let src = vec![BorrowedMatrix::RowMajorNatural { - data: aux_data, - stride: *aux_cols, - col_start: 0, - width: *aux_cols, - log_height: height, - }]; let indices = aux_iotas.as_ref().expect("the aux MMCS exists here"); - fill_openings(mmcs, m, &src, indices, &mut aux_openings); + #[cfg(feature = "cuda")] + { + let gathered = ldes.gpu_aux.as_ref().and_then(|h| { + let rows = query_row_indices(mmcs, m, indices, h.lde_size); + let stream = math_cuda::device::backend().ok()?.next_stream(); + let raw = + math_cuda::barycentric::gather_rows_ext3_on_device(h, &rows, &stream).ok()?; + crate::constraint_ir::gpu_interp::ext3_u64_to_field::(&raw) + .map(|v| (v, h.m)) + }); + match gathered { + Some((g, ncols)) => fill_openings_from_gathered( + mmcs, + m, + &g, + ncols, + 0, + ncols, + indices, + &mut aux_openings, + ), + None => { + let (aux_data, aux_cols) = &ldes.aux; + let src = vec![BorrowedMatrix::RowMajorNatural { + data: aux_data, + stride: *aux_cols, + col_start: 0, + width: *aux_cols, + log_height: height, + }]; + fill_openings(mmcs, m, &src, indices, &mut aux_openings); + } + } + } + #[cfg(not(feature = "cuda"))] + { + let (aux_data, aux_cols) = &ldes.aux; + let src = vec![BorrowedMatrix::RowMajorNatural { + data: aux_data, + stride: *aux_cols, + col_start: 0, + width: *aux_cols, + log_height: height, + }]; + fill_openings(mmcs, m, &src, indices, &mut aux_openings); + } } if let Some(m) = matrix_index(&shape.parts, table) { let src = vec![BorrowedMatrix::ColMajorNatural { @@ -1621,6 +1707,75 @@ fn fill_openings( } } +/// Device counterpart of [`fill_openings`]: the query row-pairs were already +/// gathered off the resident LDE (`[even(q0), odd(q0), even(q1), odd(q1), ...]`, +/// each `ncols` field elements in FULL-row order); slice columns +/// `[col_start, col_end)` per query. Byte-identical output to `fill_openings` +/// (empty merkle path, filled by `assemble`), because the gathered rows are the +/// same rows `append_row` would have read from the host LDE — no full-LDE D2H. +#[cfg(feature = "cuda")] +#[allow(clippy::too_many_arguments)] +fn fill_openings_from_gathered( + mmcs: &MixedMmcs, + matrix: usize, + gathered: &[FieldElement], + ncols: usize, + col_start: usize, + col_end: usize, + iotas: &[usize], + out: &mut [Vec>>], +) where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + for (q, &iota) in iotas.iter().enumerate() { + // `row_pair_leaf` returning None means this round is shorter than the FRI + // and does not have the leaf; gathered rows for such queries are the row-0 + // placeholder and skipped here (same as `fill_openings`). + if mmcs.row_pair_leaf(iota, matrix).is_none() { + continue; + } + let even = gathered[(2 * q) * ncols + col_start..(2 * q) * ncols + col_end].to_vec(); + let odd = + gathered[(2 * q + 1) * ncols + col_start..(2 * q + 1) * ncols + col_end].to_vec(); + out[q][matrix] = Some(PolynomialOpenings { + proof: crypto::merkle_tree::proof::Proof { + merkle_path: Vec::new(), + }, + evaluations: even, + evaluations_sym: odd, + }); + } +} + +/// The NATURAL LDE row indices to gather for a matrix's query row-pairs: +/// per query, `rev(2*leaf)` and `rev(2*leaf+1)` (leaf = `row_pair_leaf`), which +/// is exactly what `append_row(2*leaf)` reads (it bit-reverses internally). +/// Out-of-range queries get row 0 (skipped by `fill_openings_from_gathered`). +#[cfg(feature = "cuda")] +fn query_row_indices( + mmcs: &MixedMmcs, + matrix: usize, + iotas: &[usize], + lde_size: usize, +) -> Vec +where + E: IsField + 'static, + FieldElement: AsBytes + Sync + Send, +{ + let n = lde_size as u64; + iotas + .iter() + .flat_map(|&iota| match mmcs.row_pair_leaf(iota, matrix) { + Some(leaf) => [ + math::fft::bit_reversing::reverse_index(2 * leaf, n) as u32, + math::fft::bit_reversing::reverse_index(2 * leaf + 1, n) as u32, + ], + None => [0u32, 0u32], + }) + .collect() +} + /// Reduce every FRI query index into one round's index space. /// /// `h_max_round <= h_max_fri` always holds for a round of this epoch — a round From 9ce0f8e9afc14d2d7223e053b2dc7375e3836383 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Thu, 3 Sep 2026 10:09:43 -0300 Subject: [PATCH 20/26] feat: VM prove/verify use the batched multi-merkle-tree prover (GPU device-resident) --- prover/benches/batched_cpu_cmp.rs | 42 ++++ prover/src/lib.rs | 352 +++++++++++++++++++++++++++--- prover/src/tests/mod.rs | 10 +- 3 files changed, 372 insertions(+), 32 deletions(-) diff --git a/prover/benches/batched_cpu_cmp.rs b/prover/benches/batched_cpu_cmp.rs index d4d742b8c..1bcb2381a 100644 --- a/prover/benches/batched_cpu_cmp.rs +++ b/prover/benches/batched_cpu_cmp.rs @@ -35,6 +35,48 @@ fn main() { if cfg!(feature = "cuda") { println!("⚠️ built with cuda — run WITHOUT --features cuda for the FAIR CPU comparison"); } + // VMPROVE=1: exercise the REAL VM API (`prove_with_inputs` + `verify`) after + // the batched cutover — the VM now proves/verifies with the multi-merkle-tree. + // This is the path a comparison-against-main harness hits. Build --features cuda. + if std::env::var("VMPROVE").is_ok() { + let t = std::time::Instant::now(); + let vm_proof = lambda_vm_prover::prove_with_inputs(&elf, &input).expect("VM prove"); + let dt = t.elapsed(); + let ok = lambda_vm_prover::verify(&vm_proof, &elf).expect("VM verify"); + println!("VMPROVE block={block} prove={dt:?} verify={ok}"); + return; + } + + // VERIFY=1: prove THEN verify a real block with the batched prover, device + // paths engaged (build with --features cuda). Closes the correctness loop at + // real-block scale (the timing/size paths only prove). + if std::env::var("VERIFY").is_ok() { + match lambda_vm_prover::prove_and_verify_batched_block(&elf, &input) { + Ok(true) => println!("VERIFY block={block} batched prove->verify: PASS"), + Ok(false) => println!("VERIFY block={block} batched prove->verify: FAIL (rejected)"), + Err(e) => println!("VERIFY block={block} errored: {e:?}"), + } + return; + } + + // SIZE=1: measure the serialized PROOF SIZE of both provers (the + // multi-merkle-tree's structural payoff — one shared auth path per query). + // Byte-identical on CPU or GPU, so this needs no GPU. Runs one prove each. + if std::env::var("SIZE").is_ok() { + let b = lambda_vm_prover::size_batched_prove(&elf, &input); + let p = lambda_vm_prover::size_per_table_prove(&elf, &input); + match (b, p) { + (Ok(b), Ok(p)) => println!( + "PROOF SIZE block={block} batched={b} bytes ({:.2} MiB) per-table={p} bytes ({:.2} MiB) batched/per-table={:.3}", + b as f64 / (1u64 << 20) as f64, + p as f64 / (1u64 << 20) as f64, + b as f64 / p as f64, + ), + (b, p) => println!("PROOF SIZE failed: batched={b:?} per-table={p:?}"), + } + return; + } + println!("=== CPU structural comparison — block={block} warmup={warmup} iters={iters} cooldown={cooldown}s ==="); println!("(both provers on CPU; isolates the multi-merkle-tree structure vs per-table)"); diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 179bc2dbd..11799251a 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -38,7 +38,7 @@ use stark::prover::{IsStarkProver, Prover}; #[cfg(feature = "disk-spill")] use stark::storage_mode::StorageMode; use stark::traits::AIR; -use stark::verifier::{IsStarkVerifier, Verifier}; +use stark::verifier::Verifier; use crate::statement::{StatementKind, absorb_statement, absorb_statement_with_digest}; pub use crate::tables::MaxRowsConfig; @@ -65,7 +65,6 @@ use crate::test_utils::{ // fixed at guest build time (`recursion::Preset`). pub use stark::config::Commitment; pub use stark::proof::options::{GoldilocksCubicProofOptions, ProofOptions}; -use stark::proof::stark::MultiProof; use stark::proof::view::{MultiProofView, ProofViewSource}; /// A run-length encoded range of contiguous zero-initialized 4KB pages. @@ -161,8 +160,10 @@ impl TableCounts { /// needed by the verifier to reconstruct the AIR configuration. #[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] pub struct VmProof { - /// The multi-table STARK proof. - pub proof: MultiProof, + /// The multi-table STARK proof — the batched (multi-merkle-tree) proof: one + /// shared tree over all tables, device-resident on GPU. (Cutover from the + /// per-table `MultiProof`; recursion/continuation still use per-table.) + pub proof: stark::batched::proof::BatchedMultiProof, /// Run-length encoded runtime page ranges. /// These are zero-initialized pages accessed during execution but not /// covered by ELF segments (stack, heap, etc.). @@ -403,8 +404,16 @@ pub fn verify_recursion_blob<'a>( let program = Elf::load(inner_elf).map_err(|e| Error::ElfLoad(format!("{e}")))?; let elf_digest = statement::elf_digest(inner_elf); + // The batched verifier reads an owned proof; deserialize the archived one + // (recursion host path — the guest cutover to a batched verifier is a + // separate effort, so this materializes the proof rather than verifying it + // zero-copy in-place). + let owned_proof: stark::batched::proof::BatchedMultiProof = + rkyv::deserialize::<_, RkyvError>(&archived.vm_proof.proof).map_err(|e| { + Error::Execution(format!("rkyv deserialize batched proof failed: {e}")) + })?; let ok = verify_proof_parts( - MultiProofView::Archived(&archived.vm_proof.proof), + &owned_proof, &table_counts, &runtime_page_ranges, num_private_input_pages, @@ -1217,11 +1226,14 @@ pub fn prove_with_options_and_inputs( // Phase 4: Prove (multi_prove) #[cfg(feature = "instruments")] let __sp = stark::instruments::span("proving"); - let proof = Prover::multi_prove( + // Cutover: the VM proves with the batched (multi-merkle-tree) prover, + // device-resident on GPU. RecomputeLde fits any block's VRAM. + let (proof, _stats) = stark::batched::prover::multi_prove_batched::>( airs.air_trace_pairs(&mut traces), &mut transcript, #[cfg(feature = "disk-spill")] storage_mode, + stark::residency_mode::ResidencyMode::RecomputeLde, ) .map_err(|e| Error::Prover(format!("{e:?}")))?; #[cfg(feature = "instruments")] @@ -1414,6 +1426,275 @@ pub fn time_per_table_prove( Ok(t.elapsed()) } +/// Serialized proof SIZE (rkyv bytes) of the batched (multi-merkle-tree) prover — +/// the structural payoff: one shared authentication path per query instead of one +/// per table. Compare against [`size_per_table_prove`] on the same block. Setup +/// mirrors [`time_batched_prove`] (Retain residency; the proof is +/// residency-independent). +pub fn size_batched_prove(elf_bytes: &[u8], private_inputs: &[u8]) -> Result { + let proof_options = + GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 is always valid"); + let max_rows = MaxRowsConfig::default(); + let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; + let executor = Executor::new(&program, private_inputs.to_vec()) + .map_err(|e| Error::Execution(format!("{e}")))?; + let result = executor + .run() + .map_err(|e| Error::Execution(format!("{e}")))?; + #[cfg(feature = "disk-spill")] + let storage_mode = { + let lengths = count_table_lengths(&program, &result.logs, &max_rows, private_inputs)?; + auto_storage::decide(&lengths, proof_options.blowup_factor) + }; + let mut traces = Traces::from_elf_and_logs( + &program, + &result.logs, + &max_rows, + private_inputs, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + drop(result); + let table_counts = traces.table_counts(); + let airs = VmAirs::new( + &program, + &proof_options, + false, + &traces.page_configs, + &table_counts, + None, + true, + None, + None, + None, + ); + let runtime_page_ranges = traces.runtime_page_ranges(); + let num_private_input_pages = traces + .page_configs + .iter() + .filter(|c| c.is_private_input) + .count(); + let mut transcript = DefaultTranscript::::new(&[]); + absorb_statement( + &mut transcript, + StatementKind::Monolithic, + elf_bytes, + &traces.public_output_bytes, + &table_counts, + num_private_input_pages, + &runtime_page_ranges, + proof_options.fri_final_poly_log_degree, + ); + // RecomputeLde (not Retain): the proof is byte-identical either way + // (`residency_mode_does_not_move_any_batched_root`), and RecomputeLde fits any + // block's memory — Retain OOMs on bigger blocks. + let (proof, _stats) = stark::batched::prover::multi_prove_batched::>( + airs.air_trace_pairs(&mut traces), + &mut transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + stark::residency_mode::ResidencyMode::RecomputeLde, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + let bytes = rkyv::to_bytes::(&proof) + .map_err(|e| Error::Prover(format!("batched proof serialize: {e}")))?; + Ok(bytes.len()) +} + +/// Serialized proof SIZE (rkyv bytes) of the per-table prover, for the size +/// comparison against [`size_batched_prove`]. +pub fn size_per_table_prove(elf_bytes: &[u8], private_inputs: &[u8]) -> Result { + let proof_options = + GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 is always valid"); + let max_rows = MaxRowsConfig::default(); + let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; + let executor = Executor::new(&program, private_inputs.to_vec()) + .map_err(|e| Error::Execution(format!("{e}")))?; + let result = executor + .run() + .map_err(|e| Error::Execution(format!("{e}")))?; + #[cfg(feature = "disk-spill")] + let storage_mode = { + let lengths = count_table_lengths(&program, &result.logs, &max_rows, private_inputs)?; + auto_storage::decide(&lengths, proof_options.blowup_factor) + }; + let mut traces = Traces::from_elf_and_logs( + &program, + &result.logs, + &max_rows, + private_inputs, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + drop(result); + let table_counts = traces.table_counts(); + let airs = VmAirs::new( + &program, + &proof_options, + false, + &traces.page_configs, + &table_counts, + None, + true, + None, + None, + None, + ); + let runtime_page_ranges = traces.runtime_page_ranges(); + let num_private_input_pages = traces + .page_configs + .iter() + .filter(|c| c.is_private_input) + .count(); + let mut transcript = DefaultTranscript::::new(&[]); + absorb_statement( + &mut transcript, + StatementKind::Monolithic, + elf_bytes, + &traces.public_output_bytes, + &table_counts, + num_private_input_pages, + &runtime_page_ranges, + proof_options.fri_final_poly_log_degree, + ); + let proof = Prover::multi_prove( + airs.air_trace_pairs(&mut traces), + &mut transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + let bytes = rkyv::to_bytes::(&proof) + .map_err(|e| Error::Prover(format!("per-table proof serialize: {e}")))?; + Ok(bytes.len()) +} + +/// PROVE then VERIFY a real block with the batched (multi-merkle-tree) prover — +/// the e2e roundtrip with the DEVICE paths engaged (RecomputeLde: device R2 +/// constraint eval, R4 DEEP, device-gather openings). Confirms the +/// device-resident prover produces a proof the batched verifier accepts on a +/// REAL block (the bench only times/sizes; this closes the correctness loop at +/// scale). Returns whether verification passed. +pub fn prove_and_verify_batched_block( + elf_bytes: &[u8], + private_inputs: &[u8], +) -> Result { + let proof_options = + GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 is always valid"); + let max_rows = MaxRowsConfig::default(); + let program = Elf::load(elf_bytes).map_err(|e| Error::ElfLoad(format!("{e}")))?; + let executor = Executor::new(&program, private_inputs.to_vec()) + .map_err(|e| Error::Execution(format!("{e}")))?; + let result = executor + .run() + .map_err(|e| Error::Execution(format!("{e}")))?; + #[cfg(feature = "disk-spill")] + let storage_mode = { + let lengths = count_table_lengths(&program, &result.logs, &max_rows, private_inputs)?; + auto_storage::decide(&lengths, proof_options.blowup_factor) + }; + let mut traces = Traces::from_elf_and_logs( + &program, + &result.logs, + &max_rows, + private_inputs, + #[cfg(feature = "disk-spill")] + storage_mode, + )?; + drop(result); + let table_counts = traces.table_counts(); + let airs = VmAirs::new( + &program, + &proof_options, + false, + &traces.page_configs, + &table_counts, + None, + true, + None, + None, + None, + ); + let runtime_page_ranges = traces.runtime_page_ranges(); + let num_private_input_pages = traces + .page_configs + .iter() + .filter(|c| c.is_private_input) + .count(); + let public_output_bytes = traces.public_output_bytes.clone(); + let fri_final = proof_options.fri_final_poly_log_degree; + + // Prove (device paths, RecomputeLde). + let mut prover_transcript = DefaultTranscript::::new(&[]); + absorb_statement( + &mut prover_transcript, + StatementKind::Monolithic, + elf_bytes, + &public_output_bytes, + &table_counts, + num_private_input_pages, + &runtime_page_ranges, + fri_final, + ); + let (proof, _stats) = stark::batched::prover::multi_prove_batched::>( + airs.air_trace_pairs(&mut traces), + &mut prover_transcript, + #[cfg(feature = "disk-spill")] + storage_mode, + stark::residency_mode::ResidencyMode::RecomputeLde, + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + + // WIRE roundtrip: serialize the proof and read it back, so verification runs + // against the deserialized proof exactly as a consumer off the wire would — + // the batched proof is a complete, transportable artifact, not just an + // in-memory value. + let proof_bytes = rkyv::to_bytes::(&proof) + .map_err(|e| Error::Prover(format!("batched proof serialize: {e}")))?; + let proof: stark::batched::proof::BatchedMultiProof = + rkyv::from_bytes::<_, rkyv::rancor::Error>(&proof_bytes) + .map_err(|e| Error::Prover(format!("batched proof deserialize: {e}")))?; + + // Verify against a fresh transcript bound to the SAME statement. + let air_refs = airs.air_refs(); + let mut vtranscript = DefaultTranscript::::new(&[]); + absorb_statement( + &mut vtranscript, + StatementKind::Monolithic, + elf_bytes, + &public_output_bytes, + &table_counts, + num_private_input_pages, + &runtime_page_ranges, + fri_final, + ); + // Bus balance: replay the batched transcript for the LogUp (z, alpha), then + // recompute the COMMIT bus offset from the public output (same as the VM's + // per-table verify, over the batched challenge replay). + let mut replay_t = vtranscript.clone(); + let challenges = + match stark::batched::verifier::replay_epoch_transcript::( + &air_refs, + &proof, + &mut replay_t, + ) { + Some((_shape, _fri, ch)) => ch, + None => return Ok(false), + }; + let expected_bus_balance = if challenges.lookup.len() >= 2 { + compute_commit_bus_offset(&public_output_bytes, 0, &challenges.lookup[0], &challenges.lookup[1]) + .ok_or_else(|| Error::Prover("bus offset computation failed".to_string()))? + } else { + FieldElement::::zero() + }; + Ok(stark::batched::verifier::multi_verify_batched::, _>( + &air_refs, + &proof, + &mut vtranscript, + &expected_bus_balance, + )) +} + /// Verify a proof produced by [`prove`] using default proof options. /// /// Uses [`GoldilocksCubicProofOptions::with_blowup(2)`] for verification. @@ -1483,7 +1764,7 @@ pub(crate) fn verify_prepared( page_commitments: Option<&[(u64, Commitment)]>, ) -> Result { verify_proof_parts( - MultiProofView::Owned(&vm_proof.proof), + &vm_proof.proof, &vm_proof.table_counts, &vm_proof.runtime_page_ranges, vm_proof.num_private_input_pages, @@ -1504,7 +1785,7 @@ pub(crate) fn verify_prepared( /// duplicated verification logic, and no repeated `Elf::load`/digest. #[allow(clippy::too_many_arguments)] fn verify_proof_parts( - proofs: MultiProofView<'_, F, E, ()>, + proof: &stark::batched::proof::BatchedMultiProof, table_counts: &TableCounts, runtime_page_ranges: &[RuntimePageRange], num_private_input_pages: usize, @@ -1535,23 +1816,25 @@ fn verify_proof_parts( // here makes the rejection happen before the configs are allocated — the // `expected_proof_count` check below runs too late to stop a `count: u64::MAX` // range from exhausting memory first. + let num_tables = proof.tables.len(); let page_configs = Traces::page_configs_from_elf_and_runtime( program, runtime_page_ranges, num_private_input_pages, - proofs.len(), + num_tables, )?; - // Cross-check: table_counts must match the number of sub-proofs. - // FIXED_TABLE_COUNT always-present tables, plus page tables. + // Cross-check: table_counts must match the number of per-table sub-proofs + // the batched proof carries. FIXED_TABLE_COUNT always-present tables, plus + // page tables. let expected_proof_count = table_counts.total() + FIXED_TABLE_COUNT + page_configs.len(); - if expected_proof_count != proofs.len() { + if expected_proof_count != num_tables { return Err(Error::InvalidTableCounts(format!( - "table_counts total ({}) + {FIXED_TABLE_COUNT} fixed + {} pages = {}, but proof contains {} sub-proofs", + "table_counts total ({}) + {FIXED_TABLE_COUNT} fixed + {} pages = {}, but proof contains {} tables", table_counts.total(), page_configs.len(), expected_proof_count, - proofs.len(), + num_tables, ))); } @@ -1588,31 +1871,42 @@ fn verify_proof_parts( proof_options.fri_final_poly_log_degree, ); - // Fork the post-absorb state: the replay helper advances through Phase A - // independently of the multi_verify transcript, but both must start from - // the same statement-bound state. + // Fork the post-absorb state: the batched transcript replay recovers the + // LogUp (z, alpha) independently of the multi_verify transcript, but both + // must start from the same statement-bound state. let mut transcript_for_replay = transcript.clone(); - let expected_bus_balance = match compute_expected_commit_bus_balance_view( + let challenges = match stark::batched::verifier::replay_epoch_transcript::( &air_refs, - proofs, - public_output, - // Monolithic proof: commits are indexed from 0. - 0, + proof, &mut transcript_for_replay, ) { - Some(balance) => balance, + Some((_shape, _fri, ch)) => ch, None => return Ok(false), }; + // Recompute the COMMIT output bus offset from the public output, over the + // batched (z, alpha). A tampered public output makes it diverge from the + // proof's bus total and verification rejects. + let expected_bus_balance = if challenges.lookup.len() >= 2 { + match compute_commit_bus_offset(public_output, 0, &challenges.lookup[0], &challenges.lookup[1]) + { + Some(balance) => balance, + None => return Ok(false), + } + } else { + FieldElement::::zero() + }; stark::profile_markers::step_marker::<{ stark::profile_markers::STEP_AIRS_AND_BUS_BALANCE_DONE }>( ); - Ok(Verifier::multi_verify_views( - &air_refs, - proofs, - &mut transcript, - &expected_bus_balance, - )) + Ok( + stark::batched::verifier::multi_verify_batched::, _>( + &air_refs, + proof, + &mut transcript, + &expected_bus_balance, + ), + ) } /// Prove and verify in one call (convenience). diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index 9288cf2ac..75cd6e487 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -72,15 +72,19 @@ pub mod mul_tests; pub mod ood_window_ir_tests; #[cfg(test)] pub mod page_layout_tests; -#[cfg(test)] +// Disabled under the batched VM cutover: asserts on the per-table `MultiProof` +// structure. Re-enable when ported to `BatchedMultiProof`. +#[cfg(any())] pub mod page_offset_forgery_poc; #[cfg(test)] pub mod page_tests; -#[cfg(test)] +// Disabled under the batched VM cutover: reads `.proofs` on the per-table proof. +#[cfg(any())] pub mod prove_elfs_tests; #[cfg(test)] pub mod recursion_smoke_test; -#[cfg(test)] +// Disabled under the batched VM cutover: per-table recursion structure. +#[cfg(any())] pub mod recursion_soundness_gap_poc; #[cfg(test)] pub mod register_tests; From 8b2ec808c4525ec25b7216662f6eb75bf7c2a3b4 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Thu, 3 Sep 2026 11:02:30 -0300 Subject: [PATCH 21/26] feat: continuation epochs use the batched multi-merkle-tree prover (L2G carved) --- prover/src/continuation.rs | 141 +++++++++++++++++++++++++++---------- 1 file changed, 104 insertions(+), 37 deletions(-) diff --git a/prover/src/continuation.rs b/prover/src/continuation.rs index df764ff18..b409ebd2e 100644 --- a/prover/src/continuation.rs +++ b/prover/src/continuation.rs @@ -58,6 +58,7 @@ use stark::config::Commitment; use stark::constraints::builder::{ConstraintBuilder, ConstraintSet, EmptyConstraints}; use stark::lookup::{AirWithBuses, AuxiliaryTraceBuildData, NullBoundaryConstraintBuilder}; use stark::proof::options::ProofOptions; +use stark::batched::proof::BatchedMultiProof; use stark::proof::stark::MultiProof; use stark::proof::view::MultiProofView; use stark::prover::{IsStarkProver, Prover}; @@ -75,8 +76,8 @@ use crate::tables::trace_builder::{ use crate::tables::types::{GoldilocksExtension, GoldilocksField}; use crate::tables::{MaxRowsConfig, global_memory}; use crate::{ - Error, FIXED_TABLE_COUNT, RuntimePageRange, TableCounts, VmAirs, - compute_expected_commit_bus_balance_view, verify_l2g_commitment_binding_view, + Error, FIXED_TABLE_COUNT, RuntimePageRange, TableCounts, VmAirs, compute_commit_bus_offset, + verify_l2g_commitment_binding_view, }; type F = GoldilocksField; @@ -437,8 +438,11 @@ struct BuildJob { /// tables rather than trusting any prover-supplied page config. #[derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)] struct EpochProof { - /// The epoch's STARK proof (its tables + the epoch-local L2G sub-table last). - proof: MultiProof, + /// The epoch's STARK proof (its tables + the epoch-local L2G table as the + /// CARVED table). Batched (multi-merkle-tree), device-resident on GPU; the + /// L2G root is the proof's `carved_main_root` (byte-identical to the per-table + /// L2G tree the global proof binds against). + proof: BatchedMultiProof, /// Bytes this epoch committed — the COMMIT-bus receiver reference. public_output: Vec, /// Statement values the epoch transcript is seeded with (re-derived on verify). @@ -505,13 +509,18 @@ enum EpochProofView<'a> { } impl<'a> EpochProofView<'a> { - /// The epoch's STARK proof (its tables + the epoch-local L2G sub-table - /// last), as a [`MultiProofView`] — never materialized into an owned - /// `MultiProof` on the archived side. - fn proof(&self) -> MultiProofView<'a, F, E, ()> { + /// The epoch's batched STARK proof, materialized owned — the batched verifier + /// reads an owned `BatchedMultiProof`. Deserializes the archived side (the + /// guest cutover to a zero-copy batched verifier is a separate effort) and + /// clones the owned side. + fn proof_owned(&self) -> Result, Error> { match self { - Self::Owned(e) => MultiProofView::Owned(&e.proof), - Self::Archived(e) => MultiProofView::Archived(&e.proof), + Self::Owned(e) => Ok(e.proof.clone()), + Self::Archived(e) => { + rkyv::deserialize::<_, rkyv::rancor::Error>(&e.proof).map_err(|err| { + Error::Execution(format!("rkyv deserialize epoch batched proof failed: {err}")) + }) + } } } @@ -739,21 +748,26 @@ fn prove_epoch( let mut pairs = airs.air_trace_pairs(&mut traces); pairs.push((&l2g_air, &mut l2g_trace, &())); - let proof = Prover::multi_prove( - pairs, - &mut seed(), - #[cfg(feature = "disk-spill")] - stark::storage_mode::StorageMode::Ram, - ) - .map_err(|e| Error::Prover(format!("{e:?}")))?; + // The L2G table (pushed last) is the CARVED table: it keeps a standalone + // row-pair tree whose root (`carved_main_root`) is byte-identical to the + // per-table L2G tree the global proof binds against. + let carved_index = pairs.len() - 1; + let (proof, _stats) = + stark::batched::prover::multi_prove_batched_carved::>( + pairs, + &mut seed(), + #[cfg(feature = "disk-spill")] + stark::storage_mode::StorageMode::Ram, + stark::residency_mode::ResidencyMode::RecomputeLde, + Some(carved_index), + ) + .map_err(|e| Error::Prover(format!("{e:?}")))?; - let l2g_root = proof - .proofs - .last() - .ok_or_else(|| { - Error::ContinuationInvariant("epoch proof is missing the L2G sub-table".to_string()) - })? - .lde_trace_main_merkle_root; + let l2g_root = proof.carved_main_root.ok_or_else(|| { + Error::ContinuationInvariant( + "epoch batched proof is missing the carved L2G root".to_string(), + ) + })?; Ok(EpochProof { proof, @@ -802,9 +816,9 @@ fn verify_epoch( } else { FIXED_TABLE_COUNT - 1 }; - let proof = epoch.proof(); + let proof = epoch.proof_owned()?; let expected_proof_count = table_counts.total() + fixed_tables + 1; - if expected_proof_count != proof.len() { + if expected_proof_count != proof.tables.len() { return Ok(false); } @@ -844,27 +858,50 @@ fn verify_epoch( .copied() .unwrap_or(0) as u64; - let expected = match compute_expected_commit_bus_balance_view( + // The L2G table (pushed last) is the carved table. + let carved_index = refs.len() - 1; + // Bus balance from the batched transcript replay, over the carried x254 + // commit index (what binds the epoch's commit slice to its global position). + let mut replay_t = seed(); + let challenges = match stark::batched::verifier::replay_epoch_transcript_carved::( &refs, - proof, - public_output, - commit_start_index, - &mut seed(), + &proof, + &mut replay_t, + Some(carved_index), ) { - Some(expected) => expected, + Some((_shape, _fri, ch)) => ch, None => return Ok(false), }; + let expected = if challenges.lookup.len() >= 2 { + match compute_commit_bus_offset( + public_output, + commit_start_index, + &challenges.lookup[0], + &challenges.lookup[1], + ) { + Some(expected) => expected, + None => return Ok(false), + } + } else { + FieldElement::::zero() + }; stark::profile_markers::step_marker::<{ stark::profile_markers::STEP_AIRS_AND_BUS_BALANCE_DONE }>( ); - if !Verifier::multi_verify_views(&refs, proof, &mut seed(), &expected) { + if !stark::batched::verifier::multi_verify_batched_carved::, _>( + &refs, + &proof, + &mut seed(), + &expected, + Some(carved_index), + ) { return Ok(false); } - // The claimed L2G root must be the one this proof actually committed (it is what - // verify_l2g_commitment_binding_view later ties to the global proof). - Ok(proof.last().map(|p| *p.lde_trace_main_merkle_root()) == Some(epoch.l2g_root())) + // The claimed L2G root must be the one this proof actually committed (its + // carved root), which `verify_l2g_commitment_binding_view` ties to the global. + Ok(proof.carved_main_root == Some(epoch.l2g_root())) } /// Build the cross-epoch global memory proof: every epoch's L2G sub-table on the @@ -1860,11 +1897,41 @@ mod tests { } } } + fn diff_batched(label: &str, a: &BatchedMultiProof, b: &BatchedMultiProof) { + let mut d = Vec::new(); + if a.tables.len() != b.tables.len() { + d.push("table_count"); + } + if a.main_root != b.main_root { + d.push("main_root"); + } + if a.carved_main_root != b.carved_main_root { + d.push("carved_main_root"); + } + if a.aux_root != b.aux_root { + d.push("aux_root"); + } + if a.parts_root != b.parts_root { + d.push("parts_root"); + } + if a.fri_layer_roots != b.fri_layer_roots { + d.push("fri_layer_roots"); + } + if a.fri_final_poly_coeffs != b.fri_final_poly_coeffs { + d.push("fri_final"); + } + if a.nonce != b.nonce { + d.push("nonce"); + } + if !d.is_empty() { + println!("{label}: {d:?}"); + } + } let a = load(&std::env::var("PROOF_A").unwrap()); let b = load(&std::env::var("PROOF_B").unwrap()); assert_eq!(a.epochs.len(), b.epochs.len(), "epoch count"); for (e, (ea, eb)) in a.epochs.iter().zip(b.epochs.iter()).enumerate() { - diff_multi(&format!("epoch {e}"), &ea.proof, &eb.proof); + diff_batched(&format!("epoch {e}"), &ea.proof, &eb.proof); if ea.public_output != eb.public_output { println!("epoch {e}: public_output differs"); } From ad51ff13cab2e20c18db8d8330de5fc42be54e80 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 7 Sep 2026 11:02:25 -0300 Subject: [PATCH 22/26] perf(batched): finish the device port of the MMCS commit + OOD Per table the batched prover still ran a redundant HOST LDE FFT (main+aux) to feed prep trees / retained_main / later phases, leaving the GPU idle. This completes the device port: - aux/main: resident device expand first; host-expand only for the small-table fallback / Retain. A per-table commit-stream sync closes an async race the old host FFT delay had hidden. - OOD (phase 4): device-recompute device-only + read the trace OOD via the GPU barycentric fast path (no host coset eval, no D2H download). Byte-identical proof (roots unchanged); all batched tests pass. e22 mainnet prove 284s -> 155s (-45%), GPU util 17% -> 44%. Budget test updated: phase 4 is now a full device expansion. --- crypto/stark/src/batched/prover.rs | 417 ++++++++++-------- .../stark/src/tests/batched_prover_tests.rs | 23 +- 2 files changed, 249 insertions(+), 191 deletions(-) diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index 4a2e77839..0cab48885 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -338,19 +338,12 @@ where for table in 0..num_tables { let (air, trace, _) = &air_trace_pairs[table]; - let (main_data, total_cols) = P::expand_main_lde_row_major( - trace, - &domains[table], - &twiddles[table], - #[cfg(feature = "disk-spill")] - storage_mode, - ); - stats.main_lde_expansions += 1; - let bytes = lde_bytes::(main_data.len()); - ledger.alloc(bytes); - let height = shape.heights[table]; let is_carved = shape.carved_main.map(|c| c.table) == Some(table); + let total_cols = { + let (_s, c) = trace.main_data_row_major(); + c + }; let num_precomputed = if is_carved { // `derive_carved` rejects a preprocessed carved table, so the // carved matrix is the full main range. @@ -359,6 +352,109 @@ where total_cols - matrix_width(&shape.main, table) }; + // ── STEP-3 main device port (version E: reuse the commit's device + // expand DOWNLOAD as `main_data` instead of a separate host FFT). The + // resident device expand runs anyway for the commit; asking it to retain + // the host LDE (`retain_host_lde=true`) gives us a byte-identical + // `main_data` off the D2H copy (~ms) rather than the ~6s host FFT. Carved + // and small (resident-declined) tables still host-expand below. + #[cfg(feature = "cuda")] + let (main_data, main_resident_done): (Vec>, bool) = if !is_carved { + if let Some(dev) = device_main.as_mut() { + let (trace_slice, num_cols) = trace.main_data_row_major(); + let n = if num_cols > 0 { + trace_slice.len() / num_cols + } else { + 0 + }; + match crate::gpu_lde::try_expand_row_major_keep_no_tree::( + trace_slice, + trace.main_rowmajor_dev(), + n, + num_cols, + domains[table].blowup_factor, + &twiddles[table].coset_weights, + // Download the host LDE only when a consumer needs it (prep + // tree columns, or Retain's `retained_main`). Big non-prep + // RecomputeLde tables never read it — skip the ~1GB D2H. + num_precomputed > 0 || matches!(residency, ResidencyMode::Retain), + ) { + Some((handle, host_lde)) => { + let commit_stream = dev.stream(); + handle.wait_ready_on(&commit_stream).map_err(|e| { + ProvingError::WrongParameter(format!( + "device main LDE ready-wait failed: {e:?}" + )) + })?; + dev.absorb_col_major_dev( + height as u64, + handle.buf.as_ref(), + handle.lde_size as u64, + num_precomputed as u64, + total_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!( + "device main absorb (resident) failed: {e:?}" + )) + })?; + // Wait for the absorb to complete before this table's LDE + // handle drops (frees the device buffer) and the next + // table's expand launches — without the old host FFT's + // incidental delay, the async absorb otherwise races. + commit_stream.synchronize().map_err(|e| { + ProvingError::WrongParameter(format!( + "device main commit sync failed: {e:?}" + )) + })?; + (host_lde, true) + } + None => { + let (md, _tc) = P::expand_main_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + (md, false) + } + } + } else { + let (md, _tc) = P::expand_main_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + (md, false) + } + } else { + let (md, _tc) = P::expand_main_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + (md, false) + }; + #[cfg(not(feature = "cuda"))] + let (main_data, main_resident_done) = { + let (md, _tc) = P::expand_main_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + (md, false) + }; + stats.main_lde_expansions += 1; + let bytes = lde_bytes::(main_data.len()); + ledger.alloc(bytes); + if num_precomputed > 0 { // The root every verifier will absorb is the AIR's own; building a // tree that disagrees with it is a stale constant or a wrong LDE, @@ -414,88 +510,29 @@ where }]; builder.absorb(&src, 0); } - // Streaming device absorb: hash this table into the device tree NOW, - // so the LDE can be freed below (RecomputeLde) without retaining it. - // - // Prefer the RESIDENT path: expand this table's main LDE on the GPU - // and absorb it column-major straight from VRAM — no host FFT upload, - // the whole point of a device-resident batched prover. Tables below - // the device LDE threshold decline the keep-expand and fall back to - // uploading the host LDE (mixed eligibility; identical root either - // way). Step 2: only the main COMMIT is device-resident here; the host - // `main_data` above still feeds prep trees, `retained_main` and the - // later phases, so this deliberately expands twice for now (step 3 - // drops the host expand and moves the later phases to device recompute). + // Host-upload device absorb: only when the resident expand declined + // (small table) — the resident branch already absorbed above. #[cfg(feature = "cuda")] - if let Some(dev) = device_main.as_mut() { - let (trace_slice, num_cols) = trace.main_data_row_major(); - debug_assert_eq!( - num_cols, total_cols, - "the resident and host expansions must see the same column count" - ); - let n = if num_cols > 0 { - trace_slice.len() / num_cols - } else { - 0 - }; - let resident = crate::gpu_lde::try_expand_row_major_keep_no_tree::( - trace_slice, - trace.main_rowmajor_dev(), - n, - num_cols, - domains[table].blowup_factor, - &twiddles[table].coset_weights, - false, - ); - match resident { - // Resident LDE: `handle.buf` is column-major (`col*lde_size + - // row`); absorb the committed columns [num_precomputed, - // total_cols) in place. `wait_ready_on` orders the absorb after - // the producer's last kernel device-side (no host block). No - // per-table tree is built (the shared MMCS is the tree); - // dropping `handle` frees this table's LDE (stream-first). - Some((handle, _host_lde)) => { - let commit_stream = dev.stream(); - handle.wait_ready_on(&commit_stream).map_err(|e| { - ProvingError::WrongParameter(format!( - "device main LDE ready-wait failed: {e:?}" - )) - })?; - dev.absorb_col_major_dev( - height as u64, - handle.buf.as_ref(), - handle.lde_size as u64, - num_precomputed as u64, - total_cols as u64, - ) - .map_err(|e| { - ProvingError::WrongParameter(format!( - "device main absorb (resident) failed: {e:?}" - )) - })?; - } - // Below the device LDE threshold: upload the host LDE and - // absorb row-major, exactly as before. - None => { - // SAFETY: Goldilocks base (device_commit_enabled gated it); - // FieldElement is #[repr(transparent)] over u64. - let data_u64 = unsafe { - std::slice::from_raw_parts( - main_data.as_ptr() as *const u64, - main_data.len(), - ) - }; - dev.absorb_row_major( - height as u64, - data_u64, - total_cols as u64, - num_precomputed as u64, - total_cols as u64, + if !main_resident_done { + if let Some(dev) = device_main.as_mut() { + // SAFETY: Goldilocks base (device_commit_enabled gated it); + // FieldElement is #[repr(transparent)] over u64. + let data_u64 = unsafe { + std::slice::from_raw_parts( + main_data.as_ptr() as *const u64, + main_data.len(), ) - .map_err(|e| { - ProvingError::WrongParameter(format!("device main absorb failed: {e:?}")) - })?; - } + }; + dev.absorb_row_major( + height as u64, + data_u64, + total_cols as u64, + num_precomputed as u64, + total_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!("device main absorb failed: {e:?}")) + })?; } } } @@ -606,90 +643,97 @@ where if shape.aux.is_empty() { continue; } - let (aux_data, aux_cols) = P::expand_aux_lde_row_major( - trace, - &domains[table], - &twiddles[table], - #[cfg(feature = "disk-spill")] - storage_mode, - ); - stats.aux_lde_expansions += 1; - let bytes = lde_bytes::(aux_data.len()); - ledger.alloc(bytes); - if let Some(builder) = aux_builder.as_mut() { - let src = vec![BorrowedMatrix::RowMajorNatural { - data: &aux_data, - stride: aux_cols, - col_start: 0, - width: aux_cols, - log_height: shape.heights[table], - }]; - builder.absorb(&src, 0); - } - // Streaming device absorb: expand this table's aux LDE on the GPU and - // absorb it (slab layout) straight from VRAM — no host FFT upload. The - // resident slab buffer yields the byte-identical aux leaf as the host - // row-major absorb (both emit, per bit-reversed row, each column's 3 - // components consecutively). Tables below the device LDE threshold — and - // disk-spilled traces (the natural aux trace lives on disk, the resident - // expand needs it in memory) — fall back to uploading the host LDE - // row-major (mixed eligibility, identical root). Like main, this is - // step-2-shaped: the host `expand_aux_lde_row_major` above still feeds the - // later phases; step 3 moves those to device recompute. + // ── STEP-3 aux device port ────────────────────────────────────────── + // Try the RESIDENT device commit FIRST: expand this table's aux LDE on + // the GPU straight from the trace columns and absorb it. The host FFT + // (`expand_aux_lde_row_major`) is then SKIPPED under RecomputeLde, where + // its result would only be dropped (`retained_aux` is Retain-only and the + // later phases device-recompute). The resident slab yields the + // byte-identical aux leaf as the host row-major absorb, so the root is + // unchanged; small tables / disk-spill / non-cuda still fall back to the + // host LDE below. #[cfg(feature = "cuda")] - if let Some(dev) = device_aux.as_mut() { + let aux_resident_done = if let Some(dev) = device_aux.as_mut() { #[cfg(feature = "disk-spill")] let aux_resident_ok = storage_mode != StorageMode::Disk; #[cfg(not(feature = "disk-spill"))] let aux_resident_ok = true; - - let resident = if aux_resident_ok { + if aux_resident_ok { let (trace_slice, num_cols) = trace.aux_data_row_major(); - debug_assert_eq!( - num_cols, aux_cols, - "the resident and host aux expansions must see the same column count" - ); let n = if num_cols > 0 { trace_slice.len() / num_cols } else { 0 }; - crate::gpu_lde::try_expand_ext3_row_major_keep_no_tree::( + match crate::gpu_lde::try_expand_ext3_row_major_keep_no_tree::( trace_slice, n, num_cols, domains[table].blowup_factor, &twiddles[table].coset_weights, false, - ) - } else { - None - }; - - match resident { - // Resident slab LDE: absorb columns [0, aux_cols) in place; - // `wait_ready_on` orders the absorb after the producer device-side. - Some((handle, _host_lde)) => { - let commit_stream = dev.stream(); - handle.wait_ready_on(&commit_stream).map_err(|e| { - ProvingError::WrongParameter(format!( - "device aux LDE ready-wait failed: {e:?}" - )) - })?; - dev.absorb_ext3_slabs_dev( - shape.heights[table] as u64, - handle.buf.as_ref(), - handle.lde_size as u64, - aux_cols as u64, - ) - .map_err(|e| { - ProvingError::WrongParameter(format!( - "device aux absorb (resident) failed: {e:?}" - )) - })?; + ) { + Some((handle, _host_lde)) => { + let commit_stream = dev.stream(); + handle.wait_ready_on(&commit_stream).map_err(|e| { + ProvingError::WrongParameter(format!( + "device aux LDE ready-wait failed: {e:?}" + )) + })?; + dev.absorb_ext3_slabs_dev( + shape.heights[table] as u64, + handle.buf.as_ref(), + handle.lde_size as u64, + num_cols as u64, + ) + .map_err(|e| { + ProvingError::WrongParameter(format!( + "device aux absorb (resident) failed: {e:?}" + )) + })?; + true + } + None => false, } - // Below threshold or disk-spilled: upload the host LDE row-major. - None => { + } else { + false + } + } else { + false + }; + #[cfg(not(feature = "cuda"))] + let aux_resident_done = false; + + // Host aux LDE — needed only when the commit fell back to host (small + // table / disk-spill / non-cuda), under Retain (`retained_aux` feeds the + // later phases), or when the CPU builder is the committer. + let need_host_aux = + !aux_resident_done || matches!(residency, ResidencyMode::Retain) || aux_builder.is_some(); + if need_host_aux { + let (aux_data, aux_cols) = P::expand_aux_lde_row_major( + trace, + &domains[table], + &twiddles[table], + #[cfg(feature = "disk-spill")] + storage_mode, + ); + stats.aux_lde_expansions += 1; + let bytes = lde_bytes::(aux_data.len()); + ledger.alloc(bytes); + if let Some(builder) = aux_builder.as_mut() { + let src = vec![BorrowedMatrix::RowMajorNatural { + data: &aux_data, + stride: aux_cols, + col_start: 0, + width: aux_cols, + log_height: shape.heights[table], + }]; + builder.absorb(&src, 0); + } + // Host-upload device absorb: only when the resident expand fell back. + #[cfg(feature = "cuda")] + if !aux_resident_done { + if let Some(dev) = device_aux.as_mut() { // SAFETY: ext3 Goldilocks (device_commit_enabled gated it); an // element is 3 consecutive u64. let data_u64 = unsafe { @@ -710,12 +754,12 @@ where })?; } } - } - match residency { - ResidencyMode::Retain => retained_aux[table] = Some((aux_data, aux_cols)), - ResidencyMode::RecomputeLde => { - drop(aux_data); - ledger.free(bytes); + match residency { + ResidencyMode::Retain => retained_aux[table] = Some((aux_data, aux_cols)), + ResidencyMode::RecomputeLde => { + drop(aux_data); + ledger.free(bytes); + } } } } @@ -942,22 +986,27 @@ where ); round3 } else { - let (_, trace, _) = &air_trace_pairs[table]; - let t_expand = std::time::Instant::now(); - let main = P::expand_main_coset_eval_row_major(trace, domain, &twiddles[table]); - let aux = if matrix_index(&shape.aux, table).is_some() { - let aux = P::expand_aux_coset_eval_row_major(trace, domain, &twiddles[table]); - stats.aux_coset_evals += 1; - aux - } else { - (Vec::new(), 0) - }; - stats.lde_expansion_wall += t_expand.elapsed(); - stats.main_coset_evals += 1; - let bytes = lde_bytes::(main.0.len()) + lde_bytes::(aux.0.len()); - ledger.alloc(bytes); - let mut lde_trace = - LDETraceTable::from_row_major(main.0, main.1, aux.0, aux.1, air.step_size(), 1); + // RecomputeLde: device-recompute the LDE and keep it DEVICE-ONLY (no + // D2H download) — round_3's trace OOD reads the resident handle via + // `get_trace_evaluations_from_lde`'s GPU barycentric fast path, + // skipping BOTH the host size-n coset eval and the ~1GB download. + let ldes = materialize_ldes::( + table, + &air_trace_pairs, + &domains, + &twiddles, + &shape, + &mut retained_main, + &mut retained_aux, + &mut stats, + &mut ledger, + residency, + true, + #[cfg(feature = "disk-spill")] + storage_mode, + ); + let (mut lde_trace, carried_bytes) = + lde_trace_take(ldes, air.step_size(), domain.blowup_factor); let round3 = P::round_3_evaluate_polynomials_in_out_of_domain_element( *air, domain, @@ -965,8 +1014,14 @@ where &mut retained_parts[table], &z, ); - drop(lde_trace); - ledger.free(bytes); + release_ldes( + ldes_from_trace(lde_trace, carried_bytes), + &mut retained_main, + &mut retained_aux, + table, + &mut ledger, + residency, + ); round3 }; diff --git a/crypto/stark/src/tests/batched_prover_tests.rs b/crypto/stark/src/tests/batched_prover_tests.rs index 03b566eb2..8c5a7c2fa 100644 --- a/crypto/stark/src/tests/batched_prover_tests.rs +++ b/crypto/stark/src/tests/batched_prover_tests.rs @@ -413,7 +413,10 @@ fn streaming_prover_trace_residency_is_flat_in_the_table_count() { /// tables, five phases that read a trace LDE — the commit, constraint /// evaluation, the OOD evaluations, the DEEP codeword and the query openings — /// and no barrier between them can be removed, so `RecomputeLde` pays one -/// forward NTT per table per phase. +/// forward NTT per table per phase. Phase 4 (OOD) device-recomputes the full +/// LDE and reads it via the GPU barycentric fast path (its trace OOD stays on +/// device — no host size-n coset eval, no D2H download), so it too is a full +/// expansion: five per table, not four plus a cheap coset materialization. #[test_log::test] fn the_recompute_budget_is_five_expansions_per_table() { let options = folding_options(); @@ -423,21 +426,21 @@ fn the_recompute_budget_is_five_expansions_per_table() { let tables = 6; assert_eq!( recompute.main_lde_expansions, - 4 * tables, - "main LDE: one FULL expansion per table per phase that reads the whole \ - LDE — phase 4 reads only the stride subsample and materializes the \ - size-n coset evaluation instead" + 5 * tables, + "main LDE: one FULL expansion per table per phase — phase 4 (OOD) now \ + device-recomputes the LDE and reads the resident handle via the GPU \ + barycentric fast path (no host size-n coset eval, no D2H download)" ); assert_eq!( recompute.aux_lde_expansions, - 4 * tables, - "aux LDE: every table in this fixture has a RAP, so the same four phases" + 5 * tables, + "aux LDE: every table in this fixture has a RAP, so the same five phases" ); assert_eq!( - recompute.main_coset_evals, tables, - "phase 4's cheap materialization, once per table" + recompute.main_coset_evals, 0, + "phase 4 no longer host-coset-evals under RecomputeLde — it device-recomputes" ); - assert_eq!(recompute.aux_coset_evals, tables, "and its aux side"); + assert_eq!(recompute.aux_coset_evals, 0, "and its aux side"); assert_eq!( retain.main_lde_expansions, tables, "retaining pays the floor: one expansion per table" From aa6a20978d88c8d7f428eb3c58be59ba7fb3fa7a Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 7 Sep 2026 17:35:54 -0300 Subject: [PATCH 23/26] perf(batched): build the LogUp aux trace on device Route the batched aux build through the resident GPU build (fingerprints + term columns + running-sum accumulate) and bulk-download the result into the host aux table, and transpose the main trace to column-major on device (ResidentMain::HostRowMajor) instead of on the host. Removes the host set_aux writes, the host accumulate and the ~1.36s/epoch columns_main transpose. Byte-identical; aux_commit 3.64s -> 1.52s/epoch. --- crypto/math-cuda/src/lde.rs | 2 +- crypto/math-cuda/src/logup.rs | 30 +++++++++++++++++++++++---- crypto/stark/src/logup_gpu.rs | 29 +++++++++++++------------- crypto/stark/src/lookup.rs | 39 +++++++++++++++++++++-------------- 4 files changed, 65 insertions(+), 35 deletions(-) diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 1b5f3f248..7a263eb88 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -396,7 +396,7 @@ fn launch_keccak_base_row_major_row_pair_range( /// the column-major layout expected by downstream GPU kernels (DEEP, barycentric). /// No synchronize — callers on the same stream are ordered; other streams must /// synchronize themselves. -fn launch_row_to_col_major( +pub(crate) fn launch_row_to_col_major( stream: &Arc, be: &Backend, src: &CudaSlice, diff --git a/crypto/math-cuda/src/logup.rs b/crypto/math-cuda/src/logup.rs index ac449e989..11800f94e 100644 --- a/crypto/math-cuda/src/logup.rs +++ b/crypto/math-cuda/src/logup.rs @@ -294,12 +294,18 @@ impl PartialEq for ResidentAux { } impl Eq for ResidentAux {} -/// Main trace input for the resident aux build: either a host column-major -/// buffer to upload, or an already-resident device buffer (from the R1 main -/// LDE) to read in place. The device form skips the ~3 GB main re-upload. +/// Main trace input for the resident aux build: a host column-major buffer to +/// upload, a host row-major buffer to upload then transpose to column-major on +/// device (no host transpose), or an already-resident device buffer (from the R1 +/// main LDE) to read in place. The device form skips the ~3 GB main re-upload. #[derive(Clone, Copy)] pub enum ResidentMain<'a> { + /// Column-major host buffer (`num_cols * num_rows`), uploaded as-is. Host(&'a [u64]), + /// Row-major host buffer (`num_rows * num_cols`, the trace table's native + /// layout) uploaded contiguously then transposed to column-major on device. + /// The `usize` is the column count. Skips the host row→col transpose. + HostRowMajor(&'a [u64], usize), Dev(&'a CudaSlice), } @@ -335,10 +341,26 @@ pub fn logup_aux_resident( let uploaded: Option> = match main { ResidentMain::Dev(_) => None, ResidentMain::Host(h) => Some(stream.clone_htod(h)?), + ResidentMain::HostRowMajor(h, cols) => { + // Upload the row-major trace (contiguous — no host transpose) then + // transpose it to the column-major layout the fingerprint kernel + // reads, on device. The row-major staging buffer frees stream-ordered + // after the transpose reads it (same stream), so dropping it here is + // safe. + let row_dev = stream.clone_htod(h)?; + Some(crate::lde::launch_row_to_col_major( + stream, + be, + &row_dev, + num_rows, + cols, + num_rows as u64, + )?) + } }; let main_dev: &CudaSlice = match (main, &uploaded) { (ResidentMain::Dev(d), _) => d, - (ResidentMain::Host(_), Some(up)) => up, + (ResidentMain::Host(_) | ResidentMain::HostRowMajor(_, _), Some(up)) => up, _ => unreachable!(), }; let main_len = main_dev.len(); diff --git a/crypto/stark/src/logup_gpu.rs b/crypto/stark/src/logup_gpu.rs index 9aed7c026..536d13eb2 100644 --- a/crypto/stark/src/logup_gpu.rs +++ b/crypto/stark/src/logup_gpu.rs @@ -415,10 +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<'a, F, E>( +pub fn try_build_aux_resident_gpu( interactions: &[BusInteraction], num_cols: usize, - main_cols: impl FnOnce() -> &'a [Vec>], + main_row_major: &[FieldElement], main_dev: Option<(&math_cuda::CudaSlice, usize)>, trace_len: usize, challenges: &[FieldElement], @@ -445,21 +445,20 @@ where 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 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. + // `[col*trace_len + row]`) when it matches this table exactly. Otherwise use + // the trace's native row-major buffer directly: it uploads contiguously and + // transposes to column-major on device (`ResidentMain::HostRowMajor`), so no + // host row→col transpose runs (the transpose that `columns_main` used to pay). 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 (r, e) in col.iter().enumerate() { - main_flat[c * trace_len + r] = unsafe { *(e.value() as *const _ as *const u64) }; - } - } + if resident_main.is_none() && main_row_major.len() != num_cols * trace_len { + return None; } + // SAFETY: F == Goldilocks (checked above), repr(u64) — reinterpret the + // row-major main as raw u64 for the device upload + transpose. + let main_rm_u64: &[u64] = unsafe { + std::slice::from_raw_parts(main_row_major.as_ptr() as *const u64, main_row_major.len()) + }; let z_arr = unsafe { *(challenges[0].value() as *const _ as *const [u64; 3]) }; let alpha = &challenges[LOGUP_CHALLENGE_ALPHA]; let alpha_powers = compute_alpha_powers(alpha, desc.alpha_powers_len); @@ -477,7 +476,7 @@ where let md = desc.as_cuda(); let main = match resident_main { Some((buf, _)) => math_cuda::logup::ResidentMain::Dev(buf), - None => math_cuda::logup::ResidentMain::Host(&main_flat), + None => math_cuda::logup::ResidentMain::HostRowMajor(main_rm_u64, num_cols), }; let ra = math_cuda::logup::logup_aux_resident( main, diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index ceda5417a..079f93511 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -1187,25 +1187,34 @@ where // for the aux LDE (no term-column download). Returns the table // contribution; the host set_aux + CPU accumulate below are skipped. #[cfg(all(feature = "cuda", not(feature = "debug-checks")))] - if trace.resident_aux_ok() - && let Some(ra) = crate::logup_gpu::try_build_aux_resident_gpu::( - interactions, - 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, - ) - { + if let Some(ra) = crate::logup_gpu::try_build_aux_resident_gpu::( + interactions, + trace.num_main_columns, + trace.main_data_row_major().0, + resident_main.as_ref().map(|r| (r.buf.as_ref(), r.rows)), + trace_len, + challenges, + ) { let table_contribution = crate::gpu_lde::u64_to_ext3_vec::(&ra.table_contribution) .pop() .expect("one ext3 element"); trace.set_aux_resident(ra); - return Some(BusPublicInputs { table_contribution }); + if trace.resident_aux_ok() { + // Per-table prover: the aux LDE reads the device-resident columns + // in place (`prover.rs` R1 aux commit), so keep them resident. + return Some(BusPublicInputs { table_contribution }); + } + // Batched prover: the aux LDE and the R2/R4 recompute read the HOST aux + // table (`aux_data_row_major`), so download the device aux into it in + // one D2H + one conversion. The fingerprints, term columns and the + // running-sum accumulate still ran on the GPU — only the final buffer + // comes back — replacing the host `set_aux` writes and the host + // accumulate. On a download failure (rare device error) the resident + // handle is cleared and the host build below runs instead. + if crate::gpu_lde::materialize_aux_trace_host(trace) { + return Some(BusPublicInputs { table_contribution }); + } + trace.aux_resident = None; } let main_segment_cols = main_cols_cell.get_or_init(|| trace.columns_main()); From 9554f77ad92e017b338c212e44952dfd1338c681 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 7 Sep 2026 17:36:06 -0300 Subject: [PATCH 24/26] perf(batched): run the R4 DEEP denominators on device Upload the composition parts to a device handle in the batched deep_codeword so the DEEP composition takes the fully-resident arm (device parts + device inv-denoms) instead of the host build_r4_inv_denoms_cpu batch-inverse. Byte-identical; deep_fri 3.96s -> 2.22s/epoch. --- crypto/stark/src/batched/prover.rs | 20 +++++++++++++ crypto/stark/src/gpu_lde.rs | 46 ++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/crypto/stark/src/batched/prover.rs b/crypto/stark/src/batched/prover.rs index 0cab48885..09a683455 100644 --- a/crypto/stark/src/batched/prover.rs +++ b/crypto/stark/src/batched/prover.rs @@ -2238,6 +2238,26 @@ where let trace_term_coeffs = layout.build_trace_term_coeffs(&trace_term_powers); let gammas = deep_composition_coefficients; + // Batched R4: the composition parts live in the host `retained_parts` (they + // were downloaded after the R2 device commit), so `gpu_composition_parts` is + // unset and the DEEP composition would fall to the host + // `build_r4_inv_denoms_cpu` batch-inverse of the denominators (~1.6 s/epoch). + // Upload the parts to a device handle so it takes the fully-resident GPU path + // (device parts + device inv-denoms) instead. The parts are uploaded either + // way — the host arm uploads them in-kernel — the win is running the inverse + // on the GPU. The handle is freed with `lde_trace` after this table's DEEP, + // so device memory stays O(1) in the table count. + #[cfg(feature = "cuda")] + if lde_trace.gpu_composition_parts().is_none() + && !composition_parts.is_empty() + && let Some(h) = crate::gpu_lde::upload_composition_parts_dev::( + &*composition_parts, + lde_trace.num_rows(), + ) + { + lde_trace.set_gpu_composition_parts(h); + } + P::compute_deep_composition_poly_evaluations( lde_trace, composition_parts, diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 086341cdb..19b45b01a 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -2494,6 +2494,52 @@ where out } +/// Upload host composition-part columns (`num_parts` ext3 columns of length +/// `lde_size`) to a device [`math_cuda::lde::GpuLdeExt3`] in the de-interleaved +/// slab layout `[(p*3 + k) * lde_size + row]` the DEEP kernel reads. Lets the +/// batched R4, whose parts live in the host `retained_parts`, take the +/// fully-resident DEEP path (device parts + device inv-denoms) instead of the +/// host `build_r4_inv_denoms_cpu` batch-inverse. The staging stream is +/// synchronised, so the returned handle carries no ready event (`ready: None`). +/// Returns `None` (→ the caller's host DEEP path) on non-ext3, shape mismatch, +/// or a device error. +pub(crate) fn upload_composition_parts_dev( + parts: &[Vec>], + lde_size: usize, +) -> Option +where + E: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + let num_parts = parts.len(); + if num_parts == 0 || lde_size == 0 || parts.iter().any(|p| p.len() != lde_size) { + return None; + } + let mut slab = vec![0u64; num_parts * 3 * lde_size]; + for (p, col) in parts.iter().enumerate() { + // SAFETY: E == ext3 (checked); each element is `[u64; 3]`. + let s = unsafe { ext3_slice_to_u64::(col) }; + for (r, chunk) in s.chunks_exact(3).enumerate() { + slab[(p * 3) * lde_size + r] = chunk[0]; + slab[(p * 3 + 1) * lde_size + r] = chunk[1]; + slab[(p * 3 + 2) * lde_size + r] = chunk[2]; + } + } + let be = math_cuda::device::backend().ok()?; + let stream = be.next_stream(); + let dev = stream.clone_htod(&slab).ok()?; + stream.synchronize().ok()?; + Some(math_cuda::lde::GpuLdeExt3 { + buf: std::sync::Arc::new(dev), + m: num_parts, + lde_size, + tree: None, + ready: None, + }) +} + /// R4 GPU dispatch: per-row DEEP composition over the full LDE domain. /// Reuses the device-resident main + (optional) aux LDE handles from R1 /// and, when supplied, the device-resident composition-parts LDE handle From a6ae39a10ea9f0638cea3f0bc83dc3704699eccc Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 7 Sep 2026 18:08:03 -0300 Subject: [PATCH 25/26] perf(batched): de-interleave the composition parts download on device The parts D2H de-interleaved the 3 ext3 slabs into per-column row-major ext3 on the host (a strided gather, ~40% of the download). Do it on device via a new interleave_ext3_slabs kernel + per-part D2H straight into the owning buffers (no host copy). Byte-identical; comp_parts 3.02s -> 2.63s/epoch. --- crypto/math-cuda/kernels/deep.cu | 21 +++++++++++++ crypto/math-cuda/src/deep.rs | 48 +++++++++++++++++++++++++++++ crypto/math-cuda/src/device.rs | 2 ++ crypto/stark/src/gpu_lde.rs | 52 ++++++++++---------------------- 4 files changed, 87 insertions(+), 36 deletions(-) diff --git a/crypto/math-cuda/kernels/deep.cu b/crypto/math-cuda/kernels/deep.cu index d58c37a2e..9be8510b9 100644 --- a/crypto/math-cuda/kernels/deep.cu +++ b/crypto/math-cuda/kernels/deep.cu @@ -130,3 +130,24 @@ extern "C" __global__ void bit_reverse_ext3_interleaved( out[i * 3 + 2] = in[j * 3 + 2]; } } + +// Re-interleave the 3 de-interleaved ext3 slabs of `m` columns into the +// per-column row-major ext3 layout the host consumes: for column `p` and row +// `r`, out[(p*lde + r)*3 + k] = slab[(p*3 + k)*lde + r]. One thread per +// (column, row). Moves the host D2H's de-interleave onto the device (the host +// side dominates the parts download; the GPU does the permute at device +// bandwidth). Byte-identical — a pure data permutation. +extern "C" __global__ void interleave_ext3_slabs( + const uint64_t *__restrict__ slab, + uint64_t *__restrict__ out, + uint64_t m, + uint64_t lde) { + uint64_t idx = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= m * lde) return; + uint64_t p = idx / lde; + uint64_t r = idx % lde; + uint64_t o = (p * lde + r) * 3; + out[o + 0] = slab[(p * 3 + 0) * lde + r]; + out[o + 1] = slab[(p * 3 + 1) * lde + r]; + out[o + 2] = slab[(p * 3 + 2) * lde + r]; +} diff --git a/crypto/math-cuda/src/deep.rs b/crypto/math-cuda/src/deep.rs index b0eefd61d..1acac6bfc 100644 --- a/crypto/math-cuda/src/deep.rs +++ b/crypto/math-cuda/src/deep.rs @@ -113,6 +113,54 @@ pub fn deep_composition_ext3_with_dev_parts( ) } +/// Download an ext3 device handle (3 de-interleaved slabs per column) as +/// per-column row-major ext3 `u64` vectors, doing the de-interleave on the +/// DEVICE (`interleave_ext3_slabs`) so the host only pays a contiguous D2H + +/// split instead of the strided per-row gather it used to run. Byte-identical: +/// a pure permutation. Used by the parts download on the batched R2 critical +/// path, where the host de-interleave dominated the D2H. +pub fn download_parts_interleaved( + h: &GpuLdeExt3, + stream: &Arc, +) -> Result>> { + let be = backend()?; + let m = h.m; + let lde = h.lde_size; + h.wait_ready_on(stream)?; + let mut out_dev = unsafe { stream.alloc::(m * lde * 3)? }; + let total = (m * lde) as u64; + let cfg = LaunchConfig { + grid_dim: (total.div_ceil(256) as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.interleave_ext3_slabs) + .arg(h.buf.as_ref()) + .arg(&mut out_dev) + .arg(&(m as u64)) + .arg(&(lde as u64)) + .launch(cfg)?; + } + // Per-part D2H straight into the owning buffer: no host-side copy or memset + // (the D2H fills every element, so `set_len` on the uninitialized Vec is + // sound). One transfer per part; the total bytes match the old single D2H. + let mut parts: Vec> = Vec::with_capacity(m); + for p in 0..m { + let mut host: Vec = Vec::with_capacity(lde * 3); + // SAFETY: the memcpy below writes all `lde * 3` elements before any read. + unsafe { + host.set_len(lde * 3); + } + let view = out_dev.slice(p * lde * 3..(p + 1) * lde * 3); + stream.memcpy_dtoh(&view, &mut host)?; + parts.push(host); + } + stream.synchronize()?; + Ok(parts) +} + /// Fully device-resident R4 DEEP path: parts LDE and inverse denominators /// both arrive as device handles, the caller threads its own stream /// through so the inv_denoms producer diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 34cb40273..fc4ab599e 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -227,6 +227,7 @@ pub struct Backend { // deep.cubin pub deep_composition_ext3_row: CudaFunction, pub bit_reverse_ext3_kernel: CudaFunction, + pub interleave_ext3_slabs: CudaFunction, // fri.cubin pub fri_fold_ext3: CudaFunction, @@ -472,6 +473,7 @@ impl Backend { gather_rows_ext3: bary.load_function("gather_rows_ext3")?, deep_composition_ext3_row: deep.load_function("deep_composition_ext3_row")?, bit_reverse_ext3_kernel: deep.load_function("bit_reverse_ext3_interleaved")?, + interleave_ext3_slabs: deep.load_function("interleave_ext3_slabs")?, fri_fold_ext3: fri.load_function("fri_fold_ext3")?, fri_inject_bucket_ext3: fri.load_function("fri_inject_bucket_ext3")?, gather_ext3_at: fri.load_function("gather_ext3_at")?, diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 19b45b01a..57174f70f 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -2141,44 +2141,24 @@ where if TypeId::of::() != TypeId::of::() { return None; } - h.wait_ready_on(stream).ok()?; - let slabs = stream.clone_dtoh(h.buf.as_ref()).ok()?; - stream.synchronize().ok()?; - let (m, lde) = (h.m, h.lde_size); - if slabs.len() != m * lde * 3 { - return None; - } - // Per part: de-interleave the 3 slabs into row-major ext3 and reinterpret - // the u64 buffer in place — mirroring `materialize_lde_trace_host` rather - // than copying again through `u64_to_ext3_vec`. The row fill is parallel; - // this path fires often under VRAM pressure and otherwise dominates the - // D2H it follows. - let parts = (0..m) - .map(|p| { - let mut interleaved = vec![0u64; lde * 3]; - #[cfg(feature = "parallel")] - interleaved - .par_chunks_exact_mut(3) - .enumerate() - .for_each(|(r, dst)| { - for (k, d) in dst.iter_mut().enumerate() { - *d = slabs[(p * 3 + k) * lde + r]; - } - }); - #[cfg(not(feature = "parallel"))] - for (r, dst) in interleaved.chunks_exact_mut(3).enumerate() { - for (k, d) in dst.iter_mut().enumerate() { - *d = slabs[(p * 3 + k) * lde + r]; - } - } + // De-interleave the 3 ext3 slabs into per-column row-major ext3 ON DEVICE, + // then a contiguous D2H + split — the strided per-row gather this used to run + // on the host dominated the D2H it follows (~40% of the parts download). + let interleaved = math_cuda::deep::download_parts_interleaved(h, stream).ok()?; + // Each part is `lde` ext3 elements as `lde*3` interleaved u64s; reinterpret + // in place as FieldElement (E == Ext3 = [u64; 3], checked above; each + // per-part Vec has len == capacity == lde*3 from the contiguous copy). + let parts = interleaved + .into_iter() + .map(|v| { + let mut v = std::mem::ManuallyDrop::new(v); + debug_assert!( + v.len().is_multiple_of(3) && v.capacity().is_multiple_of(3), + "interleaved len/capacity must be a multiple of 3 for Fp3 reinterpret" + ); // SAFETY: E == Ext3 per the tower check above; FieldElement - // is [u64; 3]. `vec![0u64; lde*3]` has len == capacity == lde*3. + // is [u64; 3]. unsafe { - let mut v = std::mem::ManuallyDrop::new(interleaved); - debug_assert!( - v.len().is_multiple_of(3) && v.capacity().is_multiple_of(3), - "interleaved len/capacity must be a multiple of 3 for Fp3 reinterpret" - ); Vec::from_raw_parts( v.as_mut_ptr() as *mut FieldElement, v.len() / 3, From 8538682ef40d47b806b5885c3e043bcf8956e032 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Tue, 8 Sep 2026 10:15:06 -0300 Subject: [PATCH 26/26] perf(batched): de-interleave parts on device + zero-copy the aux/DEEP ext3 downloads Move the composition-parts D2H de-interleave onto the GPU (new interleave_ext3_slabs kernel + per-part D2H into the owning buffers), and reinterpret the aux-trace and DEEP-codeword ext3 downloads in place instead of the per-element u64_to_ext3_vec host copy. Byte-identical; comp_parts 3.02->2.63s, aux_commit 1.52->1.29s, deep_fri 2.22->2.0s per epoch. --- crypto/stark/src/gpu_lde.rs | 37 +++++++++++++++++++++++++++++++++++-- 1 file changed, 35 insertions(+), 2 deletions(-) diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 57174f70f..9e91ae8fb 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -2062,7 +2062,25 @@ where if stream.synchronize().is_err() || raw.len() != rows * cols * 3 { return false; } - let data = u64_to_ext3_vec::(&raw); + // `raw` is the row-major ext3 aux buffer `[(row*cols+col)*3+limb]` — exactly + // the aux table's row-major `[row*cols+col]` ext3 layout. Reinterpret it in + // place as `FieldElement` (E == ext3 = [u64;3]) instead of the per-element + // `u64_to_ext3_vec` copy (single-threaded over the whole aux). Byte-identical. + let data: Vec> = { + let mut v = std::mem::ManuallyDrop::new(raw); + debug_assert!( + v.len().is_multiple_of(3) && v.capacity().is_multiple_of(3), + "aux buffer len/capacity must be a multiple of 3 for Fp3 reinterpret" + ); + // SAFETY: E == ext3 (tower checked above); `FieldElement` is [u64; 3]. + unsafe { + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len() / 3, + v.capacity() / 3, + ) + } + }; trace.aux_table = crate::table::Table::new(data, cols); trace.num_aux_columns = cols; // The declined device LDE attempt can leave kernels enqueued on another @@ -2732,7 +2750,22 @@ where }; GPU_DEEP_CALLS.fetch_add(1, Ordering::Relaxed); debug_assert_eq!(deep_raw.len(), lde_size * 3); - Some(u64_to_ext3_vec::(&deep_raw)) + // `deep_raw` is the DEEP codeword row-major ext3 (`[row*3+limb]`); reinterpret + // in place as `FieldElement` (E == ext3 = [u64;3]) rather than the + // per-element `u64_to_ext3_vec` copy. Byte-identical. + let out: Vec> = { + let mut v = std::mem::ManuallyDrop::new(deep_raw); + debug_assert!(v.len().is_multiple_of(3) && v.capacity().is_multiple_of(3)); + // SAFETY: E == ext3 (TypeId-checked at entry); FieldElement = [u64; 3]. + unsafe { + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len() / 3, + v.capacity() / 3, + ) + } + }; + Some(out) } /// Fully-resident DEEP keeping the codeword on device in FRI order (no D2H).