Move vm processor (#2136)

Co-authored-by: Georg Wiese <georgwiese@gmail.com>

Author: 2 people

executor/src/witgen/data_structures/mutable_state.rs

@@ -40,6 +40,15 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
         }
     }
 
+    /// 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<String, Vec<T>> {
+        if let Some(first_machine) = self.machines.first() {
+            first_machine.try_borrow_mut().unwrap().run_timed(&self);
+        }
+        self.take_witness_col_values()
+    }
+
     /// Call the machine responsible for the right-hand-side of an identity given its ID
     /// and the row pair of the caller.
     pub fn call(&self, identity_id: u64, caller_rows: &RowPair<'_, 'a, T>) -> EvalResult<'a, T> {
@@ -59,7 +68,7 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
     }
 
     /// Extracts the witness column values from the machines.
-    pub fn take_witness_col_values(self) -> HashMap<String, Vec<T>> {
+    fn take_witness_col_values(self) -> HashMap<String, Vec<T>> {
         // We keep the already processed machines mutably borrowed so that
         // "later" machines do not try to create new rows in already processed
         // machines.

executor/src/witgen/machines/double_sorted_witness_machine_16.rs

@@ -132,6 +132,10 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses16<'a, T> {
             return None;
         }
 
+        if parts.connections.is_empty() {
+            return None;
+        }
+
         if !parts
             .connections
             .values()

executor/src/witgen/machines/double_sorted_witness_machine_32.rs

@@ -102,6 +102,10 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses32<'a, T> {
             return None;
         }
 
+        if parts.connections.is_empty() {
+            return None;
+        }
+
         if !parts.connections.values().all(|i| i.is_permutation()) {
             return None;
         }

executor/src/witgen/generator.rs executor/src/witgen/machines/dynamic_machine.rs

@@ -2,28 +2,24 @@ use powdr_ast::analyzed::AlgebraicExpression as Expression;
 use powdr_number::{DegreeType, FieldElement};
 use std::collections::{BTreeMap, HashMap};
 
+use crate::witgen::block_processor::BlockProcessor;
 use crate::witgen::data_structures::finalizable_data::FinalizableData;
+use crate::witgen::data_structures::multiplicity_counter::MultiplicityCounter;
 use crate::witgen::data_structures::mutable_state::MutableState;
-use crate::witgen::machines::profiling::{record_end, record_start};
-use crate::witgen::processor::OuterQuery;
-use crate::witgen::EvalValue;
-
-use super::affine_expression::AlgebraicVariable;
-use super::block_processor::BlockProcessor;
-use super::data_structures::multiplicity_counter::MultiplicityCounter;
-use super::machines::{Machine, MachineParts};
-use super::processor::SolverState;
-use super::rows::{Row, RowIndex, RowPair};
-use super::sequence_iterator::{DefaultSequenceIterator, ProcessingSequenceIterator};
-use super::vm_processor::VmProcessor;
-use super::{EvalResult, FixedData, QueryCallback};
+use crate::witgen::machines::{Machine, MachineParts};
+use crate::witgen::processor::{OuterQuery, SolverState};
+use crate::witgen::rows::{Row, RowIndex, RowPair};
+use crate::witgen::sequence_iterator::{DefaultSequenceIterator, ProcessingSequenceIterator};
+use crate::witgen::vm_processor::VmProcessor;
+use crate::witgen::{AlgebraicVariable, EvalResult, EvalValue, FixedData, QueryCallback};
 
 struct ProcessResult<'a, T: FieldElement> {
     eval_value: EvalValue<AlgebraicVariable<'a>, T>,
     updated_data: SolverState<'a, T>,
 }
 
-pub struct Generator<'a, T: FieldElement> {
+/// A machine is generic and can handle lookups that generate a dynamic number of rows.
+pub struct DynamicMachine<'a, T: FieldElement> {
     fixed_data: &'a FixedData<'a, T>,
     parts: MachineParts<'a, T>,
     data: FinalizableData<T>,
@@ -34,7 +30,7 @@ pub struct Generator<'a, T: FieldElement> {
     multiplicity_counter: MultiplicityCounter,
 }
 
-impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> {
+impl<'a, T: FieldElement> Machine<'a, T> for DynamicMachine<'a, T> {
     fn identity_ids(&self) -> Vec<u64> {
         self.parts.identity_ids()
     }
@@ -43,6 +39,16 @@ impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> {
         &self.name
     }
 
+    /// Runs the machine without any arguments from the first row.
+    fn run<Q: QueryCallback<T>>(&mut self, mutable_state: &MutableState<'a, 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, None, true)
+            .updated_data
+            .block;
+    }
+
     fn process_plookup<'b, Q: QueryCallback<T>>(
         &mut self,
         mutable_state: &MutableState<'a, T, Q>,
@@ -110,7 +116,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> {
     }
 }
 
-impl<'a, T: FieldElement> Generator<'a, T> {
+impl<'a, T: FieldElement> DynamicMachine<'a, T> {
     pub fn new(
         name: String,
         fixed_data: &'a FixedData<'a, T>,
@@ -132,18 +138,6 @@ impl<'a, T: FieldElement> Generator<'a, T> {
         }
     }
 
-    /// Runs the machine without any arguments from the first row.
-    pub fn run<Q: QueryCallback<T>>(&mut self, mutable_state: &MutableState<'a, T, Q>) {
-        record_start(self.name());
-        assert!(self.data.is_empty());
-        let first_row = self.compute_partial_first_row(mutable_state);
-        self.data = self
-            .process(first_row, 0, mutable_state, None, true)
-            .updated_data
-            .block;
-        record_end(self.name());
-    }
-
     fn fill_remaining_rows<Q: QueryCallback<T>>(&mut self, mutable_state: &MutableState<'a, T, Q>) {
         if self.data.len() < self.degree as usize + 1 {
             assert!(self.latch.is_some());

executor/src/witgen/machines/machine_extractor.rs

@@ -13,11 +13,9 @@ use super::fixed_lookup_machine::FixedLookup;
 use super::sorted_witness_machine::SortedWitnesses;
 use super::FixedData;
 use super::KnownMachine;
+use crate::witgen::machines::dynamic_machine::DynamicMachine;
 use crate::witgen::machines::Connection;
-use crate::witgen::{
-    generator::Generator,
-    machines::{write_once_memory::WriteOnceMemory, MachineParts},
-};
+use crate::witgen::machines::{write_once_memory::WriteOnceMemory, MachineParts};
 use crate::Identity;
 
 use powdr_ast::analyzed::{
@@ -26,11 +24,6 @@ use powdr_ast::analyzed::{
 use powdr_ast::parsed::{self, visitor::AllChildren};
 use powdr_number::FieldElement;
 
-pub struct ExtractionOutput<'a, T: FieldElement> {
-    pub machines: Vec<KnownMachine<'a, T>>,
-    pub base_parts: MachineParts<'a, T>,
-}
-
 pub struct MachineExtractor<'a, T: FieldElement> {
     fixed: &'a FixedData<'a, T>,
 }
@@ -40,12 +33,10 @@ impl<'a, T: FieldElement> MachineExtractor<'a, T> {
         Self { fixed }
     }
 
-    /// Finds machines in the witness columns and identities
-    /// and returns a list of machines and the identities
+    /// Finds machines in the witness columns and identities and returns a list of machines and the identities
     /// that are not "internal" to the machines.
-    pub fn split_out_machines(&self, identities: Vec<&'a Identity<T>>) -> ExtractionOutput<'a, T> {
-        let mut machines: Vec<KnownMachine<T>> = vec![];
-
+    /// The first returned machine is the "main machine", i.e. a machine that has no incoming connections.
+    pub fn split_out_machines(&self, identities: Vec<&'a Identity<T>>) -> Vec<KnownMachine<'a, T>> {
         // Ignore prover functions that reference columns of later stages.
         let all_witnesses = self.fixed.witness_cols.keys().collect::<HashSet<_>>();
         let current_stage_witnesses = self
@@ -68,6 +59,30 @@ impl<'a, T: FieldElement> MachineExtractor<'a, T> {
             })
             .collect::<Vec<&analyzed::Expression>>();
 
+        if self.fixed.stage() > 0 {
+            // We expect later-stage witness columns to be accumulators for lookup and permutation arguments.
+            // These don't behave like normal witness columns (e.g. in a block machine), and they might depend
+            // on witness columns of more than one machine.
+            // Therefore, we treat everything as one big machine. Also, we remove lookups and permutations,
+            // as they are assumed to be handled in stage 0.
+            let polynomial_identities = identities
+                .into_iter()
+                .filter(|identity| matches!(identity, Identity::Polynomial(_)))
+                .collect::<Vec<_>>();
+            let machine_parts = MachineParts::new(
+                self.fixed,
+                Default::default(),
+                polynomial_identities,
+                self.fixed.witness_cols.keys().collect::<HashSet<_>>(),
+                prover_functions,
+            );
+
+            return build_main_machine(self.fixed, machine_parts)
+                .into_iter()
+                .collect();
+        }
+        let mut machines: Vec<KnownMachine<T>> = vec![];
+
         let mut publics = PublicsTracker::default();
         let mut remaining_witnesses = current_stage_witnesses.clone();
         let mut base_identities = identities.clone();
@@ -197,27 +212,33 @@ impl<'a, T: FieldElement> MachineExtractor<'a, T> {
             .collect::<Vec<_>>();
 
         log::trace!(
-        "\nThe base machine contains the following witnesses:\n{}\n identities:\n{}\n and prover functions:\n{}",
-        remaining_witnesses
-            .iter()
-            .map(|s| self.fixed.column_name(s))
-            .sorted()
-            .format(", "),
-        base_identities
-            .iter()
-            .format("\n"),
-        base_prover_functions.iter().format("\n")
-    );
+            "\nThe base machine contains the following witnesses:\n{}\n identities:\n{}\n and prover functions:\n{}",
+            remaining_witnesses
+                .iter()
+                .map(|s| self.fixed.column_name(s))
+                .sorted()
+                .format(", "),
+            base_identities
+                .iter()
+                .format("\n"),
+            base_prover_functions.iter().format("\n")
+        );
 
-        ExtractionOutput {
-            machines,
-            base_parts: MachineParts::new(
-                self.fixed,
-                Default::default(),
-                base_identities,
-                remaining_witnesses,
-                base_prover_functions,
-            ),
+        let base_parts = MachineParts::new(
+            self.fixed,
+            Default::default(),
+            base_identities,
+            remaining_witnesses,
+            base_prover_functions,
+        );
+
+        if let Some(main_machine) = build_main_machine(self.fixed, base_parts) {
+            std::iter::once(main_machine).chain(machines).collect()
+        } else {
+            if !machines.is_empty() {
+                log::error!("No main machine was extracted, but secondary machines were. Does the system have a cycle?");
+            }
+            vec![]
         }
     }
 
@@ -358,6 +379,14 @@ impl<'a> PublicsTracker<'a> {
     }
 }
 
+fn build_main_machine<'a, T: FieldElement>(
+    fixed_data: &'a FixedData<'a, T>,
+    machine_parts: MachineParts<'a, T>,
+) -> Option<KnownMachine<'a, T>> {
+    (!machine_parts.witnesses.is_empty())
+        .then(|| build_machine(fixed_data, machine_parts, |t| format!("Main machine ({t})")))
+}
+
 fn build_machine<'a, T: FieldElement>(
     fixed_data: &'a FixedData<'a, T>,
     machine_parts: MachineParts<'a, T>,
@@ -395,7 +424,9 @@ fn build_machine<'a, T: FieldElement>(
         log::debug!("Detected machine: {machine}");
         KnownMachine::BlockMachine(machine)
     } else {
-        log::debug!("Detected machine: VM.");
+        log::debug!("Detected machine: Dynamic machine.");
+        // If there is a connection to this machine, all connections must have the same latch.
+        // If there is no connection to this machine, it is the main machine and there is no latch.
         let latch = machine_parts.connections
             .values()
             .fold(None, |existing_latch, identity| {
@@ -411,13 +442,12 @@ fn build_machine<'a, T: FieldElement>(
                 } else {
                     Some(current_latch.clone())
                 }
-            })
-            .unwrap();
-        KnownMachine::Vm(Generator::new(
-            name_with_type("Vm"),
+            });
+        KnownMachine::DynamicMachine(DynamicMachine::new(
+            name_with_type("Dynamic"),
             fixed_data,
             machine_parts.clone(),
-            Some(latch),
+            latch,
         ))
     }
 }