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, PUBLIC_VALUES_AIR_ID,
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::{engine::GpuBabyBearPoseidon2Engine, prover_backend::GpuBackend};
17use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
18use p3_baby_bear::BabyBear;
19
20use super::{
21    phantom::PhantomChipGPU, Poseidon2PeripheryChipGPU, SystemChipInventoryGPU, DIGEST_WIDTH,
22};
23
24/// A utility method to get the `VariableRangeCheckerChipGPU` from [ChipInventory].
25/// Note, `VariableRangeCheckerChipGPU` always will always exist in the inventory.
26pub fn get_inventory_range_checker(
27    inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
28) -> Arc<VariableRangeCheckerChipGPU> {
29    inventory
30        .find_chip::<Arc<VariableRangeCheckerChipGPU>>()
31        .next()
32        .unwrap()
33        .clone()
34}
35
36/// A utility method to find a **byte** [BitwiseOperationLookupChipGPU] or create one and add
37/// to the inventory if it does not exist.
38pub fn get_or_create_bitwise_op_lookup(
39    inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
40) -> Result<Arc<BitwiseOperationLookupChipGPU<8>>, ChipInventoryError> {
41    let bitwise_lu = {
42        let existing_chip = inventory
43            .find_chip::<Arc<BitwiseOperationLookupChipGPU<8>>>()
44            .next();
45        if let Some(chip) = existing_chip {
46            chip.clone()
47        } else {
48            let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
49
50            let chip = Arc::new(BitwiseOperationLookupChipGPU::hybrid(Arc::new(
51                BitwiseOperationLookupChip::new(air.bus),
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<GpuBabyBearPoseidon2Engine> 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    ) -> Result<
78        VmChipComplex<
79            BabyBearPoseidon2Config,
80            DenseRecordArena,
81            GpuBackend,
82            SystemChipInventoryGPU,
83        >,
84        ChipInventoryError,
85    > {
86        let range_bus = airs.range_checker().bus;
87        let range_checker = Arc::new(VariableRangeCheckerChipGPU::hybrid(Arc::new(
88            VariableRangeCheckerChip::new(range_bus),
89        )));
90
91        let mut inventory = ChipInventory::new(airs);
92        // PublicValuesChip is required when num_public_values > 0 in single segment mode.
93        if config.has_public_values_chip() {
94            assert_eq!(
95                inventory.executor_idx_to_insertion_idx.len(),
96                PV_EXECUTOR_IDX
97            );
98
99            // We set insertion_idx so that air_idx = num_airs - (insertion_idx + 1) =
100            // PUBLIC_VALUES_AIR_ID in `VmChipComplex::executor_idx_to_air_idx`. We need to do this
101            // because this chip is special and not part of the normal inventory.
102            let insertion_idx = inventory
103                .airs()
104                .num_airs()
105                .checked_sub(1 + PUBLIC_VALUES_AIR_ID)
106                .unwrap();
107            inventory.executor_idx_to_insertion_idx.push(insertion_idx);
108        }
109        inventory.next_air::<VariableRangeCheckerAir>()?;
110        inventory.add_periphery_chip(range_checker.clone());
111
112        let hasher_chip = if config.continuation_enabled {
113            let max_buffer_size = (config.segmentation_limits.max_trace_height as usize)
114                .next_power_of_two() * 2 // seems like a reliable estimate
115                * (DIGEST_WIDTH * 2); // size of one record
116            assert_eq!(inventory.chips().len(), POSEIDON2_INSERTION_IDX);
117            let sbox_registers = if config.max_constraint_degree >= 7 {
118                0
119            } else {
120                1
121            };
122            // ATTENTION: The threshold 7 here must match the one in `new_poseidon2_periphery_air`
123            let _direct_bus = if sbox_registers == 0 {
124                inventory
125                    .next_air::<Poseidon2PeripheryAir<BabyBear, 0>>()?
126                    .bus
127            } else {
128                inventory
129                    .next_air::<Poseidon2PeripheryAir<BabyBear, 1>>()?
130                    .bus
131            };
132            let chip = Arc::new(Poseidon2PeripheryChipGPU::new(
133                max_buffer_size,
134                sbox_registers,
135            ));
136            inventory.add_periphery_chip(chip.clone());
137            Some(chip)
138        } else {
139            None
140        };
141        let system = SystemChipInventoryGPU::new(
142            config,
143            &inventory.airs().system().memory,
144            range_checker,
145            hasher_chip,
146        );
147
148        let phantom_chip = PhantomChipGPU::new();
149        inventory.add_executor_chip(phantom_chip);
150
151        Ok(VmChipComplex { system, inventory })
152    }
153}
openvm_circuit/system/cuda/extensions.rs

openvm_circuit/system/cuda/
extensions.rs
