Clean up bus ids (#3420)

Author: Schaeff

openvm/src/lib.rs

@@ -228,14 +228,14 @@ impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, PowdrExtens
         // TODO: here we make assumptions about the existence of some chips in the periphery. Make this more flexible
 
         use crate::powdr_extension::trace_generator::cuda::PowdrPeripheryInstancesGpu;
-        let range_checker = inventory
-            .find_chip::<Arc<VariableRangeCheckerChipGPU>>()
-            .next()
-            .unwrap();
         let bitwise_lookup = inventory
             .find_chip::<Arc<BitwiseOperationLookupChipGPU<8>>>()
             .next()
             .cloned();
+        let range_checker = inventory
+            .find_chip::<Arc<VariableRangeCheckerChipGPU>>()
+            .next()
+            .unwrap();
         let tuple_range_checker = inventory
             .find_chip::<Arc<RangeTupleCheckerChipGPU<2>>>()
             .next()
@@ -264,6 +264,11 @@ impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, PowdrExtens
             range_checker.clone(),
             bitwise_lookup,
             tuple_range_checker,
+            PeripheryBusIds {
+                range_checker: range_checker_bus_id,
+                bitwise_lookup: bitwise_lookup_bus_id,
+                tuple_range_checker: tuple_range_checker_bus_id,
+            },
         );
 
         for precompile in &extension.precompiles {
@@ -275,11 +280,6 @@ impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, PowdrExtens
                 extension.airs.clone(),
                 extension.base_config.clone(),
                 shared_chips_pair.clone(),
-                PeripheryBusIds {
-                    range_checker: range_checker_bus_id,
-                    bitwise_lookup: bitwise_lookup_bus_id,
-                    tuple_range_checker: tuple_range_checker_bus_id,
-                },
             );
             inventory.add_executor_chip(chip);
         }
@@ -288,6 +288,7 @@ impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, PowdrExtens
     }
 }
 
+#[derive(Clone)]
 pub struct PeripheryBusIds {
     pub range_checker: u16,
     pub bitwise_lookup: Option<u16>,

openvm/src/powdr_extension/chip.rs

@@ -146,7 +146,6 @@ mod cuda {
             trace_generator::cuda::{PowdrPeripheryInstancesGpu, PowdrTraceGeneratorGpu},
             PowdrPrecompile,
         },
-        PeripheryBusIds,
     };
 
     pub struct PowdrChipGpu {
@@ -161,21 +160,15 @@ mod cuda {
             original_airs: OriginalAirs<BabyBear>,
             base_config: OriginalVmConfig,
             periphery: PowdrPeripheryInstancesGpu,
-            periphery_bus_ids: PeripheryBusIds,
         ) -> Self {
             let PowdrPrecompile {
                 name,
                 apc,
                 apc_record_arena_gpu: apc_record_arena,
                 ..
             } = precompile;
-            let trace_generator = PowdrTraceGeneratorGpu::new(
-                apc,
-                original_airs,
-                base_config,
-                periphery,
-                periphery_bus_ids,
-            );
+            let trace_generator =
+                PowdrTraceGeneratorGpu::new(apc, original_airs, base_config, periphery);
 
             Self {
                 name,

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

@@ -31,7 +31,7 @@ use crate::{
         executor::OriginalArenas,
         trace_generator::{common::create_dummy_airs, cuda::inventory::create_dummy_chip_complex},
     },
-    BabyBearSC, Instr, PeripheryBusIds,
+    BabyBearSC, Instr,
 };
 
 mod inventory;
@@ -177,7 +177,6 @@ pub struct PowdrTraceGeneratorGpu {
     pub original_airs: OriginalAirs<BabyBear>,
     pub config: OriginalVmConfig,
     pub periphery: PowdrPeripheryInstancesGpu,
-    pub periphery_bus_ids: PeripheryBusIds,
 }
 
 impl PowdrTraceGeneratorGpu {
@@ -186,14 +185,12 @@ impl PowdrTraceGeneratorGpu {
         original_airs: OriginalAirs<BabyBear>,
         config: OriginalVmConfig,
         periphery: PowdrPeripheryInstancesGpu,
-        periphery_bus_ids: PeripheryBusIds,
     ) -> Self {
         Self {
             apc,
             original_airs,
             config,
             periphery,
-            periphery_bus_ids,
         }
     }
 
@@ -346,17 +343,17 @@ impl PowdrTraceGeneratorGpu {
         let periphery = &self.periphery.real;
 
         // Range checker
-        let var_range_bus_id = self.periphery_bus_ids.range_checker as u32;
+        let var_range_bus_id = self.periphery.bus_ids.range_checker as u32;
         let var_range_count = &periphery.range_checker.count;
 
         // Tuple checker
         let tuple_range_checker_chip = periphery.tuple_range_checker.as_ref().unwrap();
-        let tuple2_bus_id = self.periphery_bus_ids.tuple_range_checker.unwrap() as u32;
+        let tuple2_bus_id = self.periphery.bus_ids.tuple_range_checker.unwrap() as u32;
         let tuple2_sizes = tuple_range_checker_chip.sizes;
         let tuple2_count_u32 = tuple_range_checker_chip.count.as_ref();
 
         // Bitwise lookup; NUM_BITS is fixed at 8 in CUDA
-        let bitwise_bus_id = self.periphery_bus_ids.bitwise_lookup.unwrap() as u32;
+        let bitwise_bus_id = self.periphery.bus_ids.bitwise_lookup.unwrap() as u32;
         let bitwise_count_u32 = periphery.bitwise_lookup_8.as_ref().unwrap().count.as_ref();
 
         // Launch GPU apply-bus to update periphery histograms on device

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

@@ -1,4 +1,4 @@
-use crate::powdr_extension::trace_generator::common::DummyExecutor;
+use crate::{powdr_extension::trace_generator::common::DummyExecutor, PeripheryBusIds};
 use openvm_circuit::arch::{
     AirInventory, AirInventoryError, ChipInventory, ChipInventoryError, ExecutorInventoryBuilder,
     ExecutorInventoryError, VmCircuitExtension, VmExecutionExtension, VmProverExtension,
@@ -27,6 +27,8 @@ pub struct PowdrPeripheryInstancesGpu {
     pub real: SharedPeripheryChipsGpu,
     /// The dummy chips used for all APCs. They share the range checker but create new instances of the bitwise lookup chip and the tuple range checker.
     pub dummy: SharedPeripheryChipsGpu,
+    /// The bus ids of the periphery
+    pub bus_ids: PeripheryBusIds,
 }
 
 #[derive(Clone)]
@@ -41,6 +43,7 @@ impl PowdrPeripheryInstancesGpu {
         range_checker: Arc<VariableRangeCheckerChipGPU>,
         bitwise_8: Option<Arc<BitwiseOperationLookupChipGPU<8>>>,
         tuple_range_checker: Option<Arc<RangeTupleCheckerChipGPU<2>>>,
+        bus_ids: PeripheryBusIds,
     ) -> Self {
         Self {
             real: SharedPeripheryChipsGpu {
@@ -73,6 +76,7 @@ impl PowdrPeripheryInstancesGpu {
                     })
                 }),
             },
+            bus_ids,
         }
     }
 }