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

openvm_bigint_circuit/extension/
mod.rs
