Skip to content

Commit

Permalink
Propagate known
Browse files Browse the repository at this point in the history
  • Loading branch information
@chriseth
chriseth committed Jan 10, 2025
1 parent c72f646 commit 5ac787c
Show file tree
Hide file tree
Showing 2 changed files with 95 additions and 44 deletions.
11 changes: 4 additions & 7 deletions executor/src/witgen/jit/processor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -176,15 +176,12 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
while progress {
progress = false;

// TODO At this point, we should call a function on `witgen`
// to propagate known concrete values across the identities
// to other known (but not concrete) variables.
witgen.process_assignments();

for (id, row_offset) in &self.identities {
if complete.contains(&(id.id(), *row_offset)) {
continue;
}
let result = witgen.process_identity(can_process.clone(), id, *row_offset);
let is_complete = complete.contains(&(id.id(), *row_offset));
let result =
witgen.process_identity(can_process.clone(), is_complete, id, *row_offset);
progress |= result.progress;
if result.complete {
complete.insert((id.id(), *row_offset));
Expand Down
128 changes: 91 additions & 37 deletions executor/src/witgen/jit/witgen_inference.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ use powdr_number::FieldElement;

use crate::witgen::{
data_structures::mutable_state::MutableState, global_constraints::RangeConstraintSet,
range_constraints::RangeConstraint, FixedData, QueryCallback,
jit::effect::format_code, range_constraints::RangeConstraint, FixedData, QueryCallback,
};

use super::{
Expand Down Expand Up @@ -145,13 +145,17 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
pub fn process_identity<CanProcess: CanProcessCall<T>>(
&mut self,
can_process: CanProcess,
is_complete: bool,
id: &'a Identity<T>,
row_offset: i32,
) -> ProcessSummary {
let result = match id {
Identity::Polynomial(PolynomialIdentity { expression, .. }) => {
self.process_equality_on_row(expression, row_offset, T::from(0).into())
}
Identity::Polynomial(PolynomialIdentity { expression, .. }) => self
.process_equality_on_row(
expression,
row_offset,
&VariableOrValue::Value(T::from(0)),
),
Identity::Lookup(LookupIdentity { id, left, .. })
| Identity::Permutation(PermutationIdentity { id, left, .. })
| Identity::PhantomPermutation(PhantomPermutationIdentity { id, left, .. })
Expand All @@ -166,7 +170,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
Identity::PhantomBusInteraction(_) => ProcessResult::empty(),
Identity::Connect(_) => ProcessResult::empty(),
};
self.ingest_effects(result)
self.ingest_effects(is_complete, result)
}

/// Process the constraint that the expression evaluated at the given offset equals the given value.
Expand Down Expand Up @@ -202,17 +206,45 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
&self,
lhs: &Expression<T>,
offset: i32,
rhs: AffineSymbolicExpression<T, Variable>,
rhs: &VariableOrValue<T, Variable>,
) -> ProcessResult<T, Variable> {
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 - rhs).solve().unwrap()
} else {
ProcessResult::empty()
// First we try to find a new assignment to a variable in the equality.

let evaluator = Evaluator::new(self);
let Some(lhs_evaluated) = evaluator.evaluate(lhs, offset) else {
return ProcessResult::empty();
};

let rhs_evaluated = match rhs {
VariableOrValue::Variable(v) => evaluator.evaluate_variable(v.clone()),
VariableOrValue::Value(v) => (*v).into(),
};

// 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.
let result = (lhs_evaluated - rhs_evaluated).solve().unwrap();
if result.complete && result.effects.is_empty() {
// A complete result without effects means that there were no unknowns
// in the constraint.
// We try again, but this time we treat all non-concrete variables
// as unknown and in that way try to find new concrete values for
// already known variables.
let evaluator = Evaluator::new(self).only_concrete_known();
if let Some(lhs_evaluated) = evaluator.evaluate(lhs, offset) {
let rhs_evaluated = match rhs {
VariableOrValue::Variable(v) => evaluator.evaluate_variable(v.clone()),
VariableOrValue::Value(v) => (*v).into(),
};
let result = (lhs_evaluated - rhs_evaluated).solve().unwrap();
if !result.effects.is_empty() {
println!("Got new effects:\n{}", format_code(&result.effects));
return result;
}
}
}
result
}

fn process_call<CanProcess: CanProcessCall<T>>(
Expand Down Expand Up @@ -280,21 +312,20 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
}
}

fn process_assignments(&mut self) {
pub fn process_assignments(&mut self) {
// TODO we need to keep old assignments because we need them
// for propagating range constraints.
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) => {
Evaluator::new(self).evaluate_variable(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);
let r = self.process_equality_on_row(
assignment.lhs,
assignment.row_offset,
&assignment.rhs,
);
let summary = self.ingest_effects(false, r);
progress |= summary.progress;
// If it is not complete, queue it again.
(!summary.complete).then_some(assignment)
Expand All @@ -307,34 +338,57 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
}
}

fn ingest_effects(&mut self, process_result: ProcessResult<T, Variable>) -> ProcessSummary {
fn ingest_effects(
&mut self,
is_complete: bool,
process_result: ProcessResult<T, Variable>,
) -> ProcessSummary {
let mut progress = false;
for e in process_result.effects {
match &e {
Effect::Assignment(variable, assignment) => {
assert!(self.known_variables.insert(variable.clone()));
// If the variable was determined to be a constant, we add this
// as a range constraint, so we can use it in future evaluations.
self.add_range_constraint(variable.clone(), assignment.range_constraint());
progress = true;
self.code.push(e);
progress |=
self.add_range_constraint(variable.clone(), assignment.range_constraint());
if !self.is_known(variable) {
self.known_variables.insert(variable.clone());
progress = true;
self.code.push(e);
}
}
Effect::RangeConstraint(variable, rc) => {
progress |= self.add_range_constraint(variable.clone(), rc.clone());
}
Effect::MachineCall(_, _, vars) => {
for v in vars {
// Inputs are already known, but it does not hurt to add all of them.
self.known_variables.insert(v.clone());
// TODO this is the only place we actually need the "is complete"
// property. Maybe we should rather record the submachine calls
// per row and ignore them if we re-call?
// Also it probably makes sense to store "completeness" in witgen_inference.

// If the machine call is already complete, it means that we have already
// processed it and created code for it. This run is just to potentially
// get better range constraints, so we ignore the actual machine call.
if !is_complete {
for v in vars {
// Inputs are already known, but it does not hurt to add all of them.
self.known_variables.insert(v.clone());
}
progress = true;

self.code.push(e);
}
progress = true;
}
Effect::Assertion(_) => {
assert!(!is_complete);
self.code.push(e);
}
Effect::Assertion(_) => self.code.push(e),
Effect::Branch(..) => unreachable!(),
}
}
if progress {
// TODO in some cases (after a branch), we might need to process assignments
// even if there was no progress in the identities above.
self.process_assignments();
}
ProcessSummary {
Expand Down Expand Up @@ -613,9 +667,9 @@ mod test {
counter += 1;
for row in rows {
for id in retained_identities.iter() {
if !complete.contains(&(id.id(), *row))
&& witgen.process_identity(&mutable_state, id, *row).complete
{
let is_complete = complete.contains(&(id.id(), *row));
let result = witgen.process_identity(&mutable_state, is_complete, id, *row);
if result.complete {
complete.insert((id.id(), *row));
}
}
Expand Down

0 comments on commit 5ac787c

Please sign in to comment.