Skip to content

Commit

Permalink
Exported shared code to temp variables in ExprEvaluator. (#893)
Browse files Browse the repository at this point in the history
  • Loading branch information
Alon-Ti authored Nov 26, 2024
1 parent af9250e commit 4d64300
Show file tree
Hide file tree
Showing 2 changed files with 91 additions and 211 deletions.
240 changes: 62 additions & 178 deletions crates/prover/src/constraint_framework/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ impl Expr {
Expr::Sub(a, b) => format!("{} - ({})", a.format_expr(), b.format_expr()),
Expr::Mul(a, b) => format!("({}) * ({})", a.format_expr(), b.format_expr()),
Expr::Neg(a) => format!("-({})", a.format_expr()),
Expr::Inv(a) => format!("1/({})", a.format_expr()),
Expr::Inv(a) => format!("1 / ({})", a.format_expr()),
}
}
}
Expand Down Expand Up @@ -248,6 +248,7 @@ pub struct ExprEvaluator {
pub cur_var_index: usize,
pub constraints: Vec<Expr>,
pub logup: FormalLogupAtRow,
pub intermediates: Vec<(String, Expr)>,
}

impl ExprEvaluator {
Expand All @@ -257,8 +258,35 @@ impl ExprEvaluator {
cur_var_index: Default::default(),
constraints: Default::default(),
logup: FormalLogupAtRow::new(INTERACTION_TRACE_IDX, has_partial_sum, log_size),
intermediates: vec![],
}
}

pub fn add_intermediate(&mut self, expr: Expr) -> Expr {
let name = format!("intermediate{}", self.intermediates.len());
let intermediate = Expr::Param(name.clone());
self.intermediates.push((name, expr));
intermediate
}

pub fn format_constraints(&self) -> String {
let lets_string = self
.intermediates
.iter()
.map(|(name, expr)| format!("let {} = {};", name, expr.format_expr()))
.collect::<Vec<String>>()
.join("\n");

let constraints_str = self
.constraints
.iter()
.enumerate()
.map(|(i, c)| format!("let constraint_{i} = ") + &c.format_expr() + ";")
.collect::<Vec<String>>()
.join("\n\n");

lets_string + "\n\n" + &constraints_str
}
}

impl EvalAtRow for ExprEvaluator {
Expand Down Expand Up @@ -286,7 +314,15 @@ impl EvalAtRow for ExprEvaluator {
where
Self::EF: std::ops::Mul<G, Output = Self::EF>,
{
self.constraints.push(Expr::one() * constraint);
match Expr::one() * constraint {
Expr::Mul(one, constraint) => {
assert_eq!(*one, Expr::one());
self.constraints.push(*constraint);
}
_ => {
unreachable!();
}
}
}

fn combine_ef(values: [Self::F; 4]) -> Self::EF {
Expand All @@ -310,7 +346,8 @@ impl EvalAtRow for ExprEvaluator {
multiplicity,
values,
}| {
Fraction::new(multiplicity.clone(), combine_formal(*relation, values))
let intermediate = self.add_intermediate(combine_formal(*relation, values));
Fraction::new(multiplicity.clone(), intermediate)
},
)
.collect();
Expand All @@ -324,187 +361,34 @@ impl EvalAtRow for ExprEvaluator {
mod tests {
use num_traits::One;

use crate::constraint_framework::expr::{ColumnExpr, Expr, ExprEvaluator};
use crate::constraint_framework::{
relation, EvalAtRow, FrameworkEval, RelationEntry, ORIGINAL_TRACE_IDX,
};
use crate::core::fields::m31::M31;
use crate::constraint_framework::expr::ExprEvaluator;
use crate::constraint_framework::{relation, EvalAtRow, FrameworkEval, RelationEntry};
use crate::core::fields::FieldExpOps;

#[test]
fn test_expr_eval() {
let test_struct = TestStruct {};
let eval = test_struct.evaluate(ExprEvaluator::new(16, false));
assert_eq!(eval.constraints.len(), 2);
assert_eq!(
eval.constraints[0],
Expr::Mul(
Box::new(Expr::one()),
Box::new(Expr::Mul(
Box::new(Expr::Mul(
Box::new(Expr::Mul(
Box::new(Expr::Col(ColumnExpr {
interaction: ORIGINAL_TRACE_IDX,
idx: 0,
offset: 0
})),
Box::new(Expr::Col(ColumnExpr {
interaction: ORIGINAL_TRACE_IDX,
idx: 1,
offset: 0
}))
)),
Box::new(Expr::Col(ColumnExpr {
interaction: ORIGINAL_TRACE_IDX,
idx: 2,
offset: 0
}))
)),
Box::new(Expr::Inv(Box::new(Expr::Add(
Box::new(Expr::Col(ColumnExpr {
interaction: ORIGINAL_TRACE_IDX,
idx: 0,
offset: 0
})),
Box::new(Expr::Col(ColumnExpr {
interaction: ORIGINAL_TRACE_IDX,
idx: 1,
offset: 0
}))
))))
))
)
);

assert_eq!(
eval.constraints[1],
Expr::Mul(
Box::new(Expr::Const(M31(1))),
Box::new(Expr::Sub(
Box::new(Expr::Mul(
Box::new(Expr::Sub(
Box::new(Expr::Sub(
Box::new(Expr::SecureCol([
Box::new(Expr::Col(ColumnExpr {
interaction: 2,
idx: 4,
offset: 0
})),
Box::new(Expr::Col(ColumnExpr {
interaction: 2,
idx: 6,
offset: 0
})),
Box::new(Expr::Col(ColumnExpr {
interaction: 2,
idx: 8,
offset: 0
})),
Box::new(Expr::Col(ColumnExpr {
interaction: 2,
idx: 10,
offset: 0
}))
])),
Box::new(Expr::Sub(
Box::new(Expr::SecureCol([
Box::new(Expr::Col(ColumnExpr {
interaction: 2,
idx: 5,
offset: -1
})),
Box::new(Expr::Col(ColumnExpr {
interaction: 2,
idx: 7,
offset: -1
})),
Box::new(Expr::Col(ColumnExpr {
interaction: 2,
idx: 9,
offset: -1
})),
Box::new(Expr::Col(ColumnExpr {
interaction: 2,
idx: 11,
offset: -1
}))
])),
Box::new(Expr::Mul(
Box::new(Expr::Col(ColumnExpr {
interaction: 0,
idx: 3,
offset: 0
})),
Box::new(Expr::Param("total_sum".into()))
))
))
)),
Box::new(Expr::Const(M31(0)))
)),
Box::new(Expr::Sub(
Box::new(Expr::Add(
Box::new(Expr::Add(
Box::new(Expr::Add(
Box::new(Expr::Const(M31(0))),
Box::new(Expr::Mul(
Box::new(Expr::Param(
"TestRelation_alpha0".to_string()
)),
Box::new(Expr::Col(ColumnExpr {
interaction: 1,
idx: 0,
offset: 0
}))
))
)),
Box::new(Expr::Mul(
Box::new(Expr::Param("TestRelation_alpha1".to_string())),
Box::new(Expr::Col(ColumnExpr {
interaction: 1,
idx: 1,
offset: 0
}))
))
)),
Box::new(Expr::Mul(
Box::new(Expr::Param("TestRelation_alpha2".to_string())),
Box::new(Expr::Col(ColumnExpr {
interaction: 1,
idx: 2,
offset: 0
}))
))
)),
Box::new(Expr::Param("TestRelation_z".to_string()))
))
)),
Box::new(Expr::Const(M31(1)))
))
)
);
}

#[test]
fn test_format_expr() {
let test_struct = TestStruct {};
let eval = test_struct.evaluate(ExprEvaluator::new(16, false));
let constraint0_str = "(1) * ((((col_1_0[0]) * (col_1_1[0])) * (col_1_2[0])) * (1/(col_1_0[0] + col_1_1[0])))";
assert_eq!(eval.constraints[0].format_expr(), constraint0_str);
let constraint1_str = "(1) \
* ((SecureCol(col_2_4[0], col_2_6[0], col_2_8[0], col_2_10[0]) \
- (SecureCol(\
col_2_5[-1], \
col_2_7[-1], \
col_2_9[-1], \
col_2_11[-1]\
) - ((col_0_3[0]) * (total_sum))) \
- (0)) \
* (0 + (TestRelation_alpha0) * (col_1_0[0]) \
+ (TestRelation_alpha1) * (col_1_1[0]) \
+ (TestRelation_alpha2) * (col_1_2[0]) \
- (TestRelation_z)) \
- (1))";
assert_eq!(eval.constraints[1].format_expr(), constraint1_str);
let expected = "let intermediate0 = 0 \
+ (TestRelation_alpha0) * (col_1_0[0]) \
+ (TestRelation_alpha1) * (col_1_1[0]) \
+ (TestRelation_alpha2) * (col_1_2[0]) \
- (TestRelation_z);
\
let constraint_0 = \
(((col_1_0[0]) * (col_1_1[0])) * (col_1_2[0])) * (1 / (col_1_0[0] + col_1_1[0]));
\
let constraint_1 = (SecureCol(col_2_4[0], col_2_6[0], col_2_8[0], col_2_10[0]) \
- (SecureCol(col_2_5[-1], col_2_7[-1], col_2_9[-1], col_2_11[-1]) \
- ((col_0_3[0]) * (total_sum))) \
- (0)) \
* (intermediate0) \
- (1);"
.to_string();

assert_eq!(eval.format_constraints(), expected);
}

relation!(TestRelation, 3);
Expand Down
62 changes: 29 additions & 33 deletions crates/prover/src/examples/state_machine/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -300,38 +300,34 @@ mod tests {
);

let eval = component.evaluate(ExprEvaluator::new(log_n_rows, true));

assert_eq!(eval.constraints.len(), 2);
let constraint0_str = "(1) \
* ((SecureCol(\
col_2_5[claimed_sum_offset], \
col_2_8[claimed_sum_offset], \
col_2_11[claimed_sum_offset], \
col_2_14[claimed_sum_offset]\
) - (claimed_sum)) \
* (col_0_2[0]))";
assert_eq!(eval.constraints[0].format_expr(), constraint0_str);
let constraint1_str = "(1) \
* ((SecureCol(col_2_3[0], col_2_6[0], col_2_9[0], col_2_12[0]) \
- (SecureCol(col_2_4[-1], col_2_7[-1], col_2_10[-1], col_2_13[-1]) \
- ((col_0_2[0]) * (total_sum))) \
- (0)) \
* ((0 \
+ (StateMachineElements_alpha0) * (col_1_0[0]) \
+ (StateMachineElements_alpha1) * (col_1_1[0]) \
- (StateMachineElements_z)) \
* (0 + (StateMachineElements_alpha0) * (col_1_0[0] + 1) \
+ (StateMachineElements_alpha1) * (col_1_1[0]) \
- (StateMachineElements_z))) \
- ((0 \
+ (StateMachineElements_alpha0) * (col_1_0[0] + 1) \
+ (StateMachineElements_alpha1) * (col_1_1[0]) \
- (StateMachineElements_z)) \
* (1) \
+ (0 + (StateMachineElements_alpha0) * (col_1_0[0]) \
+ (StateMachineElements_alpha1) * (col_1_1[0]) \
- (StateMachineElements_z)) \
* (-(1))))";
assert_eq!(eval.constraints[1].format_expr(), constraint1_str);
let expected = "let intermediate0 = 0 \
+ (StateMachineElements_alpha0) * (col_1_0[0]) \
+ (StateMachineElements_alpha1) * (col_1_1[0]) \
- (StateMachineElements_z);
\
let intermediate1 = 0 \
+ (StateMachineElements_alpha0) * (col_1_0[0] + 1) \
+ (StateMachineElements_alpha1) * (col_1_1[0]) \
- (StateMachineElements_z);
\
let constraint_0 = (SecureCol(\
col_2_5[claimed_sum_offset], \
col_2_8[claimed_sum_offset], \
col_2_11[claimed_sum_offset], \
col_2_14[claimed_sum_offset]\
) - (claimed_sum)) \
* (col_0_2[0]);
\
let constraint_1 = (SecureCol(col_2_3[0], col_2_6[0], col_2_9[0], col_2_12[0]) \
- (SecureCol(col_2_4[-1], col_2_7[-1], col_2_10[-1], col_2_13[-1]) \
- ((col_0_2[0]) * (total_sum))) \
- (0)) \
* ((intermediate0) * (intermediate1)) \
- ((intermediate1) * (1) + (intermediate0) * (-(1)));"
.to_string();

assert_eq!(eval.format_constraints(), expected);
}
}

1 comment on commit 4d64300

@github-actions
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 2.

Benchmark suite Current: 4d64300 Previous: cd8b37b Ratio
iffts/simd ifft/22 13217317 ns/iter (± 165573) 6306399 ns/iter (± 210024) 2.10
merkle throughput/simd merkle 28750753 ns/iter (± 369354) 13712527 ns/iter (± 579195) 2.10

This comment was automatically generated by workflow using github-action-benchmark.

CC: @shaharsamocha7

Please sign in to comment.