Merge lifted into dev. (#1307)

* Add lifted Merkle verifier. (#1234)

* Add SIMD MerkleOpsLifted. (#1235)

* Add auxiliary structs to vcs_lifted. (#1247)

* Add generic blake for lifted vcs. (#1252)

* Change queries from vec to slice. (#1250)

* Add lift_and_accumulate. (#1245)

* Modify finalize of DomainEvaluationAccumulator. (#1246)

* Add skeleton of lifted quotient ops. (#1248)

All the code related to the previous architecture of quotient ops has
been removed.

* Add empty poseidon merkle hasher lifted. (#1253)

* Change channel's assoc type to MerkleHasherLifted. (#1254)

* Add CPU QuotientOps impl. (#1249)

* Make pcs prove lifted values. (#1251)

* Add pcs verifier lifted. (#1255)

* Modify trace step in masks and eval at point. (#1257)

* Add CpuBackend impl to ComponentProver. (#1258)

* Add mixed wide fibonacci test. (#1259)

* Add m31 blake hasher to simd. (#1260)

* Allow small traces with cpu, fix bug. (#1261)

* Add lift and accumulate in simd. (#1262)

* Add simd impl of accumulate_numerators. (#1263)

* Add `compute_quotients_and_combine`. (#1266)

* Change fri decommit to return a vec of query positions. (#1269)

* Commit to a single column in fri. (#1270)

* Update docs and rename. (#1271)

* Remove iterations in fri verifier decommit. (#1272)

* Remove unneeded structure. (#1273)

* Update docs. (#1274)

* Add poseidon hasher. (#1283)

* Add SIMD implementation of poseidon MerkleOpsLifted. (#1280)

* Adapt the queries for the preprocessed trees. (#1281)

Co-authored-by: Alon F <alonf@starkware.co>

* Fix non-parallel bugs. (#1287)

* Adapt barycentric weights to the lifted protocol. (#1288)

* Change order of the alphas to the samples order. (#1292)

* Change queried values layout. (#1293)

* Add periodicity samples. (#1294)

* Improve performance of cpu merkle. (#1284)

* Fix lifting size in fri_answers and add test. (#1296)

* Optimize blake simd merkle lifted. (#1291)

* Add maximal size periodicity checks. (#1301)

* Add dummy serde and reexport MerkleHasherLifted. (#1297)

These changes make it easier to integrate with stwo_cairo_prover.

---------

Co-authored-by: Alon F <alonf@starkware.co>

Author: leo-starkware

Cargo.lock

@@ -220,8 +220,9 @@ impl<E: FrameworkEval> Component for FrameworkComponent<E> {
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
+        max_log_degree_bound: u32,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-        let trace_step = CanonicCoset::new(self.eval.log_size()).step();
+        let trace_step = CanonicCoset::new(max_log_degree_bound).step();
         self.info.mask_offsets.as_ref().map_cols(|col_offsets| {
             col_offsets
                 .iter()
@@ -239,6 +240,7 @@ impl<E: FrameworkEval> Component for FrameworkComponent<E> {
         point: CirclePoint<SecureField>,
         mask: &TreeVec<ColumnVec<Vec<SecureField>>>,
         evaluation_accumulator: &mut PointEvaluationAccumulator,
+        max_log_degree_bound: u32,
     ) {
         let preprocessed_mask = self
             .preprocessed_column_indices
@@ -252,7 +254,7 @@ impl<E: FrameworkEval> Component for FrameworkComponent<E> {
         self.eval.evaluate(PointEvaluator::new(
             mask_points,
             evaluation_accumulator,
-            coset_vanishing(CanonicCoset::new(self.eval.log_size()).coset, point).inverse(),
+            coset_vanishing(CanonicCoset::new(max_log_degree_bound).coset, point).inverse(),
             self.eval.log_size(),
             self.claimed_sum,
         ));

crates/constraint-framework/src/component.rs

@@ -6,13 +6,15 @@ use rayon::prelude::*;
 use stwo::core::air::Component;
 use stwo::core::constraints::coset_vanishing;
 use stwo::core::fields::m31::BaseField;
+use stwo::core::fields::qm31::SecureField;
 use stwo::core::pcs::TreeVec;
 use stwo::core::poly::circle::CanonicCoset;
 use stwo::core::utils::bit_reverse;
 use stwo::prover::backend::simd::column::VeryPackedSecureColumnByCoords;
 use stwo::prover::backend::simd::m31::LOG_N_LANES;
 use stwo::prover::backend::simd::very_packed_m31::{VeryPackedBaseField, LOG_N_VERY_PACKED_ELEMS};
 use stwo::prover::backend::simd::SimdBackend;
+use stwo::prover::backend::CpuBackend;
 use stwo::prover::poly::circle::{CircleEvaluation, PolyOps};
 use stwo::prover::poly::BitReversedOrder;
 use stwo::prover::secure_column::SecureColumnByCoords;
@@ -70,7 +72,7 @@ impl<E: FrameworkEval + Sync> ComponentProver<SimdBackend> for FrameworkComponen
             .collect_vec();
         bit_reverse(&mut denom_inv);
 
-        // Accumulator.
+        // Note that `accum` is a mutable reference to a column in `evaluation_accumulator`.
         let [mut accum] =
             evaluation_accumulator.columns([(eval_domain.log_size(), self.n_constraints())]);
         accum.random_coeff_powers.reverse();
@@ -82,32 +84,19 @@ impl<E: FrameworkEval + Sync> ComponentProver<SimdBackend> for FrameworkComponen
         )
         .entered();
 
+        // Fall back to CPU if the trace is too small.
         if trace_domain.log_size() < LOG_N_LANES + LOG_N_VERY_PACKED_ELEMS {
-            // Fall back to CPU if the trace is too small.
-            let mut col = accum.col.to_cpu();
-
-            for row in 0..(1 << eval_domain.log_size()) {
-                let trace_cols = trace.as_cols_ref().map_cols(|c| c.to_cpu());
-                let trace_cols = trace_cols.as_cols_ref();
-
-                // Evaluate constrains at row.
-                let eval = CpuDomainEvaluator::new(
-                    &trace_cols,
-                    row,
-                    &accum.random_coeff_powers,
-                    trace_domain.log_size(),
-                    eval_domain.log_size(),
-                    self.eval.log_size(),
-                    self.claimed_sum,
-                );
-                let row_res = self.eval.evaluate(eval).row_res;
-
-                // Finalize row.
-                let denom_inv = denom_inv[row >> trace_domain.log_size()];
-                col.set(row, col.at(row) + row_res * denom_inv)
-            }
-            let col = SecureColumnByCoords::from_cpu(col);
-            *accum.col = col;
+            let trace_cols = trace.as_cols_ref().map_cols(|c| c.to_cpu());
+            let trace_cols = trace_cols.as_cols_ref();
+            *accum.col = SecureColumnByCoords::from_cpu(accumulate_pointwise_cpu(
+                self,
+                trace_cols,
+                eval_domain.log_size(),
+                trace_domain.log_size(),
+                denom_inv,
+                &accum.random_coeff_powers,
+                &accum.col.to_cpu(),
+            ));
             return;
         }
 
@@ -148,16 +137,117 @@ impl<E: FrameworkEval + Sync> ComponentProver<SimdBackend> for FrameworkComponen
 
                 // Finalize row.
                 unsafe {
-                    let denom_inv = VeryPackedBaseField::broadcast(
+                    let row_denom_inv = VeryPackedBaseField::broadcast(
                         denom_inv[vec_row
                             >> (trace_domain.log_size() - LOG_N_LANES - LOG_N_VERY_PACKED_ELEMS)],
                     );
                     chunk.set_packed(
                         idx_in_chunk,
-                        chunk.packed_at(idx_in_chunk) + row_res * denom_inv,
+                        chunk.packed_at(idx_in_chunk) + row_res * row_denom_inv,
                     )
                 }
             }
         });
     }
 }
+
+impl<E: FrameworkEval + Sync> ComponentProver<CpuBackend> for FrameworkComponent<E> {
+    /// Almost all this implementation is equal to the one above for `SimdBackend`.
+    fn evaluate_constraint_quotients_on_domain(
+        &self,
+        trace: &Trace<'_, CpuBackend>,
+        evaluation_accumulator: &mut DomainEvaluationAccumulator<CpuBackend>,
+    ) {
+        if self.n_constraints() == 0 {
+            return;
+        }
+        let eval_domain = CanonicCoset::new(self.max_constraint_log_degree_bound()).circle_domain();
+        let trace_domain = CanonicCoset::new(self.eval.log_size());
+
+        let mut component_polys = trace.polys.sub_tree(&self.trace_locations);
+        component_polys[PREPROCESSED_TRACE_IDX] = self
+            .preprocessed_column_indices
+            .iter()
+            .map(|idx| &trace.polys[PREPROCESSED_TRACE_IDX][*idx])
+            .collect();
+
+        // Extend trace if necessary.
+        // TODO: Don't extend when eval_size < committed_size. Instead, pick a good
+        // subdomain. (For larger blowup factors).
+        let need_to_extend = component_polys
+            .iter()
+            .flatten()
+            .any(|c| c.evals.domain.log_size() != eval_domain.log_size());
+        let trace: TreeVec<
+            Vec<Cow<'_, CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>>,
+        > = if need_to_extend {
+            let _span = span!(Level::INFO, "Constraint Extension").entered();
+            let twiddles = CpuBackend::precompute_twiddles(eval_domain.half_coset);
+            component_polys
+                .as_cols_ref()
+                .map_cols(|col| Cow::Owned(col.get_evaluation_on_domain(eval_domain, &twiddles)))
+        } else {
+            component_polys.map_cols(|c| Cow::Borrowed(&c.evals))
+        };
+
+        // Denom inverses.
+        let log_expand = eval_domain.log_size() - trace_domain.log_size();
+        let mut denom_inv = (0..1 << log_expand)
+            .map(|i| coset_vanishing(trace_domain.coset(), eval_domain.at(i)).inverse())
+            .collect_vec();
+        bit_reverse(&mut denom_inv);
+
+        // Accumulator.
+        let [mut accum] =
+            evaluation_accumulator.columns([(eval_domain.log_size(), self.n_constraints())]);
+        accum.random_coeff_powers.reverse();
+
+        let _span = span!(
+            Level::INFO,
+            "Constraint point-wise eval",
+            class = "ConstraintEval"
+        )
+        .entered();
+        let trace_cols = trace.as_cols_ref().map_cols(|c| c.as_ref());
+
+        *accum.col = accumulate_pointwise_cpu(
+            self,
+            trace_cols,
+            eval_domain.log_size(),
+            trace_domain.log_size(),
+            denom_inv,
+            &accum.random_coeff_powers,
+            accum.col,
+        );
+    }
+}
+
+fn accumulate_pointwise_cpu<E: FrameworkEval>(
+    component: &FrameworkComponent<E>,
+    trace_cols: TreeVec<Vec<&CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>>,
+    eval_log_size: u32,
+    trace_log_size: u32,
+    denom_inv: Vec<BaseField>,
+    random_coeff_powers: &[SecureField],
+    accum: &SecureColumnByCoords<CpuBackend>,
+) -> SecureColumnByCoords<CpuBackend> {
+    let mut res = SecureColumnByCoords::zeros(1 << eval_log_size);
+    for row in 0..(1 << eval_log_size) {
+        // Evaluate constrains at row.
+        let eval = CpuDomainEvaluator::new(
+            &trace_cols,
+            row,
+            random_coeff_powers,
+            trace_log_size,
+            eval_log_size,
+            component.eval.log_size(),
+            component.claimed_sum,
+        );
+        let row_res = component.eval.evaluate(eval).row_res;
+
+        // Finalize row.
+        let row_denom_inv = denom_inv[row >> trace_log_size];
+        res.set(row, accum.at(row) + row_res * row_denom_inv)
+    }
+    res
+}

crates/constraint-framework/src/prover/component_prover.rs

@@ -9,7 +9,7 @@ use stwo::core::fields::qm31::SecureField;
 use stwo::core::pcs::{CommitmentSchemeVerifier, PcsConfig, TreeVec};
 use stwo::core::poly::circle::CanonicCoset;
 use stwo::core::proof::StarkProof;
-use stwo::core::vcs::MerkleHasher;
+use stwo::core::vcs_lifted::merkle_hasher::MerkleHasherLifted;
 use stwo::core::verifier::{verify, VerificationError};
 use stwo::prover::backend::simd::m31::LOG_N_LANES;
 use stwo::prover::backend::simd::SimdBackend;
@@ -149,7 +149,7 @@ impl BlakeStatement1 {
     }
 }
 
-pub struct BlakeProof<H: MerkleHasher> {
+pub struct BlakeProof<H: MerkleHasherLifted> {
     stmt0: BlakeStatement0,
     stmt1: BlakeStatement1,
     stark_proof: StarkProof<H>,
@@ -306,7 +306,6 @@ where
     // Setup protocol.
     let channel = &mut MC::C::default();
     let mut commitment_scheme = CommitmentSchemeProver::new(config, &twiddles);
-
     // Preprocessed trace.
     // TODO(ShaharS): share is_first column between components when constant columns support this.
     let span = span!(Level::INFO, "Preprocessed Trace").entered();
@@ -521,7 +520,7 @@ mod tests {
     use std::env;
 
     use stwo::core::pcs::PcsConfig;
-    use stwo::core::vcs::blake2_merkle::Blake2sMerkleChannel;
+    use stwo::core::vcs_lifted::blake2_merkle::Blake2sMerkleChannel;
 
     use crate::blake::air::{prove_blake, verify_blake};
 

crates/examples/src/blake/air.rs

@@ -6,7 +6,7 @@ use stwo::core::fields::qm31::SecureField;
 use stwo::core::pcs::{PcsConfig, TreeSubspan};
 use stwo::core::poly::circle::CanonicCoset;
 use stwo::core::proof::StarkProof;
-use stwo::core::vcs::blake2_merkle::{Blake2sMerkleChannel, Blake2sMerkleHasher};
+use stwo::core::vcs_lifted::blake2_merkle::{Blake2sMerkleChannel, Blake2sMerkleHasher};
 use stwo::core::ColumnVec;
 use stwo::prover::backend::simd::column::BaseColumn;
 use stwo::prover::backend::simd::m31::LOG_N_LANES;
@@ -293,7 +293,7 @@ mod tests {
     use stwo::core::channel::Blake2sChannel;
     use stwo::core::fri::FriConfig;
     use stwo::core::pcs::{CommitmentSchemeVerifier, PcsConfig};
-    use stwo::core::vcs::blake2_merkle::Blake2sMerkleChannel;
+    use stwo::core::vcs_lifted::blake2_merkle::Blake2sMerkleChannel;
     use stwo::core::verifier::verify;
 
     use crate::plonk::{prove_fibonacci_plonk, PlonkLookupElements};

crates/examples/src/plonk/mod.rs

@@ -13,7 +13,7 @@ use stwo::core::fields::FieldExpOps;
 use stwo::core::pcs::PcsConfig;
 use stwo::core::poly::circle::CanonicCoset;
 use stwo::core::proof::StarkProof;
-use stwo::core::vcs::blake2_merkle::{Blake2sMerkleChannel, Blake2sMerkleHasher};
+use stwo::core::vcs_lifted::blake2_merkle::{Blake2sMerkleChannel, Blake2sMerkleHasher};
 use stwo::core::ColumnVec;
 use stwo::prover::backend::simd::column::BaseColumn;
 use stwo::prover::backend::simd::m31::{PackedBaseField, LOG_N_LANES};
@@ -404,7 +404,7 @@ mod tests {
     use stwo::core::fri::FriConfig;
     use stwo::core::pcs::{CommitmentSchemeVerifier, PcsConfig, TreeVec};
     use stwo::core::poly::circle::CanonicCoset;
-    use stwo::core::vcs::blake2_merkle::Blake2sMerkleChannel;
+    use stwo::core::vcs_lifted::blake2_merkle::Blake2sMerkleChannel;
     use stwo::core::verifier::verify;
     use stwo_constraint_framework::assert_constraints_on_polys;
 
@@ -486,6 +486,7 @@ mod tests {
         );
     }
 
+    #[ignore = "AIRs with constraint degree >= 2 are not supported yet in the lifted protocol."]
     #[test_log::test]
     fn test_simd_poseidon_prove() {
         // Note: To see time measurement, run test with
@@ -529,6 +530,7 @@ mod tests {
         verify(&[&component], channel, commitment_scheme, proof).unwrap();
     }
 
+    #[ignore = "AIRs with constraint degree >= 2 are not supported yet in the lifted protocol."]
     #[cfg(feature = "tracing")]
     #[test]
     fn trace_simd_poseidon_prove() {

crates/examples/src/poseidon/mod.rs

@@ -6,7 +6,7 @@ use stwo::core::fields::m31::{BaseField, M31};
 use stwo::core::fields::qm31::{SecureField, QM31};
 use stwo::core::pcs::TreeVec;
 use stwo::core::proof::StarkProof;
-use stwo::core::vcs::MerkleHasher;
+use stwo::core::vcs_lifted::merkle_hasher::MerkleHasherLifted;
 use stwo::prover::backend::simd::SimdBackend;
 use stwo::prover::poly::circle::CircleEvaluation;
 use stwo::prover::poly::BitReversedOrder;
@@ -148,7 +148,7 @@ pub fn track_state_machine_relations(
     .collect()
 }
 
-pub struct StateMachineProof<H: MerkleHasher> {
+pub struct StateMachineProof<H: MerkleHasherLifted> {
     pub public_input: [State; 2], // Initial and final state.
     pub stmt0: StateMachineStatement0,
     pub stmt1: StateMachineStatement1,