-
Notifications
You must be signed in to change notification settings - Fork 116
Apc stacking #2917
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Apc stacking #2917
Changes from 57 commits
23002ae
b05703d
7a35e47
323e706
f09fc34
a2d3683
bdff7c9
5d060c5
c6443f7
c950d3f
cb4733d
e28585d
d142918
df67d47
920edba
50a1105
dedee09
09b512a
4b3c5d4
a7f5ff1
e48bf1b
13850f6
669b974
9dcff42
c8918f5
6ff8258
0f5448e
b8b0beb
4de95d2
4dd19ef
a7537c2
dfb387f
1e83820
f363529
6f6ca78
05af4ae
1075ba2
428550b
2ab7843
5193df6
1ddff4f
7ea0748
4473e9a
35c8ba0
c139514
44c8c58
a5483e9
b85b763
c01c5ca
54dd766
3b8dd68
3ada389
f176a0f
857997a
9dc6760
1905935
be0c6b6
ed27779
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,6 +1,6 @@ | ||
use crate::bus_map::{BusMap, BusType}; | ||
use crate::expression_conversion::algebraic_to_quadratic_symbolic_expression; | ||
use crate::optimizer::simplify_expression; | ||
pub use crate::optimizer::simplify_expression; | ||
use constraint_optimizer::IsBusStateful; | ||
use expression::{AlgebraicExpression, AlgebraicReference}; | ||
use itertools::Itertools; | ||
|
@@ -36,7 +36,7 @@ pub struct SymbolicInstructionStatement<T> { | |
pub args: Vec<T>, | ||
} | ||
|
||
#[derive(Debug, Clone, Serialize, Deserialize)] | ||
#[derive(Debug, Clone, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq)] | ||
pub struct SymbolicConstraint<T> { | ||
pub expr: AlgebraicExpression<T>, | ||
} | ||
|
@@ -278,6 +278,9 @@ pub fn build< | |
vm_config: VmConfig<M, B>, | ||
degree_bound: DegreeBound, | ||
opcode: u32, | ||
// If true, disables the optimization of multipliying only constants when guarding with is_valid. | ||
// Chip stacking needs variables to be guarded too. | ||
strict_is_valid_guards: bool, | ||
) -> Result<Apc<T>, crate::constraint_optimizer::Error> { | ||
let (machine, subs) = statements_to_symbolic_machine( | ||
&program, | ||
|
@@ -294,7 +297,7 @@ pub fn build< | |
)?; | ||
|
||
// add guards to constraints that are not satisfied by zeroes | ||
let machine = add_guards(machine, vm_config.bus_map); | ||
let machine = add_guards(machine, vm_config.bus_map, strict_is_valid_guards); | ||
|
||
Ok(Apc { machine, subs }) | ||
} | ||
|
@@ -346,6 +349,7 @@ fn add_guards_constraint<T: FieldElement>( | |
fn add_guards<T: FieldElement>( | ||
mut machine: SymbolicMachine<T>, | ||
bus_map: BusMap, | ||
is_strict: bool, | ||
) -> SymbolicMachine<T> { | ||
let pre_degree = machine.degree(); | ||
let exec_bus_id = bus_map.get_bus_id(&BusType::ExecutionBridge).unwrap(); | ||
|
@@ -358,11 +362,22 @@ fn add_guards<T: FieldElement>( | |
id: max_id, | ||
}); | ||
|
||
machine.constraints = machine | ||
.constraints | ||
.into_iter() | ||
.map(|c| add_guards_constraint(c.expr, &is_valid).into()) | ||
.collect(); | ||
machine.constraints = if is_strict { | ||
machine | ||
.constraints | ||
.into_iter() | ||
.map(|mut c| { | ||
c.expr = is_valid.clone() * c.expr.clone(); | ||
c | ||
}) | ||
.collect() | ||
} else { | ||
Comment on lines
+366
to
+374
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Does this disable some cheaper stuff that There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
machine | ||
.constraints | ||
.into_iter() | ||
.map(|c| add_guards_constraint(c.expr, &is_valid).into()) | ||
.collect() | ||
}; | ||
|
||
let [execution_bus_receive, execution_bus_send] = machine | ||
.bus_interactions | ||
|
@@ -389,7 +404,7 @@ fn add_guards<T: FieldElement>( | |
for b in &mut machine.bus_interactions { | ||
// already handled exec and pc lookup bus types | ||
if b.id != exec_bus_id && b.id != pc_lookup_bus_id { | ||
if !satisfies_zero_witness(&b.mult) { | ||
if is_strict || !satisfies_zero_witness(&b.mult) { | ||
// guard the multiplicity by `is_valid` | ||
b.mult = is_valid.clone() * b.mult.clone(); | ||
// TODO this would not have to be cloned if we had *= | ||
|
@@ -405,11 +420,13 @@ fn add_guards<T: FieldElement>( | |
|
||
machine.constraints.extend(is_valid_mults); | ||
|
||
assert_eq!( | ||
pre_degree, | ||
machine.degree(), | ||
"Degree should not change after adding guards" | ||
); | ||
if !is_strict { | ||
assert_eq!( | ||
pre_degree, | ||
machine.degree(), | ||
"Degree should not change after adding guards" | ||
); | ||
} | ||
|
||
// This needs to be added after the assertion above because it's a quadratic constraint | ||
// so it may increase the degree of the machine. | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -38,6 +38,9 @@ enum Commands { | |
|
||
#[arg(long)] | ||
input: Option<u32>, | ||
|
||
#[arg(long)] | ||
chip_stacking: Option<f32>, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could you add a description here, with a note that the feature is experimental? |
||
}, | ||
|
||
Execute { | ||
|
@@ -58,6 +61,9 @@ enum Commands { | |
|
||
#[arg(long)] | ||
input: Option<u32>, | ||
|
||
#[arg(long)] | ||
chip_stacking: Option<f32>, | ||
}, | ||
|
||
Prove { | ||
|
@@ -89,6 +95,9 @@ enum Commands { | |
|
||
#[arg(long)] | ||
metrics: Option<PathBuf>, | ||
|
||
#[arg(long)] | ||
chip_stacking: Option<f32>, | ||
}, | ||
} | ||
|
||
|
@@ -115,8 +124,12 @@ fn run_command(command: Commands) { | |
pgo, | ||
max_columns, | ||
input, | ||
chip_stacking, | ||
} => { | ||
let powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64); | ||
let mut powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64); | ||
if let Some(log) = chip_stacking { | ||
powdr_config = powdr_config.with_chip_stacking(log); | ||
} | ||
let pgo_config = pgo_config(guest.clone(), guest_opts.clone(), pgo, max_columns, input); | ||
let program = | ||
powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap(); | ||
|
@@ -130,8 +143,12 @@ fn run_command(command: Commands) { | |
pgo, | ||
max_columns, | ||
input, | ||
chip_stacking, | ||
} => { | ||
let powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64); | ||
let mut powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64); | ||
if let Some(log) = chip_stacking { | ||
powdr_config = powdr_config.with_chip_stacking(log); | ||
} | ||
let pgo_config = pgo_config(guest.clone(), guest_opts.clone(), pgo, max_columns, input); | ||
let program = | ||
powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap(); | ||
|
@@ -148,8 +165,12 @@ fn run_command(command: Commands) { | |
max_columns, | ||
input, | ||
metrics, | ||
chip_stacking, | ||
} => { | ||
let powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64); | ||
let mut powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64); | ||
if let Some(log) = chip_stacking { | ||
powdr_config = powdr_config.with_chip_stacking(log); | ||
} | ||
let pgo_config = pgo_config(guest.clone(), guest_opts.clone(), pgo, max_columns, input); | ||
let program = | ||
powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap(); | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -130,6 +130,13 @@ impl<T, R> AlgebraicExpression<T, R> { | |
} | ||
} | ||
|
||
impl<T: Clone + Ord, R: Clone + Ord> AlgebraicExpression<T, R> { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think the changes to this file could be nice individual PR :) |
||
/// Get a canonical form by expanding the expression and reordering operands. | ||
pub fn normalize(self) -> Self { | ||
normalize(expand(self)) | ||
} | ||
} | ||
Comment on lines
+133
to
+138
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could this be done by going through quadratic expression? cc @georgwiese There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Not quite. We could do this by introducing a If self.to_expression(
&|number| CanonicalExpression { coefficients:
[(BTreeSet::new(), *number)].into_iter().collect()
}, |reference| CanonicalExpression { coefficients:
([*reference].into_iter().collect(), T::one())].into_iter().collect()
}).into() Not sure if it's worth it, but seems a bit cleaner. |
||
|
||
impl<T, R> ops::Add for AlgebraicExpression<T, R> { | ||
type Output = Self; | ||
|
||
|
@@ -168,6 +175,192 @@ impl<T, R> From<T> for AlgebraicExpression<T, R> { | |
} | ||
} | ||
|
||
// Helper to flatten Add trees | ||
fn flatten_add<T: Ord + Clone, R: Ord + Clone>( | ||
expr: AlgebraicExpression<T, R>, | ||
acc: &mut Vec<AlgebraicExpression<T, R>>, | ||
) { | ||
match expr { | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left, | ||
op: AlgebraicBinaryOperator::Add, | ||
right, | ||
}) => { | ||
flatten_add(*left, acc); | ||
flatten_add(*right, acc); | ||
} | ||
_ => acc.push(normalize(expr)), | ||
} | ||
} | ||
|
||
// Helper to flatten Mul trees | ||
fn flatten_mul<T: Ord + Clone, R: Ord + Clone>( | ||
expr: AlgebraicExpression<T, R>, | ||
acc: &mut Vec<AlgebraicExpression<T, R>>, | ||
) { | ||
match expr { | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left, | ||
op: AlgebraicBinaryOperator::Mul, | ||
right, | ||
}) => { | ||
flatten_mul(*left, acc); | ||
flatten_mul(*right, acc); | ||
} | ||
_ => acc.push(normalize(expr)), | ||
} | ||
} | ||
|
||
fn normalize<T: Ord + Clone, R: Ord + Clone>( | ||
expr: AlgebraicExpression<T, R>, | ||
) -> AlgebraicExpression<T, R> { | ||
match expr { | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left, | ||
op: AlgebraicBinaryOperator::Add, | ||
right, | ||
}) => { | ||
let mut terms = Vec::new(); | ||
flatten_add(*left, &mut terms); | ||
flatten_add(*right, &mut terms); | ||
terms.sort(); | ||
terms.into_iter().reduce(|a, b| a + b).unwrap() | ||
} | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left, | ||
op: AlgebraicBinaryOperator::Mul, | ||
right, | ||
}) => { | ||
let mut factors = Vec::new(); | ||
flatten_mul(*left, &mut factors); | ||
flatten_mul(*right, &mut factors); | ||
factors.sort(); | ||
factors.into_iter().reduce(|a, b| a * b).unwrap() | ||
} | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left, | ||
op: AlgebraicBinaryOperator::Sub, | ||
right, | ||
}) => normalize(*left) - normalize(*right), | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This seems sus, why would it be simpler than addition? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. its simpler because there's no need to handle commutativity (no need to flatten and reorder) |
||
_ => expr, | ||
} | ||
} | ||
|
||
fn expand<T: Ord + Clone, R: Ord + Clone>( | ||
expr: AlgebraicExpression<T, R>, | ||
) -> AlgebraicExpression<T, R> { | ||
match expr { | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left, | ||
op: AlgebraicBinaryOperator::Mul, | ||
right, | ||
}) => { | ||
let left = expand(*left); | ||
let right = expand(*right); | ||
match (left, right) { | ||
// (a + b)*c ==> a*c + b*c | ||
( | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left: a, | ||
op: AlgebraicBinaryOperator::Add, | ||
right: b, | ||
}), | ||
c, | ||
) => expand(*a * c.clone()) + expand(*b * c), | ||
// a*(b + c) ==> a*b + a*c | ||
( | ||
a, | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left: b, | ||
op: AlgebraicBinaryOperator::Add, | ||
right: c, | ||
}), | ||
) => expand(a.clone() * *b) + expand(a.clone() * *c), | ||
(left, right) => left * right, | ||
} | ||
} | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left, | ||
op: AlgebraicBinaryOperator::Add, | ||
right, | ||
}) => expand(*left) + expand(*right), | ||
AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { | ||
left, | ||
op: AlgebraicBinaryOperator::Sub, | ||
right, | ||
}) => expand(*left) - expand(*right), | ||
_ => expr, | ||
} | ||
} | ||
|
||
#[cfg(test)] | ||
mod tests { | ||
use super::*; | ||
type Expr = AlgebraicExpression<u32, String>; | ||
|
||
fn v(name: &str) -> Expr { | ||
Expr::Reference(name.to_string()) | ||
} | ||
|
||
#[test] | ||
fn test_normalize() { | ||
let expr1 = normalize(v("c") + v("b") + v("a")); | ||
let expr2 = normalize(v("a") + v("b") + v("c")); | ||
assert_eq!(expr1, expr2); | ||
|
||
let expr1 = normalize(v("c") * v("b") * v("a")); | ||
let expr2 = normalize(v("a") * v("b") * v("c")); | ||
assert_eq!(expr1, expr2); | ||
|
||
let expr1 = normalize(v("a") + v("b") * v("c")); | ||
let expr2 = normalize(v("c") * v("b") + v("a")); | ||
assert_eq!(expr1, expr2); | ||
} | ||
|
||
#[test] | ||
fn test_expand() { | ||
let e1 = expand(v("a") * (v("b") + v("c"))); | ||
let e2 = v("a") * v("b") + v("a") * v("c"); | ||
assert_eq!(e1, e2); | ||
|
||
let e1 = expand((v("a") + v("b")) * v("c")); | ||
let e2 = v("a") * v("c") + v("b") * v("c"); | ||
assert_eq!(e1, e2); | ||
|
||
// (a + b) * (c + d) | ||
let e1 = expand((v("a") + v("b")) * (v("c") + v("d"))); | ||
let e2 = v("a") * v("c") + v("a") * v("d") + v("b") * v("c") + v("b") * v("d"); | ||
assert_eq!(normalize(e1), normalize(e2)); | ||
|
||
// (a + b) * (c + d) * (e + f) | ||
let e1 = expand((v("a") + v("b")) * (v("c") + v("d")) * (v("e") + v("f"))); | ||
let e2 = v("a") * v("c") * v("e") | ||
+ v("a") * v("c") * v("f") | ||
+ v("a") * v("d") * v("e") | ||
+ v("a") * v("d") * v("f") | ||
+ v("b") * v("c") * v("e") | ||
+ v("b") * v("c") * v("f") | ||
+ v("b") * v("d") * v("e") | ||
+ v("b") * v("d") * v("f"); | ||
assert_eq!(normalize(e1), normalize(e2)); | ||
|
||
// a*(b + c) + d + b*(a + b) | ||
let expr1 = normalize(expand( | ||
v("a") * (v("b") + v("c")) + v("d") + v("b") * (v("a") + v("b")), | ||
)); | ||
// a*b + a*c + (d + b*a + b*b) | ||
let expr2 = normalize(expand( | ||
v("a") * v("b") + v("a") * v("c") + (v("d") + v("b") * v("a") + v("b") * v("b")), | ||
)); | ||
// (d + b*a) + b*b + (a*b + a*c) | ||
let expr3 = normalize(expand( | ||
(v("d") + v("b") * v("a")) + v("b") * v("b") + (v("a") * v("b") + v("a") * v("c")), | ||
)); | ||
assert_eq!(expr1, expr2); | ||
assert_eq!(expr1, expr3); | ||
} | ||
} | ||
|
||
impl<T, R> ExpressionConvertible<T, R> for AlgebraicExpression<T, R> { | ||
fn to_expression< | ||
E: Add<E, Output = E> + Sub<E, Output = E> + Mul<E, Output = E> + Neg<Output = E>, | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you add some documentation on this parameter?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What about this?