From 8c50681fc19fc72eafa474ebcedf0814cb2c124a Mon Sep 17 00:00:00 2001 From: Georg Wiese Date: Fri, 24 Nov 2023 16:07:49 +0100 Subject: [PATCH] Write-Once Memory --- compiler/src/verify.rs | 9 +- compiler/tests/asm.rs | 28 ++- compiler/tests/powdr_std.rs | 2 +- executor/src/witgen/machines/block_machine.rs | 12 + .../src/witgen/machines/machine_extractor.rs | 6 + executor/src/witgen/machines/mod.rs | 7 + .../src/witgen/machines/write_once_memory.rs | 234 ++++++++++++++++++ riscv/tests/riscv.rs | 2 +- test_data/asm/mem_write_once.asm | 47 ++++ .../asm/mem_write_once_external_write.asm | 33 +++ 10 files changed, 375 insertions(+), 5 deletions(-) create mode 100644 executor/src/witgen/machines/write_once_memory.rs create mode 100644 test_data/asm/mem_write_once.asm create mode 100644 test_data/asm/mem_write_once_external_write.asm diff --git a/compiler/src/verify.rs b/compiler/src/verify.rs index 12cca0fbe8..218f1e6323 100644 --- a/compiler/src/verify.rs +++ b/compiler/src/verify.rs @@ -10,7 +10,12 @@ use std::{ use crate::compile_asm_string; -pub fn verify_asm_string(file_name: &str, contents: &str, inputs: Vec) { +pub fn verify_asm_string( + file_name: &str, + contents: &str, + inputs: Vec, + external_witness_values: Vec<(&str, Vec)>, +) { let temp_dir = mktemp::Temp::new_dir().unwrap(); let (_, result) = compile_asm_string( file_name, @@ -20,7 +25,7 @@ pub fn verify_asm_string(file_name: &str, contents: &str, input &temp_dir, true, Some(BackendType::PilStarkCli), - vec![], + external_witness_values, ) .unwrap(); diff --git a/compiler/tests/asm.rs b/compiler/tests/asm.rs index 4d7b3d6701..8cd6205e06 100644 --- a/compiler/tests/asm.rs +++ b/compiler/tests/asm.rs @@ -4,6 +4,14 @@ use std::fs; use test_log::test; fn verify_asm(file_name: &str, inputs: Vec) { + verify_asm_with_external_witness(file_name, inputs, vec![]); +} + +fn verify_asm_with_external_witness( + file_name: &str, + inputs: Vec, + external_witness_values: Vec<(&str, Vec)>, +) { let file_name = format!( "{}/../test_data/asm/{file_name}", env!("CARGO_MANIFEST_DIR") @@ -11,7 +19,7 @@ fn verify_asm(file_name: &str, inputs: Vec) { let contents = fs::read_to_string(&file_name).unwrap(); - verify_asm_string(&file_name, &contents, inputs) + verify_asm_string(&file_name, &contents, inputs, external_witness_values); } fn gen_estark_proof(file_name: &str, inputs: Vec) { @@ -71,6 +79,24 @@ fn secondary_block_machine_add2() { gen_estark_proof(f, vec![]); } +#[test] +fn mem_write_once() { + let f = "mem_write_once.asm"; + verify_asm::(f, vec![]); + gen_halo2_proof(f, vec![]); + gen_estark_proof(f, vec![]); +} + +#[test] +fn mem_write_once_external_write() { + let f = "mem_write_once_external_write.asm"; + let mut mem = vec![GoldilocksField::from(0); 256]; + mem[17] = GoldilocksField::from(42); + mem[62] = GoldilocksField::from(123); + mem[255] = GoldilocksField::from(-1); + verify_asm_with_external_witness::(f, vec![], vec![("main.v", mem)]); +} + #[test] fn block_machine_cache_miss() { let f = "block_machine_cache_miss.asm"; diff --git a/compiler/tests/powdr_std.rs b/compiler/tests/powdr_std.rs index dce760aae3..7f46d118aa 100644 --- a/compiler/tests/powdr_std.rs +++ b/compiler/tests/powdr_std.rs @@ -11,7 +11,7 @@ fn verify_asm(file_name: &str, inputs: Vec) { let contents = fs::read_to_string(&file_name).unwrap(); - verify_asm_string(&file_name, &contents, inputs) + verify_asm_string(&file_name, &contents, inputs, vec![]) } fn gen_estark_proof(file_name: &str, inputs: Vec) { diff --git a/executor/src/witgen/machines/block_machine.rs b/executor/src/witgen/machines/block_machine.rs index 92a0a60f3c..32f8f78039 100644 --- a/executor/src/witgen/machines/block_machine.rs +++ b/executor/src/witgen/machines/block_machine.rs @@ -15,6 +15,7 @@ use crate::witgen::{machines::Machine, EvalError, EvalValue, IncompleteCause}; use crate::witgen::{MutableState, QueryCallback}; use ast::analyzed::{ AlgebraicExpression as Expression, AlgebraicReference, Identity, IdentityKind, PolyID, + PolynomialType, }; use ast::parsed::SelectedExpressions; use number::{DegreeType, FieldElement}; @@ -72,6 +73,17 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> { global_range_constraints: &GlobalConstraints, ) -> Option { for id in connecting_identities { + for r in id.right.expressions.iter() { + if let Some(poly) = try_to_simple_poly(r) { + if poly.poly_id.ptype == PolynomialType::Constant { + // It does not really make sense to have constant polynomials on the RHS + // of a block machine lookup, as all constant polynomials are periodic, so + // it would always return the same value. + return None; + } + } + } + // TODO we should check that the other constraints/fixed columns are also periodic. if let Some(block_size) = try_to_period(&id.right.selector, fixed_data) { assert!(block_size <= fixed_data.degree as usize); diff --git a/executor/src/witgen/machines/machine_extractor.rs b/executor/src/witgen/machines/machine_extractor.rs index 4b4f3990e9..dd2c20d86e 100644 --- a/executor/src/witgen/machines/machine_extractor.rs +++ b/executor/src/witgen/machines/machine_extractor.rs @@ -8,6 +8,7 @@ use super::FixedData; use super::KnownMachine; use crate::witgen::generator::Generator; use crate::witgen::global_constraints::GlobalConstraints; +use crate::witgen::machines::write_once_memory::WriteOnceMemory; use ast::analyzed::{AlgebraicExpression as Expression, Identity, IdentityKind, PolyID}; use ast::parsed::visitor::ExpressionVisitable; use ast::parsed::SelectedExpressions; @@ -107,6 +108,11 @@ pub fn split_out_machines<'a, T: FieldElement>( { log::info!("Detected machine: memory"); machines.push(KnownMachine::DoubleSortedWitnesses(machine)); + } else if let Some(machine) = + WriteOnceMemory::try_new(fixed, &connecting_identities, &machine_identities) + { + log::info!("Detected machine: write-once memory"); + machines.push(KnownMachine::WriteOnceMemory(machine)); } else if let Some(machine) = BlockMachine::try_new( fixed, &connecting_identities, diff --git a/executor/src/witgen/machines/mod.rs b/executor/src/witgen/machines/mod.rs index 60bf964662..768a1b5baf 100644 --- a/executor/src/witgen/machines/mod.rs +++ b/executor/src/witgen/machines/mod.rs @@ -9,6 +9,7 @@ use self::block_machine::BlockMachine; use self::double_sorted_witness_machine::DoubleSortedWitnesses; pub use self::fixed_lookup_machine::FixedLookup; use self::sorted_witness_machine::SortedWitnesses; +use self::write_once_memory::WriteOnceMemory; use ast::analyzed::IdentityKind; use super::affine_expression::AffineExpression; @@ -23,6 +24,7 @@ mod double_sorted_witness_machine; mod fixed_lookup_machine; pub mod machine_extractor; mod sorted_witness_machine; +mod write_once_memory; /// A machine is a set of witness columns and identities where the columns /// are used on the right-hand-side of lookups. It can process plookups. @@ -54,6 +56,7 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync { pub enum KnownMachine<'a, T: FieldElement> { SortedWitnesses(SortedWitnesses<'a, T>), DoubleSortedWitnesses(DoubleSortedWitnesses), + WriteOnceMemory(WriteOnceMemory<'a, T>), BlockMachine(BlockMachine<'a, T>), Vm(Generator<'a, T>), } @@ -71,6 +74,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> { KnownMachine::DoubleSortedWitnesses(m) => { m.process_plookup(mutable_state, kind, left, right) } + KnownMachine::WriteOnceMemory(m) => m.process_plookup(mutable_state, kind, left, right), KnownMachine::BlockMachine(m) => m.process_plookup(mutable_state, kind, left, right), KnownMachine::Vm(m) => m.process_plookup(mutable_state, kind, left, right), } @@ -88,6 +92,9 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> { KnownMachine::DoubleSortedWitnesses(m) => { m.take_witness_col_values(fixed_lookup, query_callback) } + KnownMachine::WriteOnceMemory(m) => { + m.take_witness_col_values(fixed_lookup, query_callback) + } KnownMachine::BlockMachine(m) => { m.take_witness_col_values(fixed_lookup, query_callback) } diff --git a/executor/src/witgen/machines/write_once_memory.rs b/executor/src/witgen/machines/write_once_memory.rs new file mode 100644 index 0000000000..f3aa8f63bb --- /dev/null +++ b/executor/src/witgen/machines/write_once_memory.rs @@ -0,0 +1,234 @@ +use std::collections::{BTreeMap, HashMap}; + +use ast::{ + analyzed::{ + AlgebraicExpression as Expression, AlgebraicReference, Identity, IdentityKind, PolyID, + PolynomialType, + }, + parsed::SelectedExpressions, +}; +use number::{DegreeType, FieldElement}; + +use crate::witgen::{ + affine_expression::AffineExpression, util::try_to_simple_poly, EvalError, EvalResult, + EvalValue, FixedData, IncompleteCause, MutableState, QueryCallback, +}; + +use super::{FixedLookup, Machine}; + +/// A memory machine with a fixed address space, and each address can only have one +/// value during the lifetime of the program. +/// In the simplest case, it looks like this: +/// ```pil +/// let ADDR = |i| i; +/// let v; +/// // Stores a value, fails if the cell already has a value that's different +/// instr mstore X, Y -> { {X, Y} in {ADDR, v} } +/// // Loads a value. If the cell is empty, the prover can choose a value. +/// // Note that this is the same lookup, only Y is considered an output instead +/// // of an input. +/// instr mload X -> Y { {X, Y} in {ADDR, v} } +/// ``` +pub struct WriteOnceMemory<'a, T: FieldElement> { + /// The fixed data + fixed_data: &'a FixedData<'a, T>, + /// The right-hand side of the connecting identity + /// (if there are several, they must all be the same) + rhs: &'a SelectedExpressions>, + /// The polynomials that are used as values (witness polynomials on the RHS) + value_polys: Vec, + /// A map from keys to row indices + key_to_index: BTreeMap, DegreeType>, + /// The memory content + data: BTreeMap>>, +} + +impl<'a, T: FieldElement> WriteOnceMemory<'a, T> { + pub fn try_new( + fixed_data: &'a FixedData<'a, T>, + connecting_identities: &[&'a Identity>], + identities: &[&Identity>], + ) -> Option { + if !identities.is_empty() { + return None; + } + + let rhs = &connecting_identities[0].right; + if !connecting_identities.iter().all(|i| i.right == *rhs) { + return None; + } + + if rhs.selector.is_some() { + return None; + } + + let rhs_polys = rhs + .expressions + .iter() + .map(|e| try_to_simple_poly(e)) + .collect::>>(); + + // Returns None if any RHS polynomial is a complex expression + let rhs_polys = rhs_polys?; + + // Build a Vec for the key and value polynomials + let (key_polys, value_polys): (Vec<_>, Vec<_>) = rhs_polys + .into_iter() + .partition(|p| p.poly_id.ptype == PolynomialType::Constant); + let key_polys = key_polys + .into_iter() + .map(|p| { + assert!(!p.next); + p.poly_id + }) + .collect::>(); + let value_polys = value_polys + .into_iter() + .map(|p| { + assert!(!p.next); + p.poly_id + }) + .collect::>(); + + let mut key_to_index = BTreeMap::new(); + for row in 0..fixed_data.degree { + let key = key_polys + .iter() + .map(|k| fixed_data.fixed_cols[k].values[row as usize]) + .collect::>(); + if key_to_index.insert(key, row).is_some() { + // Duplicate keys, can't be a write-once memory + return None; + } + } + + Some(Self { + fixed_data, + rhs, + value_polys, + key_to_index, + data: BTreeMap::new(), + }) + } + + fn process_plookup_internal( + &mut self, + left: &[AffineExpression<&'a AlgebraicReference, T>], + right: &'a SelectedExpressions>, + ) -> EvalResult<'a, T> { + let (key_expressions, value_expressions): (Vec<_>, Vec<_>) = left + .iter() + .zip(right.expressions.iter()) + .partition(|(_, r)| { + try_to_simple_poly(r).unwrap().poly_id.ptype == PolynomialType::Constant + }); + let key = key_expressions + .into_iter() + .map(|(k, _)| k.constant_value()) + .collect::>>(); + let value_expressions = value_expressions + .into_iter() + .map(|(v, _)| v) + .collect::>(); + let value = value_expressions + .iter() + .map(|v| v.constant_value()) + .collect::>(); + + log::trace!("Key: {:?}", key); + log::trace!("Value: {:?}", value); + + let Some(key) = key else { + return Ok(EvalValue::incomplete( + IncompleteCause::NonConstantRequiredArgument("key"), + )); + }; + + let index = self.key_to_index.get(&key).cloned().ok_or_else(|| { + EvalError::from(format!("Key {:?} not found in write-once memory", key)) + })?; + + // If there is an externally provided memory value, use it + let external_witness_value = self + .value_polys + .iter() + .map(|p| self.fixed_data.external_witness(index, p)) + .collect::>(); + let stored_value = match self.data.get(&index) { + Some(values) => values + .iter() + .zip(external_witness_value.iter()) + .map(|(&stored, &external)| external.or(stored)) + .collect(), + None => external_witness_value, + }; + + let mut updates = vec![]; + let values = value_expressions + .into_iter() + .zip(stored_value.iter()) + .map(|(l, r)| { + match (l.constant_value(), r) { + // No value provided and none stored -> keep value as None + (None, None) => Ok::, EvalError>(None), + // Value provided but none stored -> write, no updates + (Some(l), None) => Ok(Some(l)), + // Value stored -> keep stored value & either update LHS or assert equality + (_, Some(r)) => { + updates.extend((l.clone() - (*r).into()).solve()?.constraints); + Ok(Some(*r)) + } + } + }) + .collect::, _>>()?; + + // Write values + let is_complete = !values.contains(&None); + self.data.insert(index, values); + + match is_complete { + true => Ok(EvalValue::complete(updates)), + false => Ok(EvalValue::incomplete_with_constraints( + updates, + IncompleteCause::NonConstantRequiredArgument("value"), + )), + } + } +} + +impl<'a, T: FieldElement> Machine<'a, T> for WriteOnceMemory<'a, T> { + fn process_plookup<'b, Q: QueryCallback>( + &mut self, + _mutable_state: &'b mut MutableState<'a, 'b, T, Q>, + kind: IdentityKind, + left: &[AffineExpression<&'a AlgebraicReference, T>], + right: &'a SelectedExpressions>, + ) -> Option> { + (right == self.rhs && kind == IdentityKind::Plookup) + .then(|| self.process_plookup_internal(left, right)) + } + + fn take_witness_col_values<'b, Q: QueryCallback>( + &mut self, + _fixed_lookup: &'b mut FixedLookup, + _query_callback: &'b mut Q, + ) -> HashMap> { + self.value_polys + .iter() + .enumerate() + .map(|(value_index, poly)| { + let column = self.fixed_data.witness_cols[poly] + .external_values + .clone() + .unwrap_or_else(|| { + let mut column = vec![T::zero(); self.fixed_data.degree as usize]; + for (row, values) in self.data.iter() { + column[*row as usize] = values[value_index].unwrap_or_default(); + } + column + }); + (self.fixed_data.column_name(poly).to_string(), column) + }) + .collect() + } +} diff --git a/riscv/tests/riscv.rs b/riscv/tests/riscv.rs index f88fe9cf71..1b6de9b6a0 100644 --- a/riscv/tests/riscv.rs +++ b/riscv/tests/riscv.rs @@ -145,7 +145,7 @@ fn verify_file(case: &str, inputs: Vec, coprocessors: &CoProces riscv::compile_rust_to_riscv_asm(&format!("tests/riscv_data/{case}"), &temp_dir); let powdr_asm = riscv::compiler::compile(riscv_asm, coprocessors); - verify_asm_string(&format!("{case}.asm"), &powdr_asm, inputs); + verify_asm_string(&format!("{case}.asm"), &powdr_asm, inputs, vec![]); } #[test] diff --git a/test_data/asm/mem_write_once.asm b/test_data/asm/mem_write_once.asm new file mode 100644 index 0000000000..accc89fb3b --- /dev/null +++ b/test_data/asm/mem_write_once.asm @@ -0,0 +1,47 @@ +machine MemReadWrite { + + degree 256; + + reg pc[@pc]; + reg X1[<=]; + reg X2[<=]; + reg Y1[<=]; + reg Y2[<=]; + reg A; + reg B; + + // Write-once memory with key (ADDR1, ADDR2) and value (v1, v2) + let ADDR1 = |i| i; + let ADDR2 = |i| i + 1; + let v1; + let v2; + // Stores a value, fails if the cell already has a value that's different + instr mstore X1, X2, Y1, Y2 -> { {X1, X2, Y1, Y2} in {ADDR1, ADDR2, v1, v2} } + // Loads a value. If the cell is empty, the prover can choose a value. + instr mload X1, X2 -> Y1, Y2 { {X1, X2, Y1, Y2} in {ADDR1, ADDR2, v1, v2} } + + instr assert_eq X1, Y1 { X1 = Y1 } + + function main { + mstore 17, 18, 42, 43; + mstore 62, 63, 123, 1234; + mstore 255, 256, -1, -2; + + // Setting the same value twice is not a problem + mstore 17, 18, 42, 43; + + A, B <== mload(17, 18); + assert_eq A, 42; + assert_eq B, 43; + + A, B <== mload(62, 63); + assert_eq A, 123; + assert_eq B, 1234; + + A, B <== mload(255, 256); + assert_eq A, -1; + assert_eq B, -2; + + return; + } +} \ No newline at end of file diff --git a/test_data/asm/mem_write_once_external_write.asm b/test_data/asm/mem_write_once_external_write.asm new file mode 100644 index 0000000000..314d8ab4c5 --- /dev/null +++ b/test_data/asm/mem_write_once_external_write.asm @@ -0,0 +1,33 @@ +// Very simple write-once memory, but without an mstore operation. +// As a result, this only works if the content of the `v` column has +// been provided externally. +machine MemReadWrite { + + degree 256; + + reg pc[@pc]; + reg X[<=]; + reg Y[<=]; + reg A; + + // Write-once memory + let ADDR = |i| i; + let v; + // Loads a value. If the cell is empty, the prover can choose a value. + instr mload X -> Y { {X, Y} in {ADDR, v} } + + instr assert_eq X, Y { X = Y } + + function main { + A <== mload(17); + assert_eq A, 42; + + A <== mload(62); + assert_eq A, 123; + + A <== mload(255); + assert_eq A, -1; + + return; + } +} \ No newline at end of file