clean up

Author: Schaeff

openvm/src/powdr_extension/trace_generator/cpu/mod.rs

@@ -1,21 +1,21 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::HashMap,
+    iter::{once, repeat},
+    sync::Arc,
+};
 
 use itertools::Itertools;
 use openvm_circuit::{
     arch::{AirInventory, MatrixRecordArena},
     utils::next_power_of_two_or_zero,
 };
 use openvm_stark_backend::{
-    p3_field::FieldAlgebra,
-    p3_matrix::dense::{DenseMatrix, RowMajorMatrix},
-    prover::{
-        hal::ProverBackend,
-        types::{AirProvingContext, AirProvingContexts},
-    },
-    Chip,
+    Chip, p3_field::FieldAlgebra, p3_matrix::dense::{DenseMatrix, RowMajorMatrix}, prover::{
+        hal::ProverBackend, types::{AirProvingContext, AirProvingContexts, Rejected}
+    }
 };
 use openvm_stark_sdk::p3_baby_bear::BabyBear;
-use powdr_autoprecompiles::{trace_handler::TraceTrait, Apc};
+use powdr_autoprecompiles::{expression::RowEvaluator, trace_handler::TraceTrait, Apc};
 use powdr_constraint_solver::constraint_system::ComputationMethod;
 
 use crate::{
@@ -62,7 +62,7 @@ impl<F> From<Arc<RowMajorMatrix<F>>> for SharedCpuTrace<F> {
 }
 
 impl<R, PB: ProverBackend<Matrix = Arc<RowMajorMatrix<BabyBear>>>> Chip<R, PB> for PowdrChipCpu {
-    fn generate_proving_ctx(&self, records: R) -> AirProvingContext<PB> {
+    fn generate_proving_ctx(&self, _records: R) -> AirProvingContext<PB> {
         unreachable!()
     }
 
@@ -73,14 +73,9 @@ impl<R, PB: ProverBackend<Matrix = Arc<RowMajorMatrix<BabyBear>>>> Chip<R, PB> f
             .trace_generator
             .generate_witness(self.record_arena_by_air_name.take());
 
-        let rejected = rejected
-            .into_iter()
-            .map(|(key, (rows, values))| (key, (rows, AirProvingContext::simple(values, vec![]))))
-            .collect();
-
         AirProvingContexts {
-            main: AirProvingContext::simple(Arc::new(trace), vec![]),
-            rejected,
+            main: AirProvingContext::simple(trace, vec![]),
+            rejected
         }
     }
 }
@@ -111,16 +106,21 @@ impl PowdrTraceGeneratorCpu {
         &self,
         mut original_arenas: OriginalArenas<MatrixRecordArena<BabyBear>>,
     ) -> (
-        DenseMatrix<BabyBear>,
-        HashMap<String, (Vec<usize>, Arc<DenseMatrix<BabyBear>>)>,
+        Arc<DenseMatrix<BabyBear>>,
+        Rejected<Arc<DenseMatrix<BabyBear>>>,
     ) {
+        let mut rejected_pcs = vec![];
+
         use powdr_autoprecompiles::trace_handler::{generate_trace, TraceData};
 
         let num_apc_calls = original_arenas.number_of_calls();
         if num_apc_calls == 0 {
             // If the APC isn't called, early return with an empty trace.
             let width = self.apc.machine().main_columns().count();
-            return (RowMajorMatrix::new(vec![], width), HashMap::default());
+            return (
+                Arc::new(RowMajorMatrix::new(vec![], width)),
+                Default::default()
+            );
         }
 
         let chip_inventory = {
@@ -175,7 +175,7 @@ impl PowdrTraceGeneratorCpu {
         let height = next_power_of_two_or_zero(num_apc_calls);
         let mut values = <BabyBear as FieldAlgebra>::zero_vec(height * width);
 
-        let mut rejected: HashMap<String, Vec<usize>> = dummy_trace_by_air_name
+        let mut rejected_rows_per_air: HashMap<String, Vec<usize>> = dummy_trace_by_air_name
             .keys()
             .cloned()
             .map(|key| (key, vec![]))
@@ -187,7 +187,9 @@ impl PowdrTraceGeneratorCpu {
             // TODO: optimize by parallelizing on chunks of rows, currently fails because `dyn AnyChip<MatrixRecordArena<Val<SC>>>` is not `Send`
             .chunks_mut(width)
             .zip(dummy_values)
-            .for_each(|(row_slice, dummy_values)| {
+            // Just for testing, reject the first call of each apc
+            .zip(once(false).chain(repeat(true)))
+            .for_each(|((row_slice, dummy_values), row_is_valid)| {
                 // map the dummy rows to the autoprecompile row
 
                 use powdr_autoprecompiles::expression::MappingRowEvaluator;
@@ -225,7 +227,6 @@ impl PowdrTraceGeneratorCpu {
                 let evaluator = MappingRowEvaluator::new(row_slice, &apc_poly_id_to_index);
 
                 // check the constraints and bus interactions
-                let row_is_valid = true;
 
                 if row_is_valid {
                     // replay the side effects of this row on the main periphery
@@ -248,12 +249,40 @@ impl PowdrTraceGeneratorCpu {
                             );
                         });
                 } else {
+                    // set the whole row to zero
+                    // TODO: this generates a gap in the table. Instead, reuse the row in the next iteration.
+                    for cell in row_slice {
+                        *cell = BabyBear::ZERO;
+                    }
+
                     // for each original row
                     for original_row_reference in dummy_values {
-                        // TODO replay the side effects of this row on the real periphery
+                        // replay the side effects of this row on the real periphery
+                        let original_row_data = &original_row_reference.data[original_row_reference
+                            .start()
+                            ..original_row_reference.start() + original_row_reference.length];
+                        let evaluator = RowEvaluator::new(original_row_data);
+                        let (machine, _) =
+                            &self.original_airs.air_name_to_machine[original_row_reference.air_id];
+
+                        for interaction in &machine.bus_interactions {
+                            use powdr_autoprecompiles::expression::{
+                                AlgebraicEvaluator, ConcreteBusInteraction,
+                            };
+
+                            let ConcreteBusInteraction { id, mult, args } =
+                                evaluator.eval_bus_interaction(interaction);
+
+                            rejected_pcs.extend(self.periphery.real.apply(
+                                id as u16,
+                                mult.as_canonical_u32(),
+                                args.map(|arg| arg.as_canonical_u32()),
+                                &self.periphery.bus_ids,
+                            ));
+                        }
 
                         // add the row index to the rejected set
-                        rejected
+                        rejected_rows_per_air
                             .get_mut(original_row_reference.air_id)
                             .unwrap()
                             .push(original_row_reference.row_index);
@@ -262,19 +291,22 @@ impl PowdrTraceGeneratorCpu {
             });
 
         // merge the rejected indices with the traces
-        let rejected = rejected
-            .into_iter()
-            .map(|(name, indices)| {
-                (
-                    name.clone(),
+        let rejected = Rejected {
+            pcs: rejected_pcs,
+            rows_per_air: rejected_rows_per_air
+                .into_iter()
+                .map(|(name, indices)| {
                     (
-                        indices,
-                        dummy_trace_by_air_name.remove(&name).unwrap().matrix,
-                    ),
-                )
-            })
-            .collect();
-
-        (RowMajorMatrix::new(values, width), rejected)
+                        name.clone(),
+                        (
+                            dummy_trace_by_air_name.remove(&name).unwrap().matrix,
+                            indices,
+                        ),
+                    )
+                })
+                .collect(),
+            };
+
+        (Arc::new(RowMajorMatrix::new(values, width)), rejected)
     }
 }

openvm/src/powdr_extension/trace_generator/cpu/periphery.rs

@@ -165,13 +165,20 @@ where
 impl SharedPeripheryChipsCpu {
     /// Sends concrete values to the shared chips using a given bus id.
     /// Panics if the bus id doesn't match any of the chips' bus ids.
+    /// Returns the `pc` of the program row that was received, if any
     pub fn apply(
         &self,
         bus_id: u16,
         mult: u32,
         mut args: impl Iterator<Item = u32>,
         periphery_bus_ids: &PeripheryBusIds,
-    ) {
+    ) -> Option<u32> {
+        // early return is the interaction is with the program chip
+        if bus_id == 2 {
+            // The passed data starts with the pc
+            return Some(args.next().unwrap());
+        }
+
         match bus_id {
             id if Some(id) == periphery_bus_ids.bitwise_lookup => {
                 // bitwise operation lookup
@@ -214,13 +221,18 @@ impl SharedPeripheryChipsCpu {
                     self.tuple_range_checker.as_ref().unwrap().add_count(&args);
                 }
             }
-            0..=2 => {
-                // execution bridge, memory, pc lookup
+            0..=1 => {
+                // execution bridge, memory
                 // do nothing
             }
+            2 => {
+                unreachable!("Program bus already handled at the beginnning of this function")
+            }
             _ => {
                 unreachable!("Bus interaction {} not implemented", bus_id);
             }
-        }
+        };
+
+        None
     }
 }