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

Linear constraint remover (#2425)

Browse Source 
# what

Replace certain multilinear constraints* with intermediate polynomials
in pilopt.
*with no next references, not touching public inputs

# why

[This
comment](https://github.com/powdr-labs/powdr/issues/2337#issuecomment-2621634278)
highlights a use-case where materialising a column which is only
constrained to a non-shifted multilinear polynomial is wasteful. Fixing
this when we create the code is tricky, since we cannot know the degree
of the constraint from pil. This is however something pilopt can do,
which is what we attempt to do here

# unknowns

- [x] We have similar optimizers for `x = 42` and `x = y` which are both
multilinear expressions. They do no yield intermediate polynomials,
since in that case it is fine to "inline" them. This implementation
currently makes sure all three optimizer apply to distinct cases, but
maybe it would be better to have a single optimizer which handles all
three cases, only introducing intermediate polynomials in the non
trivial cases (`42` and `y`)
      - Implemented in https://github.com/powdr-labs/powdr/pull/2532

---------

Co-authored-by: chriseth <chris@ethereum.org>
Co-authored-by: Gaston Zanitti <gzanitti@gmail.com>
This commit is contained in:
Thibaut Schaeffer
2025-04-10 18:56:30 +02:00
committed by GitHub
parent edaad20fd5
commit 2cb4548283
24 changed files with 536 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
ast/src/analyzed/mod.rs
View File

@@ -18,7 +18,7 @@ use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use crate::parsed::types::{ArrayType, Type, TypeBounds, TypeScheme};
use crate::parsed::visitor::{Children, ExpressionVisitable};
use crate::parsed::visitor::{AllChildren, Children, ExpressionVisitable};
pub use crate::parsed::BinaryOperator;
pub use crate::parsed::UnaryOperator;
use crate::parsed::{
@@ -87,6 +87,16 @@ impl<T> Analyzed<T> {
.collect::<HashSet<_>>()
}
/// Returns the set of all referenced challenges in this [`Analyzed<T>`].
pub fn challenges(&self) -> BTreeSet<&Challenge> {
self.all_children()
.filter_map(|expr| match expr {
AlgebraicExpression::Challenge(challenge) => Some(challenge),
_ => None,
})
.collect()
}
/// Returns the number of stages based on the maximum stage number of all definitions
pub fn stage_count(&self) -> usize {
self.definitions
@@ -97,6 +107,30 @@ impl<T> Analyzed<T> {
+ 1
}
/// Returns the next available ID for a polynomial of the specified type.
pub fn next_id_for_kind(&self, poly_type: PolynomialType) -> u64 {
match poly_type {
PolynomialType::Committed => self
.committed_polys_in_source_order()
.flat_map(|(s, _)| s.array_elements())
.map(|(_, poly)| poly.id)
.max(),
PolynomialType::Constant => self
.constant_polys_in_source_order()
.flat_map(|(s, _)| s.array_elements())
.map(|(_, poly)| poly.id)
.max(),
PolynomialType::Intermediate => self
.intermediate_columns
.values()
.flat_map(|(s, _)| s.array_elements())
.map(|(_, poly_id)| poly_id.id)
.max(),
}
.unwrap_or(0)
+ 1
}
/// @returns the number of committed polynomials (with multiplicities for arrays)
pub fn commitment_count(&self) -> usize {
self.declaration_type_count(PolynomialType::Committed)
@@ -724,6 +758,7 @@ impl DegreeRange {
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Hash)]
pub struct Symbol {
/// Unique ID of the symbol within the set of symbols of this kind.
pub id: u64,
pub source: SourceRef,
pub absolute_name: String,

28
backend/src/mock/mod.rs
View File

@@ -8,12 +8,10 @@ use std::{
use bus_checker::{BusChecker, BusInteraction};
use connection_constraint_checker::{Connection, ConnectionConstraintChecker};
use itertools::Itertools;
use machine::Machine;
use polynomial_constraint_checker::PolynomialConstraintChecker;
use powdr_ast::{
analyzed::{AlgebraicExpression, Analyzed},
parsed::visitor::AllChildren,
};
use powdr_ast::analyzed::Analyzed;
use powdr_executor::{constant_evaluator::VariablySizedColumn, witgen::WitgenCallback};
use powdr_number::{DegreeType, FieldElement};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
@@ -87,18 +85,16 @@ impl<F: FieldElement> Backend<F> for MockBackend<F> {
let challenges = self
.machine_to_pil
.values()
.flat_map(|pil| pil.identities.iter())
.flat_map(|identity| identity.all_children())
.filter_map(|expr| match expr {
AlgebraicExpression::Challenge(challenge) => {
// Use the hash of the ID as the challenge.
// This way, if the same challenge is used by different machines, they will
// have the same value.
let mut hasher = DefaultHasher::new();
challenge.id.hash(&mut hasher);
Some((challenge.id, F::from(hasher.finish())))
}
_ => None,
.flat_map(|pil| pil.challenges())
// deduplicate challenges across machines
.unique()
.map(|challenge| {
// Use the hash of the ID as the challenge.
// This way, if the same challenge is used by different machines, they will
// have the same value.
let mut hasher = DefaultHasher::new();
challenge.id.hash(&mut hasher);
(challenge.id, F::from(hasher.finish()))
})
.collect::<BTreeMap<_, _>>();

8
backend/src/plonky3/stark.rs
View File

@@ -425,13 +425,13 @@ mod tests {
}
#[test]
fn add() {
fn mul() {
let content = r#"
namespace Add(8);
col witness x;
col witness y;
col witness z;
x + y = z;
x * y = z;
"#;
run_test(content);
}
@@ -468,7 +468,7 @@ mod tests {
let beta: expr = std::prelude::challenge(0, 42);
col witness x;
col witness stage(1) y;
x = y + beta * alpha;
y = x * x + beta * alpha;
"#;
run_test(content);
}
@@ -492,7 +492,7 @@ mod tests {
#[test]
fn polynomial_identity() {
let content = "namespace Global(8); pol fixed z = [1, 2]*; pol witness a; a = z + 1;";
let content = "namespace Global(8); pol fixed z = [1, 2]*; pol witness a; a = z * z + 1;";
run_test(content);
}

74
backend/src/stwo/circuit_builder.rs
View File

@@ -5,8 +5,10 @@ use std::collections::HashSet;
extern crate alloc;
use crate::stwo::prover::into_stwo_field;
use alloc::collections::btree_map::BTreeMap;
use powdr_ast::analyzed::{AlgebraicExpression, AlgebraicReference, Analyzed, Challenge, Identity};
use alloc::collections::{btree_map::BTreeMap, btree_set::BTreeSet};
use powdr_ast::analyzed::{
AlgebraicExpression, AlgebraicReference, AlgebraicReferenceThin, Analyzed, Challenge, Identity,
};
use powdr_ast::analyzed::{PolyID, PolynomialType};
use powdr_number::Mersenne31Field as M31;
use stwo_prover::constraint_framework::preprocessed_columns::PreProcessedColumnId;
@@ -310,21 +312,55 @@ impl FrameworkEval for PowdrEval {
}
}
// This function creates a list of the names of the constant polynomials that have next references constraint
pub fn get_constant_with_next_list(analyzed: &Analyzed<M31>) -> HashSet<&String> {
let mut constant_with_next_list: HashSet<&String> = HashSet::new();
analyzed.all_children().for_each(|e| {
if let AlgebraicExpression::Reference(AlgebraicReference {
name,
poly_id,
next,
}) = e
{
if matches!(poly_id.ptype, PolynomialType::Constant) && *next {
// add the name of the constant polynomial to the list
constant_with_next_list.insert(name);
}
};
});
constant_with_next_list
/// This function creates a list of the names of the constant polynomials that have next references
/// Note that the analysis should also dereference next references to intermediate polynomials
pub fn get_constant_with_next_list(analyzed: &Analyzed<M31>) -> HashSet<String> {
let intermediate_definitions = analyzed.intermediate_definitions();
let cache = &mut BTreeMap::new();
// get all constant references
get_constant_refs_with_intermediates(analyzed, &intermediate_definitions, cache)
.into_iter()
// only keep the names of next references
.filter_map(|r| r.next.then_some(r.name))
.collect()
}
/// Auxiliary function to collect all references and next references to constant polynomials
/// References to intermediate values are resolved recursively
fn get_constant_refs_with_intermediates<T, E: AllChildren<AlgebraicExpression<T>>>(
e: &E,
intermediate_definitions: &BTreeMap<AlgebraicReferenceThin, AlgebraicExpression<T>>,
intermediates_cache: &mut BTreeMap<AlgebraicReferenceThin, BTreeSet<AlgebraicReference>>,
) -> BTreeSet<AlgebraicReference> {
e.all_children()
.filter_map(|e| {
if let AlgebraicExpression::Reference(reference) = e {
Some(reference)
} else {
None
}
})
.map(|reference| match reference.poly_id.ptype {
PolynomialType::Constant => std::iter::once(reference.clone()).collect(),
PolynomialType::Intermediate => {
let reference = reference.to_thin();
intermediates_cache
.get(&reference)
.cloned()
.unwrap_or_else(|| {
let result = get_constant_refs_with_intermediates(
&intermediate_definitions[&reference],
intermediate_definitions,
intermediates_cache,
);
intermediates_cache.insert(reference, result.clone());
result
})
}
PolynomialType::Committed => std::iter::empty().collect(),
})
.fold(BTreeSet::new(), |mut acc, set| {
acc.extend(set);
acc
})
}

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

@@ -1,7 +1,6 @@
use itertools::Itertools;
use num_traits::{One, Zero};
use powdr_ast::analyzed::{AlgebraicExpression, Analyzed, DegreeRange};
use powdr_ast::parsed::visitor::AllChildren;
use powdr_ast::analyzed::{Analyzed, DegreeRange};
use powdr_backend_utils::{machine_fixed_columns, machine_witness_columns};
use powdr_executor::constant_evaluator::VariablySizedColumn;
use powdr_executor::witgen::WitgenCallback;
@@ -670,16 +669,7 @@ fn get_challenges<MC: MerkleChannel>(
analyzed: &Analyzed<M31>,
challenge_channel: &mut MC::C,
) -> BTreeMap<u64, M31> {
let identities = &analyzed.identities;
let challenges_stage0 = identities
.iter()
.flat_map(|identity| {
identity.all_children().filter_map(|expr| match expr {
AlgebraicExpression::Challenge(challenge) => Some(challenge.id),
_ => None,
})
})
.collect::<BTreeSet<_>>();
let challenges_stage0 = analyzed.challenges();
// Stwo provides a function to draw challenges from the secure field `QM31`,
// which consists of 4 `M31` elements.
@@ -696,6 +686,7 @@ fn get_challenges<MC: MerkleChannel>(
challenges_stage0
.into_iter()
.map(|challenge| challenge.id)
.zip(draw_challenges.map(|challenge| from_stwo_field(&challenge)))
.collect::<BTreeMap<_, _>>()
}

8
pil-analyzer/src/expressionizer.rs
View File

@@ -58,6 +58,14 @@ pub fn try_value_to_expression<T: FieldElement>(
Expressionizer { poly_id_to_name }.try_value_to_expression(value)
}
/// Tries to convert an algebraic expression to an expression that evaluates to the same value.
pub fn try_algebraic_expression_to_expression<T: FieldElement>(
poly_id_to_name: &BTreeMap<(PolynomialType, u64), String>,
e: &AlgebraicExpression<T>,
) -> Result<Expression, EvalError> {
Expressionizer { poly_id_to_name }.try_algebraic_expression_to_expression(e)
}
struct Expressionizer<'a> {
/// Maps polynomial IDs to their names.
/// For arrays, this does not include the array elements, just the ID

1
pil-analyzer/src/lib.rs
View File

@@ -24,6 +24,7 @@ use powdr_ast::{
},
};
pub use expressionizer::try_algebraic_expression_to_expression;
pub use pil_analyzer::{analyze_ast, analyze_file, analyze_string};
pub trait AnalysisDriver: Clone + Copy {

4
pilopt/Cargo.toml
View File

@@ -10,14 +10,12 @@ repository = { workspace = true }
[dependencies]
powdr-ast.workspace = true
powdr-number.workspace = true
powdr-pil-analyzer.workspace = true
log = "0.4.27"
pretty_assertions = "1.4.0"
itertools = "0.13.0"
[dev-dependencies]
powdr-pil-analyzer.workspace = true
[lints]
workspace = true

409
pilopt/src/lib.rs
View File

@@ -7,18 +7,19 @@ use std::hash::{DefaultHasher, Hash, Hasher};
use itertools::Itertools;
use powdr_ast::analyzed::{
AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression, AlgebraicReference,
AlgebraicUnaryOperation, AlgebraicUnaryOperator, Analyzed, ConnectIdentity, Expression,
FunctionValueDefinition, Identity, LookupIdentity, PermutationIdentity, PhantomLookupIdentity,
PhantomPermutationIdentity, PolyID, PolynomialIdentity, PolynomialReference, PolynomialType,
Reference, StatementIdentifier, Symbol, SymbolKind,
AlgebraicReferenceThin, AlgebraicUnaryOperation, AlgebraicUnaryOperator, Analyzed,
ConnectIdentity, ContainsNextRef, Expression, FunctionValueDefinition, Identity,
LookupIdentity, PermutationIdentity, PhantomLookupIdentity, PhantomPermutationIdentity, PolyID,
PolynomialIdentity, PolynomialReference, PolynomialType, Reference, StatementIdentifier,
Symbol, SymbolKind,
};
use powdr_ast::parsed::types::Type;
use powdr_ast::parsed::visitor::{AllChildren, Children, ExpressionVisitable};
use powdr_ast::parsed::Number;
use powdr_number::{BigUint, FieldElement};
pub mod referenced_symbols;
use powdr_pil_analyzer::try_algebraic_expression_to_expression;
use referenced_symbols::{ReferencedSymbols, SymbolReference};
pub fn optimize<T: FieldElement>(mut pil_file: Analyzed<T>) -> Analyzed<T> {
@@ -30,10 +31,10 @@ pub fn optimize<T: FieldElement>(mut pil_file: Analyzed<T>) -> Analyzed<T> {
deduplicate_fixed_columns(&mut pil_file);
simplify_identities(&mut pil_file);
extract_constant_lookups(&mut pil_file);
replace_linear_witness_columns(&mut pil_file);
remove_constant_witness_columns(&mut pil_file);
remove_constant_intermediate_columns(&mut pil_file);
simplify_identities(&mut pil_file);
remove_equal_constrained_witness_columns(&mut pil_file);
inline_trivial_intermediate_polynomials(&mut pil_file);
remove_trivial_identities(&mut pil_file);
remove_duplicate_identities(&mut pil_file);
@@ -125,7 +126,11 @@ fn remove_unreferenced_definitions<T: FieldElement>(pil_file: &mut Analyzed<T>)
)
}
} else if let Some((_, value)) = pil_file.intermediate_columns.get(n.name.as_ref()) {
assert!(n.type_args.is_none());
assert!(n
.type_args
.as_ref()
.map(|args| args.is_empty())
.unwrap_or(true));
Box::new(value.iter().flat_map(|v| {
v.all_children().flat_map(|e| {
if let AlgebraicExpression::Reference(AlgebraicReference { poly_id, .. }) = e {
@@ -164,8 +169,6 @@ fn remove_unreferenced_definitions<T: FieldElement>(pil_file: &mut Analyzed<T>)
pil_file.remove_trait_impls(&impls_to_remove);
}
/// Builds a lookup-table that can be used to turn all poly ids into the names of the symbols that define them.
/// For array elements, this contains the array name and the index of the element in the array.
fn build_poly_id_to_definition_name_lookup(
pil_file: &Analyzed<impl FieldElement>,
) -> BTreeMap<PolyID, (&String, Option<usize>)> {
@@ -191,6 +194,7 @@ fn build_poly_id_to_definition_name_lookup(
})
.collect()
}
/// Collect all names that are referenced in identities and public declarations.
fn collect_required_symbols<'a, T: FieldElement>(
pil_file: &'a Analyzed<T>,
@@ -232,7 +236,10 @@ fn remove_constant_fixed_columns<T: FieldElement>(pil_file: &mut Analyzed<T>) {
"Determined fixed column {} to be constant {value}. Removing.",
poly.absolute_name
);
Some(((poly.absolute_name.clone(), poly.into()), value))
Some((
(poly.absolute_name.clone(), poly.into()),
T::from(value).into(),
))
})
.collect::<Vec<((String, PolyID), _)>>();
@@ -554,6 +561,121 @@ fn extract_constant_lookups<T: FieldElement>(pil_file: &mut Analyzed<T>) {
}
}
/// Identifies witness columns that are constrained to a non-shifted (multi)linear expression, replaces the witness column by an intermediate column defined to be that expression.
/// The pattern is the following:
/// ```pil
/// col witness x;
/// x = lin;
/// ```
/// Where `lin` is a non-shifted multi-linear expression.
/// This optimization makes sense because it saves the cost of committing to `x` and does not increase the degree of the constraints `x` is involved in.
/// TODO: the optimization is *NOT* applied to witness columns which are boolean constrained: intermediate polynomials get inlined in witness generation and
/// the boolean constraint gets lost. Remove this limitation once intermediate polynomials are handled correctly in witness generation.
fn replace_linear_witness_columns<T: FieldElement>(pil_file: &mut Analyzed<T>) {
let intermediate_definitions = pil_file.intermediate_definitions();
// We cannot remove arrays or array elements, so we filter them out.
// Also, we filter out columns that are used in public declarations.
// Also, we filter out columns that are boolean constrained.
let in_publics: HashSet<_> = pil_file
.public_declarations_in_source_order()
.map(|(_, pubd)| pubd.referenced_poly().name.clone())
.collect();
// pattern match identities looking for `w * (1 - w) = 0` and `(1 - w) * w = 0` constraints
let boolean_constrained_witnesses = pil_file
.identities
.iter()
.filter_map(|id| try_to_boolean_constrained(id))
.collect::<HashSet<_>>();
let keep = pil_file
.committed_polys_in_source_order()
.filter(|&(s, _)| {
s.is_array()
|| in_publics.contains(&s.absolute_name)
|| boolean_constrained_witnesses
.intersection(&s.array_elements().map(|(_, poly_id)| poly_id).collect())
.count()
> 0
})
.map(|(s, _)| s.into())
.collect::<HashSet<PolyID>>();
let linear_polys = pil_file
.identities
.iter()
.enumerate()
.filter_map(|(index, id)| {
if let Identity::Polynomial(i) = id {
Some((index, i))
} else {
None
}
})
.filter_map(|(index, identity)| {
try_to_linearly_constrained(identity, &intermediate_definitions).map(|e| (index, e))
})
.filter(|(_, ((_, poly_id), _))| !keep.contains(poly_id))
.collect::<Vec<(usize, ((String, PolyID), _))>>();
let (identities_to_remove, intermediate_polys): (BTreeSet<_>, Vec<_>) = linear_polys
.iter()
.filter_map(|(index, ((name, old_poly_id), expression))| {
// Remove the definition. If this fails, we have already replaced it by an intermediate column.
pil_file.definitions.remove(name).map(|(symbol, value)| {
// sanity checks that we are looking at a single witness column
assert!(symbol.kind == SymbolKind::Poly(PolynomialType::Committed));
assert!(value.is_none());
assert!(symbol.length.is_none());
// we introduce a new intermediate symbol, so we need a new unique id
let id = pil_file.next_id_for_kind(PolynomialType::Intermediate);
// the symbol is not an array, so the poly_id uses the same id as the symbol
let poly_id = PolyID {
ptype: PolynomialType::Intermediate,
id,
};
let kind = SymbolKind::Poly(PolynomialType::Intermediate);
let stage = None;
let symbol = Symbol {
id,
kind,
stage,
..symbol
};
// Add the definition to the intermediate columns
assert!(pil_file
.intermediate_columns
.insert(name.clone(), (symbol, vec![expression.clone()]))
.is_none());
// Note: the `statement_order` does not need to be updated, since we are still declaring the same symbol
let r = AlgebraicReference {
name: name.clone(),
poly_id,
next: false,
};
(
index,
(
(name.clone(), *old_poly_id),
AlgebraicExpression::Reference(r),
),
)
})
})
.unzip();
// we remove the identities by hand here, because they do *not* become trivial by applying this optimization
pil_file.remove_identities(&identities_to_remove);
// we substitute the references to the witnesses by references to the intermediate columns
substitute_polynomial_references(pil_file, intermediate_polys);
}
/// Identifies witness columns that are constrained to a single value, replaces every
/// reference to this column by the value and deletes the column.
fn remove_constant_witness_columns<T: FieldElement>(pil_file: &mut Analyzed<T>) {
@@ -568,6 +690,7 @@ fn remove_constant_witness_columns<T: FieldElement>(pil_file: &mut Analyzed<T>)
}
})
.filter_map(constrained_to_constant)
.map(|(k, v)| (k, T::from(v).into()))
.collect::<Vec<((String, PolyID), _)>>();
let in_publics: HashSet<_> = pil_file
@@ -586,9 +709,10 @@ fn remove_constant_witness_columns<T: FieldElement>(pil_file: &mut Analyzed<T>)
substitute_polynomial_references(pil_file, constant_polys);
}
/// Identifies intermediate columns that are constrained to a single value, replaces every
/// reference to this column by the value and deletes the column.
fn remove_constant_intermediate_columns<T: FieldElement>(pil_file: &mut Analyzed<T>) {
/// Inlines `col i = e` into the references to `i` where `e` is an expression with no operations.
/// The reasoning is that intermediate columns are useful to remember intermediate computation results, but in this case
/// the intermediate results are already known.
fn inline_trivial_intermediate_polynomials<T: FieldElement>(pil_file: &mut Analyzed<T>) {
let intermediate_polys = pil_file
.intermediate_polys_in_source_order()
.filter_map(|(symbol, definitions)| {
@@ -596,15 +720,17 @@ fn remove_constant_intermediate_columns<T: FieldElement>(pil_file: &mut Analyzed
match symbol.is_array() {
true => None,
false => {
let ((name, poly_id), definition) = symbols_and_definitions.next().unwrap();
match definition {
AlgebraicExpression::Number(value) => {
log::debug!(
"Determined intermediate column {name} to be constant {value}. Removing.",
);
Some(((name.clone(), poly_id), value.to_arbitrary_integer()))
let ((name, poly_id), value) = symbols_and_definitions.next().unwrap();
match value {
AlgebraicExpression::BinaryOperation(_) | AlgebraicExpression::UnaryOperation(_) => {
None
}
AlgebraicExpression::Reference(_) | AlgebraicExpression::PublicReference(_) | AlgebraicExpression::Number(_) | AlgebraicExpression::Challenge(_)=>{
log::debug!(
"Determined intermediate column {name} to be trivial value `{value}`. Removing.",
);
Some(((name, poly_id), value.clone()))
}
_ => None,
}
}
}
@@ -614,11 +740,16 @@ fn remove_constant_intermediate_columns<T: FieldElement>(pil_file: &mut Analyzed
substitute_polynomial_references(pil_file, intermediate_polys);
}
/// Substitutes all references to certain polynomials by the given field elements.
/// Substitutes all references to certain polynomials by the given non-shifted expressions.
fn substitute_polynomial_references<T: FieldElement>(
pil_file: &mut Analyzed<T>,
substitutions: Vec<((String, PolyID), BigUint)>,
substitutions: Vec<((String, PolyID), AlgebraicExpression<T>)>,
) {
let poly_id_to_name = pil_file
.name_to_poly_id()
.into_iter()
.map(|(name, poly)| ((poly.ptype, poly.id), name))
.collect();
let substitutions_by_id = substitutions
.iter()
.map(|((_, id), value)| (*id, value.clone()))
@@ -634,18 +765,19 @@ fn substitute_polynomial_references<T: FieldElement>(
) = e
{
if let Some(value) = substitutions_by_name.get(name) {
*e = Number {
value: (*value).clone(),
type_: Some(Type::Expr),
}
.into();
*e = try_algebraic_expression_to_expression(&poly_id_to_name, value).unwrap();
}
}
});
pil_file.post_visit_expressions_in_identities_mut(&mut |e: &mut AlgebraicExpression<_>| {
if let AlgebraicExpression::Reference(AlgebraicReference { poly_id, .. }) = e {
if let AlgebraicExpression::Reference(AlgebraicReference { poly_id, next, .. }) = e {
if let Some(value) = substitutions_by_id.get(poly_id) {
*e = AlgebraicExpression::Number(T::checked_from((*value).clone()).unwrap());
let value = if *next {
value.clone().next().unwrap()
} else {
value.clone()
};
*e = value;
}
}
});
@@ -681,6 +813,66 @@ fn constrained_to_constant<T: FieldElement>(
None
}
/// Tries to extract a witness column which is constrained to a non-shifted multi-linear expression.
/// Inputs:
/// - `identity`: a polynomial identity
/// - `intermediate_columns`: the intermediate column definitions
///
/// Outputs:
/// - `None` if the identity does not match the pattern
/// - `Some(((name, poly_id), expression))` if the identity matches the pattern, where:
/// - `name` is the name of the witness column
/// - `poly_id` is the id of the witness column
/// - `expression` is the (multi)linear expression that the witness column is constrained to
fn try_to_linearly_constrained<T: FieldElement>(
identity: &PolynomialIdentity<T>,
intermediate_columns: &BTreeMap<AlgebraicReferenceThin, AlgebraicExpression<T>>,
) -> Option<((String, PolyID), AlgebraicExpression<T>)> {
// We require the constraint to be of the form `left = right`, represented as `left - right`
let (left, right) = if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
left,
op: AlgebraicBinaryOperator::Sub,
right,
}) = &identity.expression
{
(left, right)
} else {
return None;
};
// Helper function to try to extract a witness column from two sides of an identity
let try_from_sides = |left: &AlgebraicExpression<T>, right: &AlgebraicExpression<T>| {
// We require `left` to be a single, non-shifted, witness column `w`
let w = if let AlgebraicExpression::Reference(
r @ AlgebraicReference {
poly_id:
PolyID {
ptype: PolynomialType::Committed,
..
},
next: false,
..
},
) = left
{
r
} else {
return None;
};
// we require `right` to be a multi-linear expression over non-shifted columns
if right.contains_next_ref(intermediate_columns)
|| right.degree_with_cache(intermediate_columns, &mut Default::default()) != 1
{
return None;
}
Some(((w.name.clone(), w.poly_id), right.clone()))
};
try_from_sides(left, right).or_else(|| try_from_sides(right, left))
}
/// Removes identities that evaluate to zero (including constraints of the form "X = X") and lookups with empty columns.
fn remove_trivial_identities<T: FieldElement>(pil_file: &mut Analyzed<T>) {
let to_remove = pil_file
@@ -845,108 +1037,61 @@ fn remove_duplicate_identities<T: FieldElement>(pil_file: &mut Analyzed<T>) {
pil_file.remove_identities(&to_remove);
}
/// Identifies witness columns that are directly constrained to be equal to other witness columns
/// through polynomial identities of the form "x = y" and returns a tuple ((name, id), (name, id))
/// for each pair of identified columns
fn equal_constrained<T: FieldElement>(
expression: &AlgebraicExpression<T>,
poly_id_to_array_elem: &BTreeMap<PolyID, (&String, Option<usize>)>,
) -> Option<((String, PolyID), (String, PolyID))> {
match expression {
/// Tries to extract a boolean constrained witness column from a polynomial identity.
/// The pattern used is `x * (1 - x) = 0` or `(1 - x) * x = 0` where `x` is a witness column.
fn try_to_boolean_constrained<T: FieldElement>(id: &Identity<T>) -> Option<PolyID> {
// we are only interested in polynomial identities
let expression = if let Identity::Polynomial(PolynomialIdentity { expression, .. }) = id {
expression
} else {
return None;
};
// we are only interested in `a * b = 0` constraints
let (a, b) = match expression {
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
left,
op: AlgebraicBinaryOperator::Sub,
right,
}) => match (left.as_ref(), right.as_ref()) {
(AlgebraicExpression::Reference(l), AlgebraicExpression::Reference(r)) => {
let is_valid = |x: &AlgebraicReference, left: bool| {
x.is_witness()
&& if left {
// We don't allow the left-hand side to be an array element
// to preserve array integrity (e.g. `x = y` is valid, but `x[0] = y` is not)
poly_id_to_array_elem.get(&x.poly_id).unwrap().1.is_none()
} else {
true
}
};
left: a,
op: AlgebraicBinaryOperator::Mul,
right: b,
}) => (a, b),
_ => return None,
};
if is_valid(l, true) && is_valid(r, false) && r.next == l.next {
Some(if l.poly_id > r.poly_id {
((l.name.clone(), l.poly_id), (r.name.clone(), r.poly_id))
} else {
((r.name.clone(), r.poly_id), (l.name.clone(), l.poly_id))
})
} else {
None
}
}
_ => None,
},
_ => None,
}
}
fn remove_equal_constrained_witness_columns<T: FieldElement>(pil_file: &mut Analyzed<T>) {
let poly_id_to_array_elem = build_poly_id_to_definition_name_lookup(pil_file);
let mut substitutions: BTreeMap<(String, PolyID), (String, PolyID)> = pil_file
.identities
.iter()
.filter_map(|id| {
if let Identity::Polynomial(PolynomialIdentity { expression, .. }) = id {
equal_constrained(expression, &poly_id_to_array_elem)
} else {
None
}
})
.collect();
resolve_transitive_substitutions(&mut substitutions);
let (subs_by_id, subs_by_name): (HashMap<_, _>, HashMap<_, _>) = substitutions
.iter()
.map(|(k, v)| ((k.1, v), (&k.0, v)))
.unzip();
pil_file.post_visit_expressions_in_identities_mut(&mut |e: &mut AlgebraicExpression<_>| {
if let AlgebraicExpression::Reference(ref mut reference) = e {
if let Some((replacement_name, replacement_id)) = subs_by_id.get(&reference.poly_id) {
reference.poly_id = *replacement_id;
reference.name = replacement_name.clone();
}
}
});
pil_file.post_visit_expressions_mut(&mut |e: &mut Expression| {
if let Expression::Reference(_, Reference::Poly(reference)) = e {
if let Some((replacement_name, _)) = subs_by_name.get(&reference.name) {
reference.name = replacement_name.clone();
}
}
});
}
fn resolve_transitive_substitutions(subs: &mut BTreeMap<(String, PolyID), (String, PolyID)>) {
let mut changed = true;
while changed {
changed = false;
let keys: Vec<_> = subs
.keys()
.map(|(name, id)| (name.to_string(), *id))
.collect();
for key in keys {
let Some(target_key) = subs.get(&key) else {
continue;
};
let Some(new_target) = subs.get(target_key) else {
continue;
};
if subs.get(&key).unwrap() != new_target {
subs.insert(key, new_target.clone());
changed = true;
}
}
// we are only interested in `b := (1 - a)` or `a := (1 - b)`
let a = match (a.as_ref(), b.as_ref()) {
(
a_0,
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
left: one,
op: AlgebraicBinaryOperator::Sub,
right: a_1,
}),
)
| (
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
left: one,
op: AlgebraicBinaryOperator::Sub,
right: a_1,
}),
a_0,
) if **one == AlgebraicExpression::Number(T::one()) && *a_0 == **a_1 => a_0,
_ => return None,
};
// we are only interested in `a` being a column
let a = match a {
AlgebraicExpression::Reference(AlgebraicReference {
poly_id,
next: false,
..
}) => poly_id,
_ => return None,
};
// we are only interested in `a` being a witness column
if matches!(a.ptype, PolynomialType::Committed) {
Some(*a)
} else {
None
}
}

133
pilopt/tests/optimizer.rs
View File

@@ -15,14 +15,14 @@ fn replace_fixed() {
let _ = one;
};
X * one = X * zero - zero + Y;
one * Y = zero * Y + 7 * X;
one * Y = zero * Y + 7 * X * X;
"#;
let expectation = r#"namespace N(65536);
col witness X;
col witness Y;
query |i| {
let _: expr = 1_expr;
};
N::X = 7 * N::X;
N::Y = 7 * N::Y * N::Y;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
assert_eq!(optimized, expectation);
@@ -120,10 +120,8 @@ fn intermediate() {
"#;
let expectation = r#"namespace N(65536);
col witness x;
col intermediate = N::x;
col int2 = N::intermediate * N::x;
col int3 = N::int2;
N::int3 = 3 * N::x + N::x;
col int2 = N::x * N::x;
N::int2 = 3 * N::x + N::x;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
assert_eq!(optimized, expectation);
@@ -235,7 +233,7 @@ fn remove_unreferenced_keep_enums() {
// Y is not mentioned anywhere.
let f: col = |i| if i == 0 { t([]) } else { (|x| 1)(Y::F([])) };
let x;
x = f;
x = f * f;
"#;
let expectation = r#"namespace N(65536);
enum X {
@@ -254,7 +252,7 @@ fn remove_unreferenced_keep_enums() {
let t: N::X[] -> int = |r| 1_int;
col fixed f(i) { if i == 0_int { N::t([]) } else { (|x| 1_int)(N::Y::F([])) } };
col witness x;
N::x = N::f;
N::x = N::f * N::f;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
assert_eq!(optimized, expectation);
@@ -275,7 +273,7 @@ fn test_trait_impl() {
impl Default<int> { f: || 1, g: |x| x }
let x: col = |_| Default::f();
let w;
w = x;
w = x * x;
"#;
let expectation = r#"namespace N(65536);
trait Default<T> {
@@ -289,7 +287,7 @@ fn test_trait_impl() {
let dep: fe -> fe = |x| x + 1_fe;
col fixed x(_) { N::Default::f::<fe>() };
col witness w;
N::w = N::x;
N::w = N::x * N::x;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
assert_eq!(optimized, expectation);
@@ -304,7 +302,7 @@ fn enum_ref_by_trait() {
impl X<fe> { f: |_| O::Y(1), g: || { let r = Q::B(1_int); 1 } }
let x: col = |i| { match X::f(1_fe) { O::Y(y) => y, _ => 0 } };
let w;
w = x;
w = x * x;
"#;
let expectation = r#"namespace N(65536);
enum O<T> {
@@ -331,7 +329,7 @@ fn enum_ref_by_trait() {
_ => 0_fe,
} };
col witness w;
N::w = N::x;
N::w = N::x * N::x;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
assert_eq!(optimized, expectation);
@@ -459,12 +457,56 @@ fn equal_constrained_transitive() {
a + b + c = 5;
"#;
let expectation = r#"namespace N(65536);
col witness a;
N::a + N::a + N::a = 5;
col witness c;
N::c + N::c + N::c = 5;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
assert_eq!(optimized, expectation);
}
#[test]
fn replace_witness_by_intermediate() {
let input = r#"namespace N(65536);
col witness w;
col fixed f = [1, 0]*;
col witness can_be_replaced;
can_be_replaced = 2 * w + 3 * f + 5;
can_be_replaced + w = 5;
// Constraining to a shifted expression should not replace the witness.
col witness linear_with_next_ref;
linear_with_next_ref = 2 * w + 3 * f' + 5;
linear_with_next_ref + w = 5;
// Constraining to a quadratic expression should not replace the witness.
col witness quadratic;
quadratic = 2 * w * w + 3 * f + 5;
quadratic + w = 5;
// The first constraint is removed, the second one is kept.
col witness constrained_twice;
constrained_twice = 2 * w + 3 * f + 5;
constrained_twice = w + f;
"#;
let expectation = r#"namespace N(65536);
col witness w;
col fixed f = [1_fe, 0_fe]*;
col can_be_replaced = 2 * N::w + 3 * N::f + 5;
N::can_be_replaced + N::w = 5;
col witness linear_with_next_ref;
N::linear_with_next_ref = 2 * N::w + 3 * N::f' + 5;
N::linear_with_next_ref + N::w = 5;
col witness quadratic;
N::quadratic = 2 * N::w * N::w + 3 * N::f + 5;
N::quadratic + N::w = 5;
col constrained_twice = 2 * N::w + 3 * N::f + 5;
N::constrained_twice = N::w + N::f;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
assert_eq!(optimized, expectation);
}
#[test]
fn simplify_associative_operations() {
let input = r#"namespace N(150);
@@ -475,27 +517,60 @@ fn simplify_associative_operations() {
col fixed c2 = [2]*;
col fixed c3 = [3]*;
(x + c2) + c1 = y;
(c2 + x) + c3 = y;
(x - c2) + c1 = y;
(x + c2) + c1 = y * y;
(c2 + x) + c3 = y * y;
(x - c2) + c1 = y * y;
((x + 3) - y) - 9 = z;
(c3 + (x + 3)) - y = z;
((-x + 3) + y) + 9 = z;
((-x + 3) + c3) + 12 = z;
((x + 3) - y) - 9 = z * z;
(c3 + (x + 3)) - y = z * z;
((-x + 3) + y) + 9 = z * z;
((-x + 3) + c3) + 12 = z * z;
"#;
let expectation = r#"namespace N(150);
col witness x;
col witness y;
col witness z;
N::x + 3 = N::y;
N::x + 5 = N::y;
N::x - 2 + 1 = N::y;
N::x + 3 - N::y - 9 = N::z;
N::x + 6 - N::y = N::z;
-N::x + N::y + 12 = N::z;
-N::x + 18 = N::z;
N::x + 3 = N::y * N::y;
N::x + 5 = N::y * N::y;
N::x - 2 + 1 = N::y * N::y;
N::x + 3 - N::y - 9 = N::z * N::z;
N::x + 6 - N::y = N::z * N::z;
-N::x + N::y + 12 = N::z * N::z;
-N::x + 18 = N::z * N::z;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();
assert_eq!(optimized, expectation);
}
#[test]
fn inline_chain_of_substitutions() {
let input = r#"namespace N(65536);
col witness a;
col witness b;
col witness x;
x = a + y;
col witness y;
y = x + b;
col witness m;
m = x - y;
a * b = 10;
m * a = 1;
"#;
let expectation = r#"namespace N(65536);
col witness a;
col witness b;
col x = N::a + N::y;
col y = N::x + N::b;
col m = N::x - N::y;
N::a * N::b = 10;
N::m * N::a = 1;
"#;
let optimized = optimize(analyze_string::<GoldilocksField>(input).unwrap()).to_string();

16
pipeline/tests/pil.rs
View File

@@ -169,20 +169,12 @@ fn external_witgen_a_provided() {
test_mock_backend(pipeline);
}
#[test]
fn external_witgen_b_provided() {
let f = "pil/external_witgen.pil";
let external_witness = vec![("main::b".to_string(), vec![GoldilocksField::from(4); 16])];
let pipeline = make_prepared_pipeline(f, Default::default(), external_witness, LinkerMode::Bus);
test_mock_backend(pipeline);
}
#[test]
fn external_witgen_both_provided() {
let f = "pil/external_witgen.pil";
let external_witness = vec![
("main::a".to_string(), vec![GoldilocksField::from(3); 16]),
("main::b".to_string(), vec![GoldilocksField::from(4); 16]),
("main::b".to_string(), vec![GoldilocksField::from(16); 16]),
];
let pipeline = make_prepared_pipeline(f, Default::default(), external_witness, LinkerMode::Bus);
test_mock_backend(pipeline);
@@ -194,8 +186,8 @@ fn external_witgen_fails_on_conflicting_external_witness() {
let f = "pil/external_witgen.pil";
let external_witness = vec![
("main::a".to_string(), vec![GoldilocksField::from(3); 16]),
// Does not satisfy b = a + 1
("main::b".to_string(), vec![GoldilocksField::from(3); 16]),
// Does not satisfy b = (a + 1) * (a + 1)
("main::b".to_string(), vec![GoldilocksField::from(15); 16]),
];
let pipeline = make_prepared_pipeline(f, Default::default(), external_witness, LinkerMode::Bus);
test_mock_backend(pipeline);
@@ -260,7 +252,7 @@ fn halo_without_lookup() {
#[test]
fn add() {
let f = "pil/add.pil";
let f = "pil/mul.pil";
regular_test_gl(f, Default::default());
}

10
plonky3/src/circuit_builder.rs
View File

@@ -28,7 +28,6 @@ use powdr_ast::analyzed::{
};
use crate::{CallbackResult, MultiStageAir, MultistageAirBuilder};
use powdr_ast::parsed::visitor::ExpressionVisitable;
use powdr_executor_utils::{
expression_evaluator::{ExpressionEvaluator, TerminalAccess},
@@ -97,12 +96,9 @@ impl<T: FieldElement> From<&Analyzed<T>> for ConstraintSystem<T> {
// we use a set to collect all used challenges
let mut challenges_by_stage = vec![BTreeSet::new(); analyzed.stage_count()];
for identity in &identities {
identity.pre_visit_expressions(&mut |expr| {
if let AlgebraicExpression::Challenge(challenge) = expr {
challenges_by_stage[challenge.stage as usize].insert(challenge.id);
}
});
let challenges = analyzed.challenges();
for challenge in challenges {
challenges_by_stage[challenge.stage as usize].insert(challenge.id);
}
// finally, we convert the set to a vector

4
test_data/asm/challenges.asm
View File

@@ -6,10 +6,10 @@ machine Main with degree: 4 {
// Stage-0 witness column, does not depend on any challenges:
col witness bar;
bar = foo + 3;
bar = foo * foo + 3;
// Stage-1 witness column, depends on after-stage-0 challenge #1:
col witness stage(1) bar2;
let alpha: expr = challenge(0, 1);
bar2 = bar + alpha;
bar2 = bar * bar + alpha;
}

2
test_data/asm/expand_fixed.asm
View File

@@ -140,5 +140,5 @@ machine Main with degree: 2**22 {
Result::Ok(r) => |i| std::convert::fe(r(i)),
Result::Err(e) => panic(e)
};
w = LAST;
w = LAST * LAST;
}

2
test_data/asm/full_pil_constant.asm
View File

@@ -1,5 +1,5 @@
machine FullConstant with degree: 4 {
pol constant C = [0, 1]*;
col commit w;
w = C;
w = C * C;
}

2
test_data/asm/set_hint.asm
View File

@@ -7,5 +7,5 @@ let new_col_with_hint: -> expr = constr || {
machine Main with degree: 4 {
let x;
let w = new_col_with_hint();
x = w + 1;
x = w * (w + 1);
}

2
test_data/asm/type_vars_in_local_decl.asm
View File

@@ -16,5 +16,5 @@ let q: int -> int = |i| {
machine Main with degree: 4 {
col witness w;
let f: col = q;
w = f;
w = f * f;
}

4
test_data/pil/external_witgen.pil
View File

@@ -10,8 +10,8 @@ namespace main(N);
// satisfied by filling everything with zero.
// However, providing values for a or b externally should lead witgen to find a
// unique witness.
b = a + 1;
b = (a + 1) * (a + 1);
// The pipeline skips witgen if the full witness is provided.
// The `c` column is here only so that the provided external witness becomes partial, making witgen run.
c = b + 1;
c = (b + 1) * (b + 1);

2
test_data/pil/fixed_columns.pil
View File

@@ -22,4 +22,4 @@ col fixed C(i) {match i {
}};
// ANCHOR_END: mapping
let w;
w = B;
w = B * B;

2
test_data/pil/halo_without_lookup.pil
View File

@@ -2,4 +2,4 @@
namespace Assembly(4);
col fixed C = [0, 1, 2, 3];
col witness w;
w = C;
w = C * C;

4
test_data/pil/add.pil → test_data/pil/mul.pil
View File

@@ -1,5 +1,5 @@
namespace Add(8);
namespace Mul(8);
let a;
let b;
let c;
a + b = c;
a * b = c;

13
test_data/pil/simple_div.pil
View File

@@ -1,5 +1,8 @@
let N: int = 0x10000;
namespace std::prover;
let provide_value: expr, int, fe -> () = [];
namespace SimpleDiv(N);
col fixed BYTE(i) { i & 0xff };
@@ -7,6 +10,12 @@ namespace SimpleDiv(N);
col witness Y;
col witness Z;
col witness R;
// We introduce an underconstrained column which we set to 1 at runtime.
// This is hacky and unsafe, but it enables blocking the linear constraint optimizer from removing the entire program.
col witness one;
query |i| {
std::prover::provide_value(one, i, 1);
};
// Compute X = Y / Z, i.e.
// X * Z + R = Y, where 0 <= R < Z
@@ -16,8 +25,8 @@ namespace SimpleDiv(N);
X * Z + R = Y;
Z - R - 1 = Y_b1 + Y_b2 * 0x100;
X = X_b1 + X_b2 * 0x100;
R = R_b1 + R_b2 * 0x100;
X = one * (X_b1 + X_b2 * 0x100);
R = one * (R_b1 + R_b2 * 0x100);
col witness Y_b1;
col witness Y_b2;

4
test_data/pil/trivial.pil
View File

@@ -1,5 +1,5 @@
// The the simplest PIL that doesn't get optimized away completely.
// The simplest PIL that doesn't get optimized away completely.
namespace main(4);
let index: col = |i| i;
let w;
w = index;
w = index * index;

5
test_data/std/fingerprint_test.asm
View File

@@ -34,7 +34,8 @@ machine Main with degree: 2048 {
col witness stage(1) fingerprint_value1_hint(i) query Query::Hint(fingerprint_hint()[1]);
// Assert consistency between `fingerprint` and `fingerprint_inter`
fingerprint_value0 = fingerprint_value0_hint;
fingerprint_value1 = fingerprint_value1_hint;
// We compare squares to avoid the linear constraint being optimized away, which currently leads to an issue because `fingerprint_value0_hint` has a hint
fingerprint_value0 * fingerprint_value0 = fingerprint_value0_hint * fingerprint_value0_hint;
fingerprint_value1 * fingerprint_value1 = fingerprint_value1_hint * fingerprint_value1_hint;
}