From ec8c9943f855bf87045ca948922353bc4c1d25cf Mon Sep 17 00:00:00 2001 From: Georg Wiese <georgwiese@gmail.com> Date: Thu, 2 Nov 2023 09:56:09 +0000 Subject: [PATCH] VM-to-VM Witness Generation --- compiler/tests/asm.rs | 28 +- compiler/tests/pil.rs | 13 +- executor/src/witgen/block_processor.rs | 10 +- .../data_structures/finalizable_data.rs | 12 + executor/src/witgen/generator.rs | 114 ++++++-- executor/src/witgen/machines/block_machine.rs | 19 +- .../machines/double_sorted_witness_machine.rs | 8 +- executor/src/witgen/machines/mod.rs | 28 +- .../witgen/machines/sorted_witness_machine.rs | 9 +- executor/src/witgen/mod.rs | 56 ++-- executor/src/witgen/query_processor.rs | 15 +- executor/src/witgen/rows.rs | 24 +- executor/src/witgen/vm_processor.rs | 276 ++++++++++++++---- test_data/asm/vm_to_vm.asm | 7 +- .../asm/vm_to_vm_dynamic_trace_length.asm | 70 +++++ test_data/asm/vm_to_vm_to_block.asm | 66 +++++ test_data/asm/vm_to_vm_to_vm.asm | 73 +++++ 17 files changed, 660 insertions(+), 168 deletions(-) create mode 100644 test_data/asm/vm_to_vm_dynamic_trace_length.asm create mode 100644 test_data/asm/vm_to_vm_to_block.asm create mode 100644 test_data/asm/vm_to_vm_to_vm.asm diff --git a/compiler/tests/asm.rs b/compiler/tests/asm.rs index 8b0a6ea819..12565e4eca 100644 --- a/compiler/tests/asm.rs +++ b/compiler/tests/asm.rs @@ -165,7 +165,6 @@ fn vm_to_block_multiple_interfaces() { } #[test] -#[should_panic = "not implemented"] fn vm_to_vm() { let f = "vm_to_vm.asm"; let i = []; @@ -174,6 +173,33 @@ fn vm_to_vm() { gen_estark_proof(f, slice_to_vec(&i)); } +#[test] +fn vm_to_vm_dynamic_trace_length() { + let f = "vm_to_vm_dynamic_trace_length.asm"; + let i = []; + verify_asm::<GoldilocksField>(f, slice_to_vec(&i)); + gen_halo2_proof(f, slice_to_vec(&i)); + gen_estark_proof(f, slice_to_vec(&i)); +} + +#[test] +fn vm_to_vm_to_block() { + let f = "vm_to_vm_to_block.asm"; + let i = []; + verify_asm::<GoldilocksField>(f, slice_to_vec(&i)); + gen_halo2_proof(f, slice_to_vec(&i)); + gen_estark_proof(f, slice_to_vec(&i)); +} + +#[test] +fn vm_to_vm_to_vm() { + let f = "vm_to_vm_to_vm.asm"; + let i = []; + verify_asm::<GoldilocksField>(f, slice_to_vec(&i)); + gen_halo2_proof(f, slice_to_vec(&i)); + gen_estark_proof(f, slice_to_vec(&i)); +} + #[test] fn test_mem_read_write() { let f = "mem_read_write.asm"; diff --git a/compiler/tests/pil.rs b/compiler/tests/pil.rs index a7c5befd8c..c0ecffde18 100644 --- a/compiler/tests/pil.rs +++ b/compiler/tests/pil.rs @@ -221,12 +221,13 @@ fn test_single_line_blocks() { gen_estark_proof(f, Default::default()); } -#[test] -fn test_two_block_machine_functions() { - let f = "two_block_machine_functions.pil"; - verify_pil(f, None); - gen_estark_proof(f, Default::default()); -} +// TODO: Add back once #693 is merged +// #[test] +// fn test_two_block_machine_functions() { +// let f = "two_block_machine_functions.pil"; +// verify_pil(f, None); +// gen_estark_proof(f, Default::default()); +// } #[test] fn test_fixed_columns() { diff --git a/executor/src/witgen/block_processor.rs b/executor/src/witgen/block_processor.rs index bc0bcb63c8..b7ce31d598 100644 --- a/executor/src/witgen/block_processor.rs +++ b/executor/src/witgen/block_processor.rs @@ -27,15 +27,19 @@ pub struct WithoutCalldata; pub struct OuterQuery<'a, T: FieldElement> { /// A local copy of the left-hand side of the outer query. /// This will be mutated while processing the block. - left: Left<'a, T>, + pub left: Left<'a, T>, /// The right-hand side of the outer query. - right: &'a SelectedExpressions<Expression<T>>, + pub right: &'a SelectedExpressions<Expression<T>>, } impl<'a, T: FieldElement> OuterQuery<'a, T> { pub fn new(left: Left<'a, T>, right: &'a SelectedExpressions<Expression<T>>) -> Self { Self { left, right } } + + pub fn is_complete(&self) -> bool { + self.left.iter().all(|l| l.is_constant()) + } } /// A basic processor that knows how to determine a unique satisfying witness @@ -343,7 +347,7 @@ mod tests { fn name_to_poly_id<T: FieldElement>(fixed_data: &FixedData<T>) -> BTreeMap<String, PolyID> { let mut name_to_poly_id = BTreeMap::new(); for (poly_id, col) in fixed_data.witness_cols.iter() { - name_to_poly_id.insert(col.name.clone(), poly_id); + name_to_poly_id.insert(col.poly.name.clone(), poly_id); } for (poly_id, col) in fixed_data.fixed_cols.iter() { name_to_poly_id.insert(col.name.clone(), poly_id); diff --git a/executor/src/witgen/data_structures/finalizable_data.rs b/executor/src/witgen/data_structures/finalizable_data.rs index 0a050ef5d1..a63340315e 100644 --- a/executor/src/witgen/data_structures/finalizable_data.rs +++ b/executor/src/witgen/data_structures/finalizable_data.rs @@ -52,6 +52,10 @@ impl<'a, T: FieldElement> FinalizableData<'a, T> { self.data.len() } + pub fn is_empty(&self) -> bool { + self.data.is_empty() + } + pub fn push(&mut self, row: Row<'a, T>) { self.data.push(Entry::InProgress(row)); } @@ -86,6 +90,14 @@ impl<'a, T: FieldElement> FinalizableData<'a, T> { } } + pub fn last(&self) -> Option<&Row<'a, T>> { + match self.data.last() { + Some(Entry::InProgress(row)) => Some(row), + Some(Entry::Finalized(_, _)) => panic!("Last row already finalized."), + None => None, + } + } + pub fn mutable_row_pair(&mut self, i: usize) -> (&mut Row<'a, T>, &mut Row<'a, T>) { let (before, after) = self.data.split_at_mut(i + 1); let current = before.last_mut().unwrap(); diff --git a/executor/src/witgen/generator.rs b/executor/src/witgen/generator.rs index 5b94a00ad2..5632e017c9 100644 --- a/executor/src/witgen/generator.rs +++ b/executor/src/witgen/generator.rs @@ -7,18 +7,23 @@ use std::collections::{HashMap, HashSet}; use crate::witgen::data_structures::finalizable_data::FinalizableData; use crate::witgen::rows::CellValue; +use crate::witgen::EvalValue; use super::affine_expression::AffineExpression; -use super::block_processor::BlockProcessor; +use super::block_processor::{BlockProcessor, OuterQuery}; use super::data_structures::column_map::WitnessColumnMap; use super::global_constraints::GlobalConstraints; -use super::machines::Machine; - +use super::machines::{FixedLookup, Machine}; use super::rows::{Row, RowFactory}; use super::sequence_iterator::{DefaultSequenceIterator, ProcessingSequenceIterator}; use super::vm_processor::VmProcessor; use super::{EvalResult, FixedData, MutableState, QueryCallback}; +struct ProcessResult<'a, T: FieldElement> { + eval_value: EvalValue<&'a AlgebraicReference, T>, + block: FinalizableData<'a, T>, +} + pub struct Generator<'a, T: FieldElement> { fixed_data: &'a FixedData<'a, T>, identities: Vec<&'a Identity<Expression<T>>>, @@ -31,15 +36,63 @@ pub struct Generator<'a, T: FieldElement> { impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> { fn process_plookup<Q: QueryCallback<T>>( &mut self, - _mutable_state: &mut MutableState<'a, '_, T, Q>, + mutable_state: &mut MutableState<'a, '_, T, Q>, _kind: IdentityKind, - _left: &[AffineExpression<&'a AlgebraicReference, T>], - _right: &'a SelectedExpressions<Expression<T>>, + left: &[AffineExpression<&'a AlgebraicReference, T>], + right: &'a SelectedExpressions<Expression<T>>, ) -> Option<EvalResult<'a, T>> { - unimplemented!() + if right.selector != self.latch { + None + } else { + log::trace!("Start processing secondary VM '{}'", self.name()); + log::trace!("Arguments:"); + for (r, l) in right.expressions.iter().zip(left) { + log::trace!(" {r} = {l}"); + } + + let first_row = self + .data + .last() + .cloned() + .unwrap_or_else(|| self.compute_partial_first_row(mutable_state)); + + let outer_query = OuterQuery { + left: left.to_vec(), + right, + }; + let ProcessResult { eval_value, block } = + self.process(first_row, 0, mutable_state, Some(outer_query)); + + if eval_value.is_complete() { + log::trace!("End processing VM '{}' (successfully)", self.name()); + // Remove the last row of the previous block, as it is the first row of the current + // block. + self.data.pop(); + self.data.extend(block); + } else { + log::trace!("End processing VM '{}' (incomplete)", self.name()); + } + Some(Ok(eval_value)) + } } - fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>> { + fn take_witness_col_values<'b, Q: QueryCallback<T>>( + &mut self, + fixed_lookup: &'b mut FixedLookup<T>, + query_callback: &'b mut Q, + ) -> HashMap<String, Vec<T>> { + log::debug!("Finalizing VM: {}", self.name()); + + // In this stage, we don't have access to other machines, as they might already be finalized. + let mut mutable_state_no_machines = MutableState { + fixed_lookup, + machines: [].into_iter().into(), + query_callback, + }; + + self.fill_remaining_rows(&mut mutable_state_no_machines); + self.fix_first_row(); + self.data .take_transposed() .map(|(id, (values, _))| (self.fixed_data.column_name(&id).to_string(), values)) @@ -66,11 +119,23 @@ impl<'a, T: FieldElement> Generator<'a, T> { } } - pub fn run<Q: QueryCallback<T>>(&mut self, mutable_state: &mut MutableState<'a, '_, T, Q>) { + pub fn name(&self) -> &str { + let first_witness = self.witnesses.iter().next().unwrap(); + let first_witness_name = self.fixed_data.column_name(first_witness); + let namespace = first_witness_name + .rfind('.') + .map(|idx| &first_witness_name[..idx]); + + // For machines compiled using Powdr ASM we'll always have a namespace, but as a last + // resort we'll use the first witness name. + namespace.unwrap_or(first_witness_name) + } + + /// Runs the machine without any arguments from the first row. + pub fn run<'b, Q: QueryCallback<T>>(&mut self, mutable_state: &mut MutableState<'a, 'b, T, Q>) { + assert!(self.data.is_empty()); let first_row = self.compute_partial_first_row(mutable_state); - self.data = self.process(first_row, 0, mutable_state); - self.fill_remaining_rows(mutable_state); - self.fix_first_row(); + self.data = self.process(first_row, 0, mutable_state, None).block; } fn fill_remaining_rows<Q>(&mut self, mutable_state: &mut MutableState<'a, '_, T, Q>) @@ -81,9 +146,15 @@ impl<'a, T: FieldElement> Generator<'a, T> { assert!(self.latch.is_some()); let first_row = self.data.pop().unwrap(); - - self.data - .extend(self.process(first_row, self.data.len() as DegreeType, mutable_state)); + let ProcessResult { block, eval_value } = self.process( + first_row, + self.data.len() as DegreeType, + mutable_state, + None, + ); + assert!(eval_value.is_complete()); + + self.data.extend(block); } } @@ -131,7 +202,8 @@ impl<'a, T: FieldElement> Generator<'a, T> { first_row: Row<'a, T>, row_offset: DegreeType, mutable_state: &mut MutableState<'a, '_, T, Q>, - ) -> FinalizableData<'a, T> { + outer_query: Option<OuterQuery<'a, T>>, + ) -> ProcessResult<'a, T> { log::trace!( "Running main machine from row {row_offset} with the following initial values in the first row:\n{}", first_row.render_values(false, None) ); @@ -147,11 +219,13 @@ impl<'a, T: FieldElement> Generator<'a, T> { self.witnesses.clone(), data, row_factory, - self.latch.clone(), ); - let result = processor.run(mutable_state); - assert!(result.is_complete(), "Main machine did not complete"); - processor.finish() + if let Some(outer_query) = outer_query { + processor = processor.with_outer_query(outer_query); + } + let eval_value = processor.run(mutable_state); + let block = processor.finish(); + ProcessResult { eval_value, block } } /// At the end of the solving algorithm, we'll have computed the first row twice diff --git a/executor/src/witgen/machines/block_machine.rs b/executor/src/witgen/machines/block_machine.rs index 18f8c020c6..6c4d175ad0 100644 --- a/executor/src/witgen/machines/block_machine.rs +++ b/executor/src/witgen/machines/block_machine.rs @@ -1,6 +1,6 @@ use std::collections::{HashMap, HashSet}; -use super::{EvalResult, FixedData}; +use super::{EvalResult, FixedData, FixedLookup}; use crate::witgen::affine_expression::AffineExpression; use crate::witgen::block_processor::{BlockProcessor, OuterQuery}; @@ -164,7 +164,11 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> { }) } - fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>> { + fn take_witness_col_values<'b, Q: QueryCallback<T>>( + &mut self, + _fixed_lookup: &'b mut FixedLookup<T>, + _query_callback: &'b mut Q, + ) -> HashMap<String, Vec<T>> { if self.data.len() < 2 * self.block_size { log::warn!( "Filling empty blocks with zeros, because the block machine is never used. \ @@ -253,13 +257,10 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> { let namespace = first_witness_name .rfind('.') .map(|idx| &first_witness_name[..idx]); - if let Some(namespace) = namespace { - namespace - } else { - // For machines compiled using Powdr ASM we'll always have a namespace, but as a last - // resort we'll use the first witness name. - first_witness_name - } + + // For machines compiled using Powdr ASM we'll always have a namespace, but as a last + // resort we'll use the first witness name. + namespace.unwrap_or(first_witness_name) } fn process_plookup_internal<'b, Q: QueryCallback<T>>( diff --git a/executor/src/witgen/machines/double_sorted_witness_machine.rs b/executor/src/witgen/machines/double_sorted_witness_machine.rs index b040290ab0..d9ffb8efd4 100644 --- a/executor/src/witgen/machines/double_sorted_witness_machine.rs +++ b/executor/src/witgen/machines/double_sorted_witness_machine.rs @@ -5,7 +5,7 @@ use ast::parsed::SelectedExpressions; use itertools::Itertools; use num_traits::Zero; -use super::Machine; +use super::{FixedLookup, Machine}; use crate::witgen::affine_expression::AffineExpression; use crate::witgen::util::is_simple_poly_of_name; use crate::witgen::{EvalResult, FixedData, MutableState, QueryCallback}; @@ -111,7 +111,11 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<T> { Some(self.process_plookup_internal(left, right)) } - fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>> { + fn take_witness_col_values<'b, Q: QueryCallback<T>>( + &mut self, + _fixed_lookup: &'b mut FixedLookup<T>, + _query_callback: &'b mut Q, + ) -> HashMap<String, Vec<T>> { let mut addr = vec![]; let mut step = vec![]; let mut value = vec![]; diff --git a/executor/src/witgen/machines/mod.rs b/executor/src/witgen/machines/mod.rs index 3e4ed07d90..60bf964662 100644 --- a/executor/src/witgen/machines/mod.rs +++ b/executor/src/witgen/machines/mod.rs @@ -41,7 +41,11 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync { ) -> Option<EvalResult<'a, T>>; /// Returns the final values of the witness columns. - fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>>; + fn take_witness_col_values<'b, Q: QueryCallback<T>>( + &mut self, + fixed_lookup: &'b mut FixedLookup<T>, + query_callback: &'b mut Q, + ) -> HashMap<String, Vec<T>>; } /// All known implementations of [Machine]. @@ -61,7 +65,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> { kind: IdentityKind, left: &[AffineExpression<&'a AlgebraicReference, T>], right: &'a SelectedExpressions<Expression<T>>, - ) -> Option<crate::witgen::EvalResult<'a, T>> { + ) -> Option<EvalResult<'a, T>> { match self { KnownMachine::SortedWitnesses(m) => m.process_plookup(mutable_state, kind, left, right), KnownMachine::DoubleSortedWitnesses(m) => { @@ -72,12 +76,22 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> { } } - fn take_witness_col_values(&mut self) -> std::collections::HashMap<String, Vec<T>> { + fn take_witness_col_values<'b, Q: QueryCallback<T>>( + &mut self, + fixed_lookup: &'b mut FixedLookup<T>, + query_callback: &'b mut Q, + ) -> HashMap<String, Vec<T>> { match self { - KnownMachine::SortedWitnesses(m) => m.take_witness_col_values(), - KnownMachine::DoubleSortedWitnesses(m) => m.take_witness_col_values(), - KnownMachine::BlockMachine(m) => m.take_witness_col_values(), - KnownMachine::Vm(m) => m.take_witness_col_values(), + KnownMachine::SortedWitnesses(m) => { + m.take_witness_col_values(fixed_lookup, query_callback) + } + KnownMachine::DoubleSortedWitnesses(m) => { + m.take_witness_col_values(fixed_lookup, query_callback) + } + KnownMachine::BlockMachine(m) => { + m.take_witness_col_values(fixed_lookup, query_callback) + } + KnownMachine::Vm(m) => m.take_witness_col_values(fixed_lookup, query_callback), } } } diff --git a/executor/src/witgen/machines/sorted_witness_machine.rs b/executor/src/witgen/machines/sorted_witness_machine.rs index 81cd8c2086..40b41af53a 100644 --- a/executor/src/witgen/machines/sorted_witness_machine.rs +++ b/executor/src/witgen/machines/sorted_witness_machine.rs @@ -4,8 +4,8 @@ use ast::parsed::SelectedExpressions; use itertools::Itertools; use super::super::affine_expression::AffineExpression; -use super::Machine; use super::{EvalResult, FixedData}; +use super::{FixedLookup, Machine}; use crate::witgen::{ expression_evaluator::ExpressionEvaluator, fixed_evaluator::FixedEvaluator, symbolic_evaluator::SymbolicEvaluator, @@ -152,7 +152,12 @@ impl<'a, T: FieldElement> Machine<'a, T> for SortedWitnesses<'a, T> { Some(self.process_plookup_internal(left, right, rhs)) } - fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>> { + + fn take_witness_col_values<'b, Q: QueryCallback<T>>( + &mut self, + _fixed_lookup: &'b mut FixedLookup<T>, + _query_callback: &'b mut Q, + ) -> HashMap<String, Vec<T>> { let mut result = HashMap::new(); let (mut keys, mut values): (Vec<_>, Vec<_>) = diff --git a/executor/src/witgen/mod.rs b/executor/src/witgen/mod.rs index 0d343e08c3..a3f6ab1ea0 100644 --- a/executor/src/witgen/mod.rs +++ b/executor/src/witgen/mod.rs @@ -131,11 +131,15 @@ impl<'a, 'b, T: FieldElement, Q: QueryCallback<T>> WitnessGenerator<'a, 'b, T, Q generator.run(&mut mutable_state); // Get columns from machines - let main_columns = generator.take_witness_col_values(); + let main_columns = generator + .take_witness_col_values(mutable_state.fixed_lookup, mutable_state.query_callback); let mut columns = mutable_state .machines .iter_mut() - .flat_map(|m| m.take_witness_col_values().into_iter()) + .flat_map(|m| { + m.take_witness_col_values(mutable_state.fixed_lookup, mutable_state.query_callback) + .into_iter() + }) .chain(main_columns) .collect::<BTreeMap<_, _>>(); @@ -211,7 +215,7 @@ impl<'a, T: FieldElement> FixedData<'a, T> { fn column_name(&self, poly_id: &PolyID) -> &str { match poly_id.ptype { - PolynomialType::Committed => &self.witness_cols[poly_id].name, + PolynomialType::Committed => &self.witness_cols[poly_id].poly.name, PolynomialType::Constant => &self.fixed_cols[poly_id].name, PolynomialType::Intermediate => unimplemented!(), } @@ -238,18 +242,17 @@ impl<'a, T> FixedColumn<'a, T> { } } -#[derive(Debug)] -pub struct Query<'a, T> { - /// The query expression - expr: &'a Expression<T>, - /// The polynomial that is referenced by the query - poly: AlgebraicReference, -} - #[derive(Debug)] pub struct WitnessColumn<'a, T> { - name: String, - query: Option<Query<'a, T>>, + /// A polynomial reference that points to this column in the "current" row + /// (i.e., the "next" flag is set to false). + /// This is needed in situations where we want to update a cell when the + /// update does not come from an identity (which also has an AlgebraicReference). + poly: AlgebraicReference, + /// The prover query expression, if any. + query: Option<&'a Expression<T>>, + /// A list of externally computed witness values, if any. + /// The length of this list must be equal to the degree. external_values: Option<Vec<T>>, } @@ -265,24 +268,17 @@ impl<'a, T> WitnessColumn<'a, T> { } else { None }; - let name = name.to_string(); - let query = query.as_ref().map(|callback| { - let poly = AlgebraicReference { - poly_id: PolyID { - id: id as u64, - ptype: PolynomialType::Committed, - }, - name: name.clone(), - next: false, - index: None, - }; - Query { - poly, - expr: callback, - } - }); + let poly = AlgebraicReference { + poly_id: PolyID { + id: id as u64, + ptype: PolynomialType::Committed, + }, + name: name.to_string(), + next: false, + index: None, + }; WitnessColumn { - name, + poly, query, external_values, } diff --git a/executor/src/witgen/query_processor.rs b/executor/src/witgen/query_processor.rs index 9fd97af1cc..c9ea88cbd5 100644 --- a/executor/src/witgen/query_processor.rs +++ b/executor/src/witgen/query_processor.rs @@ -5,7 +5,7 @@ use ast::{ }; use number::FieldElement; -use super::{rows::RowPair, Constraint, EvalValue, FixedData, IncompleteCause, Query}; +use super::{rows::RowPair, Constraint, EvalValue, FixedData, IncompleteCause}; /// Computes value updates that result from a query. pub struct QueryProcessor<'a, 'b, T: FieldElement, QueryCallback: Send + Sync> { @@ -32,8 +32,8 @@ where let column = &self.fixed_data.witness_cols[poly_id]; if let Some(query) = column.query.as_ref() { - if rows.get_value(&query.poly).is_none() { - return self.process_witness_query(query, rows); + if rows.get_value(&column.poly).is_none() { + return self.process_witness_query(query, &column.poly, rows); } } // Either no query or the value is already known. @@ -42,19 +42,20 @@ where fn process_witness_query( &mut self, - query: &'a Query<'_, T>, + query: &'a Expression<T>, + poly: &'a AlgebraicReference, rows: &RowPair<T>, ) -> EvalValue<&'a AlgebraicReference, T> { - let query_str = match self.interpolate_query(query.expr, rows) { + let query_str = match self.interpolate_query(query, rows) { Ok(query) => query, Err(incomplete) => return EvalValue::incomplete(incomplete), }; if let Some(value) = (self.query_callback)(&query_str) { - EvalValue::complete(vec![(&query.poly, Constraint::Assignment(value))]) + EvalValue::complete(vec![(poly, Constraint::Assignment(value))]) } else { EvalValue::incomplete(IncompleteCause::NoQueryAnswer( query_str, - query.poly.name.to_string(), + poly.name.to_string(), )) } } diff --git a/executor/src/witgen/rows.rs b/executor/src/witgen/rows.rs index 068b88c29b..34aa44e67e 100644 --- a/executor/src/witgen/rows.rs +++ b/executor/src/witgen/rows.rs @@ -13,7 +13,7 @@ use super::{ global_constraints::{GlobalConstraints, RangeConstraintSet}, range_constraints::RangeConstraint, symbolic_witness_evaluator::{SymoblicWitnessEvaluator, WitnessColumnEvaluator}, - EvalValue, FixedData, + FixedData, }; #[derive(Clone, PartialEq, Debug)] @@ -240,28 +240,6 @@ impl<'row, 'a, T: FieldElement> RowUpdater<'row, 'a, T> { self.get_cell_mut(poly).apply_update(c); } - /// Applies the updates to the underlying rows. Returns true if any updates - /// were applied. - /// - /// # Panics - /// Panics if any updates are redundant, as this indicates a bug that would - /// potentially cause infinite loops otherwise. - pub fn apply_updates( - &mut self, - updates: &EvalValue<&AlgebraicReference, T>, - source_name: impl Fn() -> String, - ) -> bool { - if updates.constraints.is_empty() { - return false; - } - - log::trace!(" Updates from: {}", source_name()); - for (poly, c) in &updates.constraints { - self.apply_update(poly, c) - } - true - } - fn get_cell_mut<'b>(&'b mut self, poly: &AlgebraicReference) -> &'b mut Cell<'a, T> { match poly.next { false => &mut self.current[&poly.poly_id], diff --git a/executor/src/witgen/vm_processor.rs b/executor/src/witgen/vm_processor.rs index 943d47f751..1abd29301a 100644 --- a/executor/src/witgen/vm_processor.rs +++ b/executor/src/witgen/vm_processor.rs @@ -5,19 +5,24 @@ use itertools::Itertools; use number::{DegreeType, FieldElement}; use parser_util::lines::indent; use std::cmp::max; -use std::collections::HashSet; +use std::collections::{BTreeMap, HashSet}; use std::time::Instant; use crate::witgen::identity_processor::{self, IdentityProcessor}; use crate::witgen::rows::RowUpdater; +use crate::witgen::util::try_to_simple_poly; +use crate::witgen::Constraint; use crate::witgen::IncompleteCause; +use super::block_processor::OuterQuery; use super::data_structures::column_map::WitnessColumnMap; use super::data_structures::finalizable_data::FinalizableData; use super::query_processor::QueryProcessor; -use super::rows::{Row, RowFactory, RowPair, UnknownStrategy}; -use super::{EvalError, EvalResult, EvalValue, FixedData, MutableState, QueryCallback}; +use super::rows::{CellValue, Row, RowFactory, RowPair, UnknownStrategy}; +use super::{ + Constraints, EvalError, EvalResult, EvalValue, FixedData, MutableState, QueryCallback, +}; /// Maximal period checked during loop detection. const MAX_PERIOD: usize = 4; @@ -60,7 +65,9 @@ pub struct VmProcessor<'a, T: FieldElement> { last_report: DegreeType, last_report_time: Instant, row_factory: RowFactory<'a, T>, - latch: Option<Expression<T>>, + outer_query: Option<OuterQuery<'a, T>>, + inputs: BTreeMap<PolyID, T>, + previously_set_inputs: BTreeMap<PolyID, DegreeType>, } impl<'a, T: FieldElement> VmProcessor<'a, T> { @@ -71,7 +78,6 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { witnesses: HashSet<PolyID>, data: FinalizableData<'a, T>, row_factory: RowFactory<'a, T>, - latch: Option<Expression<T>>, ) -> Self { let (identities_with_next, identities_without_next): (Vec<_>, Vec<_>) = identities .iter() @@ -95,7 +101,27 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { row_factory, last_report: 0, last_report_time: Instant::now(), - latch, + outer_query: None, + inputs: BTreeMap::new(), + previously_set_inputs: BTreeMap::new(), + } + } + + pub fn with_outer_query(self, outer_query: OuterQuery<'a, T>) -> Self { + log::trace!(" Extracting inputs:"); + let mut inputs = BTreeMap::new(); + for (l, r) in outer_query.left.iter().zip(&outer_query.right.expressions) { + if let Some(right_poly) = try_to_simple_poly(r).map(|p| p.poly_id) { + if let Some(l) = l.constant_value() { + log::trace!(" {} = {}", r, l); + inputs.insert(right_poly, l); + } + } + } + Self { + outer_query: Some(outer_query), + inputs, + ..self } } @@ -111,6 +137,8 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { ) -> EvalValue<&'a AlgebraicReference, T> { assert!(self.data.len() == 1); + let mut outer_assignments = vec![]; + // Are we in an infinite loop and can just re-use the old values? let mut looping_period = None; let mut loop_detection_log_level = log::Level::Info; @@ -155,29 +183,26 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { // add and compute some values for the next row as well. if looping_period.is_none() && row_index != rows_left - 1 { self.ensure_has_next_row(row_index); - self.compute_row(row_index, mutable_state); - - if let Some(latch) = self.latch.as_ref() { - // Evaluate latch expression and return if it evaluates to 1. - let row_pair = RowPair::from_single_row( - self.row(row_index), - row_index + self.row_offset, - self.fixed_data, - UnknownStrategy::Unknown, - ); - let latch_value = row_pair - .evaluate(latch) - .ok() - .and_then(|l| l.constant_value()); - - if let Some(latch_value) = latch_value { - if latch_value.is_one() { - log::trace!("Machine returns!"); - return EvalValue::complete(vec![]); + outer_assignments.extend(self.compute_row(row_index, mutable_state).into_iter()); + + // Evaluate latch expression and return if it evaluates to 1. + if let Some(latch) = self.latch_value(row_index) { + if latch { + log::trace!("Machine returns!"); + if self.outer_query.as_ref().unwrap().is_complete() { + return EvalValue::complete(outer_assignments); + } else { + return EvalValue::incomplete_with_constraints( + outer_assignments, + IncompleteCause::BlockMachineLookupIncomplete, + ); } - } else { - return EvalValue::incomplete(IncompleteCause::UnknownLatch); } + } else if self.outer_query.is_some() { + // If we have an outer query (and therefore a latch expression), + // its value should be known at this point. + // Probably, we don't have all the necessary inputs. + return EvalValue::incomplete(IncompleteCause::UnknownLatch); } }; } @@ -187,7 +212,7 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { self.fixed_data.degree + 1 ); - EvalValue::complete(vec![]) + EvalValue::complete(outer_assignments) } /// Checks if the last rows are repeating and returns the period. @@ -208,6 +233,21 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { }) } + fn latch_value(&self, row_index: DegreeType) -> Option<bool> { + let row_pair = RowPair::from_single_row( + self.row(row_index), + row_index + self.row_offset, + self.fixed_data, + UnknownStrategy::Unknown, + ); + self.outer_query + .as_ref() + .and_then(|outer_query| outer_query.right.selector.as_ref()) + .and_then(|latch| row_pair.evaluate(latch).ok()) + .and_then(|l| l.constant_value()) + .map(|l| l.is_one()) + } + // Returns a reference to a given row fn row(&self, row_index: DegreeType) -> &Row<'a, T> { &self.data[row_index as usize] @@ -224,7 +264,7 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { &mut self, row_index: DegreeType, mutable_state: &mut MutableState<'a, '_, T, Q>, - ) { + ) -> Constraints<&'a AlgebraicReference, T> { log::trace!("Row: {}", row_index + self.row_offset); log::trace!(" Going over all identities until no more progress is made"); @@ -234,52 +274,69 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { CompletableIdentities::new(self.identities_without_next_ref.iter().cloned()); let mut identities_with_next_ref = CompletableIdentities::new(self.identities_with_next_ref.iter().cloned()); - self.loop_until_no_progress(row_index, &mut identities_without_next_ref, mutable_state) - .and_then(|_| { - self.loop_until_no_progress(row_index, &mut identities_with_next_ref, mutable_state) + let outer_assignments = self + .loop_until_no_progress(row_index, &mut identities_without_next_ref, mutable_state) + .and_then(|outer_assignments| { + Ok(outer_assignments + .into_iter() + .chain(self.loop_until_no_progress( + row_index, + &mut identities_with_next_ref, + mutable_state, + )?) + .collect::<Vec<_>>()) }) .map_err(|e| self.report_failure_and_panic_unsatisfiable(row_index, e)) .unwrap(); // Check that the computed row is "final" by asserting that all unknown values can // be set to 0. - log::trace!(" Checking that remaining identities hold when unknown values are set to 0"); - let mut identity_processor = IdentityProcessor::new(self.fixed_data, mutable_state); - self.process_identities( - row_index, - &mut identities_without_next_ref, - UnknownStrategy::Zero, - &mut identity_processor, - ) - .and_then(|_| { + // This check is only done for the primary machine, as secondary machines might simply + // not have all the inputs yet and therefore be underconstrained. + if self.outer_query.is_none() { + log::trace!( + " Checking that remaining identities hold when unknown values are set to 0" + ); + let mut identity_processor = IdentityProcessor::new(self.fixed_data, mutable_state); self.process_identities( row_index, - &mut identities_with_next_ref, + &mut identities_without_next_ref, UnknownStrategy::Zero, &mut identity_processor, ) - }) - .map_err(|e| self.report_failure_and_panic_underconstrained(row_index, e)) - .unwrap(); + .and_then(|_| { + self.process_identities( + row_index, + &mut identities_with_next_ref, + UnknownStrategy::Zero, + &mut identity_processor, + ) + }) + .map_err(|e| self.report_failure_and_panic_underconstrained(row_index, e)) + .unwrap(); + } log::trace!( "{}", - self.row(row_index) - .render(&format!("===== Row {}", row_index), true, &self.witnesses) + self.row(row_index).render( + &format!("===== Row {}", row_index as DegreeType + self.row_offset), + true, + &self.witnesses + ) ); + + outer_assignments } /// Loops over all identities and queries, until no further progress is made. /// @returns the "incomplete" identities, i.e. identities that contain unknown values. - fn loop_until_no_progress<Q>( + fn loop_until_no_progress<Q: QueryCallback<T>>( &mut self, row_index: DegreeType, identities: &mut CompletableIdentities<'a, T>, mutable_state: &mut MutableState<'a, '_, T, Q>, - ) -> Result<(), Vec<EvalError<T>>> - where - Q: FnMut(&str) -> Option<T> + Send + Sync, - { + ) -> Result<Constraints<&'a AlgebraicReference, T>, Vec<EvalError<T>>> { + let mut outer_assignments = vec![]; loop { let mut identity_processor = IdentityProcessor::new(self.fixed_data, mutable_state); let mut progress = self.process_identities( @@ -288,6 +345,16 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { UnknownStrategy::Unknown, &mut identity_processor, )?; + if let Some(true) = self.latch_value(row_index) { + let (outer_query_progress, new_outer_assignments) = self + .process_outer_query(row_index, mutable_state) + .map_err(|e| vec![e])?; + progress |= outer_query_progress; + outer_assignments.extend(new_outer_assignments); + } + + progress |= self.set_inputs_if_unset(row_index); + let mut updates = EvalValue::complete(vec![]); let mut query_processor = QueryProcessor::new(self.fixed_data, mutable_state.query_callback); @@ -309,7 +376,7 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { break; } } - Ok(()) + Ok(outer_assignments) } /// Loops over all identities once and updates the current row and next row. @@ -383,18 +450,117 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> { } } + fn process_outer_query<Q: QueryCallback<T>>( + &mut self, + row_index: DegreeType, + mutable_state: &mut MutableState<'a, '_, T, Q>, + ) -> Result<(bool, Constraints<&'a AlgebraicReference, T>), EvalError<T>> { + let OuterQuery { left, right } = self + .outer_query + .as_ref() + .expect("Asked to process outer query, but it was not set!"); + + let row_pair = RowPair::new( + &self.data[row_index as usize], + &self.data[row_index as usize + 1], + (row_index + self.row_offset) % self.fixed_data.degree, + self.fixed_data, + UnknownStrategy::Unknown, + ); + + let mut identity_processor = IdentityProcessor::new(self.fixed_data, mutable_state); + let updates = identity_processor + .process_link(left, right, &row_pair) + .map_err(|e| { + log::warn!("Error in outer query: {e}"); + log::warn!("Some of the following entries could not be matched:"); + for (l, r) in left.iter().zip(right.expressions.iter()) { + if let Ok(r) = row_pair.evaluate(r) { + log::warn!(" => {} = {}", l, r); + } + } + e + })?; + + let progress = self.apply_updates(row_index, &updates, || "outer query".to_string()); + + let outer_assignments = updates + .constraints + .into_iter() + .filter(|(poly, _)| !self.witnesses.contains(&poly.poly_id)) + .collect::<Vec<_>>(); + + Ok((progress, outer_assignments)) + } + + /// Sets the inputs to the values given in [VmProcessor::inputs] if they are not already set. + /// Typically, inputs will have a constraint of the form: `((1 - instr__reset) * (_input' - _input)) = 0;` + /// So, once the value of `_input` is set, this function will do nothing until the next reset instruction. + /// However, if `_input` does become unconstrained, we need to undo all changes we've done so far. + /// For this reason, we keep track of all changes we've done to inputs in [VmProcessor::previously_set_inputs]. + fn set_inputs_if_unset(&mut self, row_index: DegreeType) -> bool { + let mut input_updates = EvalValue::complete(vec![]); + for (poly_id, value) in self.inputs.iter() { + match &self.data[row_index as usize][poly_id].value { + CellValue::Known(_) => {} + CellValue::RangeConstraint(_) | CellValue::Unknown => { + input_updates.combine(EvalValue::complete([( + &self.fixed_data.witness_cols[poly_id].poly, + Constraint::Assignment(*value), + )])); + } + }; + } + + for (poly, _) in &input_updates.constraints { + let poly_id = poly.poly_id; + if let Some(start_row) = self.previously_set_inputs.remove(&poly_id) { + log::trace!( + " Resetting previously set inputs for column: {}", + self.fixed_data.column_name(&poly_id) + ); + for row_index in start_row..row_index { + self.data[row_index as usize][&poly_id].value = CellValue::Unknown; + } + } + } + for (poly, _) in &input_updates.constraints { + self.previously_set_inputs.insert(poly.poly_id, row_index); + } + self.apply_updates(row_index, &input_updates, || "inputs".to_string()) + } + fn apply_updates( &mut self, row_index: DegreeType, - updates: &EvalValue<&AlgebraicReference, T>, + updates: &EvalValue<&'a AlgebraicReference, T>, source_name: impl Fn() -> String, ) -> bool { if updates.constraints.is_empty() { return false; } + + log::trace!(" Updates from: {}", source_name()); + + // Build RowUpdater + // (a bit complicated, because we need two mutable + // references to elements of the same vector) let (current, next) = self.data.mutable_row_pair(row_index as usize); let mut row_updater = RowUpdater::new(current, next, row_index + self.row_offset); - row_updater.apply_updates(updates, source_name) + + for (poly, c) in &updates.constraints { + if self.witnesses.contains(&poly.poly_id) { + row_updater.apply_update(poly, c); + } else if let Constraint::Assignment(v) = c { + let left = &mut self.outer_query.as_mut().unwrap().left; + log::trace!(" => {} (outer) = {}", poly, v); + for l in left.iter_mut() { + l.assign(poly, *v); + } + }; + } + + true } fn report_failure_and_panic_unsatisfiable( diff --git a/test_data/asm/vm_to_vm.asm b/test_data/asm/vm_to_vm.asm index 4962cc2761..9cd7e217ae 100644 --- a/test_data/asm/vm_to_vm.asm +++ b/test_data/asm/vm_to_vm.asm @@ -1,5 +1,7 @@ machine Main { + degree 256; + VM vm; reg pc[@pc]; @@ -14,9 +16,8 @@ machine Main { function main { A <== add(1, 1); - // TODO: uncomment the following two lines once we support having many calls to a machine - // A <== add(A, 1); - // A <== sub(A, 1); + A <== add(A, 1); + A <== sub(A, 1); assert_eq A, 2; return; } diff --git a/test_data/asm/vm_to_vm_dynamic_trace_length.asm b/test_data/asm/vm_to_vm_dynamic_trace_length.asm new file mode 100644 index 0000000000..31e1b51520 --- /dev/null +++ b/test_data/asm/vm_to_vm_dynamic_trace_length.asm @@ -0,0 +1,70 @@ +machine Main { + + degree 256; + + Pow pow; + + reg pc[@pc]; + reg X[<=]; + reg Y[<=]; + reg Z[<=]; + reg A; + + instr pow X, Y -> Z = pow.pow + instr assert_eq X, Y { X = Y } + + function main { + A <== pow(7, 0); + assert_eq A, 1; + + A <== pow(7, 1); + assert_eq A, 7; + + A <== pow(7, 2); + assert_eq A, 49; + + A <== pow(2, 2); + assert_eq A, 4; + + A <== pow(2, 10); + assert_eq A, 1024; + + return; + } +} + +// Computes X^Y by multiplying X by itself Y times +machine Pow { + + reg pc[@pc]; + reg X[<=]; + reg Y[<=]; + reg Z[<=]; + reg A; + reg CNT; + + col witness XInv; + col witness XIsZero; + XIsZero = 1 - X * XInv; + XIsZero * X = 0; + XIsZero * (1 - XIsZero) = 0; + + instr mul X, Y -> Z { X * Y = Z } + instr jmpz X, l: label { pc' = XIsZero * l + (1 - XIsZero) * (pc + 1) } + instr jmp l: label { pc' = l } + + function pow x: field, y: field -> field { + A <=X= 1; + CNT <=X= y; + + start:: + jmpz CNT, done; + A <== mul(A, x); + CNT <=X= CNT - 1; + jmp start; + + done:: + return A; + } + +} \ No newline at end of file diff --git a/test_data/asm/vm_to_vm_to_block.asm b/test_data/asm/vm_to_vm_to_block.asm new file mode 100644 index 0000000000..214a70f052 --- /dev/null +++ b/test_data/asm/vm_to_vm_to_block.asm @@ -0,0 +1,66 @@ +machine Main { + + degree 256; + + Pythagoras pythagoras; + + reg pc[@pc]; + reg X[<=]; + reg Y[<=]; + reg Z[<=]; + reg A; + + instr pythagoras X, Y -> Z = pythagoras.pythagoras + instr assert_eq X, Y { X = Y } + + function main { + A <== pythagoras(3, 4); + assert_eq A, 25; + + A <== pythagoras(4, 3); + assert_eq A, 25; + + A <== pythagoras(1, 2); + assert_eq A, 5; + + return; + } +} + + +machine Pythagoras { + + Arith arith; + + reg pc[@pc]; + reg X[<=]; + reg Y[<=]; + reg Z[<=]; + reg A; + reg B; + + + instr add X, Y -> Z = arith.add + instr mul X, Y -> Z = arith.mul + + function pythagoras a: field, b: field -> field { + A <== mul(a, a); + B <== mul(b, b); + A <== add(A, B); + return A; + } +} + +machine Arith(latch, operation_id) { + + operation add<0> x1, x2 -> y; + operation mul<1> x1, x2 -> y; + + col fixed latch = [1]*; + col witness operation_id; + col witness x1; + col witness x2; + col witness y; + + y = operation_id * (x1 * x2) + (1 - operation_id) * (x1 + x2); +} diff --git a/test_data/asm/vm_to_vm_to_vm.asm b/test_data/asm/vm_to_vm_to_vm.asm new file mode 100644 index 0000000000..7033c8aae6 --- /dev/null +++ b/test_data/asm/vm_to_vm_to_vm.asm @@ -0,0 +1,73 @@ +machine Main { + + degree 256; + + Pythagoras pythagoras; + + reg pc[@pc]; + reg X[<=]; + reg Y[<=]; + reg Z[<=]; + reg A; + + instr pythagoras X, Y -> Z = pythagoras.pythagoras + instr assert_eq X, Y { X = Y } + + function main { + A <== pythagoras(3, 4); + assert_eq A, 25; + + A <== pythagoras(4, 3); + assert_eq A, 25; + + A <== pythagoras(1, 2); + assert_eq A, 5; + + return; + } +} + +machine Pythagoras { + + Arith arith; + + reg pc[@pc]; + reg X[<=]; + reg Y[<=]; + reg Z[<=]; + reg A; + reg B; + + + instr add X, Y -> Z = arith.add + instr mul X, Y -> Z = arith.mul + + function pythagoras a: field, b: field -> field { + A <== mul(a, a); + B <== mul(b, b); + A <== add(A, B); + return A; + } +} + +machine Arith { + + reg pc[@pc]; + reg X[<=]; + reg Y[<=]; + reg Z[<=]; + reg A; + + instr add X, Y -> Z { X + Y = Z } + instr mul X, Y -> Z { X * Y = Z } + + function add x: field, y: field -> field { + A <== add(x, y); + return A; + } + + function mul x: field, y: field -> field { + A <== mul(x, y); + return A; + } +} \ No newline at end of file