Merge branch 'starkware-libs:dev' into dev

README.md

@@ -1,15 +1,62 @@
+<div align="center">
+
+![STWO](resources/img/logo.png)
+
+<a href="https://github.com/starkware-libs/stwo/actions/workflows/ci.yaml"><img alt="GitHub Workflow Status (with event)" src="https://img.shields.io/github/actions/workflow/status/starkware-libs/stwo/ci.yaml?style=for-the-badge" height=30></a>
+<a href="https://codecov.io/gh/starkware-libs/stwo" >
+<img src="https://img.shields.io/codecov/c/github/starkware-libs/stwo?style=for-the-badge&logo=codecov" height=30/>
+</a>
+<a href="https://github.com/starkware-libs/stwo/blob/main/LICENSE"><img src="https://img.shields.io/github/license/starkware-libs/stwo.svg?style=for-the-badge" alt="Project license" height="30"></a>
+<a href="https://starkware.co/"><img src="https://img.shields.io/badge/By StarkWare-29296E.svg?&style=for-the-badge&logo=data:image/svg%2bxml;base64,PD94bWwgdmVyc2lvbj0iMS4wIiBlbmNvZGluZz0iVVRGLTgiPz48c3ZnIGlkPSJhIiB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAxODEgMTgxIj48ZGVmcz48c3R5bGU+LmJ7ZmlsbDojZmZmO308L3N0eWxlPjwvZGVmcz48cGF0aCBjbGFzcz0iYiIgZD0iTTE3Ni43Niw4OC4xOGwtMzYtMzcuNDNjLTEuMzMtMS40OC0zLjQxLTIuMDQtNS4zMS0xLjQybC0xMC42MiwyLjk4LTEyLjk1LDMuNjNoLjc4YzUuMTQtNC41Nyw5LjktOS41NSwxNC4yNS0xNC44OSwxLjY4LTEuNjgsMS44MS0yLjcyLDAtNC4yN0w5Mi40NSwuNzZxLTEuOTQtMS4wNC00LjAxLC4xM2MtMTIuMDQsMTIuNDMtMjMuODMsMjQuNzQtMzYsMzcuNjktMS4yLDEuNDUtMS41LDMuNDQtLjc4LDUuMThsNC4yNywxNi41OGMwLDIuNzIsMS40Miw1LjU3LDIuMDcsOC4yOS00LjczLTUuNjEtOS43NC0xMC45Ny0xNS4wMi0xNi4wNi0xLjY4LTEuODEtMi41OS0xLjgxLTQuNCwwTDQuMzksODguMDVjLTEuNjgsMi4zMy0xLjgxLDIuMzMsMCw0LjUzbDM1Ljg3LDM3LjNjMS4zNiwxLjUzLDMuNSwyLjEsNS40NCwxLjQybDExLjQtMy4xMSwxMi45NS0zLjYzdi45MWMtNS4yOSw0LjE3LTEwLjIyLDguNzYtMTQuNzYsMTMuNzNxLTMuNjMsMi45OC0uNzgsNS4zMWwzMy40MSwzNC44NGMyLjIsMi4yLDIuOTgsMi4yLDUuMTgsMGwzNS40OC0zNy4xN2MxLjU5LTEuMzgsMi4xNi0zLjYsMS40Mi01LjU3LTEuNjgtNi4wOS0zLjI0LTEyLjMtNC43OS0xOC4zOS0uNzQtMi4yNy0xLjIyLTQuNjItMS40Mi02Ljk5LDQuMyw1LjkzLDkuMDcsMTEuNTIsMTQuMjUsMTYuNzEsMS42OCwxLjY4LDIuNzIsMS42OCw0LjQsMGwzNC4zMi0zNS43NHExLjU1LTEuODEsMC00LjAxWm0tNzIuMjYsMTUuMTVjLTMuMTEtLjc4LTYuMDktMS41NS05LjE5LTIuNTktMS43OC0uMzQtMy42MSwuMy00Ljc5LDEuNjhsLTEyLjk1LDEzLjg2Yy0uNzYsLjg1LTEuNDUsMS43Ni0yLjA3LDIuNzJoLS42NWMxLjMtNS4zMSwyLjcyLTEwLjYyLDQuMDEtMTUuOGwxLjY4LTYuNzNjLjg0LTIuMTgsLjE1LTQuNjUtMS42OC02LjA5bC0xMi45NS0xNC4xMmMtLjY0LS40NS0xLjE0LTEuMDgtMS40Mi0xLjgxbDE5LjA0LDUuMTgsMi41OSwuNzhjMi4wNCwuNzYsNC4zMywuMTQsNS43LTEuNTVsMTIuOTUtMTQuMzhzLjc4LTEuMDQsMS42OC0xLjE3Yy0xLjgxLDYuNi0yLjk4LDE0LjEyLTUuNDQsMjAuNDYtMS4wOCwyLjk2LS4wOCw2LjI4LDIuNDYsOC4xNiw0LjI3LDQuMTQsOC4yOSw4LjU1LDEyLjk1LDEyLjk1LDAsMCwxLjMsLjkxLDEuNDIsMi4wN2wtMTMuMzQtMy42M1oiLz48L3N2Zz4=" alt="StarkWare" height="30"></a>
+</div>
+
+<div align="center">
+  <h3>
+    <a href="https://eprint.iacr.org/2024/278">
+      Paper
+    </a>
+    <span> | </span>
+    <a href="https://github.com/starkware-libs/stwo">
+      Documentation
+    </a>
+    <span> | </span>
+    <a href="https://starkware-libs.github.io/stwo/dev/bench/index.html">
+      Benchmarks
+    </a>
+  </h3>
+</div>
+
 # Stwo
 
-## About
+## 🌟 About
+
+Stwo is a next generation implementation of a [CSTARK](https://eprint.iacr.org/2024/278) prover and verifier, written in Rust 🦀.
+
+> **Stwo is a work in progress.**
+>
+> It is not recommended to use it in a production setting yet.
 
-Stwo is a Rust implementation of a [CSTARK](https://eprint.iacr.org/2024/278) prover and verifier.
+## 🚀 Key Features
 
-## Disclaimer
+- **Circle STARKs:** Based on the latest cryptographic research and innovations in the ZK field.
+- **High performance:** Stwo is designed to be extremely fast and efficient.
+- **Flexible:** Adaptable for various validity proof applications.
 
-Stwo is a work in progress.
+## 📊 Benchmarks
 
-## License
+Run `poseidon_benchmark.sh` to run a single-threaded poseidon2 hash proof benchmark.
+
+Further benchmarks can be run using `cargo bench`.
+
+Visual representation of benchmarks can be found [here](https://starkware-libs.github.io/stwo/dev/bench/index.html).
+
+## 📜 License
 
 This project is licensed under the **Apache 2.0 license**.
 
 See [LICENSE](LICENSE) for more information.
+
+<!-- markdownlint-restore -->
+<!-- prettier-ignore-end -->
+
+<!-- ALL-CONTRIBUTORS-LIST:END -->

crates/prover/Cargo.toml

@@ -3,20 +3,27 @@ name = "stwo-prover"
 version.workspace = true
 edition.workspace = true
 
+[features]
+parallel = ["rayon"]
+
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
 blake2.workspace = true
 blake3.workspace = true
 bytemuck = { workspace = true, features = ["derive"] }
 cfg-if = "1.0.0"
+downcast-rs = "1.2"
 educe.workspace = true
 hex.workspace = true
 itertools.workspace = true
 num-traits.workspace = true
 rand = { version = "0.8.5", default-features = false, features = ["small_rng"] }
+starknet-crypto = "0.6.2"
+starknet-ff = "0.3.7"
 thiserror.workspace = true
 tracing.workspace = true
+rayon = { version = "1.10.0", optional = true }
 
 [dev-dependencies]
 aligned = "0.4.2"
@@ -75,3 +82,7 @@ harness = false
 [[bench]]
 name = "quotients"
 harness = false
+
+[[bench]]
+name = "poseidon"
+harness = false

crates/prover/benches/poseidon.rs

@@ -0,0 +1,32 @@
+use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+use stwo_prover::core::backend::simd::SimdBackend;
+use stwo_prover::core::channel::{Blake2sChannel, Channel};
+use stwo_prover::core::fields::m31::BaseField;
+use stwo_prover::core::fields::IntoSlice;
+use stwo_prover::core::prover::prove;
+use stwo_prover::core::vcs::blake2_hash::Blake2sHasher;
+use stwo_prover::core::vcs::hasher::Hasher;
+use stwo_prover::examples::poseidon::{gen_trace, PoseidonAir, PoseidonComponent};
+
+pub fn simd_poseidon(c: &mut Criterion) {
+    const LOG_N_ROWS: u32 = 15;
+    let mut group = c.benchmark_group("poseidon2");
+    group.throughput(Throughput::Elements(1u64 << (LOG_N_ROWS + 3)));
+    group.bench_function(format!("poseidon2 2^{} instances", LOG_N_ROWS + 3), |b| {
+        b.iter(|| {
+            let component = PoseidonComponent {
+                log_n_instances: LOG_N_ROWS,
+            };
+            let trace = gen_trace(component.log_column_size());
+            let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[])));
+            let air = PoseidonAir { component };
+            prove::<SimdBackend>(&air, channel, trace).unwrap()
+        });
+    });
+}
+
+criterion_group!(
+    name = bit_rev;
+    config = Criterion::default().sample_size(10);
+    targets = simd_poseidon);
+criterion_main!(bit_rev);

crates/prover/src/core/air/air_ext.rs

@@ -1,19 +1,16 @@
-use std::collections::{BTreeMap, BTreeSet};
-use std::iter::zip;
-
 use itertools::{zip_eq, Itertools};
 
 use super::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use super::{Air, AirProver, ComponentTrace};
 use crate::core::backend::Backend;
-use crate::core::channel::{Blake2sChannel, Channel};
 use crate::core::circle::CirclePoint;
-use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
-use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, CirclePoly, SecureCirclePoly};
-use crate::core::poly::BitReversedOrder;
-use crate::core::prover::LOG_BLOWUP_FACTOR;
-use crate::core::{ColumnVec, ComponentVec, InteractionElements};
+use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
+use crate::core::pcs::{CommitmentTreeProver, TreeVec};
+use crate::core::poly::circle::SecureCirclePoly;
+use crate::core::vcs::blake2_merkle::Blake2sMerkleHasher;
+use crate::core::vcs::ops::MerkleOps;
+use crate::core::{ColumnVec, InteractionElements, LookupValues};
 
 pub trait AirExt: Air {
     fn composition_log_degree_bound(&self) -> u32 {
@@ -24,101 +21,121 @@ pub trait AirExt: Air {
             .unwrap()
     }
 
-    fn trace_commitment_domains(&self) -> Vec<CanonicCoset> {
-        self.column_log_sizes()
+    fn n_interaction_phases(&self) -> u32 {
+        self.components()
             .iter()
-            .map(|&log_size| CanonicCoset::new(log_size + LOG_BLOWUP_FACTOR))
-            .collect_vec()
+            .map(|component| component.n_interaction_phases())
+            .max()
+            .unwrap()
     }
 
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
-    ) -> ComponentVec<Vec<CirclePoint<SecureField>>> {
-        let mut component_points = ComponentVec(Vec::new());
-        for component in self.components() {
-            let points = component.mask_points(point);
-            component_points.push(points);
-        }
-        component_points
-    }
-
-    fn interaction_elements(&self, channel: &mut Blake2sChannel) -> InteractionElements {
-        let mut ids = BTreeSet::new();
+    ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
+        let mut air_points = TreeVec::default();
         for component in self.components() {
-            ids.extend(component.interaction_element_ids());
+            let component_points = component.mask_points(point);
+            if air_points.len() < component_points.len() {
+                air_points.resize(component_points.len(), vec![]);
+            }
+            air_points.as_mut().zip_eq(component_points).map(
+                |(air_tree_points, component_tree_points)| {
+                    air_tree_points.extend(component_tree_points);
+                },
+            );
         }
-        let elements = channel.draw_felts(ids.len()).into_iter().map(|e| e.0 .0);
-        InteractionElements(BTreeMap::from_iter(zip_eq(ids, elements)))
+        // Add the composition polynomial mask points.
+        air_points.push(vec![vec![point]; SECURE_EXTENSION_DEGREE]);
+        air_points
     }
 
     fn eval_composition_polynomial_at_point(
         &self,
         point: CirclePoint<SecureField>,
-        mask_values: &ComponentVec<Vec<SecureField>>,
+        mask_values: &Vec<TreeVec<Vec<Vec<SecureField>>>>,
         random_coeff: SecureField,
+        interaction_elements: &InteractionElements,
+        lookup_values: &LookupValues,
     ) -> SecureField {
         let mut evaluation_accumulator = PointEvaluationAccumulator::new(random_coeff);
-        zip(self.components(), &mask_values.0).for_each(|(component, mask)| {
+        zip_eq(self.components(), mask_values).for_each(|(component, mask)| {
             component.evaluate_constraint_quotients_at_point(
                 point,
                 mask,
                 &mut evaluation_accumulator,
+                interaction_elements,
+                lookup_values,
             )
         });
         evaluation_accumulator.finalize()
     }
 
-    fn column_log_sizes(&self) -> Vec<u32> {
-        self.components()
-            .iter()
-            .flat_map(|component| component.trace_log_degree_bounds())
-            .collect()
+    fn column_log_sizes(&self) -> TreeVec<ColumnVec<u32>> {
+        let mut air_sizes = TreeVec::default();
+        self.components().iter().for_each(|component| {
+            let component_sizes = component.trace_log_degree_bounds();
+            if air_sizes.len() < component_sizes.len() {
+                air_sizes.resize(component_sizes.len(), vec![]);
+            }
+            air_sizes.as_mut().zip_eq(component_sizes).map(
+                |(air_tree_sizes, component_tree_sizes)| {
+                    air_tree_sizes.extend(component_tree_sizes)
+                },
+            );
+        });
+        air_sizes
     }
 
-    fn component_traces<'a, B: Backend>(
+    fn component_traces<'a, B: Backend + MerkleOps<Blake2sMerkleHasher>>(
         &'a self,
-        polynomials: &'a [CirclePoly<B>],
-        evals: &'a [CircleEvaluation<B, BaseField, BitReversedOrder>],
+        trees: &'a [CommitmentTreeProver<B>],
     ) -> Vec<ComponentTrace<'_, B>> {
-        let poly_iter = &mut polynomials.iter();
-        let eval_iter = &mut evals.iter();
+        let mut poly_iters = trees
+            .iter()
+            .map(|tree| tree.polynomials.iter())
+            .collect_vec();
+        let mut eval_iters = trees
+            .iter()
+            .map(|tree| tree.evaluations.iter())
+            .collect_vec();
+
         self.components()
             .iter()
             .map(|component| {
-                let n_columns = component.trace_log_degree_bounds().len();
-                let polys = poly_iter.take(n_columns).collect();
-                let evals = eval_iter.take(n_columns).collect();
-                ComponentTrace::new(polys, evals)
+                let col_sizes_per_tree = component
+                    .trace_log_degree_bounds()
+                    .iter()
+                    .map(|col_sizes| col_sizes.len())
+                    .collect_vec();
+                let polys = col_sizes_per_tree
+                    .iter()
+                    .zip(poly_iters.iter_mut())
+                    .map(|(n_columns, iter)| iter.take(*n_columns).collect_vec())
+                    .collect_vec();
+                let evals = col_sizes_per_tree
+                    .iter()
+                    .zip(eval_iters.iter_mut())
+                    .map(|(n_columns, iter)| iter.take(*n_columns).collect_vec())
+                    .collect_vec();
+                ComponentTrace {
+                    polys: TreeVec::new(polys),
+                    evals: TreeVec::new(evals),
+                }
             })
-            .collect()
+            .collect_vec()
     }
 }
+
 impl<A: Air + ?Sized> AirExt for A {}
 
 pub trait AirProverExt<B: Backend>: AirProver<B> {
-    fn interact(
-        &self,
-        trace: &ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
-        elements: &InteractionElements,
-    ) -> ComponentVec<CircleEvaluation<B, BaseField, BitReversedOrder>> {
-        let trace_iter = &mut trace.iter();
-        ComponentVec(
-            self.prover_components()
-                .iter()
-                .map(|component| {
-                    let n_columns = component.trace_log_degree_bounds().len();
-                    let trace_columns = trace_iter.take(n_columns).collect_vec();
-                    component.write_interaction_trace(&trace_columns, elements)
-                })
-                .collect(),
-        )
-    }
-
     fn compute_composition_polynomial(
         &self,
         random_coeff: SecureField,
         component_traces: &[ComponentTrace<'_, B>],
+        interaction_elements: &InteractionElements,
+        lookup_values: &LookupValues,
     ) -> SecureCirclePoly<B> {
         let total_constraints: usize = self
             .prover_components()
@@ -130,11 +147,23 @@ pub trait AirProverExt<B: Backend>: AirProver<B> {
             self.composition_log_degree_bound(),
             total_constraints,
         );
-        zip(self.prover_components(), component_traces).for_each(|(component, trace)| {
-            component.evaluate_constraint_quotients_on_domain(trace, &mut accumulator)
+        zip_eq(self.prover_components(), component_traces).for_each(|(component, trace)| {
+            component.evaluate_constraint_quotients_on_domain(
+                trace,
+                &mut accumulator,
+                interaction_elements,
+                lookup_values,
+            )
         });
         accumulator.finalize()
     }
+
+    fn lookup_values(&self, component_traces: &[ComponentTrace<'_, B>]) -> LookupValues {
+        let mut values = LookupValues::default();
+        zip_eq(self.prover_components(), component_traces)
+            .for_each(|(component, trace)| values.extend(component.lookup_values(trace)));
+        values
+    }
 }
 
 impl<B: Backend, A: AirProver<B>> AirProverExt<B> for A {}