Skip to content
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

Determine range constraints from fixed lookup. #2300

Merged
merged 11 commits into from
Jan 13, 2025
9 changes: 3 additions & 6 deletions executor/src/witgen/data_structures/mutable_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,15 +56,12 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
identity_id: u64,
known_inputs: &BitVec,
range_constraints: &[RangeConstraint<T>],
) -> bool {
) -> Option<Vec<RangeConstraint<T>>> {
// TODO We are currently ignoring bus interaction (also, but not only because there is no
// unique machine responsible for handling a bus send), so just answer "false" if the identity
// has no responsible machine.
self.responsible_machine(identity_id)
.ok()
.is_some_and(|mut machine| {
machine.can_process_call_fully(identity_id, known_inputs, range_constraints)
})
let mut machine = self.responsible_machine(identity_id).ok()?;
machine.can_process_call_fully(identity_id, known_inputs, range_constraints)
}

/// Call the machine responsible for the right-hand-side of an identity given its ID
Expand Down
21 changes: 21 additions & 0 deletions executor/src/witgen/jit/affine_symbolic_expression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,27 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
}
}

/// Returns the range constraint of the whole expression.
/// This only works for simple expressions since all coefficients
/// must be known numbers.
pub fn range_constraint(&self) -> RangeConstraint<T> {
self.coefficients
.iter()
.map(|(var, coeff)| {
let coeff = coeff.try_to_number()?;
let rc = self.range_constraints.get(var)?;
Some(rc.multiple(coeff))
})
.collect::<Option<Vec<_>>>()
.and_then(|summands| {
summands
.into_iter()
.chain(std::iter::once(self.offset.range_constraint()))
.reduce(|c1, c2| c1.combine_sum(&c2))
})
.unwrap_or_default()
}

/// If this expression contains a single unknown variable, returns it.
pub fn single_unknown_variable(&self) -> Option<&V> {
if self.coefficients.len() == 1 {
Expand Down
36 changes: 36 additions & 0 deletions executor/src/witgen/jit/single_step_processor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -205,4 +205,40 @@ if (VM::instr_add[0] == 1) {
}"
);
}

#[test]
fn range_constraints_from_lookup() {
let input = "
namespace VM(256);
let instr_add: col;
let instr_mul: col;
let pc: col;

col fixed LINE = [0, 1] + [2]*;
col fixed INSTR_ADD = [0, 1] + [0]*;
col fixed INSTR_MUL = [1, 0] + [1]*;

pc' = pc + 1;
instr_add = 0;
[ pc, instr_add, instr_mul ] in [ LINE, INSTR_ADD, INSTR_MUL ];

";
let code = generate_single_step(input, "Main").unwrap();
// After the machine call, we should have a direct assignment `VM::instr_mul[1] = 1`,
// instead of just an assignment from the call variable.
// This is because the fixed lookup machine can already provide a range constraint.
// For reasons of processing order, the call variable will also be assigned
// right before the call.
assert_eq!(
format_code(&code),
"\
VM::pc[1] = (VM::pc[0] + 1);
VM::instr_add[1] = 0;
call_var(2, 1, 0) = VM::pc[1];
call_var(2, 1, 1) = 0;
call_var(2, 1, 2) = 1;
machine_call(2, [Known(call_var(2, 1, 0)), Known(call_var(2, 1, 1)), Unknown(call_var(2, 1, 2))]);
VM::instr_mul[1] = 1;"
);
}
}
31 changes: 19 additions & 12 deletions executor/src/witgen/jit/witgen_inference.rs
Original file line number Diff line number Diff line change
Expand Up @@ -233,43 +233,49 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
}
let evaluated = arguments
.iter()
.map(|a| {
self.evaluate(a, row_offset)
.and_then(|a| a.try_to_known().cloned())
})
.map(|a| self.evaluate(a, row_offset))
.collect::<Vec<_>>();
let range_constraints = evaluated
.iter()
.map(|e| e.as_ref().map(|e| e.range_constraint()).unwrap_or_default())
.collect_vec();
let known = evaluated.iter().map(|e| e.is_some()).collect();
let known: BitVec = evaluated
.iter()
.map(|e| e.as_ref().and_then(|e| e.try_to_known()).is_some())
.collect();

if !can_process_call.can_process_call_fully(lookup_id, &known, &range_constraints) {
let Some(new_range_constraints) =
can_process_call.can_process_call_fully(lookup_id, &known, &range_constraints)
else {
log::trace!(
"Sub-machine cannot process call fully (will retry later): {lookup_id}, arguments: {}",
arguments.iter().zip(known).map(|(arg, known)| {
format!("{arg} [{}]", if known { "known" } else { "unknown" })
}).format(", "));
return ProcessResult::empty();
}
};
let mut effects = vec![];
let vars = arguments
.iter()
.zip_eq(new_range_constraints)
.enumerate()
.map(|(index, arg)| {
.map(|(index, (arg, new_rc))| {
let var = Variable::MachineCallParam(MachineCallVariable {
identity_id: lookup_id,
row_offset,
index,
});
self.assign_variable(arg, row_offset, var.clone());
effects.push(Effect::RangeConstraint(var.clone(), new_rc.clone()));
if known[index] {
assert!(self.is_known(&var));
}
var
})
.collect_vec();
effects.push(Effect::MachineCall(lookup_id, known, vars.clone()));
ProcessResult {
effects: vec![Effect::MachineCall(lookup_id, known, vars)],
effects,
complete: true,
}
}
Expand Down Expand Up @@ -514,16 +520,17 @@ pub trait FixedEvaluator<T: FieldElement>: Clone {
}

pub trait CanProcessCall<T: FieldElement> {
/// Returns true if a call to the machine that handles the given identity
/// Returns Some(..) if a call to the machine that handles the given identity
/// can always be processed with the given known inputs and range constraints
/// on the parameters.
/// The value in the Option is a vector of new range constraints.
/// @see Machine::can_process_call
fn can_process_call_fully(
&self,
_identity_id: u64,
_known_inputs: &BitVec,
_range_constraints: &[RangeConstraint<T>],
) -> bool;
) -> Option<Vec<RangeConstraint<T>>>;
}

impl<T: FieldElement, Q: QueryCallback<T>> CanProcessCall<T> for &MutableState<'_, T, Q> {
Expand All @@ -532,7 +539,7 @@ impl<T: FieldElement, Q: QueryCallback<T>> CanProcessCall<T> for &MutableState<'
identity_id: u64,
known_inputs: &BitVec,
range_constraints: &[RangeConstraint<T>],
) -> bool {
) -> Option<Vec<RangeConstraint<T>>> {
MutableState::can_process_call_fully(self, identity_id, known_inputs, range_constraints)
}
}
Expand Down
11 changes: 6 additions & 5 deletions executor/src/witgen/machines/double_sorted_witness_machine_32.rs
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses32<'a, T> {
identity_id: u64,
known_arguments: &BitVec,
range_constraints: &[RangeConstraint<T>],
) -> bool {
) -> Option<Vec<RangeConstraint<T>>> {
assert!(self.parts.connections.contains_key(&identity_id));
assert_eq!(known_arguments.len(), 4);
assert_eq!(range_constraints.len(), 4);
Expand All @@ -200,17 +200,18 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses32<'a, T> {

// We need to known operation_id, step and address for all calls.
if !known_arguments[0] || !known_arguments[1] || !known_arguments[2] {
return false;
return None;
}

// For the value, it depends: If we write, we need to know it, if we read we do not need to know it.
if known_arguments[3] {
// It is known, so we are good anyway.
true
Some(vec![RangeConstraint::unconstrained(); 4])
} else {
// It is not known, so we can only process if we do not write.
!range_constraints[0].allows_value(T::from(OPERATION_ID_BOOTLOADER_WRITE))
&& !range_constraints[0].allows_value(T::from(OPERATION_ID_WRITE))
(!range_constraints[0].allows_value(T::from(OPERATION_ID_BOOTLOADER_WRITE))
&& !range_constraints[0].allows_value(T::from(OPERATION_ID_WRITE)))
.then(|| vec![RangeConstraint::unconstrained(); 4])
}
}

Expand Down
85 changes: 80 additions & 5 deletions executor/src/witgen/machines/fixed_lookup_machine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -260,10 +260,10 @@ impl<'a, T: FieldElement> Machine<'a, T> for FixedLookup<'a, T> {
&mut self,
identity_id: u64,
known_arguments: &BitVec,
_range_constraints: &[RangeConstraint<T>],
) -> bool {
range_constraints: &[RangeConstraint<T>],
) -> Option<Vec<RangeConstraint<T>>> {
if !Self::is_responsible(&self.connections[&identity_id]) {
return false;
return None;
}
let index = self
.indices
Expand All @@ -274,8 +274,72 @@ impl<'a, T: FieldElement> Machine<'a, T> for FixedLookup<'a, T> {
.or_insert_with_key(|application| {
create_index(self.fixed_data, application, &self.connections)
});
// Check that if there is a match, it is unique.
index.values().all(|value| value.0.is_some())
let input_range_constraints = known_arguments
.iter()
.zip_eq(range_constraints)
.filter_map(|(is_known, range_constraint)| is_known.then_some(range_constraint.clone()))
.collect_vec();

// Now only consider the index entries that match the input range constraints,
// see that the result is unique and determine new output range constraints.
let values_matching_input_constraints = index
.iter()
.filter(|(key, _)| matches_range_constraint(key, &input_range_constraints))
.map(|(_, value)| {
let (_, values) = value.get()?;
Some(values)
});
let mut new_range_constraints: Option<Vec<(T, T, T::Integer)>> = None;
for values in values_matching_input_constraints {
// If any value is None, it means the lookup does not have a unique answer,
// and thus we cannot process the call.
let values = values?;
new_range_constraints = Some(match new_range_constraints {
// First item, turn each value into (min, max, mask).
None => values
.iter()
.map(|v| (*v, *v, v.to_integer()))
.collect_vec(),
// Reduce range constraint by disjunction.
Some(mut acc) => {
chriseth marked this conversation as resolved.
Show resolved Hide resolved
for ((min, max, mask), v) in acc.iter_mut().zip_eq(values) {
*min = (*min).min(*v);
*max = (*max).max(*v);
*mask |= v.to_integer();
}
acc
}
})
}
Some(match new_range_constraints {
None => {
// The iterator was empty, i.e. there are no inputs in the index matching the
Copy link
Collaborator

Choose a reason for hiding this comment

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

I still don't understand whether this is a case that would happen in practice (for example because some branches are actually impossible), or whether hitting this path would mean we can't complete the code generation. If the latter, maybe we can at least log an error?

Copy link
Member Author

Choose a reason for hiding this comment

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

I think this can happen if we decide to branch, then derive further concrete values, do a lookup again and then discover that we do not have these values in the lookup table at all.

// range constraints.
// This means that every call like this will lead to a fatal error, but there is
// enough information in the inputs to hanlde the call. Unfortunately, there is
// no way to signal this in the return type, yet.
// TODO(#2324): change this.
// We just return the input range constraints to signal "everything allright".
range_constraints.to_vec()
}
Some(new_output_range_constraints) => {
let mut new_output_range_constraints = new_output_range_constraints.into_iter();
known_arguments
.iter()
.enumerate()
.map(|(i, is_known)| {
if is_known {
// Just copy the input range constraint.
range_constraints[i].clone()
} else {
let (min, max, mask) = new_output_range_constraints.next().unwrap();
RangeConstraint::from_range(min, max)
.conjunction(&RangeConstraint::from_mask(mask))
}
})
.collect()
}
})
}

fn process_plookup<Q: crate::witgen::QueryCallback<T>>(
Expand Down Expand Up @@ -398,3 +462,14 @@ impl<'a, T: FieldElement> RangeConstraintSet<AlgebraicVariable<'a>, T>
}
}
}

/// Checks that an array of values satisfies a set of range constraints.
fn matches_range_constraint<T: FieldElement>(
values: &[T],
range_constraints: &[RangeConstraint<T>],
) -> bool {
values
.iter()
.zip_eq(range_constraints)
.all(|(value, range_constraint)| range_constraint.allows_value(*value))
}
12 changes: 7 additions & 5 deletions executor/src/witgen/machines/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,21 +55,23 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync {
);
}

/// Returns true if this machine can alway fully process a call via the given
/// Returns Some(..) if this machine can alway fully process a call via the given
/// identity, the set of known arguments and a list of range constraints
/// on the parameters. Note that the range constraints can be imposed both
/// on inputs and on outputs.
/// If this returns true, then corresponding calls to `process_lookup_direct`
/// If this returns Some(..), then corresponding calls to `process_lookup_direct`
/// are safe.
/// The value returned inside the option is a vector of range constraints on the arguments,
/// which again can be imposed both on inputs and on outputs.
/// The function requires `&mut self` because it usually builds an index structure
/// or something similar.
fn can_process_call_fully(
&mut self,
_identity_id: u64,
_known_arguments: &BitVec,
_range_constraints: &[RangeConstraint<T>],
) -> bool {
false
) -> Option<Vec<RangeConstraint<T>>> {
None
}

/// Like `process_plookup`, but also records the time spent in this machine.
Expand Down Expand Up @@ -188,7 +190,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
identity_id: u64,
known_arguments: &BitVec,
range_constraints: &[RangeConstraint<T>],
) -> bool {
) -> Option<Vec<RangeConstraint<T>>> {
match self {
KnownMachine::SecondStageMachine(m) => {
m.can_process_call_fully(identity_id, known_arguments, range_constraints)
Expand Down
2 changes: 1 addition & 1 deletion pipeline/tests/asm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -467,7 +467,7 @@ fn permutation_to_block() {
}

#[test]
#[should_panic = "called `Result::unwrap()` on an `Err` value: Linear constraint is not satisfiable: 18446744069414584320 != 0"]
#[should_panic = "Column main_bin::pc is not stackable in a 1-row block, conflict in rows 0 and 1"]
fn permutation_to_vm() {
// TODO: witgen issue: Machine incorrectly detected as block machine.
let f = "asm/permutations/vm_to_vm.asm";
Expand Down
Loading