Skip to content

Commit

Permalink
Introduce concept of assignment. (#2244)
Browse files Browse the repository at this point in the history
  • Loading branch information
@chriseth
chriseth authored Dec 18, 2024
1 parent 18f6da0 commit 1b58d1e
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 52 deletions.
49 changes: 30 additions & 19 deletions executor/src/witgen/jit/block_machine_processor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,34 +47,22 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
let mut witgen = WitgenInference::new(self.fixed_data, self, known_variables);

// In the latch row, set the RHS selector to 1.
witgen.assign(
&connection_rhs.selector,
self.latch_row as i32,
T::one().into(),
)?;
witgen.assign_constant(&connection_rhs.selector, self.latch_row as i32, T::one());

// For each known argument, transfer the value to the expression in the connection's RHS.
// For each argument, connect the expression on the RHS with the formal parameter.
for (index, expr) in connection_rhs.expressions.iter().enumerate() {
if known_args[index] {
let param_i =
AffineSymbolicExpression::from_known_symbol(Variable::Param(index), None);
witgen.assign(expr, self.latch_row as i32, param_i)?;
}
witgen.assign_variable(expr, self.latch_row as i32, Variable::Param(index));
}

// Solve for the block witness.
// Fails if any machine call cannot be completed.
self.solve_block(&mut witgen)?;

// For each unknown argument, get the value from the expression in the connection's RHS.
// If assign() fails, it means that we weren't able to to solve for the operation's output.
for (index, expr) in connection_rhs.expressions.iter().enumerate() {
if !known_args[index] {
let param_i =
AffineSymbolicExpression::from_unknown_variable(Variable::Param(index), None);
witgen
.assign(expr, self.latch_row as i32, param_i)
.map_err(|_| format!("Could not solve for params[{index}]"))?;
if !witgen.is_known(&Variable::Param(index)) {
return Err(format!(
"Unable to derive algorithm to compute output value \"{expr}\""
));
}
}

Expand Down Expand Up @@ -265,6 +253,29 @@ params[2] = Add::c[0];"
);
}

#[test]
fn unconstrained_output() {
let input = "
namespace Unconstrained(256);
col witness sel, a, b, c;
a + b = 0;
";
let err_str = generate_for_block_machine(
input,
1,
0,
"Unconstrained::sel",
&["Unconstrained::a", "Unconstrained::b"],
&["Unconstrained::c"],
)
.err()
.unwrap();
assert_eq!(
err_str,
"Unable to derive algorithm to compute output value \"Unconstrained::c\""
);
}

#[test]
// TODO: Currently fails, because the machine has a non-rectangular block shape.
#[should_panic = "Incomplete machine calls"]
Expand Down
130 changes: 97 additions & 33 deletions executor/src/witgen/jit/witgen_inference.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@ pub struct WitgenInference<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> {
fixed_evaluator: FixedEval,
derived_range_constraints: HashMap<Variable, RangeConstraint<T>>,
known_variables: HashSet<Variable>,
/// Internal equalities we were not able to solve yet.
assignments: Vec<Assignment<'a, T>>,
code: Vec<Effect<T, Variable>>,
}

Expand All @@ -53,6 +55,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
fixed_evaluator,
derived_range_constraints: Default::default(),
known_variables: known_variables.into_iter().collect(),
assignments: Default::default(),
code: Default::default(),
}
}
Expand All @@ -61,11 +64,15 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
self.code
}

pub fn is_known(&self, variable: &Variable) -> bool {
self.known_variables.contains(variable)
}

/// Process an identity on a certain row.
pub fn process_identity(&mut self, id: &Identity<T>, row_offset: i32) -> ProcessSummary {
let result = match id {
Identity::Polynomial(PolynomialIdentity { expression, .. }) => {
self.process_polynomial_identity(expression, row_offset)
self.process_equality_on_row(expression, row_offset, T::from(0).into())
}
Identity::Lookup(LookupIdentity {
id, left, right, ..
Expand All @@ -89,38 +96,63 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
self.ingest_effects(result)
}

/// Process the constraint that the expression evaluated at the given offset equals the given affine expression.
/// Note that either the expression or the value might contain unknown variables, but if we are not able to
/// solve the equation, we return an error.
pub fn assign(
&mut self,
expression: &Expression<T>,
offset: i32,
value: AffineSymbolicExpression<T, Variable>,
) -> Result<(), String> {
let affine_expression = self
.evaluate(expression, offset)
.ok_or_else(|| format!("Expression is not affine: {expression}"))?;
let result = (affine_expression - value.clone())
.solve()
.map_err(|err| format!("Could not solve ({expression} - {value}): {err}"))?;
match self.ingest_effects(result).complete {
true => Ok(()),
false => Err("Wasn't able to complete the assignment".to_string()),
/// Turns the given variable either to a known symbolic value or an unknown symbolic value
/// depending on if it is known or not.
/// If it is known to be range-constrained to a single value, that value is used.
fn variable_to_expression(&self, variable: Variable) -> AffineSymbolicExpression<T, Variable> {
// If a variable is known and has a compile-time constant value,
// that value is stored in the range constraints.
let rc = self.range_constraint(variable.clone());
if let Some(val) = rc.as_ref().and_then(|rc| rc.try_to_single_value()) {
val.into()
} else if self.known_variables.contains(&variable) {
AffineSymbolicExpression::from_known_symbol(variable, rc)
} else {
AffineSymbolicExpression::from_unknown_variable(variable, rc)
}
}

fn process_polynomial_identity(
/// Process the constraint that the expression evaluated at the given offset equals the given value.
/// This does not have to be solvable right away, but is always processed as soon as we have progress.
/// Note that all variables in the expression can be unknown and their status can also change over time.
pub fn assign_constant(&mut self, expression: &'a Expression<T>, row_offset: i32, value: T) {
self.assignments.push(Assignment {
lhs: expression,
row_offset,
rhs: VariableOrValue::Value(value),
});
self.process_assignments();
}

/// Process the constraint that the expression evaluated at the given offset equals the given formal variable.
/// This does not have to be solvable right away, but is always processed as soon as we have progress.
/// Note that all variables in the expression can be unknown and their status can also change over time.
pub fn assign_variable(
&mut self,
expression: &'a Expression<T>,
row_offset: i32,
variable: Variable,
) {
self.assignments.push(Assignment {
lhs: expression,
row_offset,
rhs: VariableOrValue::Variable(variable),
});
self.process_assignments();
}

fn process_equality_on_row(
&self,
expression: &Expression<T>,
lhs: &Expression<T>,
offset: i32,
rhs: AffineSymbolicExpression<T, Variable>,
) -> ProcessResult<T, Variable> {
if let Some(r) = self.evaluate(expression, offset) {
if let Some(r) = self.evaluate(lhs, offset) {
// TODO propagate or report error properly.
// If solve returns an error, it means that the constraint is conflicting.
// In the future, we might run this in a runtime-conditional, so an error
// could just mean that this case cannot happen in practice.
r.solve().unwrap()
(r - rhs).solve().unwrap()
} else {
ProcessResult::empty()
}
Expand Down Expand Up @@ -180,6 +212,31 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
ProcessResult::empty()
}

fn process_assignments(&mut self) {
loop {
let mut progress = false;
let new_assignments = std::mem::take(&mut self.assignments)
.into_iter()
.flat_map(|assignment| {
let rhs = match &assignment.rhs {
VariableOrValue::Variable(v) => self.variable_to_expression(v.clone()),
VariableOrValue::Value(v) => (*v).into(),
};
let r =
self.process_equality_on_row(assignment.lhs, assignment.row_offset, rhs);
let summary = self.ingest_effects(r);
progress |= summary.progress;
// If it is not complete, queue it again.
(!summary.complete).then_some(assignment)
})
.collect_vec();
self.assignments.extend(new_assignments);
if !progress {
break;
}
}
}

fn ingest_effects(&mut self, process_result: ProcessResult<T, Variable>) -> ProcessSummary {
let mut progress = false;
for e in process_result.effects {
Expand Down Expand Up @@ -210,6 +267,9 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
Effect::Assertion(_) => self.code.push(e),
}
}
if progress {
self.process_assignments();
}
ProcessSummary {
complete: process_result.complete,
progress,
Expand Down Expand Up @@ -248,16 +308,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
self.fixed_evaluator.evaluate(r, offset)?.into()
} else {
let variable = Variable::from_reference(r, offset);
// If a variable is known and has a compile-time constant value,
// that value is stored in the range constraints.
let rc = self.range_constraint(variable.clone());
if let Some(val) = rc.as_ref().and_then(|rc| rc.try_to_single_value()) {
val.into()
} else if self.known_variables.contains(&variable) {
AffineSymbolicExpression::from_known_symbol(variable, rc)
} else {
AffineSymbolicExpression::from_unknown_variable(variable, rc)
}
self.variable_to_expression(variable)
}
}
Expression::PublicReference(_) | Expression::Challenge(_) => {
Expand Down Expand Up @@ -323,6 +374,19 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
}
}

/// An equality constraint between an algebraic expression evaluated
/// on a certain row offset and a variable or fixed constant value.
struct Assignment<'a, T: FieldElement> {
lhs: &'a Expression<T>,
row_offset: i32,
rhs: VariableOrValue<T, Variable>,
}

enum VariableOrValue<T, V> {
Variable(V),
Value(T),
}

pub trait FixedEvaluator<T: FieldElement> {
fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option<T> {
None
Expand Down

0 comments on commit 1b58d1e

Please sign in to comment.