From 01fda8648647dc09cf63b90db18ce1aa021483a5 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 14:10:41 +0000 Subject: [PATCH 01/21] evaluate eagerly --- executor/src/witgen/jit/witgen_inference.rs | 63 ++++++++++++--------- 1 file changed, 35 insertions(+), 28 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index cf936ef476..2459a15b64 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -328,6 +328,30 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F selector: &Expression, arguments: &'a [Expression], row_offset: i32, + ) -> ProcessResult { + let vars = arguments + .iter() + .enumerate() + .map(|(index, arg)| { + let var = Variable::MachineCallParam(MachineCallVariable { + identity_id: lookup_id, + row_offset, + index, + }); + self.assign_variable(arg, row_offset, var.clone()); + var + }) + .collect_vec(); + self.process_call_(can_process_call, lookup_id, selector, &vars, row_offset) + } + + fn process_call_( + &mut self, + can_process_call: impl CanProcessCall, + lookup_id: u64, + selector: &Expression, + arguments: &[Variable], + row_offset: i32, ) -> ProcessResult { // We need to know the selector. let Some(selector) = self @@ -346,44 +370,27 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F assert_eq!(selector, 1.into(), "Selector is non-binary"); } - let evaluated = arguments - .iter() - .map(|a| self.evaluate(a, row_offset)) - .collect::>(); - let range_constraints = evaluated + let range_constraints = arguments .iter() - .map(|e| e.as_ref().map(|e| e.range_constraint()).unwrap_or_default()) + .map(|v| self.range_constraint(v)) .collect_vec(); - let known: BitVec = evaluated - .iter() - .map(|e| e.as_ref().and_then(|e| e.try_to_known()).is_some()) - .collect(); + let known: BitVec = arguments.iter().map(|v| self.is_known(v)).collect(); let Some(new_range_constraints) = can_process_call.can_process_call_fully(lookup_id, &known, &range_constraints) else { return ProcessResult::empty(); }; - let mut effects = vec![]; - let vars = arguments + let effects = arguments .iter() .zip_eq(new_range_constraints) - .enumerate() - .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())); + .map(|(var, new_rc)| Effect::RangeConstraint(var.clone(), new_rc.clone())) + .chain(std::iter::once(Effect::MachineCall( + lookup_id, + known, + arguments.to_vec(), + ))) + .collect(); ProcessResult { effects, complete: true, From 144e0666e3e342726dd34fbdcdd2edbfbe110ccb Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:05:20 +0000 Subject: [PATCH 02/21] oe --- executor/src/witgen/jit/witgen_inference.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 2459a15b64..b71768c1bf 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -197,6 +197,19 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F self.ingest_effects(result, Some((id.id(), row_offset))) } + pub fn process_polynomial_identity( + &mut self, + expression: &'a Expression, + row_offset: i32, + ) -> Result, Error> { + let result = self.process_equality_on_row( + expression, + row_offset, + &VariableOrValue::Value(T::from(0)), + )?; + self.ingest_effects(result, None) + } + /// Process a prover function on a row, i.e. determine if we can execute it and if it will /// help us to compute the value of previously unknown variables. /// Returns the list of updated variables. From eb79b4250c2e740ad43e608b93b14783e507d0f7 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:08:20 +0000 Subject: [PATCH 03/21] x --- executor/src/witgen/jit/processor.rs | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 195d0c2de7..86675ee3cf 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -5,12 +5,14 @@ use std::{ }; use itertools::Itertools; -use powdr_ast::analyzed::{PolyID, PolynomialType}; +use powdr_ast::analyzed::{PolyID, PolynomialIdentity, PolynomialType}; use powdr_number::FieldElement; use crate::witgen::{ - data_structures::identity::Identity, jit::debug_formatter::format_identities, - range_constraints::RangeConstraint, FixedData, + data_structures::identity::{BusSend, Identity}, + jit::debug_formatter::format_identities, + range_constraints::RangeConstraint, + FixedData, }; use super::{ @@ -283,9 +285,12 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv loop { let identity = identity_queue.next(); let updated_vars = match identity { - Some((identity, row_offset)) => { - witgen.process_identity(can_process.clone(), identity, row_offset) - } + Some((identity, row_offset)) => match identity { + Identity::Polynomial(PolynomialIdentity { expression, .. }) => { + witgen.process_polynomial_identity(expression, row_offset) + } + _ => witgen.process_identity(can_process.clone(), identity, row_offset), + }, None => self.process_prover_functions(witgen), }?; if updated_vars.is_empty() && identity.is_none() { From 7f9d358f557c3af5f43c17b06da11b65ecf47337 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:15:11 +0000 Subject: [PATCH 04/21] fix --- executor/src/witgen/jit/processor.rs | 4 ++-- executor/src/witgen/jit/witgen_inference.rs | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 86675ee3cf..24359ac2a8 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -286,8 +286,8 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv let identity = identity_queue.next(); let updated_vars = match identity { Some((identity, row_offset)) => match identity { - Identity::Polynomial(PolynomialIdentity { expression, .. }) => { - witgen.process_polynomial_identity(expression, row_offset) + Identity::Polynomial(PolynomialIdentity { id, expression, .. }) => { + witgen.process_polynomial_identity(*id, expression, row_offset) } _ => witgen.process_identity(can_process.clone(), identity, row_offset), }, diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index b71768c1bf..8ec5c3a8eb 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -199,6 +199,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F pub fn process_polynomial_identity( &mut self, + identity_id: u64, expression: &'a Expression, row_offset: i32, ) -> Result, Error> { @@ -207,7 +208,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F row_offset, &VariableOrValue::Value(T::from(0)), )?; - self.ingest_effects(result, None) + self.ingest_effects(result, Some((identity_id, row_offset))) } /// Process a prover function on a row, i.e. determine if we can execute it and if it will From f549ab0754cde0fceb8429468ddc40be69605560 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:17:25 +0000 Subject: [PATCH 05/21] x --- executor/src/witgen/jit/witgen_inference.rs | 23 ++++++++++++--------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 8ec5c3a8eb..328d51403c 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -185,13 +185,15 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F row_offset, &VariableOrValue::Value(T::from(0)), )?, - Identity::BusSend(bus_interaction) => self.process_call( - can_process, - bus_interaction.identity_id, - &bus_interaction.selected_payload.selector, - &bus_interaction.selected_payload.expressions, - row_offset, - ), + Identity::BusSend(bus_interaction) => { + return self.process_call( + can_process, + bus_interaction.identity_id, + &bus_interaction.selected_payload.selector, + &bus_interaction.selected_payload.expressions, + row_offset, + ) + } Identity::Connect(_) => ProcessResult::empty(), }; self.ingest_effects(result, Some((id.id(), row_offset))) @@ -335,14 +337,14 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F (lhs_evaluated - rhs_evaluated).solve() } - fn process_call( + pub fn process_call( &mut self, can_process_call: impl CanProcessCall, lookup_id: u64, selector: &Expression, arguments: &'a [Expression], row_offset: i32, - ) -> ProcessResult { + ) -> Result, Error> { let vars = arguments .iter() .enumerate() @@ -356,7 +358,8 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F var }) .collect_vec(); - self.process_call_(can_process_call, lookup_id, selector, &vars, row_offset) + let result = self.process_call_(can_process_call, lookup_id, selector, &vars, row_offset); + self.ingest_effects(result, Some((lookup_id, row_offset))) } fn process_call_( From f3c0c108161270858d5977fdfe9308feb06c71da Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:18:48 +0000 Subject: [PATCH 06/21] x --- executor/src/witgen/jit/processor.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 24359ac2a8..96f6e3b887 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -289,6 +289,17 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv Identity::Polynomial(PolynomialIdentity { id, expression, .. }) => { witgen.process_polynomial_identity(*id, expression, row_offset) } + Identity::BusSend(BusSend { + bus_id: _, + identity_id, + selected_payload, + }) => witgen.process_call( + can_process.clone(), + *identity_id, + &selected_payload.selector, + &selected_payload.expressions, + row_offset, + ), _ => witgen.process_identity(can_process.clone(), identity, row_offset), }, None => self.process_prover_functions(witgen), From d1bbde698012a9c369167c3be742d27e590a6444 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:23:10 +0000 Subject: [PATCH 07/21] x --- executor/src/witgen/jit/processor.rs | 4 +- executor/src/witgen/jit/witgen_inference.rs | 55 ++++++++------------- 2 files changed, 24 insertions(+), 35 deletions(-) diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 96f6e3b887..39f43cdc3f 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -300,8 +300,10 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv &selected_payload.expressions, row_offset, ), - _ => witgen.process_identity(can_process.clone(), identity, row_offset), + Identity::Connect(..) => Ok(vec![]), }, + // TODO Also add prover functions to the queue (activated by their variables) + // and sort them so that they are always last. None => self.process_prover_functions(witgen), }?; if updated_vars.is_empty() && identity.is_none() { diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 328d51403c..b8bd0bf7a5 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -169,36 +169,6 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F self.add_range_constraint(var.clone(), range_constraint); } - /// Process an identity on a certain row. - /// Returns Ok(true) if there was progress and Ok(false) if there was no progress. - /// If this returns an error, it means we have conflicting constraints. - pub fn process_identity( - &mut self, - can_process: impl CanProcessCall, - id: &'a Identity, - row_offset: i32, - ) -> Result, Error> { - let result = match id { - Identity::Polynomial(PolynomialIdentity { expression, .. }) => self - .process_equality_on_row( - expression, - row_offset, - &VariableOrValue::Value(T::from(0)), - )?, - Identity::BusSend(bus_interaction) => { - return self.process_call( - can_process, - bus_interaction.identity_id, - &bus_interaction.selected_payload.selector, - &bus_interaction.selected_payload.expressions, - row_offset, - ) - } - Identity::Connect(_) => ProcessResult::empty(), - }; - self.ingest_effects(result, Some((id.id(), row_offset))) - } - pub fn process_polynomial_identity( &mut self, identity_id: u64, @@ -762,6 +732,7 @@ mod test { use test_log::test; use crate::witgen::{ + data_structures::identity::BusSend, global_constraints, jit::{effect::format_code, test_util::read_pil, variable::Cell}, machines::{Connection, FixedLookup, KnownMachine}, @@ -821,10 +792,26 @@ mod test { counter += 1; for row in rows { for id in fixed_data.identities.iter() { - progress |= !witgen - .process_identity(&mutable_state, id, *row) - .unwrap() - .is_empty(); + let updated_vars = match id { + Identity::Polynomial(PolynomialIdentity { id, expression, .. }) => witgen + .process_polynomial_identity(*id, expression, *row) + .unwrap(), + Identity::BusSend(BusSend { + bus_id: _, + identity_id, + selected_payload, + }) => witgen + .process_call( + &mutable_state, + *identity_id, + &selected_payload.selector, + &selected_payload.expressions, + *row, + ) + .unwrap(), + Identity::Connect(..) => vec![], + }; + progress |= !updated_vars.is_empty(); } } if !progress { From 47604d087aa423c5b7f91809e9d551cc779ada1c Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:27:02 +0000 Subject: [PATCH 08/21] x --- executor/src/witgen/jit/processor.rs | 17 +++++++++++++++-- executor/src/witgen/jit/witgen_inference.rs | 3 ++- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 39f43cdc3f..6796212e67 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -20,7 +20,7 @@ use super::{ effect::{format_code, Effect}, identity_queue::IdentityQueue, prover_function_heuristics::ProverFunction, - variable::{Cell, Variable}, + variable::{Cell, MachineCallVariable, Variable}, witgen_inference::{BranchResult, CanProcessCall, FixedEvaluator, Value, WitgenInference}, }; @@ -111,8 +111,21 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv pub fn generate_code( self, can_process: impl CanProcessCall, - witgen: WitgenInference<'a, T, FixedEval>, + mut witgen: WitgenInference<'a, T, FixedEval>, ) -> Result, Error<'a, T, FixedEval>> { + // Create variables for bus send arguments. + for (id, row_offset) in &self.identities { + if let Identity::BusSend(bus_send) = id { + for (index, arg) in bus_send.selected_payload.expressions.iter().enumerate() { + let var = Variable::MachineCallParam(MachineCallVariable { + identity_id: bus_send.identity_id, + row_offset: *row_offset, + index, + }); + witgen.assign_variable(arg, *row_offset, var.clone()); + } + } + } let branch_depth = 0; let identity_queue = IdentityQueue::new(self.fixed_data, &self.identities); self.generate_code_for_branch(can_process, witgen, identity_queue, branch_depth) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index b8bd0bf7a5..98ac0729fe 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -324,10 +324,11 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F row_offset, index, }); - self.assign_variable(arg, row_offset, var.clone()); + //self.assign_variable(arg, row_offset, var.clone()); var }) .collect_vec(); + self.process_assignments().unwrap(); let result = self.process_call_(can_process_call, lookup_id, selector, &vars, row_offset); self.ingest_effects(result, Some((lookup_id, row_offset))) } From 13ea71a0c43901adf5632f68cae549d25a6d744d Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:28:09 +0000 Subject: [PATCH 09/21] x --- executor/src/witgen/jit/witgen_inference.rs | 22 ++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 98ac0729fe..118f5d7183 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -328,8 +328,14 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F var }) .collect_vec(); - self.process_assignments().unwrap(); - let result = self.process_call_(can_process_call, lookup_id, selector, &vars, row_offset); + + let result = self.process_call_( + can_process_call, + lookup_id, + selector, + vars.len(), + row_offset, + ); self.ingest_effects(result, Some((lookup_id, row_offset))) } @@ -338,9 +344,10 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F can_process_call: impl CanProcessCall, lookup_id: u64, selector: &Expression, - arguments: &[Variable], + argument_count: usize, row_offset: i32, ) -> ProcessResult { + self.process_assignments().unwrap(); // We need to know the selector. let Some(selector) = self .evaluate(selector, row_offset) @@ -358,6 +365,15 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F assert_eq!(selector, 1.into(), "Selector is non-binary"); } + let arguments = (0..argument_count) + .map(|index| { + Variable::MachineCallParam(MachineCallVariable { + identity_id: lookup_id, + row_offset, + index, + }) + }) + .collect_vec(); let range_constraints = arguments .iter() .map(|v| self.range_constraint(v)) From 7183a491969b5d4272a430bfb520aadf718e6b0d Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:28:26 +0000 Subject: [PATCH 10/21] x --- executor/src/witgen/jit/witgen_inference.rs | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 118f5d7183..5ee123ac4a 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -315,25 +315,11 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F arguments: &'a [Expression], row_offset: i32, ) -> Result, Error> { - let vars = arguments - .iter() - .enumerate() - .map(|(index, arg)| { - let var = Variable::MachineCallParam(MachineCallVariable { - identity_id: lookup_id, - row_offset, - index, - }); - //self.assign_variable(arg, row_offset, var.clone()); - var - }) - .collect_vec(); - let result = self.process_call_( can_process_call, lookup_id, selector, - vars.len(), + arguments.len(), row_offset, ); self.ingest_effects(result, Some((lookup_id, row_offset))) From dc087511cfb24de71fec62e65dd406a6ae6228c4 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:31:21 +0000 Subject: [PATCH 11/21] x --- executor/src/witgen/jit/processor.rs | 2 +- executor/src/witgen/jit/witgen_inference.rs | 24 ++++++++++++++++++--- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 6796212e67..74728127b6 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -310,7 +310,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv can_process.clone(), *identity_id, &selected_payload.selector, - &selected_payload.expressions, + selected_payload.expressions.len(), row_offset, ), Identity::Connect(..) => Ok(vec![]), diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 5ee123ac4a..fec6e517d6 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -312,14 +312,14 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F can_process_call: impl CanProcessCall, lookup_id: u64, selector: &Expression, - arguments: &'a [Expression], + argument_count: usize, row_offset: i32, ) -> Result, Error> { let result = self.process_call_( can_process_call, lookup_id, selector, - arguments.len(), + argument_count, row_offset, ); self.ingest_effects(result, Some((lookup_id, row_offset))) @@ -489,6 +489,8 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F if let Some(identity_id) = identity_id { // We actually only need to store completeness for submachine calls, // but we do it for all identities. + // TODO this is not correct, somehow we get problems if we do not provide + // IDs for poly identities. Need to revisit this. self.complete_identities.insert(identity_id); } } @@ -790,6 +792,22 @@ mod test { let ref_eval = FixedEvaluatorForFixedData(&fixed_data); let mut witgen = WitgenInference::new(&fixed_data, ref_eval, known_cells, []); let mut counter = 0; + // Create variables for bus send arguments. + for row in rows { + for id in &fixed_data.identities { + if let Identity::BusSend(bus_send) = id { + for (index, arg) in bus_send.selected_payload.expressions.iter().enumerate() { + let var = Variable::MachineCallParam(MachineCallVariable { + identity_id: bus_send.identity_id, + row_offset: *row, + index, + }); + witgen.assign_variable(arg, *row, var.clone()); + } + } + } + } + loop { let mut progress = false; counter += 1; @@ -808,7 +826,7 @@ mod test { &mutable_state, *identity_id, &selected_payload.selector, - &selected_payload.expressions, + selected_payload.expressions.len(), *row, ) .unwrap(), From ebf4612667f7fc14c4ddaf1f4ca0c64d1fb2e254 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:42:11 +0000 Subject: [PATCH 12/21] x --- .../src/witgen/jit/block_machine_processor.rs | 20 +++++++++---------- .../src/witgen/jit/single_step_processor.rs | 4 ++-- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs index fa02647ccd..4002b1cdc3 100644 --- a/executor/src/witgen/jit/block_machine_processor.rs +++ b/executor/src/witgen/jit/block_machine_processor.rs @@ -321,50 +321,50 @@ main_binary::operation_id[2] = main_binary::operation_id[3]; main_binary::operation_id[1] = main_binary::operation_id[2]; main_binary::operation_id[0] = main_binary::operation_id[1]; main_binary::operation_id_next[-1] = main_binary::operation_id[0]; +call_var(9, -1, 0) = main_binary::operation_id_next[-1]; main_binary::operation_id_next[0] = main_binary::operation_id[1]; +call_var(9, 0, 0) = main_binary::operation_id_next[0]; main_binary::operation_id_next[1] = main_binary::operation_id[2]; +call_var(9, 1, 0) = main_binary::operation_id_next[1]; main_binary::A_byte[2] = ((main_binary::A[3] & 0xff000000) // 16777216); main_binary::A[2] = (main_binary::A[3] & 0xffffff); assert (main_binary::A[3] & 0xffffffff00000000) == 0; +call_var(9, 2, 1) = main_binary::A_byte[2]; main_binary::A_byte[1] = ((main_binary::A[2] & 0xff0000) // 65536); main_binary::A[1] = (main_binary::A[2] & 0xffff); assert (main_binary::A[2] & 0xffffffffff000000) == 0; +call_var(9, 1, 1) = main_binary::A_byte[1]; main_binary::A_byte[0] = ((main_binary::A[1] & 0xff00) // 256); main_binary::A[0] = (main_binary::A[1] & 0xff); assert (main_binary::A[1] & 0xffffffffffff0000) == 0; +call_var(9, 0, 1) = main_binary::A_byte[0]; main_binary::A_byte[-1] = main_binary::A[0]; +call_var(9, -1, 1) = main_binary::A_byte[-1]; main_binary::B_byte[2] = ((main_binary::B[3] & 0xff000000) // 16777216); main_binary::B[2] = (main_binary::B[3] & 0xffffff); assert (main_binary::B[3] & 0xffffffff00000000) == 0; +call_var(9, 2, 2) = main_binary::B_byte[2]; main_binary::B_byte[1] = ((main_binary::B[2] & 0xff0000) // 65536); main_binary::B[1] = (main_binary::B[2] & 0xffff); assert (main_binary::B[2] & 0xffffffffff000000) == 0; +call_var(9, 1, 2) = main_binary::B_byte[1]; main_binary::B_byte[0] = ((main_binary::B[1] & 0xff00) // 256); main_binary::B[0] = (main_binary::B[1] & 0xff); assert (main_binary::B[1] & 0xffffffffffff0000) == 0; +call_var(9, 0, 2) = main_binary::B_byte[0]; main_binary::B_byte[-1] = main_binary::B[0]; -call_var(9, -1, 0) = main_binary::operation_id_next[-1]; -call_var(9, -1, 1) = main_binary::A_byte[-1]; call_var(9, -1, 2) = main_binary::B_byte[-1]; machine_call(9, [Known(call_var(9, -1, 0)), Known(call_var(9, -1, 1)), Known(call_var(9, -1, 2)), Unknown(call_var(9, -1, 3))]); main_binary::C_byte[-1] = call_var(9, -1, 3); main_binary::C[0] = main_binary::C_byte[-1]; -call_var(9, 0, 0) = main_binary::operation_id_next[0]; -call_var(9, 0, 1) = main_binary::A_byte[0]; -call_var(9, 0, 2) = main_binary::B_byte[0]; machine_call(9, [Known(call_var(9, 0, 0)), Known(call_var(9, 0, 1)), Known(call_var(9, 0, 2)), Unknown(call_var(9, 0, 3))]); main_binary::C_byte[0] = call_var(9, 0, 3); main_binary::C[1] = (main_binary::C[0] + (main_binary::C_byte[0] * 256)); -call_var(9, 1, 0) = main_binary::operation_id_next[1]; -call_var(9, 1, 1) = main_binary::A_byte[1]; -call_var(9, 1, 2) = main_binary::B_byte[1]; machine_call(9, [Known(call_var(9, 1, 0)), Known(call_var(9, 1, 1)), Known(call_var(9, 1, 2)), Unknown(call_var(9, 1, 3))]); main_binary::C_byte[1] = call_var(9, 1, 3); main_binary::C[2] = (main_binary::C[1] + (main_binary::C_byte[1] * 65536)); main_binary::operation_id_next[2] = main_binary::operation_id[3]; call_var(9, 2, 0) = main_binary::operation_id_next[2]; -call_var(9, 2, 1) = main_binary::A_byte[2]; -call_var(9, 2, 2) = main_binary::B_byte[2]; machine_call(9, [Known(call_var(9, 2, 0)), Known(call_var(9, 2, 1)), Known(call_var(9, 2, 2)), Unknown(call_var(9, 2, 3))]); main_binary::C_byte[2] = call_var(9, 2, 3); main_binary::C[3] = (main_binary::C[2] + (main_binary::C_byte[2] * 16777216)); diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index 289566474f..dbc763889e 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -237,12 +237,12 @@ namespace M(256); assert_eq!( format_code(&code), "\ -VM::pc[1] = (VM::pc[0] + 1); call_var(1, 0, 0) = VM::pc[0]; call_var(1, 0, 1) = VM::instr_add[0]; call_var(1, 0, 2) = VM::instr_mul[0]; -VM::B[1] = VM::B[0]; +VM::pc[1] = (VM::pc[0] + 1); call_var(1, 1, 0) = VM::pc[1]; +VM::B[1] = VM::B[0]; machine_call(1, [Known(call_var(1, 1, 0)), Unknown(call_var(1, 1, 1)), Unknown(call_var(1, 1, 2))]); VM::instr_add[1] = call_var(1, 1, 1); VM::instr_mul[1] = call_var(1, 1, 2); From eb78a9e196dd395bb1ed130ee9264afba4ab8f1f Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:42:34 +0000 Subject: [PATCH 13/21] clippy --- executor/src/witgen/jit/witgen_inference.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index fec6e517d6..978d139956 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -7,8 +7,7 @@ use bit_vec::BitVec; use itertools::Itertools; use powdr_ast::analyzed::{ AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression as Expression, - AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, PolynomialIdentity, - PolynomialType, + AlgebraicReference, AlgebraicUnaryOperation, AlgebraicUnaryOperator, PolynomialType, }; use powdr_number::FieldElement; From 33fc52e41e7d81b0a6801a6bd7ebf071200d9ba2 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:44:37 +0000 Subject: [PATCH 14/21] move --- executor/src/witgen/jit/witgen_inference.rs | 38 ++++++++++----------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 978d139956..abaff0433b 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -182,6 +182,24 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F self.ingest_effects(result, Some((identity_id, row_offset))) } + pub fn process_call( + &mut self, + can_process_call: impl CanProcessCall, + lookup_id: u64, + selector: &Expression, + argument_count: usize, + row_offset: i32, + ) -> Result, Error> { + let result = self.process_call_inner( + can_process_call, + lookup_id, + selector, + argument_count, + row_offset, + ); + self.ingest_effects(result, Some((lookup_id, row_offset))) + } + /// Process a prover function on a row, i.e. determine if we can execute it and if it will /// help us to compute the value of previously unknown variables. /// Returns the list of updated variables. @@ -306,25 +324,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F (lhs_evaluated - rhs_evaluated).solve() } - pub fn process_call( - &mut self, - can_process_call: impl CanProcessCall, - lookup_id: u64, - selector: &Expression, - argument_count: usize, - row_offset: i32, - ) -> Result, Error> { - let result = self.process_call_( - can_process_call, - lookup_id, - selector, - argument_count, - row_offset, - ); - self.ingest_effects(result, Some((lookup_id, row_offset))) - } - - fn process_call_( + fn process_call_inner( &mut self, can_process_call: impl CanProcessCall, lookup_id: u64, From 4583672487512ecb41f2b4d98a8a495ea2ada40b Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 15:49:57 +0000 Subject: [PATCH 15/21] fix --- executor/src/witgen/jit/witgen_inference.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index abaff0433b..519543638f 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -730,7 +730,7 @@ impl> CanProcessCall for &MutableState<' #[cfg(test)] mod test { - use powdr_ast::analyzed::PolyID; + use powdr_ast::analyzed::{PolyID, PolynomialIdentity}; use powdr_number::GoldilocksField; use pretty_assertions::assert_eq; use test_log::test; From 04e9a9a8dcb2e92bc3fc276e69c192c58020b06d Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 16:23:13 +0000 Subject: [PATCH 16/21] Handle assignments in identity queue as well --- ast/src/analyzed/mod.rs | 10 ++ executor/src/witgen/jit/identity_queue.rs | 142 +++++++++++++------- executor/src/witgen/jit/processor.rs | 9 +- executor/src/witgen/jit/witgen_inference.rs | 10 +- 4 files changed, 114 insertions(+), 57 deletions(-) diff --git a/ast/src/analyzed/mod.rs b/ast/src/analyzed/mod.rs index 6bf3210254..558f22f666 100644 --- a/ast/src/analyzed/mod.rs +++ b/ast/src/analyzed/mod.rs @@ -1209,6 +1209,16 @@ impl From<&AlgebraicReference> for AlgebraicReferenceThin { } } +impl AlgebraicReferenceThin { + pub fn with_name(&self, name: String) -> AlgebraicReference { + AlgebraicReference { + name, + poly_id: self.poly_id, + next: self.next, + } + } +} + #[derive(Debug, Clone, Eq, Serialize, Deserialize, JsonSchema)] pub struct AlgebraicReference { /// Name of the polynomial - just for informational purposes. diff --git a/executor/src/witgen/jit/identity_queue.rs b/executor/src/witgen/jit/identity_queue.rs index eca9ac779e..9df3dc7c16 100644 --- a/executor/src/witgen/jit/identity_queue.rs +++ b/executor/src/witgen/jit/identity_queue.rs @@ -15,7 +15,10 @@ use powdr_number::FieldElement; use crate::witgen::{data_structures::identity::Identity, FixedData}; -use super::variable::Variable; +use super::{ + variable::Variable, + witgen_inference::{Assignment, VariableOrValue}, +}; /// Keeps track of identities that still need to be processed and /// updates this list based on the occurrence of updated variables @@ -27,98 +30,120 @@ pub struct IdentityQueue<'a, T: FieldElement> { } impl<'a, T: FieldElement> IdentityQueue<'a, T> { - pub fn new(fixed_data: &'a FixedData<'a, T>, identities: &[(&'a Identity, i32)]) -> Self { - let occurrences = compute_occurrences_map(fixed_data, identities).into(); - Self { - queue: identities - .iter() - .map(|(id, row)| QueueItem(id, *row)) - .collect(), - occurrences, - } + pub fn new( + fixed_data: &'a FixedData<'a, T>, + identities: &[(&'a Identity, i32)], + assignments: &[Assignment<'a, T>], + ) -> Self { + let queue: BTreeSet<_> = identities + .iter() + .map(|(id, row)| QueueItem::Identity(id, *row)) + .chain(assignments.iter().map(|a| QueueItem::Assignment(a.clone()))) + .collect(); + let occurrences = compute_occurrences_map(fixed_data, &queue).into(); + Self { queue, occurrences } } /// Returns the next identity to be processed and its row and /// removes it from the queue. - pub fn next(&mut self) -> Option<(&'a Identity, i32)> { - self.queue.pop_first().map(|QueueItem(id, row)| (id, row)) + pub fn next(&mut self) -> Option> { + self.queue.pop_first() } pub fn variables_updated( &mut self, variables: impl IntoIterator, - skip_identity: Option<(&'a Identity, i32)>, + skip_item: Option>, ) { self.queue.extend( variables .into_iter() .flat_map(|var| self.occurrences.get(&var)) .flatten() - .filter(|QueueItem(id, row)| match skip_identity { - Some((id2, row2)) => (id.id(), *row) != (id2.id(), row2), + .filter(|item| match &skip_item { + Some(it) => *item != it, None => true, - }), + }) + .cloned(), ) } } -/// Sorts identities by row and then by ID. -#[derive(Clone, Copy)] -struct QueueItem<'a, T>(&'a Identity, i32); - -impl QueueItem<'_, T> { - fn key(&self) -> (i32, u64) { - let QueueItem(id, row) = self; - (*row, id.id()) - } +#[derive(Clone)] +pub enum QueueItem<'a, T: FieldElement> { + Identity(&'a Identity, i32), + Assignment(Assignment<'a, T>), } -impl Ord for QueueItem<'_, T> { +/// Sorts identities by row and then by ID, followed by assignments. +impl Ord for QueueItem<'_, T> { fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.key().cmp(&other.key()) + match (self, other) { + (QueueItem::Identity(id1, row1), QueueItem::Identity(id2, row2)) => { + (row1, id1.id()).cmp(&(row2, id2.id())) + } + (QueueItem::Assignment(a1), QueueItem::Assignment(a2)) => a1.cmp(a2), + (QueueItem::Identity(_, _), QueueItem::Assignment(_)) => std::cmp::Ordering::Less, + (QueueItem::Assignment(_), QueueItem::Identity(_, _)) => std::cmp::Ordering::Greater, + } } } -impl PartialOrd for QueueItem<'_, T> { +impl PartialOrd for QueueItem<'_, T> { fn partial_cmp(&self, other: &Self) -> Option { Some(self.cmp(other)) } } -impl PartialEq for QueueItem<'_, T> { +impl PartialEq for QueueItem<'_, T> { fn eq(&self, other: &Self) -> bool { - self.key() == other.key() + self.cmp(other) == std::cmp::Ordering::Equal } } -impl Eq for QueueItem<'_, T> {} +impl Eq for QueueItem<'_, T> {} -/// Computes a map from each variable to the identity-row-offset pairs it occurs in. -fn compute_occurrences_map<'a, T: FieldElement>( +/// Computes a map from each variable to the queue items it occurs in. +fn compute_occurrences_map<'b, 'a: 'b, T: FieldElement>( fixed_data: &'a FixedData<'a, T>, - identities: &[(&'a Identity, i32)], + items: &BTreeSet>, ) -> HashMap>> { - let mut references_per_identity = HashMap::new(); let mut intermediate_cache = HashMap::new(); - for id in identities.iter().map(|(id, _)| *id).unique_by(|id| id.id()) { + + // Compute references only once per identity. + let mut references_per_identity = HashMap::new(); + for id in items + .iter() + .filter_map(|item| match item { + QueueItem::Identity(id, _) => Some(id), + _ => None, + }) + .unique_by(|id| id.id()) + { references_per_identity.insert( - id, + id.id(), references_in_identity(id, fixed_data, &mut intermediate_cache), ); } - identities + + items .iter() - .flat_map(|(id, row)| { - references_per_identity[id].iter().map(move |reference| { - let name = fixed_data.column_name(&reference.poly_id).to_string(); - let fat_ref = AlgebraicReference { - name, - poly_id: reference.poly_id, - next: reference.next, - }; - let var = Variable::from_reference(&fat_ref, *row); - (var, QueueItem(*id, *row)) - }) + .flat_map(|item| { + let variables = match item { + QueueItem::Identity(id, row) => { + references_in_identity(id, fixed_data, &mut intermediate_cache) + .into_iter() + .map(|r| { + let name = fixed_data.column_name(&r.poly_id).to_string(); + Variable::from_reference(&r.with_name(name), *row) + }) + .collect_vec() + } + QueueItem::Assignment(a) => { + variables_in_assignment(a, fixed_data, &mut intermediate_cache) + } + }; + variables.into_iter().map(move |v| (v, item.clone())) }) .into_group_map() } @@ -184,3 +209,22 @@ fn references_in_expression<'a, T: FieldElement>( ) .unique() } + +/// Returns a vector of all variables that occur in the assignment. +fn variables_in_assignment<'a, T: FieldElement>( + assignment: &Assignment<'a, T>, + fixed_data: &'a FixedData<'a, T>, + intermediate_cache: &mut HashMap>, +) -> Vec { + let rhs_var = match &assignment.rhs { + VariableOrValue::Variable(v) => Some(v.clone()), + VariableOrValue::Value(_) => None, + }; + references_in_expression(assignment.lhs, fixed_data, intermediate_cache) + .map(|r| { + let name = fixed_data.column_name(&r.poly_id).to_string(); + Variable::from_reference(&r.with_name(name), assignment.row_offset) + }) + .chain(rhs_var) + .collect() +} diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 195d0c2de7..9c1edec6c1 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -16,7 +16,7 @@ use crate::witgen::{ use super::{ affine_symbolic_expression, effect::{format_code, Effect}, - identity_queue::IdentityQueue, + identity_queue::{IdentityQueue, QueueItem}, prover_function_heuristics::ProverFunction, variable::{Cell, Variable}, witgen_inference::{BranchResult, CanProcessCall, FixedEvaluator, Value, WitgenInference}, @@ -112,7 +112,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv witgen: WitgenInference<'a, T, FixedEval>, ) -> Result, Error<'a, T, FixedEval>> { let branch_depth = 0; - let identity_queue = IdentityQueue::new(self.fixed_data, &self.identities); + let identity_queue = IdentityQueue::new(self.fixed_data, &self.identities, &[]); self.generate_code_for_branch(can_process, witgen, identity_queue, branch_depth) } @@ -283,9 +283,12 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv loop { let identity = identity_queue.next(); let updated_vars = match identity { - Some((identity, row_offset)) => { + Some(QueueItem::Identity(identity, row_offset)) => { witgen.process_identity(can_process.clone(), identity, row_offset) } + Some(QueueItem::Assignment(assignment)) => { + todo!() + } None => self.process_prover_functions(witgen), }?; if updated_vars.is_empty() && identity.is_none() { diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index cf936ef476..4b32b0a656 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -682,14 +682,14 @@ fn is_known_zero(x: &Option { - lhs: &'a Expression, - row_offset: i32, - rhs: VariableOrValue, +pub struct Assignment<'a, T: FieldElement> { + pub lhs: &'a Expression, + pub row_offset: i32, + pub rhs: VariableOrValue, } #[derive(Clone, derive_more::Display, Ord, PartialOrd, Eq, PartialEq, Debug)] -enum VariableOrValue { +pub enum VariableOrValue { Variable(V), Value(T), } From df24e9019aad8e2910f0dd70ab6ed91eeb1c38bf Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 16:33:40 +0000 Subject: [PATCH 17/21] fix --- executor/src/witgen/jit/single_step_processor.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index dbc763889e..aae13fb543 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -280,12 +280,12 @@ if (VM::instr_add[0] == 1) { assert_eq!( format_code(&code), "\ -VM::pc[1] = VM::pc[0]; call_var(2, 0, 0) = VM::pc[0]; -call_var(2, 0, 1) = 0; +call_var(2, 0, 1) = VM::instr_add[0]; call_var(2, 0, 2) = VM::instr_mul[0]; -VM::instr_add[1] = 0; +VM::pc[1] = VM::pc[0]; call_var(2, 1, 0) = VM::pc[1]; +VM::instr_add[1] = 0; 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))]); From b7052230dd563c703e4c4561c724781dfd05aaf5 Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 6 Feb 2025 16:36:47 +0000 Subject: [PATCH 18/21] clippy --- executor/src/witgen/jit/identity_queue.rs | 5 +---- executor/src/witgen/jit/processor.rs | 2 +- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/executor/src/witgen/jit/identity_queue.rs b/executor/src/witgen/jit/identity_queue.rs index 9df3dc7c16..8ccc9e60f2 100644 --- a/executor/src/witgen/jit/identity_queue.rs +++ b/executor/src/witgen/jit/identity_queue.rs @@ -5,10 +5,7 @@ use std::{ use itertools::Itertools; use powdr_ast::{ - analyzed::{ - AlgebraicExpression as Expression, AlgebraicReference, AlgebraicReferenceThin, - PolynomialType, - }, + analyzed::{AlgebraicExpression as Expression, AlgebraicReferenceThin, PolynomialType}, parsed::visitor::{AllChildren, Children}, }; use powdr_number::FieldElement; diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 9c1edec6c1..6860163d39 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -286,7 +286,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv Some(QueueItem::Identity(identity, row_offset)) => { witgen.process_identity(can_process.clone(), identity, row_offset) } - Some(QueueItem::Assignment(assignment)) => { + Some(QueueItem::Assignment(_assignment)) => { todo!() } None => self.process_prover_functions(witgen), From 160cf109e2cd0eb2aba869c5e7e32d34c0750fa7 Mon Sep 17 00:00:00 2001 From: chriseth Date: Fri, 7 Feb 2025 08:54:36 +0000 Subject: [PATCH 19/21] fix tests. --- .../src/witgen/jit/block_machine_processor.rs | 7 ++-- executor/src/witgen/jit/identity_queue.rs | 37 ++++++++++++++++--- executor/src/witgen/jit/processor.rs | 5 +-- .../src/witgen/jit/single_step_processor.rs | 6 +-- executor/src/witgen/jit/witgen_inference.rs | 17 +++++++++ 5 files changed, 56 insertions(+), 16 deletions(-) diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs index bec3d35ef0..fed1109522 100644 --- a/executor/src/witgen/jit/block_machine_processor.rs +++ b/executor/src/witgen/jit/block_machine_processor.rs @@ -7,9 +7,8 @@ use powdr_number::FieldElement; use crate::witgen::{ jit::{ - processor::Processor, - prover_function_heuristics::decode_prover_functions, - witgen_inference::{Assignment, VariableOrValue}, + processor::Processor, prover_function_heuristics::decode_prover_functions, + witgen_inference::Assignment, }, machines::MachineParts, FixedData, @@ -65,7 +64,7 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> { .enumerate() .filter_map(|(i, is_input)| is_input.then_some(Variable::Param(i))) .collect::>(); - let mut witgen = WitgenInference::new(self.fixed_data, self, known_variables, []); + let witgen = WitgenInference::new(self.fixed_data, self, known_variables, []); let prover_functions = decode_prover_functions(&self.machine_parts, self.fixed_data)?; diff --git a/executor/src/witgen/jit/identity_queue.rs b/executor/src/witgen/jit/identity_queue.rs index 8ccc9e60f2..3f8fe88eb7 100644 --- a/executor/src/witgen/jit/identity_queue.rs +++ b/executor/src/witgen/jit/identity_queue.rs @@ -10,7 +10,9 @@ use powdr_ast::{ }; use powdr_number::FieldElement; -use crate::witgen::{data_structures::identity::Identity, FixedData}; +use crate::witgen::{ + data_structures::identity::Identity, jit::variable::MachineCallVariable, FixedData, +}; use super::{ variable::Variable, @@ -128,13 +130,25 @@ fn compute_occurrences_map<'b, 'a: 'b, T: FieldElement>( .flat_map(|item| { let variables = match item { QueueItem::Identity(id, row) => { - references_in_identity(id, fixed_data, &mut intermediate_cache) - .into_iter() + let mut variables = references_per_identity[&id.id()] + .iter() .map(|r| { let name = fixed_data.column_name(&r.poly_id).to_string(); Variable::from_reference(&r.with_name(name), *row) }) - .collect_vec() + .collect_vec(); + if let Identity::BusSend(bus_send) = id { + variables.extend((0..bus_send.selected_payload.expressions.len()).map( + |index| { + Variable::MachineCallParam(MachineCallVariable { + identity_id: id.id(), + row_offset: *row, + index, + }) + }, + )); + }; + variables } QueueItem::Assignment(a) => { variables_in_assignment(a, fixed_data, &mut intermediate_cache) @@ -152,9 +166,20 @@ fn references_in_identity( intermediate_cache: &mut HashMap>, ) -> Vec { let mut result = BTreeSet::new(); - for e in identity.children() { - result.extend(references_in_expression(e, fixed_data, intermediate_cache)); + + match identity { + Identity::BusSend(bus_send) => result.extend(references_in_expression( + &bus_send.selected_payload.selector, + fixed_data, + intermediate_cache, + )), + _ => { + for e in identity.children() { + result.extend(references_in_expression(e, fixed_data, intermediate_cache)); + } + } } + result.into_iter().collect() } diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 2c19dc19e2..0ac4724bec 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -22,8 +22,7 @@ use super::{ prover_function_heuristics::ProverFunction, variable::{Cell, MachineCallVariable, Variable}, witgen_inference::{ - Assignment, BranchResult, CanProcessCall, FixedEvaluator, Value, VariableOrValue, - WitgenInference, + Assignment, BranchResult, CanProcessCall, FixedEvaluator, Value, WitgenInference, }, }; @@ -336,7 +335,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv ), Identity::Connect(..) => Ok(vec![]), }, - Some(QueueItem::Assignment(assignment)) => witgen.process_assignment(&assignment), + Some(QueueItem::Assignment(assignment)) => witgen.process_assignment(assignment), // TODO Also add prover functions to the queue (activated by their variables) // and sort them so that they are always last. None => self.process_prover_functions(witgen), diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index 3eb577bf03..e07bc2ea76 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -238,12 +238,12 @@ namespace M(256); assert_eq!( format_code(&code), "\ -call_var(1, 0, 0) = VM::pc[0]; +VM::pc[1] = (VM::pc[0] + 1); +VM::B[1] = VM::B[0]; call_var(1, 0, 1) = VM::instr_add[0]; call_var(1, 0, 2) = VM::instr_mul[0]; -VM::pc[1] = (VM::pc[0] + 1); +call_var(1, 0, 0) = VM::pc[0]; call_var(1, 1, 0) = VM::pc[1]; -VM::B[1] = VM::B[0]; machine_call(1, [Known(call_var(1, 1, 0)), Unknown(call_var(1, 1, 1)), Unknown(call_var(1, 1, 2))]); VM::instr_add[1] = call_var(1, 1, 1); VM::instr_mul[1] = call_var(1, 1, 2); diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 3a86d25e98..1b7dba0527 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -627,6 +627,8 @@ fn is_known_zero(x: &Option { Value(T), } +impl Display for Assignment<'_, T> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{} = {} [row {}]", + self.lhs, + match &self.rhs { + VariableOrValue::Variable(v) => v.to_string(), + VariableOrValue::Value(v) => v.to_string(), + }, + self.row_offset + ) + } +} + pub trait FixedEvaluator: Clone { /// Evaluate a fixed column cell and returns its value if it is /// compile-time constant, otherwise return None. From e4d29eede5977d55bc2a17dc34491840b23f0d27 Mon Sep 17 00:00:00 2001 From: chriseth Date: Fri, 7 Feb 2025 08:55:14 +0000 Subject: [PATCH 20/21] remove assignments. --- executor/src/witgen/jit/witgen_inference.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 1b7dba0527..2594b2e980 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeSet, HashMap, HashSet}, + collections::{HashMap, HashSet}, fmt::{Display, Formatter}, }; @@ -42,8 +42,6 @@ pub struct WitgenInference<'a, T: FieldElement, FixedEval> { /// This mainly avoids generating multiple submachine calls for the same /// connection on the same row. complete_identities: HashSet<(u64, i32)>, - /// Internal equality constraints that are not identities from the constraint set. - assignments: BTreeSet>, code: Vec>, } @@ -88,7 +86,6 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F derived_range_constraints: Default::default(), known_variables: known_variables.into_iter().collect(), complete_identities: complete_identities.into_iter().collect(), - assignments: Default::default(), code: Default::default(), } } From b1a34d80fbecb1f2a22fc615e00c5f44e3253a6f Mon Sep 17 00:00:00 2001 From: chriseth Date: Fri, 7 Feb 2025 09:04:23 +0000 Subject: [PATCH 21/21] fix tests. --- executor/src/witgen/jit/block_machine_processor.rs | 4 ++-- executor/src/witgen/jit/identity_queue.rs | 6 +++--- executor/src/witgen/jit/single_step_processor.rs | 6 +++--- executor/src/witgen/jit/witgen_inference.rs | 6 +++--- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs index fed1109522..fc41ff6203 100644 --- a/executor/src/witgen/jit/block_machine_processor.rs +++ b/executor/src/witgen/jit/block_machine_processor.rs @@ -330,10 +330,10 @@ params[2] = Add::c[0];" assert_eq!(c_rc, &RangeConstraint::from_mask(0xffffffffu64)); assert_eq!( format_code(&result.code), - "main_binary::sel[0][3] = 1; -main_binary::operation_id[3] = params[0]; + "main_binary::operation_id[3] = params[0]; main_binary::A[3] = params[1]; main_binary::B[3] = params[2]; +main_binary::sel[0][3] = 1; main_binary::operation_id[2] = main_binary::operation_id[3]; main_binary::operation_id[1] = main_binary::operation_id[2]; main_binary::operation_id[0] = main_binary::operation_id[1]; diff --git a/executor/src/witgen/jit/identity_queue.rs b/executor/src/witgen/jit/identity_queue.rs index 3f8fe88eb7..dbe012f2d2 100644 --- a/executor/src/witgen/jit/identity_queue.rs +++ b/executor/src/witgen/jit/identity_queue.rs @@ -74,7 +74,7 @@ pub enum QueueItem<'a, T: FieldElement> { Assignment(Assignment<'a, T>), } -/// Sorts identities by row and then by ID, followed by assignments. +/// Sorts identities by row and then by ID, preceded by assignments. impl Ord for QueueItem<'_, T> { fn cmp(&self, other: &Self) -> std::cmp::Ordering { match (self, other) { @@ -82,8 +82,8 @@ impl Ord for QueueItem<'_, T> { (row1, id1.id()).cmp(&(row2, id2.id())) } (QueueItem::Assignment(a1), QueueItem::Assignment(a2)) => a1.cmp(a2), - (QueueItem::Identity(_, _), QueueItem::Assignment(_)) => std::cmp::Ordering::Less, - (QueueItem::Assignment(_), QueueItem::Identity(_, _)) => std::cmp::Ordering::Greater, + (QueueItem::Identity(_, _), QueueItem::Assignment(_)) => std::cmp::Ordering::Greater, + (QueueItem::Assignment(_), QueueItem::Identity(_, _)) => std::cmp::Ordering::Less, } } } diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index e07bc2ea76..a5cac75dd8 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -238,12 +238,12 @@ namespace M(256); assert_eq!( format_code(&code), "\ -VM::pc[1] = (VM::pc[0] + 1); -VM::B[1] = VM::B[0]; call_var(1, 0, 1) = VM::instr_add[0]; call_var(1, 0, 2) = VM::instr_mul[0]; call_var(1, 0, 0) = VM::pc[0]; +VM::pc[1] = (VM::pc[0] + 1); call_var(1, 1, 0) = VM::pc[1]; +VM::B[1] = VM::B[0]; machine_call(1, [Known(call_var(1, 1, 0)), Unknown(call_var(1, 1, 1)), Unknown(call_var(1, 1, 2))]); VM::instr_add[1] = call_var(1, 1, 1); VM::instr_mul[1] = call_var(1, 1, 2); @@ -281,9 +281,9 @@ if (VM::instr_add[0] == 1) { assert_eq!( format_code(&code), "\ -call_var(2, 0, 0) = VM::pc[0]; call_var(2, 0, 1) = VM::instr_add[0]; call_var(2, 0, 2) = VM::instr_mul[0]; +call_var(2, 0, 0) = VM::pc[0]; VM::pc[1] = VM::pc[0]; call_var(2, 1, 0) = VM::pc[1]; VM::instr_add[1] = 0; diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 2594b2e980..fdc1770318 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -931,7 +931,6 @@ assert (Xor::C[6] & 0xffffffffff000000) == 0; call_var(0, 6, 0) = Xor::A_byte[6]; call_var(0, 6, 2) = Xor::C_byte[6]; machine_call(0, [Known(call_var(0, 6, 0)), Unknown(call_var(0, 6, 1)), Known(call_var(0, 6, 2))]); -Xor::B_byte[6] = call_var(0, 6, 1); Xor::A_byte[4] = ((Xor::A[5] & 0xff00) // 256); Xor::A[4] = (Xor::A[5] & 0xff); assert (Xor::A[5] & 0xffffffffffff0000) == 0; @@ -941,16 +940,17 @@ assert (Xor::C[5] & 0xffffffffffff0000) == 0; call_var(0, 5, 0) = Xor::A_byte[5]; call_var(0, 5, 2) = Xor::C_byte[5]; machine_call(0, [Known(call_var(0, 5, 0)), Unknown(call_var(0, 5, 1)), Known(call_var(0, 5, 2))]); -Xor::B_byte[5] = call_var(0, 5, 1); +Xor::B_byte[6] = call_var(0, 6, 1); Xor::A_byte[3] = Xor::A[4]; Xor::C_byte[3] = Xor::C[4]; call_var(0, 4, 0) = Xor::A_byte[4]; call_var(0, 4, 2) = Xor::C_byte[4]; machine_call(0, [Known(call_var(0, 4, 0)), Unknown(call_var(0, 4, 1)), Known(call_var(0, 4, 2))]); -Xor::B_byte[4] = call_var(0, 4, 1); +Xor::B_byte[5] = call_var(0, 5, 1); call_var(0, 3, 0) = Xor::A_byte[3]; call_var(0, 3, 2) = Xor::C_byte[3]; machine_call(0, [Known(call_var(0, 3, 0)), Unknown(call_var(0, 3, 1)), Known(call_var(0, 3, 2))]); +Xor::B_byte[4] = call_var(0, 4, 1); Xor::B_byte[3] = call_var(0, 3, 1); Xor::B[4] = Xor::B_byte[3]; Xor::B[5] = (Xor::B[4] + (Xor::B_byte[4] * 256));