Move block shape check (#2471)

Fixes point 2 of #2327

Author: chriseth

executor/src/witgen/jit/block_machine_processor.rs

@@ -1,4 +1,4 @@
-use std::collections::HashSet;
+use std::collections::{BTreeMap, BTreeSet, HashSet};
 
 use bit_vec::BitVec;
 use itertools::Itertools;
@@ -15,9 +15,10 @@ use crate::witgen::{
 };
 
 use super::{
+    effect::Effect,
     processor::ProcessorResult,
     prover_function_heuristics::ProverFunction,
-    variable::{Cell, Variable},
+    variable::{Cell, MachineCallVariable, Variable},
     witgen_inference::{CanProcessCall, FixedEvaluator, WitgenInference},
 };
 
@@ -117,29 +118,20 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
             .identities
             .iter()
             .any(|id| id.contains_next_ref(&intermediate_definitions));
-        let start_row = if !have_next_ref {
+        let (start_row, end_row) = if !have_next_ref {
             // No identity contains a next reference - we do not need to consider row -1,
             // and the block has to be rectangular-shaped.
-            0
+            (0, self.block_size as i32 - 1)
         } else {
             // A machine that might have a non-rectangular shape.
             // We iterate over all rows of the block +/- one row.
-            -1
+            (-1, self.block_size as i32)
         };
-        let identities = (start_row..self.block_size as i32).flat_map(|row| {
+        let identities = (start_row..=end_row).flat_map(|row| {
             self.machine_parts
                 .identities
                 .iter()
-                .filter_map(|id| {
-                    // Filter out identities with next references on the last row.
-                    if row as usize == self.block_size - 1
-                        && id.contains_next_ref(&intermediate_definitions)
-                    {
-                        None
-                    } else {
-                        Some((*id, row))
-                    }
-                })
+                .map(|id| (*id, row))
                 .collect_vec()
         });
 
@@ -160,7 +152,6 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
             requested_known,
             BLOCK_MACHINE_MAX_BRANCH_DEPTH,
         )
-        .with_block_shape_check()
         .with_block_size(self.block_size)
         .with_requested_range_constraints((0..known_args.len()).map(Variable::Param))
         .generate_code(can_process, witgen)
@@ -183,8 +174,100 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
                 .format("\n  ");
             format!("Code generation failed: {shortened_error}\nRun with RUST_LOG=trace to see the code generated so far.")
         })?;
+        self.check_block_shape(&result.code)?;
         Ok((result, prover_functions))
     }
+
+    /// Verifies that each column and each bus send is stackable in the block.
+    /// This means that if we have a cell write or a bus send in row `i`, we cannot
+    /// have another one in row `i + block_size`.
+    fn check_block_shape(&self, code: &[Effect<T, Variable>]) -> Result<(), String> {
+        for (column_id, row_offsets) in written_rows_per_column(code) {
+            for offset in &row_offsets {
+                if row_offsets.contains(&(*offset + self.block_size as i32)) {
+                    return Err(format!(
+                        "Column {} is not stackable in a {}-row block, conflict in rows {} and {}.",
+                        self.fixed_data.column_name(&PolyID {
+                            id: column_id,
+                            ptype: PolynomialType::Committed
+                        }),
+                        self.block_size,
+                        offset,
+                        offset + self.block_size as i32
+                    ));
+                }
+            }
+        }
+        for (identity_id, row_offsets) in completed_rows_for_bus_send(code) {
+            let row_offsets: BTreeSet<_> = row_offsets.into_iter().collect();
+            for offset in &row_offsets {
+                if row_offsets.contains(&(*offset + self.block_size as i32)) {
+                    return Err(format!(
+                        "Bus send for identity {} is not stackable in a {}-row block, conflict in rows {} and {}.",
+                        identity_id,
+                        self.block_size,
+                        offset,
+                        offset + self.block_size as i32
+                    ));
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+/// Returns, for each column ID, the collection of row offsets that have a cell write.
+/// Combines writes from branches.
+fn written_rows_per_column<T: FieldElement>(
+    code: &[Effect<T, Variable>],
+) -> BTreeMap<u64, BTreeSet<i32>> {
+    code.iter()
+        .flat_map(|e| e.written_vars())
+        .filter_map(|(v, _)| match v {
+            Variable::WitnessCell(cell) => Some((cell.id, cell.row_offset)),
+            _ => None,
+        })
+        .fold(BTreeMap::new(), |mut map, (id, row)| {
+            map.entry(id).or_default().insert(row);
+            map
+        })
+}
+
+/// Returns, for each bus send ID, the collection of row offsets that have a machine call.
+/// Combines calls from branches.
+fn completed_rows_for_bus_send<T: FieldElement>(
+    code: &[Effect<T, Variable>],
+) -> BTreeMap<u64, BTreeSet<i32>> {
+    code.iter()
+        .flat_map(machine_calls)
+        .fold(BTreeMap::new(), |mut map, (id, row)| {
+            map.entry(id).or_default().insert(row);
+            map
+        })
+}
+
+/// Returns all machine calls (bus identity and row offset) found in the effect.
+/// Recurses into branches.
+fn machine_calls<T: FieldElement>(
+    e: &Effect<T, Variable>,
+) -> Box<dyn Iterator<Item = (u64, i32)> + '_> {
+    match e {
+        Effect::MachineCall(id, _, arguments) => match &arguments[0] {
+            Variable::MachineCallParam(MachineCallVariable {
+                identity_id,
+                row_offset,
+                ..
+            }) => {
+                assert_eq!(*id, *identity_id);
+                Box::new(std::iter::once((*identity_id, *row_offset)))
+            }
+            _ => panic!("Expected machine call variable."),
+        },
+        Effect::Branch(_, first, second) => {
+            Box::new(first.iter().chain(second.iter()).flat_map(machine_calls))
+        }
+        _ => Box::new(std::iter::empty()),
+    }
 }
 
 impl<T: FieldElement> FixedEvaluator<T> for &BlockMachineProcessor<'_, T> {

executor/src/witgen/jit/processor.rs

@@ -1,13 +1,8 @@
 #![allow(dead_code)]
-use std::{
-    collections::BTreeSet,
-    fmt::{self, Display, Formatter, Write},
-};
+use std::fmt::{self, Display, Formatter, Write};
 
 use itertools::Itertools;
-use powdr_ast::analyzed::{
-    AlgebraicExpression as Expression, PolyID, PolynomialIdentity, PolynomialType,
-};
+use powdr_ast::analyzed::PolynomialIdentity;
 use powdr_number::FieldElement;
 
 use crate::witgen::{
@@ -21,8 +16,8 @@ use super::{
     affine_symbolic_expression,
     effect::{format_code, Effect},
     identity_queue::{IdentityQueue, QueueItem},
-    variable::{Cell, MachineCallVariable, Variable},
-    witgen_inference::{BranchResult, CanProcessCall, FixedEvaluator, Value, WitgenInference},
+    variable::{MachineCallVariable, Variable},
+    witgen_inference::{BranchResult, CanProcessCall, FixedEvaluator, WitgenInference},
 };
 
 /// A generic processor for generating JIT code.
@@ -36,8 +31,6 @@ pub struct Processor<'a, T: FieldElement, FixedEval> {
     initial_queue: Vec<QueueItem<'a, T>>,
     /// The size of a block.
     block_size: usize,
-    /// If the processor should check for correctly stackable block shapes.
-    check_block_shape: bool,
     /// List of variables we want to be known at the end. One of them not being known
     /// is a failure.
     requested_known_vars: Vec<Variable>,
@@ -71,7 +64,6 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
             identities,
             initial_queue,
             block_size: 1,
-            check_block_shape: false,
             requested_known_vars: requested_known_vars.into_iter().collect(),
             requested_range_constraints: vec![],
             max_branch_depth,
@@ -93,13 +85,6 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
         self
     }
 
-    /// Activates the check to see if the code for two subsequently generated
-    /// blocks conflicts.
-    pub fn with_block_shape_check(mut self) -> Self {
-        self.check_block_shape = true;
-        self
-    }
-
     pub fn generate_code(
         self,
         can_process: impl CanProcessCall<T>,
@@ -143,16 +128,6 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
             ));
         }
 
-        if self.check_block_shape {
-            // Check that the "spill" into the previous block is compatible
-            // with the "missing pieces" in the next block.
-            // If this is not the case, this is a hard error
-            // (i.e. cannot be fixed by runtime witgen) and thus we panic inside.
-            // We could do this only at the end of each branch, but it's a bit
-            // more convenient to do it here.
-            self.check_block_shape(&witgen);
-        }
-
         // Check that we could derive all requested variables.
         let missing_variables = self
             .requested_known_vars
@@ -347,27 +322,23 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
             .unique()
             .flat_map(|&call| {
                 let rows = self.rows_for_identity(call);
-                let complete_rows = rows
+                if rows
                     .iter()
                     .filter(|&&row| witgen.is_complete_call(call, row))
-                    .collect::<Vec<_>>();
-                // We might process more rows than `self.block_size`, so we check
-                // that the complete calls are on consecutive rows.
-                if complete_rows.len() >= self.block_size {
-                    let (min, max) = complete_rows.iter().minmax().into_option().unwrap();
-                    // TODO instead of checking for consecutive rows, we could also check
-                    // that they "fit" the next block.
-                    // TODO actually I think that we should not allow more than block size
-                    // completed calls.
-                    let is_consecutive = *max - *min == complete_rows.len() as i32 - 1;
-                    if is_consecutive {
-                        return vec![];
-                    }
+                    .count()
+                    >= self.block_size
+                {
+                    // We might process more rows than `self.block_size`, so we check
+                    // that we have the reqired amount of calls.
+                    // The block shape check done by block_machine_processor will do a more
+                    // thorough check later on.
+                    vec![]
+                } else {
+                    rows.iter()
+                        .filter(|&row| !witgen.is_complete_call(call, *row))
+                        .map(|row| (call, *row))
+                        .collect_vec()
                 }
-                rows.iter()
-                    .filter(|&row| !witgen.is_complete_call(call, *row))
-                    .map(|row| (call, *row))
-                    .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>()
     }
@@ -386,77 +357,6 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
             .collect()
     }
 
-    /// After solving, the known cells should be such that we can stack different blocks.
-    /// If this is not the case, this function panics.
-    /// TODO the same is actually true for machine calls.
-    fn check_block_shape(&self, witgen: &WitgenInference<'a, T, FixedEval>) {
-        let known_columns: BTreeSet<_> = witgen
-            .known_variables()
-            .iter()
-            .filter_map(|var| match var {
-                Variable::WitnessCell(cell) => Some(cell.id),
-                _ => None,
-            })
-            .collect();
-        for column_id in known_columns {
-            let known_rows = witgen
-                .known_variables()
-                .iter()
-                .filter_map(|var| match var {
-                    Variable::WitnessCell(cell) if cell.id == column_id => Some(cell.row_offset),
-                    _ => None,
-                })
-                .collect::<BTreeSet<_>>();
-
-            // Two values that refer to the same row (modulo block size) are compatible if:
-            // - One of them is unknown, or
-            // - Both are concrete and equal
-            let is_compatible = |v1: Value<T>, v2: Value<T>| match (v1, v2) {
-                (Value::Unknown, _) | (_, Value::Unknown) => true,
-                (Value::Concrete(a), Value::Concrete(b)) => a == b,
-                _ => false,
-            };
-            let cell_var = |row_offset| {
-                Variable::WitnessCell(Cell {
-                    // Column name does not matter.
-                    column_name: "".to_string(),
-                    id: column_id,
-                    row_offset,
-                })
-            };
-
-            // A column is stackable if all rows equal to each other modulo
-            // the block size are compatible.
-            for row in &known_rows {
-                let this_val = witgen.value(&cell_var(*row));
-                let next_block_val = witgen.value(&cell_var(row + self.block_size as i32));
-                if !is_compatible(this_val, next_block_val) {
-                    let column_name = self.fixed_data.column_name(&PolyID {
-                        id: column_id,
-                        ptype: PolynomialType::Committed,
-                    });
-                    let row_vals = known_rows
-                        .iter()
-                        .map(|&r| format!("  row {r}: {}\n", witgen.value(&cell_var(r))))
-                        .format("");
-                    log::debug!(
-                        "Code generated so far:\n{}\n\
-                        Column {column_name} is not stackable in a {}-row block, \
-                        conflict in rows {row} and {}.\n{row_vals}",
-                        format_code(witgen.code()),
-                        self.block_size,
-                        row + self.block_size as i32
-                    );
-                    panic!(
-                        "Column {column_name} is not stackable in a {}-row block, conflict in rows {row} and {}.\n{row_vals}",
-                        self.block_size,
-                        row + self.block_size as i32
-                    );
-                }
-            }
-        }
-    }
-
     /// If the only missing sends all only have a single argument, try to set those arguments
     /// to zero.
     fn try_fix_simple_sends(
@@ -494,7 +394,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> Processor<'a, T, FixedEv
             };
             assert!(!witgen.is_known(param));
             match modified_witgen.process_equation_on_row(
-                &Expression::Number(T::from(0)),
+                &T::from(0).into(),
                 Some(param.clone()),
                 0.into(),
                 row,