openvm_rv32im_circuit/extension/
cuda.rs

1use std::sync::Arc;
2
3use openvm_circuit::{
4    arch::{ChipInventory, ChipInventoryError, DenseRecordArena, VmProverExtension},
5    system::cuda::extensions::{get_inventory_range_checker, get_or_create_bitwise_op_lookup},
6};
7use openvm_circuit_primitives::range_tuple::{RangeTupleCheckerAir, RangeTupleCheckerChipGPU};
8use openvm_cuda_backend::{BabyBearPoseidon2GpuEngine as GpuBabyBearPoseidon2Engine, GpuBackend};
9use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
10
11use crate::{
12    Rv32AuipcAir, Rv32AuipcChipGpu, Rv32BaseAluAir, Rv32BaseAluChipGpu, Rv32BranchEqualAir,
13    Rv32BranchEqualChipGpu, Rv32BranchLessThanAir, Rv32BranchLessThanChipGpu, Rv32DivRemAir,
14    Rv32DivRemChipGpu, Rv32HintStoreAir, Rv32HintStoreChipGpu, Rv32I, Rv32Io, Rv32JalLuiAir,
15    Rv32JalLuiChipGpu, Rv32JalrAir, Rv32JalrChipGpu, Rv32LessThanAir, Rv32LessThanChipGpu,
16    Rv32LoadSignExtendAir, Rv32LoadSignExtendChipGpu, Rv32LoadStoreAir, Rv32LoadStoreChipGpu,
17    Rv32M, Rv32MulHAir, Rv32MulHChipGpu, Rv32MultiplicationAir, Rv32MultiplicationChipGpu,
18    Rv32ShiftAir, Rv32ShiftChipGpu,
19};
20
21pub struct Rv32ImGpuProverExt;
22
23// This implementation is specific to GpuBackend because the lookup chips
24// (VariableRangeCheckerChipGPU, BitwiseOperationLookupChipGPU) are specific to GpuBackend.
25impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, Rv32I> for Rv32ImGpuProverExt {
26    fn extend_prover(
27        &self,
28        _: &Rv32I,
29        inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
30    ) -> Result<(), ChipInventoryError> {
31        let pointer_max_bits = inventory.airs().pointer_max_bits();
32        let timestamp_max_bits = inventory.timestamp_max_bits();
33
34        let range_checker = get_inventory_range_checker(inventory);
35        let bitwise_lu = get_or_create_bitwise_op_lookup(inventory)?;
36
37        // These calls to next_air are not strictly necessary to construct the chips, but provide a
38        // safeguard to ensure that chip construction matches the circuit definition
39        inventory.next_air::<Rv32BaseAluAir>()?;
40        let base_alu = Rv32BaseAluChipGpu::new(
41            range_checker.clone(),
42            bitwise_lu.clone(),
43            timestamp_max_bits,
44        );
45        inventory.add_executor_chip(base_alu);
46
47        inventory.next_air::<Rv32LessThanAir>()?;
48        let lt = Rv32LessThanChipGpu::new(
49            range_checker.clone(),
50            bitwise_lu.clone(),
51            timestamp_max_bits,
52        );
53        inventory.add_executor_chip(lt);
54
55        inventory.next_air::<Rv32ShiftAir>()?;
56        let shift = Rv32ShiftChipGpu::new(
57            range_checker.clone(),
58            bitwise_lu.clone(),
59            timestamp_max_bits,
60        );
61        inventory.add_executor_chip(shift);
62
63        inventory.next_air::<Rv32LoadStoreAir>()?;
64        let load_store_chip =
65            Rv32LoadStoreChipGpu::new(range_checker.clone(), pointer_max_bits, timestamp_max_bits);
66        inventory.add_executor_chip(load_store_chip);
67
68        inventory.next_air::<Rv32LoadSignExtendAir>()?;
69        let load_sign_extend = Rv32LoadSignExtendChipGpu::new(
70            range_checker.clone(),
71            pointer_max_bits,
72            timestamp_max_bits,
73        );
74        inventory.add_executor_chip(load_sign_extend);
75
76        inventory.next_air::<Rv32BranchEqualAir>()?;
77        let beq = Rv32BranchEqualChipGpu::new(range_checker.clone(), timestamp_max_bits);
78        inventory.add_executor_chip(beq);
79
80        inventory.next_air::<Rv32BranchLessThanAir>()?;
81        let blt = Rv32BranchLessThanChipGpu::new(
82            range_checker.clone(),
83            bitwise_lu.clone(),
84            timestamp_max_bits,
85        );
86        inventory.add_executor_chip(blt);
87
88        inventory.next_air::<Rv32JalLuiAir>()?;
89        let jal_lui = Rv32JalLuiChipGpu::new(
90            range_checker.clone(),
91            bitwise_lu.clone(),
92            timestamp_max_bits,
93        );
94        inventory.add_executor_chip(jal_lui);
95
96        inventory.next_air::<Rv32JalrAir>()?;
97        let jalr = Rv32JalrChipGpu::new(
98            range_checker.clone(),
99            bitwise_lu.clone(),
100            timestamp_max_bits,
101        );
102        inventory.add_executor_chip(jalr);
103
104        inventory.next_air::<Rv32AuipcAir>()?;
105        let auipc = Rv32AuipcChipGpu::new(
106            range_checker.clone(),
107            bitwise_lu.clone(),
108            timestamp_max_bits,
109        );
110        inventory.add_executor_chip(auipc);
111
112        Ok(())
113    }
114}
115
116// This implementation is specific to GpuBackend because the lookup chips
117// (VariableRangeCheckerChipGPU, BitwiseOperationLookupChipGPU) are specific to GpuBackend.
118impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, Rv32M> for Rv32ImGpuProverExt {
119    fn extend_prover(
120        &self,
121        extension: &Rv32M,
122        inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
123    ) -> Result<(), ChipInventoryError> {
124        let pointer_max_bits = inventory.airs().pointer_max_bits();
125        let timestamp_max_bits = inventory.timestamp_max_bits();
126
127        let range_checker = get_inventory_range_checker(inventory);
128        let bitwise_lu = get_or_create_bitwise_op_lookup(inventory)?;
129
130        let range_tuple_checker = {
131            let existing_chip = inventory
132                .find_chip::<Arc<RangeTupleCheckerChipGPU<2>>>()
133                .find(|c| {
134                    c.sizes[0] >= extension.range_tuple_checker_sizes[0]
135                        && c.sizes[1] >= extension.range_tuple_checker_sizes[1]
136                });
137            if let Some(chip) = existing_chip {
138                chip.clone()
139            } else {
140                inventory.next_air::<RangeTupleCheckerAir<2>>()?;
141                let chip = Arc::new(RangeTupleCheckerChipGPU::new(
142                    extension.range_tuple_checker_sizes,
143                    range_checker.device_ctx.clone(),
144                ));
145                inventory.add_periphery_chip(chip.clone());
146                chip
147            }
148        };
149
150        // These calls to next_air are not strictly necessary to construct the chips, but provide a
151        // safeguard to ensure that chip construction matches the circuit definition
152        inventory.next_air::<Rv32MultiplicationAir>()?;
153        let mult = Rv32MultiplicationChipGpu::new(
154            range_checker.clone(),
155            range_tuple_checker.clone(),
156            timestamp_max_bits,
157        );
158        inventory.add_executor_chip(mult);
159
160        inventory.next_air::<Rv32MulHAir>()?;
161        let mul_h = Rv32MulHChipGpu::new(
162            range_checker.clone(),
163            bitwise_lu.clone(),
164            range_tuple_checker.clone(),
165            timestamp_max_bits,
166        );
167        inventory.add_executor_chip(mul_h);
168
169        inventory.next_air::<Rv32DivRemAir>()?;
170        let div_rem = Rv32DivRemChipGpu::new(
171            range_checker.clone(),
172            bitwise_lu.clone(),
173            range_tuple_checker.clone(),
174            pointer_max_bits,
175            timestamp_max_bits,
176        );
177        inventory.add_executor_chip(div_rem);
178
179        Ok(())
180    }
181}
182
183// This implementation is specific to GpuBackend because the lookup chips
184// (VariableRangeCheckerChipGPU, BitwiseOperationLookupChipGPU) are specific to GpuBackend.
185impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, Rv32Io>
186    for Rv32ImGpuProverExt
187{
188    fn extend_prover(
189        &self,
190        _: &Rv32Io,
191        inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
192    ) -> Result<(), ChipInventoryError> {
193        let pointer_max_bits = inventory.airs().pointer_max_bits();
194        let timestamp_max_bits = inventory.timestamp_max_bits();
195
196        let range_checker = get_inventory_range_checker(inventory);
197        let bitwise_lu = get_or_create_bitwise_op_lookup(inventory)?;
198
199        inventory.next_air::<Rv32HintStoreAir>()?;
200        let hint_store = Rv32HintStoreChipGpu::new(
201            range_checker.clone(),
202            bitwise_lu.clone(),
203            pointer_max_bits,
204            timestamp_max_bits,
205        );
206        inventory.add_executor_chip(hint_store);
207
208        Ok(())
209    }
210}
openvm_rv32im_circuit/extension/cuda.rs

openvm_rv32im_circuit/extension/
cuda.rs
