[WIP] Bus prototype

backend/src/mock/machine.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, time::Instant};
 
 use itertools::Itertools;
 use powdr_ast::analyzed::{AlgebraicExpression, AlgebraicReferenceThin, Analyzed, PolyID};
@@ -39,9 +39,14 @@ impl<'a, F: FieldElement> Machine<'a, F> {
         }
 
         for stage in 1..pil.stage_count() {
-            log::debug!("Generating stage-{stage} witness for machine {machine_name}");
+            log::info!("Generating stage-{stage} witness for machine {machine_name}");
+            let start_time = Instant::now();
             witness =
                 witgen_callback.next_stage_witness(pil, &witness, challenges.clone(), stage as u8);
+            log::info!(
+                "Generated stage-{stage} witness for machine {machine_name} in {}s",
+                start_time.elapsed().as_secs_f64()
+            );
         }
 
         let fixed = machine_fixed_columns(fixed, pil);

executor/src/witgen/bus_accumulator.rs

@@ -0,0 +1,182 @@
+use std::collections::BTreeMap;
+
+use powdr_ast::analyzed::Analyzed;
+use powdr_number::FieldElement;
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+
+pub fn generate_bus_accumulators<T: FieldElement>(
+    _pil: &Analyzed<T>,
+    witness_columns: &[(String, Vec<T>)],
+    _fixed_columns: Vec<(String, &[T])>,
+    _challenges: BTreeMap<u64, T>,
+) -> Vec<(String, Vec<T>)> {
+    let accumulators = (0..31)
+        .into_par_iter()
+        .map(|i| interaction_columns(i, witness_columns))
+        .collect::<Vec<_>>();
+
+    witness_columns
+        .to_vec()
+        .into_iter()
+        .chain(accumulators.into_iter().flatten())
+        .collect()
+}
+
+fn interaction_columns<T: FieldElement>(
+    connection_index: usize,
+    witness_columns: &[(String, Vec<T>)],
+) -> Vec<(String, Vec<T>)> {
+    let tuple_size = if connection_index == 0 {
+        // Simulate PC lookup
+        700
+    } else {
+        1 + connection_index / 10
+    };
+
+    // Pick random indices
+    let indices = (0..tuple_size)
+        .map(|j| {
+            ((42usize
+                .wrapping_mul(connection_index * 70 + j)
+                .wrapping_add(123))
+                % 17482394)
+                % witness_columns.len()
+        })
+        .collect::<Vec<_>>();
+
+    let size = witness_columns[0].1.len();
+    let mut acc1 = vec![T::zero(); size];
+    let mut acc2 = vec![T::zero(); size];
+    let mut acc1_next = vec![T::zero(); size];
+    let mut acc2_next = vec![T::zero(); size];
+
+    for i in 0..size {
+        let current_acc = if i == 0 {
+            (T::zero(), T::zero())
+        } else {
+            (acc1[i - 1], acc2[i - 1])
+        };
+
+        let tuple = indices
+            .iter()
+            .map(|&j| witness_columns[j].1[i])
+            .collect::<Vec<_>>();
+
+        let fingerprint = add_ext(
+            fingerprint(&tuple, (T::from(1234), T::from(12345))),
+            (T::from(8764), T::from(876324)),
+        );
+        let multiplicity = T::from(42);
+
+        /*
+        add_ext(
+            current_acc,
+            mul_ext(m_ext_next, inv_ext(folded_next))
+        )
+         */
+        let update = add_ext(
+            current_acc,
+            mul_ext((multiplicity, T::from(0)), inv_ext(fingerprint)),
+        );
+
+        acc1[i] = update.0;
+        acc2[i] = update.1;
+        acc1_next[(i + size - 1) % size] = update.0;
+        acc2_next[(i + size - 1) % size] = update.1;
+    }
+
+    vec![
+        (name("main::acc", connection_index * 2), acc1),
+        (name("main::acc", connection_index * 2 + 1), acc2),
+        (name("main::acc_next", connection_index * 2), acc1_next),
+        (name("main::acc_next", connection_index * 2 + 1), acc2_next),
+    ]
+}
+
+fn name(base: &str, i: usize) -> String {
+    if i == 0 {
+        return base.to_string();
+    }
+    format!("{base}_{}", i)
+}
+
+/*
+let<T: Add + FromLiteral + Mul> mul_ext: Fp2<T>, Fp2<T> -> Fp2<T> = |a, b| match (a, b) {
+    (Fp2::Fp2(a0, a1), Fp2::Fp2(b0, b1)) => Fp2::Fp2(
+        // Multiplication modulo the polynomial x^2 - 11. We'll use the fact
+        // that x^2 == 11 (mod x^2 - 11), so:
+        // (a0 + a1 * x) * (b0 + b1 * x) = a0 * b0 + 11 * a1 * b1 + (a1 * b0 + a0 * b1) * x (mod x^2 - 11)
+        a0 * b0 + 11 * a1 * b1,
+        a1 * b0 + a0 * b1
+    )
+};
+*/
+
+fn mul_ext<T: FieldElement>(a: (T, T), b: (T, T)) -> (T, T) {
+    (a.0 * b.0 + a.1 * b.1 * T::from(11), a.1 * b.0 + a.0 * b.1)
+}
+
+fn add_ext<T: FieldElement>(a: (T, T), b: (T, T)) -> (T, T) {
+    (a.0 + b.0, a.1 + b.1)
+}
+
+fn sub_ext<T: FieldElement>(a: (T, T), b: (T, T)) -> (T, T) {
+    (a.0 - b.0, a.1 - b.1)
+}
+
+/*
+/// Maps [x_1, x_2, ..., x_n] to its Read-Solomon fingerprint, using a challenge alpha: $\sum_{i=1}^n alpha**{(n - i)} * x_i$
+/// To generate an expression that computes the fingerprint, use `fingerprint_inter` instead.
+/// Note that alpha is passed as an expressions, so that it is only evaluated if needed (i.e., if len(expr_array) > 1).
+let fingerprint: fe[], Fp2<expr> -> Fp2<fe> = query |expr_array, alpha| {
+    fingerprint_impl(expr_array, eval_ext(alpha), len(expr_array))
+};
+
+let fingerprint_impl: fe[], Fp2<fe>, int -> Fp2<fe> = query |expr_array, alpha, l| if l == 1 {
+    // Base case
+    from_base(expr_array[0])
+} else {
+
+    // Recursively compute the fingerprint as fingerprint(expr_array[:-1], alpha) * alpha + expr_array[-1]
+    let intermediate_fingerprint = fingerprint_impl(expr_array, alpha, l - 1);
+    add_ext(mul_ext(alpha, intermediate_fingerprint), from_base(expr_array[l - 1]))
+};
+*/
+
+fn fingerprint<T: FieldElement>(expr_array: &[T], alpha: (T, T)) -> (T, T) {
+    fingerprint_impl(expr_array, alpha, expr_array.len())
+}
+
+fn fingerprint_impl<T: FieldElement>(expr_array: &[T], alpha: (T, T), l: usize) -> (T, T) {
+    if l == 1 {
+        return (expr_array[0], T::zero());
+    }
+
+    let intermediate_fingerprint = fingerprint_impl(expr_array, alpha, l - 1);
+    add_ext(
+        mul_ext(alpha, intermediate_fingerprint),
+        (expr_array[l - 1], T::zero()),
+    )
+}
+
+/*
+
+/// Extension field inversion
+let inv_ext: Fp2<fe> -> Fp2<fe> = |a| match a {
+    // The inverse of (a0, a1) is a point (b0, b1) such that:
+    // (a0 + a1 * x) (b0 + b1 * x) = 1 (mod x^2 - 11)
+    // Multiplying out and plugging in x^2 = 11 yields the following system of linear equations:
+    // a0 * b0 + 11 * a1 * b1 = 1
+    // a1 * b0 + a0 * b1 = 0
+    // Solving for (b0, b1) yields:
+    Fp2::Fp2(a0, a1) => {
+        let factor = inv_field(11 * a1 * a1 - a0 * a0);
+        Fp2::Fp2(-a0 * factor, a1 * factor)
+    }
+};
+
+*/
+fn inv_ext<T: FieldElement>(a: (T, T)) -> (T, T) {
+    let factor = T::from(1) / (T::from(11) * a.1 * a.1 - a.0 * a.0);
+    (-a.0 * factor, a.1 * factor)
+}

executor/src/witgen/machines/dynamic_machine.rs

@@ -230,6 +230,8 @@ impl<'a, T: FieldElement> DynamicMachine<'a, T> {
             &self.parts,
             SolverState::new(data, self.publics.clone()),
             mutable_state,
+            self.degree,
+            true,
         );
         if let Some(outer_query) = outer_query {
             processor = processor.with_outer_query(outer_query);

executor/src/witgen/machines/machine_extractor.rs

@@ -14,6 +14,7 @@ use super::sorted_witness_machine::SortedWitnesses;
 use super::FixedData;
 use super::KnownMachine;
 use crate::witgen::machines::dynamic_machine::DynamicMachine;
+use crate::witgen::machines::second_stage_machine::SecondStageMachine;
 use crate::witgen::machines::Connection;
 use crate::witgen::machines::{write_once_memory::WriteOnceMemory, MachineParts};
 use crate::Identity;
@@ -60,26 +61,19 @@ 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,
+                identities,
                 self.fixed.witness_cols.keys().collect::<HashSet<_>>(),
                 prover_functions,
             );
 
-            return build_main_machine(self.fixed, machine_parts)
-                .into_iter()
-                .collect();
+            return vec![KnownMachine::SecondStageMachine(SecondStageMachine::new(
+                "Bus Machine".to_string(),
+                self.fixed,
+                machine_parts,
+            ))];
         }
         let mut machines: Vec<KnownMachine<T>> = vec![];
 
@@ -193,7 +187,7 @@ impl<'a, T: FieldElement> MachineExtractor<'a, T> {
                 }
             }
 
-            let name = suggest_machine_name(&machine_parts);
+            let name = machine_parts.name();
             let id = id_counter;
             id_counter += 1;
             let name_with_type = |t: &str| format!("Secondary machine {id}: {name} ({t})");
@@ -343,18 +337,6 @@ fn log_extracted_machine<T: FieldElement>(parts: &MachineParts<'_, T>) {
     );
 }
 
-fn suggest_machine_name<T: FieldElement>(parts: &MachineParts<'_, T>) -> String {
-    let first_witness = parts.witnesses.iter().next().unwrap();
-    let first_witness_name = parts.column_name(first_witness);
-    let namespace = first_witness_name
-        .rfind("::")
-        .map(|idx| &first_witness_name[..idx]);
-
-    // For machines compiled using Powdr ASM we'll always have a namespace, but as a last
-    // resort we'll use the first witness name.
-    namespace.unwrap_or(first_witness_name).to_string()
-}
-
 #[derive(Default)]
 /// Keeps track of the global set of publics that are referenced by the machine's identities.
 struct PublicsTracker<'a>(BTreeSet<&'a String>);

executor/src/witgen/machines/mod.rs

@@ -17,6 +17,7 @@ use self::double_sorted_witness_machine_16::DoubleSortedWitnesses16;
 use self::double_sorted_witness_machine_32::DoubleSortedWitnesses32;
 pub use self::fixed_lookup_machine::FixedLookup;
 use self::profiling::{record_end, record_start};
+use self::second_stage_machine::SecondStageMachine;
 use self::sorted_witness_machine::SortedWitnesses;
 use self::write_once_memory::WriteOnceMemory;
 
@@ -30,6 +31,7 @@ mod dynamic_machine;
 mod fixed_lookup_machine;
 pub mod machine_extractor;
 pub mod profiling;
+mod second_stage_machine;
 mod sorted_witness_machine;
 mod write_once_memory;
 
@@ -117,6 +119,7 @@ pub enum LookupCell<'a, T> {
 /// This allows us to treat machines uniformly without putting them into a `Box`,
 /// which requires that all lifetime parameters are 'static.
 pub enum KnownMachine<'a, T: FieldElement> {
+    SecondStageMachine(SecondStageMachine<'a, T>),
     SortedWitnesses(SortedWitnesses<'a, T>),
     DoubleSortedWitnesses16(DoubleSortedWitnesses16<'a, T>),
     DoubleSortedWitnesses32(DoubleSortedWitnesses32<'a, T>),
@@ -129,6 +132,7 @@ pub enum KnownMachine<'a, T: FieldElement> {
 impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
     fn run<Q: QueryCallback<T>>(&mut self, mutable_state: &MutableState<'a, T, Q>) {
         match self {
+            KnownMachine::SecondStageMachine(m) => m.run(mutable_state),
             KnownMachine::SortedWitnesses(m) => m.run(mutable_state),
             KnownMachine::DoubleSortedWitnesses16(m) => m.run(mutable_state),
             KnownMachine::DoubleSortedWitnesses32(m) => m.run(mutable_state),
@@ -146,6 +150,9 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
         caller_rows: &'b RowPair<'b, 'a, T>,
     ) -> EvalResult<'a, T> {
         match self {
+            KnownMachine::SecondStageMachine(m) => {
+                m.process_plookup(mutable_state, identity_id, caller_rows)
+            }
             KnownMachine::SortedWitnesses(m) => {
                 m.process_plookup(mutable_state, identity_id, caller_rows)
             }
@@ -172,6 +179,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
 
     fn name(&self) -> &str {
         match self {
+            KnownMachine::SecondStageMachine(m) => m.name(),
             KnownMachine::SortedWitnesses(m) => m.name(),
             KnownMachine::DoubleSortedWitnesses16(m) => m.name(),
             KnownMachine::DoubleSortedWitnesses32(m) => m.name(),
@@ -187,6 +195,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
         mutable_state: &'b MutableState<'a, T, Q>,
     ) -> HashMap<String, Vec<T>> {
         match self {
+            KnownMachine::SecondStageMachine(m) => m.take_witness_col_values(mutable_state),
             KnownMachine::SortedWitnesses(m) => m.take_witness_col_values(mutable_state),
             KnownMachine::DoubleSortedWitnesses16(m) => m.take_witness_col_values(mutable_state),
             KnownMachine::DoubleSortedWitnesses32(m) => m.take_witness_col_values(mutable_state),
@@ -199,6 +208,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
 
     fn identity_ids(&self) -> Vec<u64> {
         match self {
+            KnownMachine::SecondStageMachine(m) => m.identity_ids(),
             KnownMachine::SortedWitnesses(m) => m.identity_ids(),
             KnownMachine::DoubleSortedWitnesses16(m) => m.identity_ids(),
             KnownMachine::DoubleSortedWitnesses32(m) => m.identity_ids(),
@@ -323,13 +333,35 @@ impl<'a, T: FieldElement> MachineParts<'a, T> {
         witnesses: HashSet<PolyID>,
         prover_functions: Vec<&'a analyzed::Expression>,
     ) -> Self {
-        Self {
+        let parts = Self {
             fixed_data,
             connections,
             identities,
             witnesses,
             prover_functions,
-        }
+        };
+
+        log::info!(
+            "Machine '{}' has {} identities, {} witnesses, and {} prover functions.",
+            parts.name(),
+            parts.identities.len(),
+            parts.witnesses.len(),
+            parts.prover_functions.len()
+        );
+
+        parts
+    }
+
+    fn name(&self) -> String {
+        let first_witness = self.witnesses.iter().next().unwrap();
+        let first_witness_name = self.column_name(first_witness);
+        let namespace = first_witness_name
+            .rfind("::")
+            .map(|idx| &first_witness_name[..idx]);
+
+        // For machines compiled using Powdr ASM we'll always have a namespace, but as a last
+        // resort we'll use the first witness name.
+        namespace.unwrap_or(first_witness_name).to_string()
     }
 
     /// Returns a copy of the machine parts but only containing identities that