From ab1f43a1e24fe18141448935e4c96744e34cc115 Mon Sep 17 00:00:00 2001 From: chriseth Date: Mon, 11 Nov 2024 15:35:19 +0100 Subject: [PATCH] Introduce direct lookup. (#2063) This PR introduces a more "direct" way to perform a lookup during witness generation. It removes the concept of `EvalValue`, which is a list of "updates", and instead requests the called machine to directly fill in mutable pointers to field elements. The goal is to use this (hopefully) faster interface if - the lookup can be fully solved in a single call - no cell-based range constraints are needed If the LHS of the lookup consists of direct polynomial references (or next references), the caller can pass a pointer to the final table and does not need to move data around any further. Some numbers: For the "keccak test", and only looking at the PC lookup, we get: Inside `process_plookup_internal`: 40 ns: preparing the `data` and `values` arrays 290 ns: call to process_lookup_direct 1300 ns: computing the result EvalValue. --- executor/src/witgen/eval_result.rs | 1 + .../witgen/machines/fixed_lookup_machine.rs | 157 ++++++++++++------ executor/src/witgen/machines/mod.rs | 28 +++- 3 files changed, 133 insertions(+), 53 deletions(-) diff --git a/executor/src/witgen/eval_result.rs b/executor/src/witgen/eval_result.rs index c48907a086..94cece09d1 100644 --- a/executor/src/witgen/eval_result.rs +++ b/executor/src/witgen/eval_result.rs @@ -161,6 +161,7 @@ impl EvalValue { /// New assignments or constraints for witness columns identified by an ID. pub type EvalResult<'a, T, K = AlgebraicVariable<'a>> = Result, EvalError>; +/// A fatal error for witness generation. #[derive(Clone, PartialEq)] pub enum EvalError { /// We ran out of rows diff --git a/executor/src/witgen/machines/fixed_lookup_machine.rs b/executor/src/witgen/machines/fixed_lookup_machine.rs index 3e9883fc8c..8ef89102e6 100644 --- a/executor/src/witgen/machines/fixed_lookup_machine.rs +++ b/executor/src/witgen/machines/fixed_lookup_machine.rs @@ -18,7 +18,7 @@ use crate::witgen::util::try_to_simple_poly_ref; use crate::witgen::{EvalError, EvalValue, IncompleteCause, MutableState, QueryCallback}; use crate::witgen::{EvalResult, FixedData}; -use super::{Connection, Machine}; +use super::{Connection, LookupCell, Machine}; /// An Application specifies a lookup cache. #[derive(Hash, Eq, PartialEq, Ord, PartialOrd, Clone)] @@ -197,8 +197,9 @@ impl<'a, T: FieldElement> FixedLookup<'a, T> { } } - fn process_plookup_internal( + fn process_plookup_internal<'b, Q: QueryCallback>( &mut self, + mutable_state: &'b mut MutableState<'a, 'b, T, Q>, identity_id: u64, rows: &RowPair<'_, '_, T>, left: &[AffineExpression, T>], @@ -217,62 +218,42 @@ impl<'a, T: FieldElement> FixedLookup<'a, T> { } // Split the left-hand-side into known input values and unknown output expressions. - let mut input_values = vec![]; - let mut known_inputs: BitVec = Default::default(); - let mut output_expressions = vec![]; - - for l in left { - if let Some(value) = l.constant_value() { - input_values.push(value); - known_inputs.push(true); - } else { - output_expressions.push(l); - known_inputs.push(false); - } - } - - let application = Application { - identity_id, - inputs: known_inputs, - }; - - let index = self - .indices - .entry(application) - .or_insert_with_key(|application| { - create_index(self.fixed_data, application, &self.connections) - }); - let index_value = index.get(&input_values).ok_or_else(|| { - let input_assignment = left - .iter() - .zip(right) - .filter_map(|(l, r)| l.constant_value().map(|v| (r.name.clone(), v))) - .collect(); - EvalError::FixedLookupFailed(input_assignment) - })?; + let mut data = vec![T::zero(); left.len()]; + let values = left + .iter() + .zip(&mut data) + .map(|(l, d)| { + if let Some(value) = l.constant_value() { + *d = value; + LookupCell::Input(d) + } else { + LookupCell::Output(d) + } + }) + .collect::>(); - let Some((row, output)) = index_value.get() else { + if !self.process_lookup_direct(mutable_state, identity_id, values)? { // multiple matches, we stop and learnt nothing return Ok(EvalValue::incomplete( IncompleteCause::MultipleLookupMatches, )); }; - self.multiplicity_counter.increment_at_row(identity_id, row); - let mut result = EvalValue::complete(vec![]); - for (l, r) in output_expressions.into_iter().zip(output) { - let evaluated = l.clone() - (*r).into(); - // TODO we could use bit constraints here - match evaluated.solve() { - Ok(constraints) => { - result.combine(constraints); - } - Err(_) => { - // Fail the whole lookup - return Err(EvalError::ConstraintUnsatisfiable(format!( - "Constraint is invalid ({l} != {r}).", - ))); + for (l, v) in left.iter().zip(data) { + if !l.is_constant() { + let evaluated = l.clone() - v.into(); + // TODO we could use bit constraints here + match evaluated.solve() { + Ok(constraints) => { + result.combine(constraints); + } + Err(_) => { + // Fail the whole lookup + return Err(EvalError::ConstraintUnsatisfiable(format!( + "Constraint is invalid ({l} != {v}).", + ))); + } } } } @@ -347,7 +328,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for FixedLookup<'a, T> { fn process_plookup<'b, Q: crate::witgen::QueryCallback>( &mut self, - _mutable_state: &'b mut crate::witgen::MutableState<'a, 'b, T, Q>, + mutable_state: &'b mut crate::witgen::MutableState<'a, 'b, T, Q>, identity_id: u64, caller_rows: &'b RowPair<'b, 'a, T>, ) -> EvalResult<'a, T> { @@ -362,7 +343,79 @@ impl<'a, T: FieldElement> Machine<'a, T> for FixedLookup<'a, T> { .peekable(); let outer_query = OuterQuery::new(caller_rows, identity); - self.process_plookup_internal(identity_id, caller_rows, &outer_query.left, right) + self.process_plookup_internal( + mutable_state, + identity_id, + caller_rows, + &outer_query.left, + right, + ) + } + + fn process_lookup_direct<'b, 'c, Q: QueryCallback>( + &mut self, + _mutable_state: &'b mut MutableState<'a, 'b, T, Q>, + identity_id: u64, + values: Vec>, + ) -> Result> { + let mut input_values = vec![]; + + let known_inputs = values + .iter() + .map(|v| match v { + LookupCell::Input(value) => { + input_values.push(**value); + true + } + LookupCell::Output(_) => false, + }) + .collect(); + + let application = Application { + identity_id, + inputs: known_inputs, + }; + + let index = self + .indices + .entry(application) + .or_insert_with_key(|application| { + create_index(self.fixed_data, application, &self.connections) + }); + let index_value = index.get(&input_values).ok_or_else(|| { + let right = self.connections[&identity_id].right; + let input_assignment = values + .iter() + .zip(&right.expressions) + .filter_map(|(l, r)| match l { + LookupCell::Input(v) => { + let name = try_to_simple_poly_ref(r).unwrap().name.clone(); + Some((name, **v)) + } + _ => None, + }) + .collect(); + EvalError::FixedLookupFailed(input_assignment) + })?; + + let Some((row, output)) = index_value.get() else { + // multiple matches, we stop and learnt nothing + return Ok(false); + }; + + self.multiplicity_counter.increment_at_row(identity_id, row); + + values + .into_iter() + .filter_map(|v| match v { + LookupCell::Output(e) => Some(e), + _ => None, + }) + .zip(output) + .for_each(|(e, v)| { + *e = *v; + }); + Ok(true) } fn take_witness_col_values<'b, Q: QueryCallback>( diff --git a/executor/src/witgen/machines/mod.rs b/executor/src/witgen/machines/mod.rs index 81e19d8aa8..16e73735e0 100644 --- a/executor/src/witgen/machines/mod.rs +++ b/executor/src/witgen/machines/mod.rs @@ -20,7 +20,7 @@ use self::write_once_memory::WriteOnceMemory; use super::generator::Generator; use super::rows::RowPair; -use super::{EvalResult, FixedData, MutableState, QueryCallback}; +use super::{EvalError, EvalResult, FixedData, MutableState, QueryCallback}; mod block_machine; mod double_sorted_witness_machine_16; @@ -60,6 +60,25 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync { caller_rows: &'b RowPair<'b, 'a, T>, ) -> EvalResult<'a, T>; + /// Process a connection of a given ID (which must be known to the callee). + /// This is a more direct version of `process_plookup`, where the caller + /// provides values or targets to where to write the results directly. + /// The length of `values` needs to be the same as the number of expressions + /// in the LHS / RHS of the connection. + /// It does not allow to return range constraints or complex expressions. + /// The boolean return value indicates whether the lookup was successful. + /// If it returns true, all output values in `values` need to have been set. + /// If it returns false, none of them should be changed. + /// An error is always unrecoverable. + fn process_lookup_direct<'b, 'c, Q: QueryCallback>( + &mut self, + _mutable_state: &'b mut MutableState<'a, 'b, T, Q>, + _identity_id: u64, + _values: Vec>, + ) -> Result> { + unimplemented!("Direct lookup is not supported for this machine."); + } + /// Returns the final values of the witness columns. fn take_witness_col_values<'b, Q: QueryCallback>( &mut self, @@ -70,6 +89,13 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync { fn identity_ids(&self) -> Vec; } +pub enum LookupCell<'a, T> { + /// Value is known (i.e. an input) + Input(&'a T), + /// Value is not known (i.e. an output) + Output(&'a mut T), +} + /// All known implementations of [Machine]. /// This allows us to treat machines uniformly without putting them into a `Box`, /// which requires that all lifetime parameters are 'static.