openvm_bigint_circuit/extension/
mod.rs

1use std::sync::Arc;
2
3use derive_more::derive::From;
4use openvm_bigint_transpiler::{
5    Rv32BaseAlu256Opcode, Rv32BranchEqual256Opcode, Rv32BranchLessThan256Opcode,
6    Rv32LessThan256Opcode, Rv32Mul256Opcode, Rv32Shift256Opcode,
7};
8use openvm_circuit::{
9    arch::{
10        AirInventory, AirInventoryError, ChipInventory, ChipInventoryError, ExecutionBridge,
11        ExecutorInventoryBuilder, ExecutorInventoryError, MatrixRecordArena, RowMajorMatrixArena,
12        VmBuilder, VmChipComplex, VmCircuitExtension, VmExecutionExtension, VmProverExtension,
13    },
14    system::{memory::SharedMemoryHelper, SystemChipInventory, SystemCpuBuilder, SystemPort},
15};
16use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor};
17use openvm_circuit_primitives::{
18    bitwise_op_lookup::{
19        BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip,
20        SharedBitwiseOperationLookupChip,
21    },
22    range_tuple::{
23        RangeTupleCheckerAir, RangeTupleCheckerBus, RangeTupleCheckerChip,
24        SharedRangeTupleCheckerChip,
25    },
26};
27use openvm_instructions::{program::DEFAULT_PC_STEP, LocalOpcode};
28use openvm_rv32im_circuit::Rv32ImCpuProverExt;
29use openvm_stark_backend::{
30    config::{StarkGenericConfig, Val},
31    engine::StarkEngine,
32    p3_field::PrimeField32,
33    prover::cpu::{CpuBackend, CpuDevice},
34};
35use serde::{Deserialize, Serialize};
36
37use crate::*;
38
39cfg_if::cfg_if! {
40    if #[cfg(feature = "cuda")] {
41        mod cuda;
42        pub use self::cuda::*;
43        pub use self::cuda::{
44            Int256GpuProverExt as Int256ProverExt,
45            Int256Rv32GpuBuilder as Int256Rv32Builder,
46        };
47    } else {
48        pub use self::{
49            Int256CpuProverExt as Int256ProverExt,
50            Int256Rv32CpuBuilder as Int256Rv32Builder,
51        };
52    }
53}
54
55// =================================== VM Extension Implementation =================================
56#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
57pub struct Int256 {
58    #[serde(default = "default_range_tuple_checker_sizes")]
59    pub range_tuple_checker_sizes: [u32; 2],
60}
61
62impl Default for Int256 {
63    fn default() -> Self {
64        Self {
65            range_tuple_checker_sizes: default_range_tuple_checker_sizes(),
66        }
67    }
68}
69
70fn default_range_tuple_checker_sizes() -> [u32; 2] {
71    [1 << 8, 32 * (1 << 8)]
72}
73
74#[derive(Clone, From, AnyEnum, Executor, MeteredExecutor, PreflightExecutor)]
75pub enum Int256Executor {
76    BaseAlu256(Rv32BaseAlu256Executor),
77    LessThan256(Rv32LessThan256Executor),
78    BranchEqual256(Rv32BranchEqual256Executor),
79    BranchLessThan256(Rv32BranchLessThan256Executor),
80    Multiplication256(Rv32Multiplication256Executor),
81    Shift256(Rv32Shift256Executor),
82}
83
84impl<F: PrimeField32> VmExecutionExtension<F> for Int256 {
85    type Executor = Int256Executor;
86
87    fn extend_execution(
88        &self,
89        inventory: &mut ExecutorInventoryBuilder<F, Int256Executor>,
90    ) -> Result<(), ExecutorInventoryError> {
91        let pointer_max_bits = inventory.pointer_max_bits();
92
93        let alu = Rv32BaseAlu256Executor::new(
94            Rv32HeapAdapterExecutor::new(pointer_max_bits),
95            Rv32BaseAlu256Opcode::CLASS_OFFSET,
96        );
97        inventory.add_executor(alu, Rv32BaseAlu256Opcode::iter().map(|x| x.global_opcode()))?;
98
99        let lt = Rv32LessThan256Executor::new(
100            Rv32HeapAdapterExecutor::new(pointer_max_bits),
101            Rv32LessThan256Opcode::CLASS_OFFSET,
102        );
103        inventory.add_executor(lt, Rv32LessThan256Opcode::iter().map(|x| x.global_opcode()))?;
104
105        let beq = Rv32BranchEqual256Executor::new(
106            Rv32HeapBranchAdapterExecutor::new(pointer_max_bits),
107            Rv32BranchEqual256Opcode::CLASS_OFFSET,
108            DEFAULT_PC_STEP,
109        );
110        inventory.add_executor(
111            beq,
112            Rv32BranchEqual256Opcode::iter().map(|x| x.global_opcode()),
113        )?;
114
115        let blt = Rv32BranchLessThan256Executor::new(
116            Rv32HeapBranchAdapterExecutor::new(pointer_max_bits),
117            Rv32BranchLessThan256Opcode::CLASS_OFFSET,
118        );
119        inventory.add_executor(
120            blt,
121            Rv32BranchLessThan256Opcode::iter().map(|x| x.global_opcode()),
122        )?;
123
124        let mult = Rv32Multiplication256Executor::new(
125            Rv32HeapAdapterExecutor::new(pointer_max_bits),
126            Rv32Mul256Opcode::CLASS_OFFSET,
127        );
128        inventory.add_executor(mult, Rv32Mul256Opcode::iter().map(|x| x.global_opcode()))?;
129
130        let shift = Rv32Shift256Executor::new(
131            Rv32HeapAdapterExecutor::new(pointer_max_bits),
132            Rv32Shift256Opcode::CLASS_OFFSET,
133        );
134        inventory.add_executor(shift, Rv32Shift256Opcode::iter().map(|x| x.global_opcode()))?;
135
136        Ok(())
137    }
138}
139
140impl<SC: StarkGenericConfig> VmCircuitExtension<SC> for Int256 {
141    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
142        let SystemPort {
143            execution_bus,
144            program_bus,
145            memory_bridge,
146        } = inventory.system().port();
147
148        let exec_bridge = ExecutionBridge::new(execution_bus, program_bus);
149        let range_checker = inventory.range_checker().bus;
150        let pointer_max_bits = inventory.pointer_max_bits();
151
152        let bitwise_lu = {
153            // A trick to get around Rust's borrow rules
154            let existing_air = inventory.find_air::<BitwiseOperationLookupAir<8>>().next();
155            if let Some(air) = existing_air {
156                air.bus
157            } else {
158                let bus = BitwiseOperationLookupBus::new(inventory.new_bus_idx());
159                let air = BitwiseOperationLookupAir::<8>::new(bus);
160                inventory.add_air(air);
161                air.bus
162            }
163        };
164
165        let range_tuple_checker = {
166            let existing_air = inventory.find_air::<RangeTupleCheckerAir<2>>().find(|c| {
167                c.bus.sizes[0] >= self.range_tuple_checker_sizes[0]
168                    && c.bus.sizes[1] >= self.range_tuple_checker_sizes[1]
169            });
170            if let Some(air) = existing_air {
171                air.bus
172            } else {
173                let bus = RangeTupleCheckerBus::new(
174                    inventory.new_bus_idx(),
175                    self.range_tuple_checker_sizes,
176                );
177                let air = RangeTupleCheckerAir { bus };
178                inventory.add_air(air);
179                air.bus
180            }
181        };
182
183        let alu = Rv32BaseAlu256Air::new(
184            Rv32HeapAdapterAir::new(exec_bridge, memory_bridge, bitwise_lu, pointer_max_bits),
185            BaseAluCoreAir::new(bitwise_lu, Rv32BaseAlu256Opcode::CLASS_OFFSET),
186        );
187        inventory.add_air(alu);
188
189        let lt = Rv32LessThan256Air::new(
190            Rv32HeapAdapterAir::new(exec_bridge, memory_bridge, bitwise_lu, pointer_max_bits),
191            LessThanCoreAir::new(bitwise_lu, Rv32LessThan256Opcode::CLASS_OFFSET),
192        );
193        inventory.add_air(lt);
194
195        let beq = Rv32BranchEqual256Air::new(
196            Rv32HeapBranchAdapterAir::new(exec_bridge, memory_bridge, bitwise_lu, pointer_max_bits),
197            BranchEqualCoreAir::new(Rv32BranchEqual256Opcode::CLASS_OFFSET, DEFAULT_PC_STEP),
198        );
199        inventory.add_air(beq);
200
201        let blt = Rv32BranchLessThan256Air::new(
202            Rv32HeapBranchAdapterAir::new(exec_bridge, memory_bridge, bitwise_lu, pointer_max_bits),
203            BranchLessThanCoreAir::new(bitwise_lu, Rv32BranchLessThan256Opcode::CLASS_OFFSET),
204        );
205        inventory.add_air(blt);
206
207        let mult = Rv32Multiplication256Air::new(
208            Rv32HeapAdapterAir::new(exec_bridge, memory_bridge, bitwise_lu, pointer_max_bits),
209            MultiplicationCoreAir::new(range_tuple_checker, Rv32Mul256Opcode::CLASS_OFFSET),
210        );
211        inventory.add_air(mult);
212
213        let shift = Rv32Shift256Air::new(
214            Rv32HeapAdapterAir::new(exec_bridge, memory_bridge, bitwise_lu, pointer_max_bits),
215            ShiftCoreAir::new(bitwise_lu, range_checker, Rv32Shift256Opcode::CLASS_OFFSET),
216        );
217        inventory.add_air(shift);
218
219        Ok(())
220    }
221}
222
223pub struct Int256CpuProverExt;
224// This implementation is specific to CpuBackend because the lookup chips (VariableRangeChecker,
225// BitwiseOperationLookupChip) are specific to CpuBackend.
226impl<E, SC, RA> VmProverExtension<E, RA, Int256> for Int256CpuProverExt
227where
228    SC: StarkGenericConfig,
229    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
230    RA: RowMajorMatrixArena<Val<SC>>,
231    Val<SC>: PrimeField32,
232{
233    fn extend_prover(
234        &self,
235        extension: &Int256,
236        inventory: &mut ChipInventory<SC, RA, CpuBackend<SC>>,
237    ) -> Result<(), ChipInventoryError> {
238        let range_checker = inventory.range_checker()?.clone();
239        let timestamp_max_bits = inventory.timestamp_max_bits();
240        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
241        let pointer_max_bits = inventory.airs().config().memory_config.pointer_max_bits;
242
243        let bitwise_lu = {
244            let existing_chip = inventory
245                .find_chip::<SharedBitwiseOperationLookupChip<8>>()
246                .next();
247            if let Some(chip) = existing_chip {
248                chip.clone()
249            } else {
250                let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
251                let chip = Arc::new(BitwiseOperationLookupChip::new(air.bus));
252                inventory.add_periphery_chip(chip.clone());
253                chip
254            }
255        };
256
257        let range_tuple_checker = {
258            let existing_chip = inventory
259                .find_chip::<SharedRangeTupleCheckerChip<2>>()
260                .find(|c| {
261                    c.bus().sizes[0] >= extension.range_tuple_checker_sizes[0]
262                        && c.bus().sizes[1] >= extension.range_tuple_checker_sizes[1]
263                });
264            if let Some(chip) = existing_chip {
265                chip.clone()
266            } else {
267                let air: &RangeTupleCheckerAir<2> = inventory.next_air()?;
268                let chip = SharedRangeTupleCheckerChip::new(RangeTupleCheckerChip::new(air.bus));
269                inventory.add_periphery_chip(chip.clone());
270                chip
271            }
272        };
273
274        inventory.next_air::<Rv32BaseAlu256Air>()?;
275        let alu = Rv32BaseAlu256Chip::new(
276            BaseAluFiller::new(
277                Rv32HeapAdapterFiller::new(pointer_max_bits, bitwise_lu.clone()),
278                bitwise_lu.clone(),
279                Rv32BaseAlu256Opcode::CLASS_OFFSET,
280            ),
281            mem_helper.clone(),
282        );
283        inventory.add_executor_chip(alu);
284
285        inventory.next_air::<Rv32LessThan256Air>()?;
286        let lt = Rv32LessThan256Chip::new(
287            LessThanFiller::new(
288                Rv32HeapAdapterFiller::new(pointer_max_bits, bitwise_lu.clone()),
289                bitwise_lu.clone(),
290                Rv32LessThan256Opcode::CLASS_OFFSET,
291            ),
292            mem_helper.clone(),
293        );
294        inventory.add_executor_chip(lt);
295
296        inventory.next_air::<Rv32BranchEqual256Air>()?;
297        let beq = Rv32BranchEqual256Chip::new(
298            BranchEqualFiller::new(
299                Rv32HeapBranchAdapterFiller::new(pointer_max_bits, bitwise_lu.clone()),
300                Rv32BranchEqual256Opcode::CLASS_OFFSET,
301                DEFAULT_PC_STEP,
302            ),
303            mem_helper.clone(),
304        );
305        inventory.add_executor_chip(beq);
306
307        inventory.next_air::<Rv32BranchLessThan256Air>()?;
308        let blt = Rv32BranchLessThan256Chip::new(
309            BranchLessThanFiller::new(
310                Rv32HeapBranchAdapterFiller::new(pointer_max_bits, bitwise_lu.clone()),
311                bitwise_lu.clone(),
312                Rv32BranchLessThan256Opcode::CLASS_OFFSET,
313            ),
314            mem_helper.clone(),
315        );
316        inventory.add_executor_chip(blt);
317
318        inventory.next_air::<Rv32Multiplication256Air>()?;
319        let mult = Rv32Multiplication256Chip::new(
320            MultiplicationFiller::new(
321                Rv32HeapAdapterFiller::new(pointer_max_bits, bitwise_lu.clone()),
322                range_tuple_checker.clone(),
323                Rv32Mul256Opcode::CLASS_OFFSET,
324            ),
325            mem_helper.clone(),
326        );
327        inventory.add_executor_chip(mult);
328
329        inventory.next_air::<Rv32Shift256Air>()?;
330        let shift = Rv32Shift256Chip::new(
331            ShiftFiller::new(
332                Rv32HeapAdapterFiller::new(pointer_max_bits, bitwise_lu.clone()),
333                bitwise_lu.clone(),
334                range_checker.clone(),
335                Rv32Shift256Opcode::CLASS_OFFSET,
336            ),
337            mem_helper.clone(),
338        );
339        inventory.add_executor_chip(shift);
340        Ok(())
341    }
342}
343
344#[derive(Clone)]
345pub struct Int256Rv32CpuBuilder;
346
347impl<E, SC> VmBuilder<E> for Int256Rv32CpuBuilder
348where
349    SC: StarkGenericConfig,
350    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
351    Val<SC>: PrimeField32,
352{
353    type VmConfig = Int256Rv32Config;
354    type SystemChipInventory = SystemChipInventory<SC>;
355    type RecordArena = MatrixRecordArena<Val<SC>>;
356
357    fn create_chip_complex(
358        &self,
359        config: &Int256Rv32Config,
360        circuit: AirInventory<E::SC>,
361    ) -> Result<
362        VmChipComplex<E::SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
363        ChipInventoryError,
364    > {
365        let mut chip_complex =
366            VmBuilder::<E>::create_chip_complex(&SystemCpuBuilder, &config.system, circuit)?;
367        let inventory = &mut chip_complex.inventory;
368        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32i, inventory)?;
369        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32m, inventory)?;
370        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.io, inventory)?;
371        VmProverExtension::<E, _, _>::extend_prover(
372            &Int256CpuProverExt,
373            &config.bigint,
374            inventory,
375        )?;
376        Ok(chip_complex)
377    }
378}
openvm_bigint_circuit/extension/mod.rs

openvm_bigint_circuit/extension/
mod.rs
