From cce4375497acee48edaf95a1958ebafd49f77497 Mon Sep 17 00:00:00 2001 From: chriseth Date: Fri, 20 Dec 2024 11:32:35 +0000 Subject: [PATCH 1/6] Single step processor. --- executor/src/witgen/jit/mod.rs | 1 + .../src/witgen/jit/single_step_processor.rs | 154 ++++++++++++++++++ 2 files changed, 155 insertions(+) create mode 100644 executor/src/witgen/jit/single_step_processor.rs diff --git a/executor/src/witgen/jit/mod.rs b/executor/src/witgen/jit/mod.rs index 91df4a9c6f..c2b33e43f2 100644 --- a/executor/src/witgen/jit/mod.rs +++ b/executor/src/witgen/jit/mod.rs @@ -3,6 +3,7 @@ mod block_machine_processor; mod compiler; mod effect; pub(crate) mod function_cache; +mod single_step_processor; mod symbolic_expression; mod variable; pub(crate) mod witgen_inference; diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs new file mode 100644 index 0000000000..8b949921e4 --- /dev/null +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -0,0 +1,154 @@ +#![allow(dead_code)] +use std::collections::HashSet; + +use itertools::Itertools; +use powdr_ast::analyzed::AlgebraicReference; +use powdr_number::FieldElement; + +use crate::witgen::{machines::MachineParts, FixedData}; + +use super::{ + effect::Effect, + variable::{Cell, Variable}, + witgen_inference::{FixedEvaluator, WitgenInference}, +}; + +/// A processor for generating JIT code that computes the next row from the previous row. +pub struct SingleStepProcessor<'a, T: FieldElement> { + fixed_data: &'a FixedData<'a, T>, + machine_parts: MachineParts<'a, T>, +} + +impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { + pub fn new(fixed_data: &'a FixedData<'a, T>, machine_parts: MachineParts<'a, T>) -> Self { + SingleStepProcessor { + fixed_data, + machine_parts, + } + } + + pub fn generate_code(&self) -> Result>, String> { + // All witness columns in row 0 are known. + let known_variables = self.machine_parts.witnesses.iter().map(|id| { + Variable::Cell(Cell { + column_name: self.fixed_data.column_name(id).to_string(), + id: id.id, + row_offset: 0, + }) + }); + let mut witgen = WitgenInference::new(self.fixed_data, self, known_variables); + + let mut complete = HashSet::new(); + for iteration in 0.. { + let mut progress = false; + + for id in &self.machine_parts.identities { + let row_offset = if id.contains_next_ref() { 0 } else { 1 }; + if !complete.contains(&(id.id(), row_offset)) { + let result = witgen.process_identity(id, row_offset); + if result.complete { + complete.insert((id.id(), row_offset)); + } + progress |= result.progress; + } + } + if !progress { + log::debug!("Finishing after {iteration} iterations"); + break; + } + } + + // Check that we could derive all witness values in the next row. + let unknown_witnesses = self + .machine_parts + .witnesses + .iter() + .filter(|wit| { + !witgen.is_known(&Variable::Cell(Cell { + column_name: self.fixed_data.column_name(wit).to_string(), + id: wit.id, + row_offset: 1, + })) + }) + .sorted() + .collect_vec(); + + // TODO also check that we completed all machine calls? + if unknown_witnesses.is_empty() { + Ok(witgen.code()) + } else { + Err(format!( + "Unable to derive algorithm to compute values for witness columns in the next row for the following columns: {}", + unknown_witnesses.iter().map(|wit| self.fixed_data.column_name(wit)).format(", ") + )) + } + } +} + +impl FixedEvaluator for &SingleStepProcessor<'_, T> { + fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option { + // We can only return something here if the fixed column is constant + // in the region wer are considering. + // This might be the case if we know we are not evaluating the first or the last + // row, but this is not yet implemented. + None + } +} + +#[cfg(test)] +mod test { + + use powdr_ast::analyzed::Analyzed; + use powdr_number::GoldilocksField; + + use crate::{ + constant_evaluator, + witgen::{ + global_constraints, + jit::{effect::Effect, test_util::format_code}, + machines::MachineParts, + FixedData, + }, + }; + + use super::{SingleStepProcessor, Variable}; + + fn generate_single_step( + input_pil: &str, + ) -> Result>, String> { + let analyzed: Analyzed = + powdr_pil_analyzer::analyze_string(input_pil).unwrap(); + let fixed_col_vals = constant_evaluator::generate(&analyzed); + let fixed_data = FixedData::new(&analyzed, &fixed_col_vals, &[], Default::default(), 0); + let (fixed_data, retained_identities) = + global_constraints::set_global_constraints(fixed_data, &analyzed.identities); + + let witness_columns = analyzed + .committed_polys_in_source_order() + .flat_map(|(symbol, _)| symbol.array_elements().map(|(_, id)| id)) + .collect(); + + let machine_parts = MachineParts::new( + &fixed_data, + // no connections + Default::default(), + retained_identities, + witness_columns, + // No prover functions + Vec::new(), + ); + + SingleStepProcessor { + fixed_data: &fixed_data, + machine_parts, + } + .generate_code() + } + + #[test] + fn fib() { + let input = "let X; let Y; X' = Y; Y' = X + Y;"; + let code = generate_single_step(input).unwrap(); + assert_eq!(format_code(&code), "X[1] = Y[0];\nY[1] = (X[0] + Y[0]);"); + } +} From c54c0dd1c25358b93ef4786927c344d01dc3e076 Mon Sep 17 00:00:00 2001 From: chriseth Date: Fri, 20 Dec 2024 11:38:59 +0000 Subject: [PATCH 2/6] Negative test. --- executor/src/witgen/jit/single_step_processor.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index 8b949921e4..c32db27a96 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -151,4 +151,14 @@ mod test { let code = generate_single_step(input).unwrap(); assert_eq!(format_code(&code), "X[1] = Y[0];\nY[1] = (X[0] + Y[0]);"); } + + #[test] + fn no_progress() { + let input = "let X; let Y; X' = X;"; + let err = generate_single_step(input).err().unwrap(); + assert_eq!( + err.to_string(), + "Unable to derive algorithm to compute values for witness columns in the next row for the following columns: Y" + ); + } } From eb245acf9760006375a6df94239b238b8bc63ef7 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 2 Jan 2025 13:14:24 +0100 Subject: [PATCH 3/6] Apply suggestions from code review Co-authored-by: Georg Wiese --- executor/src/witgen/jit/single_step_processor.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index c32db27a96..6f93e05ec3 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -73,8 +73,8 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { .sorted() .collect_vec(); - // TODO also check that we completed all machine calls? if unknown_witnesses.is_empty() { + assert_eq!(complete, self.machine_parts.identities.len()); Ok(witgen.code()) } else { Err(format!( @@ -88,7 +88,7 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { impl FixedEvaluator for &SingleStepProcessor<'_, T> { fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option { // We can only return something here if the fixed column is constant - // in the region wer are considering. + // in the region we are considering. // This might be the case if we know we are not evaluating the first or the last // row, but this is not yet implemented. None From 1588eb2077f045cf563e99c8e031c968c6084b7a Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 2 Jan 2025 13:25:52 +0000 Subject: [PATCH 4/6] adjust tests. --- .../src/witgen/jit/single_step_processor.rs | 48 +++++++++++-------- 1 file changed, 29 insertions(+), 19 deletions(-) diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index 6f93e05ec3..af67901487 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -27,7 +27,10 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { } } - pub fn generate_code(&self) -> Result>, String> { + pub fn generate_code + Clone>( + &self, + can_process: CanProcess, + ) -> Result>, String> { // All witness columns in row 0 are known. let known_variables = self.machine_parts.witnesses.iter().map(|id| { Variable::Cell(Cell { @@ -43,11 +46,11 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { let mut progress = false; for id in &self.machine_parts.identities { - let row_offset = if id.contains_next_ref() { 0 } else { 1 }; - if !complete.contains(&(id.id(), row_offset)) { - let result = witgen.process_identity(id, row_offset); + if !complete.contains(&id.id()) { + let row_offset = if id.contains_next_ref() { 0 } else { 1 }; + let result = witgen.process_identity(can_process.clone(), id, row_offset); if result.complete { - complete.insert((id.id(), row_offset)); + complete.insert(id.id()); } progress |= result.progress; } @@ -73,13 +76,14 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { .sorted() .collect_vec(); - if unknown_witnesses.is_empty() { - assert_eq!(complete, self.machine_parts.identities.len()); + let missing_identities = self.machine_parts.len() - complete.len(); + if unknown_witnesses.is_empty() && missing_identities == 0{ Ok(witgen.code()) } else { Err(format!( - "Unable to derive algorithm to compute values for witness columns in the next row for the following columns: {}", - unknown_witnesses.iter().map(|wit| self.fixed_data.column_name(wit)).format(", ") + "Unable to derive algorithm to compute values for witness columns in the next row for the following columns:\n{}\nand {missing_identities} identities are missing.", + unknown_witnesses.iter().map(|wit| self.fixed_data.column_name(wit)).format(", "), + )) } } @@ -104,10 +108,7 @@ mod test { use crate::{ constant_evaluator, witgen::{ - global_constraints, - jit::{effect::Effect, test_util::format_code}, - machines::MachineParts, - FixedData, + data_structures::mutable_state::MutableState, global_constraints, jit::{effect::Effect, test_util::read_pil}, machines::{machine_extractor::MachineExtractor, KnownMachine, MachineParts}, FixedData }, }; @@ -116,12 +117,18 @@ mod test { fn generate_single_step( input_pil: &str, ) -> Result>, String> { - let analyzed: Analyzed = - powdr_pil_analyzer::analyze_string(input_pil).unwrap(); - let fixed_col_vals = constant_evaluator::generate(&analyzed); + let (analyzed, fixed_col_vals) = read_pil(input_pil); + let fixed_data = FixedData::new(&analyzed, &fixed_col_vals, &[], Default::default(), 0); let (fixed_data, retained_identities) = global_constraints::set_global_constraints(fixed_data, &analyzed.identities); + let machines = MachineExtractor::new(&fixed_data).split_out_machines(retained_identities); + let mutable_state = MutableState::new(machines.into_iter(), &|_| { + Err("Query not implemented".to_string()) + }); + + let KnownMachine::DynamicMachine(machine)= mutable_state.get_machine("Main") else { panic!(); } + assert_eq!(machines.len(), 1); let witness_columns = analyzed .committed_polys_in_source_order() @@ -137,24 +144,27 @@ mod test { // No prover functions Vec::new(), ); + let mutable_state = MutableState::new(machines.into_iter(), &|_| { + Err("Query not implemented".to_string()) + }); SingleStepProcessor { fixed_data: &fixed_data, machine_parts, } - .generate_code() + .generate_code(&mutable_state) } #[test] fn fib() { - let input = "let X; let Y; X' = Y; Y' = X + Y;"; + let input = "namespace Main(256); let X; let Y; X' = Y; Y' = X + Y;"; let code = generate_single_step(input).unwrap(); assert_eq!(format_code(&code), "X[1] = Y[0];\nY[1] = (X[0] + Y[0]);"); } #[test] fn no_progress() { - let input = "let X; let Y; X' = X;"; + let input = "namespace Main(256); let X; let Y; X' = X;"; let err = generate_single_step(input).err().unwrap(); assert_eq!( err.to_string(), From cc70c191b9aaa51da7cd9a3da42895d1a3df361f Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 2 Jan 2025 13:33:41 +0000 Subject: [PATCH 5/6] Remove test function. --- .../witgen/data_structures/mutable_state.rs | 14 ---------- .../src/witgen/jit/block_machine_processor.rs | 28 +++++++++++-------- 2 files changed, 16 insertions(+), 26 deletions(-) diff --git a/executor/src/witgen/data_structures/mutable_state.rs b/executor/src/witgen/data_structures/mutable_state.rs index c8777af7eb..37564084bb 100644 --- a/executor/src/witgen/data_structures/mutable_state.rs +++ b/executor/src/witgen/data_structures/mutable_state.rs @@ -42,20 +42,6 @@ impl<'a, T: FieldElement, Q: QueryCallback> MutableState<'a, T, Q> { } } - #[cfg(test)] - pub fn get_machine(&self, substring: &str) -> &RefCell> { - use itertools::Itertools; - - self.machines - .iter() - .filter(|m| m.borrow().name().contains(substring)) - .exactly_one() - .map_err(|e| { - format!("Expected exactly one machine with substring '{substring}', but found {e}.") - }) - .unwrap() - } - /// Runs the first machine (unless there are no machines) end returns the generated columns. /// The first machine might call other machines, which is handled automatically. pub fn run(self) -> HashMap> { diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs index fa8ad9effc..c0b1067625 100644 --- a/executor/src/witgen/jit/block_machine_processor.rs +++ b/executor/src/witgen/jit/block_machine_processor.rs @@ -195,31 +195,35 @@ mod test { let (fixed_data, retained_identities) = global_constraints::set_global_constraints(fixed_data, &analyzed.identities); let machines = MachineExtractor::new(&fixed_data).split_out_machines(retained_identities); - let mutable_state = MutableState::new(machines.into_iter(), &|_| { - Err("Query not implemented".to_string()) - }); - - let machine = mutable_state.get_machine(machine_name); - let ((machine_parts, block_size, latch_row), connection_ids) = match *machine.borrow() { - KnownMachine::BlockMachine(ref m) => (m.machine_info(), m.identity_ids()), - _ => panic!("Expected a block machine"), + let [KnownMachine::BlockMachine(machine)] = machines + .iter() + .filter(|m| m.name().contains(machine_name)) + .collect_vec() + .as_slice() + else { + panic!("Expected exactly one matching block machine") }; - assert_eq!(connection_ids.len(), 1); - + let (machine_parts, block_size, latch_row) = machine.machine_info(); + assert_eq!(machine_parts.connections.len(), 1); + let connection_id = *machine_parts.connections.keys().next().unwrap(); let processor = BlockMachineProcessor { fixed_data: &fixed_data, - machine_parts, + machine_parts: machine_parts.clone(), block_size, latch_row, }; + let mutable_state = MutableState::new(machines.into_iter(), &|_| { + Err("Query not implemented".to_string()) + }); + let known_values = BitVec::from_iter( (0..num_inputs) .map(|_| true) .chain((0..num_outputs).map(|_| false)), ); - processor.generate_code(&mutable_state, connection_ids[0], &known_values) + processor.generate_code(&mutable_state, connection_id, &known_values) } #[test] From 96d5506994eac037dccc6ec0da00663d85e3dae4 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 2 Jan 2025 14:04:46 +0000 Subject: [PATCH 6/6] Fix testing functions. --- .../src/witgen/jit/single_step_processor.rs | 72 +++++++++---------- .../src/witgen/machines/dynamic_machine.rs | 5 ++ 2 files changed, 40 insertions(+), 37 deletions(-) diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index af67901487..c2780c5241 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -10,7 +10,7 @@ use crate::witgen::{machines::MachineParts, FixedData}; use super::{ effect::Effect, variable::{Cell, Variable}, - witgen_inference::{FixedEvaluator, WitgenInference}, + witgen_inference::{CanProcessCall, FixedEvaluator, WitgenInference}, }; /// A processor for generating JIT code that computes the next row from the previous row. @@ -76,14 +76,13 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { .sorted() .collect_vec(); - let missing_identities = self.machine_parts.len() - complete.len(); - if unknown_witnesses.is_empty() && missing_identities == 0{ + let missing_identities = self.machine_parts.identities.len() - complete.len(); + if unknown_witnesses.is_empty() && missing_identities == 0 { Ok(witgen.code()) } else { Err(format!( "Unable to derive algorithm to compute values for witness columns in the next row for the following columns:\n{}\nand {missing_identities} identities are missing.", - unknown_witnesses.iter().map(|wit| self.fixed_data.column_name(wit)).format(", "), - + unknown_witnesses.iter().map(|wit| self.fixed_data.column_name(wit)).format(", ") )) } } @@ -102,20 +101,26 @@ impl FixedEvaluator for &SingleStepProcessor<'_, T> { #[cfg(test)] mod test { - use powdr_ast::analyzed::Analyzed; + use itertools::Itertools; + use powdr_number::GoldilocksField; - use crate::{ - constant_evaluator, - witgen::{ - data_structures::mutable_state::MutableState, global_constraints, jit::{effect::Effect, test_util::read_pil}, machines::{machine_extractor::MachineExtractor, KnownMachine, MachineParts}, FixedData + use crate::witgen::{ + data_structures::mutable_state::MutableState, + global_constraints, + jit::{ + effect::{format_code, Effect}, + test_util::read_pil, }, + machines::{machine_extractor::MachineExtractor, KnownMachine, Machine}, + FixedData, }; use super::{SingleStepProcessor, Variable}; fn generate_single_step( input_pil: &str, + machine_name: &str, ) -> Result>, String> { let (analyzed, fixed_col_vals) = read_pil(input_pil); @@ -123,27 +128,15 @@ mod test { let (fixed_data, retained_identities) = global_constraints::set_global_constraints(fixed_data, &analyzed.identities); let machines = MachineExtractor::new(&fixed_data).split_out_machines(retained_identities); - let mutable_state = MutableState::new(machines.into_iter(), &|_| { - Err("Query not implemented".to_string()) - }); - - let KnownMachine::DynamicMachine(machine)= mutable_state.get_machine("Main") else { panic!(); } - assert_eq!(machines.len(), 1); - - let witness_columns = analyzed - .committed_polys_in_source_order() - .flat_map(|(symbol, _)| symbol.array_elements().map(|(_, id)| id)) - .collect(); - - let machine_parts = MachineParts::new( - &fixed_data, - // no connections - Default::default(), - retained_identities, - witness_columns, - // No prover functions - Vec::new(), - ); + let [KnownMachine::DynamicMachine(machine)] = machines + .iter() + .filter(|m| m.name().contains(machine_name)) + .collect_vec() + .as_slice() + else { + panic!("Expected exactly one matching dynamic machine") + }; + let machine_parts = machine.machine_parts().clone(); let mutable_state = MutableState::new(machines.into_iter(), &|_| { Err("Query not implemented".to_string()) }); @@ -157,18 +150,23 @@ mod test { #[test] fn fib() { - let input = "namespace Main(256); let X; let Y; X' = Y; Y' = X + Y;"; - let code = generate_single_step(input).unwrap(); - assert_eq!(format_code(&code), "X[1] = Y[0];\nY[1] = (X[0] + Y[0]);"); + let input = "namespace M(256); let X; let Y; X' = Y; Y' = X + Y;"; + let code = generate_single_step(input, "M").unwrap(); + assert_eq!( + format_code(&code), + "M::X[1] = M::Y[0];\nM::Y[1] = (M::X[0] + M::Y[0]);" + ); } #[test] fn no_progress() { - let input = "namespace Main(256); let X; let Y; X' = X;"; - let err = generate_single_step(input).err().unwrap(); + let input = "namespace M(256); let X; let Y; X' = X;"; + let err = generate_single_step(input, "M").err().unwrap(); assert_eq!( err.to_string(), - "Unable to derive algorithm to compute values for witness columns in the next row for the following columns: Y" + "Unable to derive algorithm to compute values for witness columns in the next row for the following columns:\n\ + M::Y\n\ + and 0 identities are missing." ); } } diff --git a/executor/src/witgen/machines/dynamic_machine.rs b/executor/src/witgen/machines/dynamic_machine.rs index 26f25f1bd8..762d662cd1 100644 --- a/executor/src/witgen/machines/dynamic_machine.rs +++ b/executor/src/witgen/machines/dynamic_machine.rs @@ -278,4 +278,9 @@ impl<'a, T: FieldElement> DynamicMachine<'a, T> { ); } } + + #[cfg(test)] + pub fn machine_parts(&self) -> &MachineParts<'a, T> { + &self.parts + } }