From f6b0b7baf74390c45ed0798b7bbaa12e8584b7c7 Mon Sep 17 00:00:00 2001 From: chriseth Date: Mon, 13 Jan 2025 14:00:28 +0100 Subject: [PATCH] Determine range constraints from fixed lookup. (#2300) Co-authored-by: Georg Wiese --- .../witgen/data_structures/mutable_state.rs | 9 +- .../witgen/jit/affine_symbolic_expression.rs | 21 +++++ .../src/witgen/jit/single_step_processor.rs | 36 ++++++++ executor/src/witgen/jit/witgen_inference.rs | 31 ++++--- .../double_sorted_witness_machine_32.rs | 11 +-- .../witgen/machines/fixed_lookup_machine.rs | 85 +++++++++++++++++-- executor/src/witgen/machines/mod.rs | 12 +-- pipeline/tests/asm.rs | 2 +- 8 files changed, 173 insertions(+), 34 deletions(-) diff --git a/executor/src/witgen/data_structures/mutable_state.rs b/executor/src/witgen/data_structures/mutable_state.rs index 31b0b74fe8..add2f2915e 100644 --- a/executor/src/witgen/data_structures/mutable_state.rs +++ b/executor/src/witgen/data_structures/mutable_state.rs @@ -56,15 +56,12 @@ impl<'a, T: FieldElement, Q: QueryCallback> MutableState<'a, T, Q> { identity_id: u64, known_inputs: &BitVec, range_constraints: &[RangeConstraint], - ) -> bool { + ) -> Option>> { // 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 diff --git a/executor/src/witgen/jit/affine_symbolic_expression.rs b/executor/src/witgen/jit/affine_symbolic_expression.rs index 62acf5faef..402e4c4745 100644 --- a/executor/src/witgen/jit/affine_symbolic_expression.rs +++ b/executor/src/witgen/jit/affine_symbolic_expression.rs @@ -113,6 +113,27 @@ impl AffineSymbolicExpression { } } + /// 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 { + self.coefficients + .iter() + .map(|(var, coeff)| { + let coeff = coeff.try_to_number()?; + let rc = self.range_constraints.get(var)?; + Some(rc.multiple(coeff)) + }) + .collect::>>() + .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 { diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index 20f23f7e42..8acbe0680c 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -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;" + ); + } } diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 915e187de8..00dd92b7cb 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -233,43 +233,49 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> 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::>(); 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, } } @@ -514,16 +520,17 @@ pub trait FixedEvaluator: Clone { } pub trait CanProcessCall { - /// 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], - ) -> bool; + ) -> Option>>; } impl> CanProcessCall for &MutableState<'_, T, Q> { @@ -532,7 +539,7 @@ impl> CanProcessCall for &MutableState<' identity_id: u64, known_inputs: &BitVec, range_constraints: &[RangeConstraint], - ) -> bool { + ) -> Option>> { MutableState::can_process_call_fully(self, identity_id, known_inputs, range_constraints) } } diff --git a/executor/src/witgen/machines/double_sorted_witness_machine_32.rs b/executor/src/witgen/machines/double_sorted_witness_machine_32.rs index 1137bcb9ab..32582244da 100644 --- a/executor/src/witgen/machines/double_sorted_witness_machine_32.rs +++ b/executor/src/witgen/machines/double_sorted_witness_machine_32.rs @@ -191,7 +191,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses32<'a, T> { identity_id: u64, known_arguments: &BitVec, range_constraints: &[RangeConstraint], - ) -> bool { + ) -> Option>> { assert!(self.parts.connections.contains_key(&identity_id)); assert_eq!(known_arguments.len(), 4); assert_eq!(range_constraints.len(), 4); @@ -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]) } } diff --git a/executor/src/witgen/machines/fixed_lookup_machine.rs b/executor/src/witgen/machines/fixed_lookup_machine.rs index 72e693a9e3..6f4c2135a8 100644 --- a/executor/src/witgen/machines/fixed_lookup_machine.rs +++ b/executor/src/witgen/machines/fixed_lookup_machine.rs @@ -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], - ) -> bool { + range_constraints: &[RangeConstraint], + ) -> Option>> { if !Self::is_responsible(&self.connections[&identity_id]) { - return false; + return None; } let index = self .indices @@ -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> = 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) => { + 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 + // 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>( @@ -398,3 +462,14 @@ impl<'a, T: FieldElement> RangeConstraintSet, T> } } } + +/// Checks that an array of values satisfies a set of range constraints. +fn matches_range_constraint( + values: &[T], + range_constraints: &[RangeConstraint], +) -> bool { + values + .iter() + .zip_eq(range_constraints) + .all(|(value, range_constraint)| range_constraint.allows_value(*value)) +} diff --git a/executor/src/witgen/machines/mod.rs b/executor/src/witgen/machines/mod.rs index 969f81f81c..91f7cdd86e 100644 --- a/executor/src/witgen/machines/mod.rs +++ b/executor/src/witgen/machines/mod.rs @@ -55,12 +55,14 @@ 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( @@ -68,8 +70,8 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync { _identity_id: u64, _known_arguments: &BitVec, _range_constraints: &[RangeConstraint], - ) -> bool { - false + ) -> Option>> { + None } /// Like `process_plookup`, but also records the time spent in this machine. @@ -188,7 +190,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> { identity_id: u64, known_arguments: &BitVec, range_constraints: &[RangeConstraint], - ) -> bool { + ) -> Option>> { match self { KnownMachine::SecondStageMachine(m) => { m.can_process_call_fully(identity_id, known_arguments, range_constraints) diff --git a/pipeline/tests/asm.rs b/pipeline/tests/asm.rs index 62b420e8d7..cf6b476d90 100644 --- a/pipeline/tests/asm.rs +++ b/pipeline/tests/asm.rs @@ -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";