github/powdr
1
1
Fork 0
mirror of https://github.com/powdr-labs/powdr.git synced 2026-04-20 03:03:25 -04:00
Code Issues Packages Projects Releases Wiki Activity

Plonky3 circuit builder: Cache computations on Analyzed PIL (#1754)

Browse Source 
On the Keccak test (and with the `RUSTFLAGS='-Ctarget-cpu=native'`
flag), this brings down the proof time from 7.2 to 4.4s. The share for
computing the quotient polynomial drops from > 50% for some machines to
< 10% for all machines.
This commit is contained in:
Georg Wiese
2024-09-05 14:46:23 +02:00
committed by GitHub
parent 7fe5fb4d02
commit 80170a6e43
1 changed files with 57 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
plonky3/src/circuit_builder.rs
View File

@@ -6,7 +6,10 @@
//! everywhere save for at row j is constructed to constrain s * (pub - x) on
//! every row.
use std::{any::TypeId, collections::BTreeMap};
use std::{
any::TypeId,
collections::{BTreeMap, HashSet},
};
use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder};
use p3_field::AbstractField;
@@ -14,16 +17,46 @@ use p3_goldilocks::Goldilocks;
use p3_matrix::{dense::RowMajorMatrix, Matrix};
use powdr_ast::analyzed::{
AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression,
AlgebraicUnaryOperation, AlgebraicUnaryOperator, Analyzed, IdentityKind, PolynomialType,
AlgebraicUnaryOperation, AlgebraicUnaryOperator, Analyzed, Identity, IdentityKind,
PolynomialType, SelectedExpressions,
};
use powdr_executor::witgen::WitgenCallback;
use powdr_number::{FieldElement, GoldilocksField, LargeInt};
use powdr_number::{DegreeType, FieldElement, GoldilocksField, LargeInt};
pub type Val = p3_goldilocks::Goldilocks;
/// A description of the constraint system.
/// All of the data is derived from the analyzed PIL, but is materialized
/// here for performance reasons.
struct ConstraintSystem<T> {
identities: Vec<Identity<SelectedExpressions<AlgebraicExpression<T>>>>,
publics: Vec<(String, usize, usize)>,
commitment_count: usize,
constant_count: usize,
degrees: HashSet<DegreeType>,
}
impl<T: FieldElement> From<&Analyzed<T>> for ConstraintSystem<T> {
fn from(analyzed: &Analyzed<T>) -> Self {
let identities = analyzed.identities_with_inlined_intermediate_polynomials();
let publics = analyzed.get_publics();
let commitment_count = analyzed.commitment_count();
let constant_count = analyzed.constant_count();
let degrees = analyzed.degrees();
Self {
identities,
publics,
commitment_count,
constant_count,
degrees,
}
}
}
pub(crate) struct PowdrCircuit<'a, T> {
/// The analyzed PIL
analyzed: &'a Analyzed<T>,
/// The constraint system description
constraint_system: ConstraintSystem<T>,
/// The value of the witness columns, if set
witness: Option<&'a [(String, Vec<T>)]>,
/// Callback to augment the witness in the later stages
@@ -37,7 +70,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
pub fn generate_trace_rows(&self) -> RowMajorMatrix<Goldilocks> {
// an iterator over all columns, committed then fixed
let witness = self.witness().iter();
let degrees = self.analyzed.degrees();
let degrees = &self.constraint_system.degrees;
let values = match degrees.len() {
1 => {
@@ -78,7 +111,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
}
Self {
analyzed,
constraint_system: analyzed.into(),
witness: None,
_witgen_callback: None,
#[cfg(debug_assertions)]
@@ -100,8 +133,8 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
.map(|(name, values)| (name, values))
.collect::<BTreeMap<_, _>>();
self.analyzed
.get_publics()
self.constraint_system
.publics
.iter()
.map(|(col_name, _, idx)| {
let vals = *witness.get(&col_name).unwrap();
@@ -111,7 +144,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
}
pub(crate) fn with_witness(self, witness: &'a [(String, Vec<T>)]) -> Self {
assert_eq!(witness.len(), self.analyzed.commitment_count());
assert_eq!(witness.len(), self.constraint_system.commitment_count);
Self {
witness: Some(witness),
..self
@@ -149,7 +182,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
match poly_id.ptype {
PolynomialType::Committed => {
assert!(
r.poly_id.id < self.analyzed.commitment_count() as u64,
r.poly_id.id < self.constraint_system.commitment_count as u64,
"Plonky3 expects `poly_id` to be contiguous"
);
let row = main.row_slice(r.next as usize);
@@ -157,7 +190,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
}
PolynomialType::Constant => {
assert!(
r.poly_id.id < self.analyzed.constant_count() as u64,
r.poly_id.id < self.constraint_system.constant_count as u64,
"Plonky3 expects `poly_id` to be contiguous"
);
let row = fixed.row_slice(r.next as usize);
@@ -206,11 +239,11 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
impl<'a, T: FieldElement> BaseAir<Val> for PowdrCircuit<'a, T> {
fn width(&self) -> usize {
self.analyzed.commitment_count()
self.constraint_system.commitment_count
}
fn preprocessed_width(&self) -> usize {
self.analyzed.constant_count() + self.analyzed.publics_count()
self.constraint_system.constant_count + self.constraint_system.publics.len()
}
fn preprocessed_trace(&self) -> Option<RowMajorMatrix<Val>> {
@@ -230,38 +263,35 @@ impl<'a, T: FieldElement, AB: AirBuilderWithPublicValues<F = Val> + PairBuilder>
let main = builder.main();
let fixed = builder.preprocessed();
let pi = builder.public_values();
let publics = self.analyzed.get_publics();
assert_eq!(publics.len(), pi.len());
assert_eq!(self.constraint_system.publics.len(), pi.len());
let local = main.row_slice(0);
// public constraints
let public_vals_by_id = publics
let public_vals_by_id = self
.constraint_system
.publics
.iter()
.zip(pi.to_vec())
.map(|((id, _, _), val)| (id, val))
.collect::<BTreeMap<&String, <AB as AirBuilderWithPublicValues>::PublicVar>>();
let fixed_local = fixed.row_slice(0);
let public_offset = self.analyzed.constant_count();
let public_offset = self.constraint_system.constant_count;
publics
.iter()
.enumerate()
.for_each(|(index, (pub_id, col_id, _))| {
self.constraint_system.publics.iter().enumerate().for_each(
|(index, (pub_id, col_id, _))| {
let selector = fixed_local[public_offset + index];
let witness_col = local[*col_id];
let public_value = public_vals_by_id[pub_id];
// constraining s(i) * (pub[i] - x(i)) = 0
builder.assert_zero(selector * (public_value.into() - witness_col));
});
},
);
// circuit constraints
for identity in &self
.analyzed
.identities_with_inlined_intermediate_polynomials()
{
for identity in &self.constraint_system.identities {
match identity.kind {
IdentityKind::Polynomial => {
assert_eq!(identity.left.expressions.len(), 0);