openvm_circuit/system/cuda/
extensions.rs

1use std::sync::Arc;
2
3use openvm_circuit::{
4    arch::{
5        AirInventory, ChipInventory, ChipInventoryError, DenseRecordArena, SystemConfig, VmBuilder,
6        VmChipComplex,
7    },
8    system::poseidon2::air::Poseidon2PeripheryAir,
9};
10use openvm_circuit_primitives::{
11    bitwise_op_lookup::{
12        BitwiseOperationLookupAir, BitwiseOperationLookupChip, BitwiseOperationLookupChipGPU,
13    },
14    var_range::{VariableRangeCheckerAir, VariableRangeCheckerChip, VariableRangeCheckerChipGPU},
15};
16use openvm_cuda_backend::{BabyBearPoseidon2GpuEngine, GpuBackend};
17use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
18use p3_baby_bear::BabyBear;
19
20use super::{phantom::PhantomChipGPU, Poseidon2PeripheryChipGPU, SystemChipInventoryGPU};
21
22/// A utility method to get the `VariableRangeCheckerChipGPU` from [ChipInventory].
23/// Note, `VariableRangeCheckerChipGPU` always will always exist in the inventory.
24pub fn get_inventory_range_checker(
25    inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
26) -> Arc<VariableRangeCheckerChipGPU> {
27    inventory
28        .find_chip::<Arc<VariableRangeCheckerChipGPU>>()
29        .next()
30        .unwrap()
31        .clone()
32}
33
34/// A utility method to find a **byte** [BitwiseOperationLookupChipGPU] or create one and add
35/// to the inventory if it does not exist.
36pub fn get_or_create_bitwise_op_lookup(
37    inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
38) -> Result<Arc<BitwiseOperationLookupChipGPU<8>>, ChipInventoryError> {
39    let device_ctx = get_inventory_range_checker(inventory).device_ctx.clone();
40    let bitwise_lu = {
41        let existing_chip = inventory
42            .find_chip::<Arc<BitwiseOperationLookupChipGPU<8>>>()
43            .next();
44        if let Some(chip) = existing_chip {
45            chip.clone()
46        } else {
47            let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
48
49            let chip = Arc::new(BitwiseOperationLookupChipGPU::hybrid(
50                Arc::new(BitwiseOperationLookupChip::new(air.bus)),
51                device_ctx,
52            ));
53            inventory.add_periphery_chip(chip.clone());
54            chip
55        }
56    };
57    Ok(bitwise_lu)
58}
59
60/// **If** internal poseidon2 chip exists, then its insertion index is 1.
61const POSEIDON2_INSERTION_IDX: usize = 1;
62/// **If** public values chip exists, then its executor index is 0.
63pub const PV_EXECUTOR_IDX: usize = 0;
64
65#[derive(Clone)]
66pub struct SystemGpuBuilder;
67
68impl VmBuilder<BabyBearPoseidon2GpuEngine> for SystemGpuBuilder {
69    type VmConfig = SystemConfig;
70    type RecordArena = DenseRecordArena;
71    type SystemChipInventory = SystemChipInventoryGPU;
72
73    fn create_chip_complex(
74        &self,
75        config: &SystemConfig,
76        airs: AirInventory<BabyBearPoseidon2Config>,
77        device_ctx: &openvm_stark_backend::EngineDeviceCtx<BabyBearPoseidon2GpuEngine>,
78    ) -> Result<
79        VmChipComplex<
80            BabyBearPoseidon2Config,
81            DenseRecordArena,
82            GpuBackend,
83            SystemChipInventoryGPU,
84        >,
85        ChipInventoryError,
86    > {
87        let device_ctx = device_ctx.clone();
88        let range_bus = airs.range_checker().bus;
89        let range_checker = Arc::new(VariableRangeCheckerChipGPU::hybrid(
90            Arc::new(VariableRangeCheckerChip::new(range_bus)),
91            device_ctx.clone(),
92        ));
93
94        let mut inventory = ChipInventory::new(airs);
95        inventory.next_air::<VariableRangeCheckerAir>()?;
96        inventory.add_periphery_chip(range_checker.clone());
97
98        assert_eq!(inventory.chips().len(), POSEIDON2_INSERTION_IDX);
99        let sbox_registers = if config.max_constraint_degree >= 7 {
100            0
101        } else {
102            1
103        };
104        // ATTENTION: The threshold 7 here must match the one in `new_poseidon2_periphery_air`
105        let _direct_bus = if sbox_registers == 0 {
106            inventory
107                .next_air::<Poseidon2PeripheryAir<BabyBear, 0>>()?
108                .bus
109        } else {
110            inventory
111                .next_air::<Poseidon2PeripheryAir<BabyBear, 1>>()?
112                .bus
113        };
114        let hasher_chip = Arc::new(Poseidon2PeripheryChipGPU::new(
115            sbox_registers,
116            device_ctx.clone(),
117        ));
118        inventory.add_periphery_chip(hasher_chip.clone());
119        let system =
120            SystemChipInventoryGPU::new(config, range_checker, hasher_chip, device_ctx.clone());
121
122        let phantom_chip = PhantomChipGPU::new(device_ctx.clone());
123        inventory.add_executor_chip(phantom_chip);
124
125        Ok(VmChipComplex { system, inventory })
126    }
127}
openvm_circuit/system/cuda/extensions.rs

openvm_circuit/system/cuda/
extensions.rs
