github/powdr
1
1
Fork 0
mirror of https://github.com/powdr-labs/powdr.git synced 2026-01-10 11:27:58 -05:00
Code Issues Packages Projects Releases Wiki Activity

Stwo support "addition and equal test" (#1942)

Browse Source 
## Initial Support for `stwo` Backend

This PR introduces the initial support for the `stwo` backend in the
Powdr pipeline.

### Summary
- **Constraint Proving**: Added functionality to prove and verify
constraints between columns.
- **Test Added**: A test case has been added to validate the above
functionality.

### Remaining Tasks
- Implement the `setup` function.
- Use `logup` for the next reference (now the next reference creates a
new column).
- Check the constant columns and public values in stwo and implement to
Powdr.

### How to Test
To test the `stwo` backend, use the following command:

```bash
cargo test --features stwo --package powdr-pipeline --test pil -- stwo_add_and_equal --exact --show-output

---------

Co-authored-by: Shuang Wu <shuangwu@Shuangs-MacBook-Pro.local>
Co-authored-by: Shuang Wu <shuangwu@Shuangs-MBP.localdomain>
Co-authored-by: Thibaut Schaeffer <schaeffer.thibaut@gmail.com>
This commit is contained in:
ShuangWu121
2024-11-18 16:35:59 +08:00
committed by GitHub
parent 58624e5f5d
commit 8f3a572076
8 changed files with 391 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
backend/Cargo.toml
View File

@@ -62,7 +62,7 @@ p3-commit = { git = "https://github.com/plonky3/Plonky3.git", rev = "2192432ddf2
p3-matrix = { git = "https://github.com/plonky3/Plonky3.git", rev = "2192432ddf28e7359dd2c577447886463e6124f0", optional = true }
p3-uni-stark = { git = "https://github.com/plonky3/Plonky3.git", rev = "2192432ddf28e7359dd2c577447886463e6124f0", optional = true }
# TODO: Change this to main branch when the `andrew/dev/update-toolchain` branch is merged,the main branch is using "nightly-2024-01-04", not compatiable with plonky3
stwo-prover = { git = "https://github.com/starkware-libs/stwo.git", optional = true, rev = "52d050c18b5dbc74af40214b3b441a6f60a20d41" }
stwo-prover = { git = "https://github.com/starkware-libs/stwo.git", optional = true, rev = "e6d10bc107c11cce54bb4aa152c3afa2e15e92c1" }
strum = { version = "0.24.1", features = ["derive"] }
log = "0.4.17"

10
backend/src/lib.rs
View File

@@ -107,15 +107,7 @@ impl BackendType {
Box::new(composite::CompositeBackendFactory::new(plonky3::Factory))
}
#[cfg(feature = "stwo")]
BackendType::Stwo => Box::new(stwo::StwoProverFactory),
#[cfg(not(any(
feature = "halo2",
feature = "estark-polygon",
feature = "estark-starky",
feature = "plonky3",
feature = "stwo"
)))]
_ => panic!("Empty backend."),
BackendType::Stwo => Box::new(stwo::Factory),
}
}
}

193
backend/src/stwo/circuit_builder.rs Normal file
View File

@@ -0,0 +1,193 @@
use num_traits::Zero;
use std::fmt::Debug;
use std::ops::{Add, AddAssign, Mul, Neg, Sub};
extern crate alloc;
use alloc::{collections::btree_map::BTreeMap, string::String, vec::Vec};
use powdr_ast::analyzed::{
AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression, Analyzed, Identity,
};
use powdr_number::{FieldElement, LargeInt};
use std::sync::Arc;
use powdr_ast::analyzed::{
AlgebraicUnaryOperation, AlgebraicUnaryOperator, PolyID, PolynomialType,
};
use stwo_prover::constraint_framework::{EvalAtRow, FrameworkComponent, FrameworkEval};
use stwo_prover::core::backend::ColumnOps;
use stwo_prover::core::fields::m31::{BaseField, M31};
use stwo_prover::core::fields::{ExtensionOf, FieldExpOps, FieldOps};
use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation};
use stwo_prover::core::poly::BitReversedOrder;
use stwo_prover::core::ColumnVec;
pub type PowdrComponent<'a, F> = FrameworkComponent<PowdrEval<F>>;
pub(crate) fn gen_stwo_circuit_trace<T, B, F>(
witness: &[(String, Vec<T>)],
) -> ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>
where
T: FieldElement, //only Merenne31Field is supported, checked in runtime
B: FieldOps<M31> + ColumnOps<F>, // Ensure B implements FieldOps for M31
F: ExtensionOf<BaseField>,
{
assert!(
witness
.iter()
.all(|(_name, vec)| vec.len() == witness[0].1.len()),
"All Vec<T> in witness must have the same length. Mismatch found!"
);
let domain = CanonicCoset::new(witness[0].1.len().ilog2()).circle_domain();
witness
.iter()
.map(|(_name, values)| {
let values = values
.iter()
.map(|v| v.try_into_i32().unwrap().into())
.collect();
CircleEvaluation::new(domain, values)
})
.collect()
}
pub struct PowdrEval<T> {
analyzed: Arc<Analyzed<T>>,
witness_columns: BTreeMap<PolyID, usize>,
}
impl<T: FieldElement> PowdrEval<T> {
pub fn new(analyzed: Arc<Analyzed<T>>) -> Self {
let witness_columns: BTreeMap<PolyID, usize> = analyzed
.definitions_in_source_order(PolynomialType::Committed)
.flat_map(|(symbol, _)| symbol.array_elements())
.enumerate()
.map(|(index, (_, id))| (id, index))
.collect();
Self {
analyzed,
witness_columns,
}
}
}
impl<T: FieldElement> FrameworkEval for PowdrEval<T> {
fn log_size(&self) -> u32 {
self.analyzed.degree().ilog2()
}
fn max_constraint_log_degree_bound(&self) -> u32 {
self.analyzed.degree().ilog2() + 1
}
fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
assert!(
self.analyzed.constant_count() == 0 && self.analyzed.publics_count() == 0,
"Error: Expected no fixed columns nor public inputs, as they are not supported yet.",
);
let witness_eval: BTreeMap<PolyID, [<E as EvalAtRow>::F; 2]> = self
.witness_columns
.keys()
.map(|poly_id| (*poly_id, eval.next_interaction_mask(0, [0, 1])))
.collect();
for id in self
.analyzed
.identities_with_inlined_intermediate_polynomials()
{
match id {
Identity::Polynomial(identity) => {
let expr = to_stwo_expression(&identity.expression, &witness_eval);
eval.add_constraint(expr);
}
Identity::Connect(..) => {
unimplemented!("Connect is not implemented in stwo yet")
}
Identity::Lookup(..) => {
unimplemented!("Lookup is not implemented in stwo yet")
}
Identity::Permutation(..) => {
unimplemented!("Permutation is not implemented in stwo yet")
}
Identity::PhantomPermutation(..) => {}
Identity::PhantomLookup(..) => {}
}
}
eval
}
}
fn to_stwo_expression<T: FieldElement, F>(
expr: &AlgebraicExpression<T>,
witness_eval: &BTreeMap<PolyID, [F; 2]>,
) -> F
where
F: FieldExpOps
+ Clone
+ Debug
+ Zero
+ Neg<Output = F>
+ AddAssign
+ AddAssign<BaseField>
+ Add<F, Output = F>
+ Sub<F, Output = F>
+ Mul<BaseField, Output = F>
+ Neg<Output = F>
+ From<BaseField>,
{
use AlgebraicBinaryOperator::*;
match expr {
AlgebraicExpression::Reference(r) => {
let poly_id = r.poly_id;
match poly_id.ptype {
PolynomialType::Committed => match r.next {
false => witness_eval[&poly_id][0].clone(),
true => witness_eval[&poly_id][1].clone(),
},
PolynomialType::Constant => {
unimplemented!("Constant polynomials are not supported in stwo yet")
}
PolynomialType::Intermediate => {
unimplemented!("Intermediate polynomials are not supported in stwo yet")
}
}
}
AlgebraicExpression::PublicReference(..) => {
unimplemented!("Public references are not supported in stwo yet")
}
AlgebraicExpression::Number(n) => F::from(M31::from(n.try_into_i32().unwrap())),
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
left,
op: Pow,
right,
}) => match **right {
AlgebraicExpression::Number(n) => {
let left = to_stwo_expression(left, witness_eval);
(0u32..n.to_integer().try_into_u32().unwrap())
.fold(F::one(), |acc, _| acc * left.clone())
}
_ => unimplemented!("pow with non-constant exponent"),
},
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) => {
let left = to_stwo_expression(left, witness_eval);
let right = to_stwo_expression(right, witness_eval);
match op {
Add => left + right,
Sub => left - right,
Mul => left * right,
Pow => unreachable!("This case was handled above"),
}
}
AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op, expr }) => {
let expr = to_stwo_expression(expr, witness_eval);
match op {
AlgebraicUnaryOperator::Minus => -expr,
}
}
AlgebraicExpression::Challenge(_challenge) => {
unimplemented!("challenges are not supported in stwo yet")
}
}
}

40
backend/src/stwo/mod.rs
View File

@@ -1,20 +1,28 @@
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::io;
use std::path::PathBuf;
use std::sync::Arc;
use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
use crate::{
field_filter::generalize_factory, Backend, BackendFactory, BackendOptions, Error, Proof,
};
use powdr_ast::analyzed::Analyzed;
use powdr_executor::constant_evaluator::{get_uniquely_sized_cloned, VariablySizedColumn};
use powdr_executor::witgen::WitgenCallback;
use powdr_number::FieldElement;
use powdr_number::{FieldElement, Mersenne31Field};
use prover::StwoProver;
use stwo_prover::core::backend::{simd::SimdBackend, BackendForChannel};
use stwo_prover::core::channel::{Blake2sChannel, Channel, MerkleChannel};
use stwo_prover::core::vcs::blake2_merkle::Blake2sMerkleChannel;
mod circuit_builder;
mod prover;
#[allow(dead_code)]
pub(crate) struct StwoProverFactory;
impl<F: FieldElement> BackendFactory<F> for StwoProverFactory {
struct RestrictedFactory;
impl<F: FieldElement> BackendFactory<F> for RestrictedFactory {
#[allow(unreachable_code)]
#[allow(unused_variables)]
fn create(
@@ -37,16 +45,28 @@ impl<F: FieldElement> BackendFactory<F> for StwoProverFactory {
let fixed = Arc::new(
get_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
);
let stwo = Box::new(StwoProver::new(pil, fixed, setup)?);
let stwo: Box<StwoProver<F, SimdBackend, Blake2sMerkleChannel, Blake2sChannel>> =
Box::new(StwoProver::new(pil, fixed)?);
Ok(stwo)
}
}
impl<T: FieldElement> Backend<T> for StwoProver<T> {
generalize_factory!(Factory <- RestrictedFactory, [Mersenne31Field]);
impl<T: FieldElement, MC: MerkleChannel + Send, C: Channel + Send> Backend<T>
for StwoProver<T, SimdBackend, MC, C>
where
SimdBackend: BackendForChannel<MC>,
MC: MerkleChannel,
C: Channel,
MC::H: DeserializeOwned + Serialize,
{
#[allow(unused_variables)]
fn verify(&self, proof: &[u8], instances: &[Vec<T>]) -> Result<(), Error> {
assert!(instances.len() == 1);
unimplemented!()
assert_eq!(instances.len(), 1);
let instances = &instances[0];
Ok(self.verify(proof, instances)?)
}
#[allow(unreachable_code)]
#[allow(unused_variables)]
@@ -59,7 +79,7 @@ impl<T: FieldElement> Backend<T> for StwoProver<T> {
if prev_proof.is_some() {
return Err(Error::NoAggregationAvailable);
}
unimplemented!()
Ok(StwoProver::prove(self, witness)?)
}
#[allow(unused_variables)]
fn export_verification_key(&self, output: &mut dyn io::Write) -> Result<(), Error> {

135
backend/src/stwo/prover.rs
View File

@@ -1,32 +1,145 @@
use powdr_ast::analyzed::Analyzed;
use serde::de::DeserializeOwned;
use serde::ser::Serialize;
use std::io;
use std::marker::PhantomData;
use std::sync::Arc;
use powdr_number::FieldElement;
use crate::stwo::circuit_builder::{gen_stwo_circuit_trace, PowdrComponent, PowdrEval};
#[allow(unused_variables)]
pub struct StwoProver<F> {
_analyzed: Arc<Analyzed<F>>,
_fixed: Arc<Vec<(String, Vec<F>)>>,
use stwo_prover::constraint_framework::TraceLocationAllocator;
use stwo_prover::core::prover::StarkProof;
use powdr_number::FieldElement;
use stwo_prover::core::air::{Component, ComponentProver};
use stwo_prover::core::backend::{Backend, BackendForChannel};
use stwo_prover::core::channel::{Channel, MerkleChannel};
use stwo_prover::core::fields::m31::M31;
use stwo_prover::core::fri::FriConfig;
use stwo_prover::core::pcs::{CommitmentSchemeProver, CommitmentSchemeVerifier, PcsConfig};
use stwo_prover::core::poly::circle::CanonicCoset;
const FRI_LOG_BLOWUP: usize = 1;
const FRI_NUM_QUERIES: usize = 100;
const FRI_PROOF_OF_WORK_BITS: usize = 16;
const LOG_LAST_LAYER_DEGREE_BOUND: usize = 0;
pub struct StwoProver<T, B: Backend + Send, MC: MerkleChannel, C: Channel> {
pub analyzed: Arc<Analyzed<T>>,
_fixed: Arc<Vec<(String, Vec<T>)>>,
/// Proving key placeholder
_proving_key: Option<()>,
/// Verifying key placeholder
_verifying_key: Option<()>,
_channel_marker: PhantomData<C>,
_backend_marker: PhantomData<B>,
_merkle_channel_marker: PhantomData<MC>,
}
impl<F: FieldElement> StwoProver<F> {
#[allow(dead_code)]
#[allow(unused_variables)]
impl<'a, F: FieldElement, B, MC, C> StwoProver<F, B, MC, C>
where
B: Backend + Send + BackendForChannel<MC>, // Ensure B implements BackendForChannel<MC>
MC: MerkleChannel + Send,
C: Channel + Send,
MC::H: DeserializeOwned + Serialize,
PowdrComponent<'a, F>: ComponentProver<B>,
{
pub fn new(
_analyzed: Arc<Analyzed<F>>,
analyzed: Arc<Analyzed<F>>,
_fixed: Arc<Vec<(String, Vec<F>)>>,
setup: Option<&mut dyn io::Read>,
) -> Result<Self, io::Error> {
Ok(Self {
_analyzed,
analyzed,
_fixed,
_proving_key: None,
_verifying_key: None,
_channel_marker: PhantomData,
_backend_marker: PhantomData,
_merkle_channel_marker: PhantomData,
})
}
pub fn prove(&self, witness: &[(String, Vec<F>)]) -> Result<Vec<u8>, String> {
let config = get_config();
// twiddles are used for FFT, they are computed in a bigger group than the eval domain.
// the eval domain is the half coset G_{2n} + <G_{n/2}>
// twiddles are computed in half coset G_{4n} + <G_{n}>, double the size of eval doamin.
let twiddles = B::precompute_twiddles(
CanonicCoset::new(self.analyzed.degree().ilog2() + 1 + FRI_LOG_BLOWUP as u32)
.circle_domain()
.half_coset,
);
// Setup protocol.
let mut prover_channel = <MC as MerkleChannel>::C::default();
let commitment_scheme = &mut CommitmentSchemeProver::<B, MC>::new(config, &twiddles);
let trace = gen_stwo_circuit_trace::<F, B, M31>(witness);
let mut tree_builder = commitment_scheme.tree_builder();
tree_builder.extend_evals(trace);
tree_builder.commit(&mut prover_channel);
let component = PowdrComponent::new(
&mut TraceLocationAllocator::default(),
PowdrEval::new(self.analyzed.clone()),
);
let proof = stwo_prover::core::prover::prove::<B, MC>(
&[&component],
&mut prover_channel,
commitment_scheme,
)
.unwrap();
Ok(bincode::serialize(&proof).unwrap())
}
pub fn verify(&self, proof: &[u8], _instances: &[F]) -> Result<(), String> {
assert!(
_instances.is_empty(),
"Expected _instances slice to be empty, but it has {} elements.",
_instances.len()
);
let config = get_config();
let proof: StarkProof<MC::H> =
bincode::deserialize(proof).map_err(|e| format!("Failed to deserialize proof: {e}"))?;
let mut verifier_channel = <MC as MerkleChannel>::C::default();
let mut commitment_scheme = CommitmentSchemeVerifier::<MC>::new(config);
//Constraints that are to be proved
let component = PowdrComponent::new(
&mut TraceLocationAllocator::default(),
PowdrEval::new(self.analyzed.clone()),
);
// Retrieve the expected column sizes in each commitment interaction, from the AIR.
// TODO: When constant columns are supported, there will be more than one sizes and proof.commitments
// size[0] is for constant columns, size[1] is for witness columns, size[2] is for lookup columns
// pass size[1] for witness columns now is not doable due to this branch is outdated for the new feature of constant columns
// it will throw errors.
let sizes = component.trace_log_degree_bounds();
assert_eq!(sizes.len(), 1);
commitment_scheme.commit(proof.commitments[0], &sizes[0], &mut verifier_channel);
stwo_prover::core::prover::verify(
&[&component],
&mut verifier_channel,
&mut commitment_scheme,
proof,
)
.map_err(|e| e.to_string())
}
}
fn get_config() -> PcsConfig {
PcsConfig {
pow_bits: FRI_PROOF_OF_WORK_BITS as u32,
fri_config: FriConfig::new(
LOG_LAST_LAYER_DEGREE_BOUND as u32,
FRI_LOG_BLOWUP as u32,
FRI_NUM_QUERIES,
),
}
}

26
pipeline/src/test_util.rs
View File

@@ -541,3 +541,29 @@ pub fn run_reparse_test_with_blacklist(file: &str, blacklist: &[&str]) {
.compute_optimized_pil()
.unwrap();
}
#[cfg(feature = "stwo")]
use powdr_number::Mersenne31Field;
#[cfg(feature = "stwo")]
pub fn test_stwo(file_name: &str, inputs: Vec<Mersenne31Field>) {
let backend = powdr_backend::BackendType::Stwo;
let mut pipeline = Pipeline::default()
.with_tmp_output()
.from_file(resolve_test_file(file_name))
.with_prover_inputs(inputs)
.with_backend(backend, None);
let proof = pipeline.compute_proof().cloned().unwrap();
let publics: Vec<Mersenne31Field> = pipeline
.publics()
.clone()
.unwrap()
.iter()
.map(|(_name, v)| v.expect("all publics should be known since we created a proof"))
.collect();
pipeline.verify(&proof, &[publics]).unwrap();
}
#[cfg(not(feature = "stwo"))]
pub fn test_stwo(_file_name: &str, _inputs: Vec<u32>) {}

7
pipeline/tests/pil.rs
View File

@@ -8,7 +8,7 @@ use powdr_pipeline::{
assert_proofs_fail_for_invalid_witnesses_pilcom, gen_estark_proof,
gen_estark_proof_with_backend_variant, make_prepared_pipeline,
make_simple_prepared_pipeline, regular_test, run_pilcom_with_backend_variant, test_halo2,
test_halo2_with_backend_variant, test_pilcom, test_plonky3_with_backend_variant,
test_halo2_with_backend_variant, test_pilcom, test_plonky3_with_backend_variant, test_stwo,
BackendVariant,
},
Pipeline,
@@ -264,6 +264,11 @@ fn add() {
);
}
#[test]
fn stwo_add_and_equal() {
let f = "pil/add_and_equal.pil";
test_stwo(f, Default::default());
}
#[test]
fn simple_div() {
let f = "pil/simple_div.pil";

10
test_data/pil/add_and_equal.pil Normal file
View File

@@ -0,0 +1,10 @@
namespace Add(4);
let a;
std::prelude::set_hint(a, |i| if i==0||i==2 {std::prelude::Query::Hint(2147483646)}else{std::prelude::Query::Hint(1)});
let b;
let c;
b = c;
a+a'=0;