From d61307dc188a7cb9441115d7df7a589924958a3c Mon Sep 17 00:00:00 2001 From: Sai Deng Date: Wed, 31 Jan 2024 15:16:59 -0800 Subject: [PATCH] add mul air test --- starky/Cargo.toml | 3 + starky/src/lib.rs | 1 + starky/src/mul_air.rs | 141 ++++++++++++++++++++++++++++++++++++++++++ starky/src/proof.rs | 7 ++- 4 files changed, 149 insertions(+), 3 deletions(-) create mode 100644 starky/src/mul_air.rs diff --git a/starky/Cargo.toml b/starky/Cargo.toml index 645f03035b..8423f99b74 100644 --- a/starky/Cargo.toml +++ b/starky/Cargo.toml @@ -22,9 +22,12 @@ itertools = { version = "0.11.0", default-features = false } log = { version = "0.4.14", default-features = false } plonky2_maybe_rayon = { path = "../maybe_rayon", default-features = false } plonky2 = { path = "../plonky2", default-features = false } +rand = "0.8.5" +serde = { version = "1.0.193", features = ["derive"] } [dev-dependencies] env_logger = { version = "0.10.0", default-features = false } +postcard = { version = "1.0.0", default-features = false, features = ["alloc"] } # Display math equations properly in documentation [package.metadata.docs.rs] diff --git a/starky/src/lib.rs b/starky/src/lib.rs index 635e57bd0b..1c078fa14e 100644 --- a/starky/src/lib.rs +++ b/starky/src/lib.rs @@ -21,3 +21,4 @@ pub mod verifier; #[cfg(test)] pub mod fibonacci_stark; +pub mod mul_air; diff --git a/starky/src/mul_air.rs b/starky/src/mul_air.rs new file mode 100644 index 0000000000..0cc65b8ac7 --- /dev/null +++ b/starky/src/mul_air.rs @@ -0,0 +1,141 @@ +use alloc::vec::Vec; +use core::marker::PhantomData; + +use log::debug; +use plonky2::field::extension::{Extendable, FieldExtension}; +use plonky2::field::packed::PackedField; +use plonky2::field::polynomial::PolynomialValues; +use plonky2::hash::hash_types::RichField; +use plonky2::iop::ext_target::ExtensionTarget; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use rand::{thread_rng, Rng}; + +use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; +use crate::evaluation_frame::{StarkEvaluationFrame, StarkFrame}; +use crate::stark::Stark; +use crate::util::trace_rows_to_poly_values; + +/// How many `a * b = c` operations to do per row in the AIR. +const REPETITIONS: usize = 911; +const TRACE_WIDTH: usize = REPETITIONS * 3; + +#[derive(Copy, Clone)] +struct MulAirStark, const D: usize> { + num_rows: usize, + _phantom: PhantomData, +} + +impl, const D: usize> MulAirStark { + const fn new(num_rows: usize) -> Self { + Self { + num_rows, + _phantom: PhantomData, + } + } + + fn generate_trace(&self) -> Vec> { + let mut rng = thread_rng(); + let mut trace_rows = Vec::with_capacity(self.num_rows); + for _ in 0..self.num_rows { + let mut row: [F; REPETITIONS * 3] = [F::ZERO; REPETITIONS * 3]; + for i in 0..REPETITIONS { + row[i * 3] = F::from_canonical_u64(rng.gen::()); + row[i * 3 + 1] = F::from_canonical_u64(rng.gen::()); + row[i * 3 + 2] = row[i * 3] * row[i * 3 + 1]; + } + trace_rows.push(row); + } + debug!( + "trace height: {}, trace width: {}", + trace_rows.len(), + trace_rows[0].len() + ); + + trace_rows_to_poly_values(trace_rows) + } +} + +const COLUMNS: usize = TRACE_WIDTH; +const PUBLIC_INPUTS: usize = 0; + +impl, const D: usize> Stark for MulAirStark { + type EvaluationFrame = StarkFrame + where + FE: FieldExtension, + P: PackedField; + + type EvaluationFrameTarget = + StarkFrame, ExtensionTarget, COLUMNS, PUBLIC_INPUTS>; + + fn eval_packed_generic( + &self, + vars: &Self::EvaluationFrame, + yield_constr: &mut ConstraintConsumer

, + ) where + FE: FieldExtension, + P: PackedField, + { + let local_values = vars.get_local_values(); + + for i in 0..REPETITIONS { + yield_constr.constraint_transition( + local_values[i * 3 + 2] - local_values[i * 3] * local_values[i * 3 + 1], + ); + } + } + + fn eval_ext_circuit( + &self, + _builder: &mut CircuitBuilder, + _vars: &Self::EvaluationFrameTarget, + _yield_constr: &mut RecursiveConstraintConsumer, + ) { + unimplemented!() + } + + fn constraint_degree(&self) -> usize { + 2 + } +} + +#[cfg(test)] +mod tests { + use anyhow::Result; + use log::{debug, Level}; + use plonky2::plonk::config::{GenericConfig, Poseidon2GoldilocksConfig}; + use plonky2::util::timing::TimingTree; + + use crate::config::StarkConfig; + use crate::mul_air::MulAirStark; + use crate::prover::prove; + use crate::verifier::verify_stark_proof; + + #[test] + fn test_mulair_stark() -> Result<()> { + init_logger(); + + const D: usize = 2; + type C = Poseidon2GoldilocksConfig; + type F = >::F; + type S = MulAirStark; + + let mut config = StarkConfig::standard_fast_config(); + config.fri_config.num_query_rounds = 100; + let num_rows = 1 << 14; + let public_inputs = []; + let stark = S::new(num_rows); + let trace = stark.generate_trace(); + let mut timing = TimingTree::new("prove mul air stark", Level::Debug); + let proof = prove::(stark, &config, trace, &public_inputs, &mut timing)?; + timing.print(); + let serialized_proof = postcard::to_allocvec(&proof).expect("unable to serialize proof"); + + debug!("serialized_proof len: {} bytes", serialized_proof.len()); + + verify_stark_proof(stark, proof, &config) + } + + fn init_logger() { + let _ = env_logger::builder().format_timestamp(None).try_init(); + } +} diff --git a/starky/src/proof.rs b/starky/src/proof.rs index 1a04772d86..bc24f2c3a5 100644 --- a/starky/src/proof.rs +++ b/starky/src/proof.rs @@ -16,11 +16,12 @@ use plonky2::iop::ext_target::ExtensionTarget; use plonky2::iop::target::Target; use plonky2::plonk::config::GenericConfig; use plonky2_maybe_rayon::*; +use serde::Serialize; use crate::config::StarkConfig; use crate::permutation::PermutationChallengeSet; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize)] pub struct StarkProof, C: GenericConfig, const D: usize> { /// Merkle cap of LDEs of trace values. pub trace_cap: MerkleCap, @@ -66,7 +67,7 @@ impl StarkProofTarget { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize)] pub struct StarkProofWithPublicInputs< F: RichField + Extendable, C: GenericConfig, @@ -125,7 +126,7 @@ pub(crate) struct StarkProofChallengesTarget { } /// Purported values of each polynomial at the challenge point. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize)] pub struct StarkOpeningSet, const D: usize> { pub local_values: Vec, pub next_values: Vec,