openvm_algebra_circuit/extension/
hybrid.rs

1//! Prover extension for the GPU backend which still does trace generation on CPU.
2
3use openvm_algebra_transpiler::Rv32ModularArithmeticOpcode;
4use openvm_circuit::{
5    arch::{DEFAULT_BLOCK_SIZE, *},
6    system::{
7        cuda::{
8            extensions::{
9                get_inventory_range_checker, get_or_create_bitwise_op_lookup, SystemGpuBuilder,
10            },
11            SystemChipInventoryGPU,
12        },
13        memory::SharedMemoryHelper,
14    },
15};
16use openvm_circuit_primitives::{
17    bigint::utils::big_uint_to_limbs, hybrid_chip::cpu_proving_ctx_to_gpu, Chip,
18};
19use openvm_cpu_backend::CpuBackend;
20use openvm_cuda_backend::{
21    base::DeviceMatrix,
22    prelude::{F, SC},
23    BabyBearPoseidon2GpuEngine as GpuBabyBearPoseidon2Engine, GpuBackend,
24};
25use openvm_cuda_common::stream::GpuDeviceCtx;
26use openvm_instructions::LocalOpcode;
27use openvm_mod_circuit_builder::{ExprBuilderConfig, FieldExpressionMetadata};
28use openvm_rv32_adapters::{
29    Rv32IsEqualModAdapterCols, Rv32IsEqualModAdapterExecutor, Rv32IsEqualModAdapterFiller,
30    Rv32IsEqualModAdapterRecord, Rv32VecHeapAdapterCols, Rv32VecHeapAdapterExecutor,
31};
32use openvm_rv32im_circuit::Rv32ImGpuProverExt;
33use openvm_stark_backend::{p3_air::BaseAir, prover::AirProvingContext};
34use strum::EnumCount;
35
36use crate::{
37    fp2_chip::{get_fp2_addsub_chip, get_fp2_muldiv_chip, Fp2Air, Fp2Chip},
38    modular_chip::*,
39    AlgebraRecord, Fp2Extension, ModularExtension, Rv32ModularConfig, Rv32ModularWithFp2Config,
40    FP2_BLOCKS_32, FP2_BLOCKS_48, MODULAR_BLOCKS_32, MODULAR_BLOCKS_48, NUM_LIMBS_32, NUM_LIMBS_48,
41};
42
43#[derive(derive_new::new)]
44pub struct HybridModularChip<F, const BLOCKS: usize, const BLOCK_SIZE: usize> {
45    cpu: ModularChip<F, BLOCKS, BLOCK_SIZE>,
46    device_ctx: GpuDeviceCtx,
47}
48
49// Auto-implementation of Chip for GpuBackend for a Cpu Chip by doing conversion
50// of Dense->Matrix Record Arena, cpu tracegen, and then H2D transfer of the trace matrix.
51impl<const BLOCKS: usize, const BLOCK_SIZE: usize> Chip<DenseRecordArena, GpuBackend>
52    for HybridModularChip<F, BLOCKS, BLOCK_SIZE>
53{
54    fn generate_proving_ctx(&self, mut arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
55        let total_input_limbs =
56            self.cpu.inner.num_inputs() * self.cpu.inner.expr.canonical_num_limbs();
57        let layout = AdapterCoreLayout::with_metadata(FieldExpressionMetadata::<
58            F,
59            Rv32VecHeapAdapterExecutor<2, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>,
60        >::new(total_input_limbs));
61
62        let record_size = RecordSeeker::<
63            DenseRecordArena,
64            AlgebraRecord<2, BLOCKS, BLOCK_SIZE>,
65            _,
66        >::get_aligned_record_size(&layout);
67
68        let records = arena.allocated();
69        if records.is_empty() {
70            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
71        }
72        debug_assert_eq!(records.len() % record_size, 0);
73
74        let num_records = records.len() / record_size;
75
76        let height = num_records.next_power_of_two();
77        let mut seeker = arena
78            .get_record_seeker::<AlgebraRecord<2, BLOCKS, BLOCK_SIZE>, AdapterCoreLayout<
79                FieldExpressionMetadata<
80                    F,
81                    Rv32VecHeapAdapterExecutor<2, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>,
82                >,
83            >>();
84        let adapter_width =
85            Rv32VecHeapAdapterCols::<F, 2, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>::width();
86        let width = adapter_width + BaseAir::<F>::width(&self.cpu.inner.expr);
87        let mut matrix_arena = MatrixRecordArena::<F>::with_capacity(height, width);
88        seeker.transfer_to_matrix_arena(&mut matrix_arena, layout);
89        let cpu_ctx = Chip::<_, CpuBackend<SC>>::generate_proving_ctx(&self.cpu, matrix_arena);
90        cpu_proving_ctx_to_gpu(cpu_ctx, &self.device_ctx)
91    }
92}
93
94#[derive(derive_new::new)]
95pub struct HybridModularIsEqualChip<
96    F,
97    const NUM_LANES: usize,
98    const LANE_SIZE: usize,
99    const TOTAL_LIMBS: usize,
100> {
101    cpu: ModularIsEqualChip<F, NUM_LANES, LANE_SIZE, TOTAL_LIMBS>,
102    device_ctx: GpuDeviceCtx,
103}
104
105impl<const NUM_LANES: usize, const LANE_SIZE: usize, const TOTAL_LIMBS: usize>
106    Chip<DenseRecordArena, GpuBackend>
107    for HybridModularIsEqualChip<F, NUM_LANES, LANE_SIZE, TOTAL_LIMBS>
108{
109    fn generate_proving_ctx(&self, mut arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
110        let record_size = size_of::<(
111            Rv32IsEqualModAdapterRecord<2, NUM_LANES, LANE_SIZE, TOTAL_LIMBS>,
112            ModularIsEqualRecord<TOTAL_LIMBS>,
113        )>();
114        let trace_width = Rv32IsEqualModAdapterCols::<F, 2, NUM_LANES, LANE_SIZE>::width()
115            + ModularIsEqualCoreCols::<F, TOTAL_LIMBS>::width();
116        let records = arena.allocated();
117        if records.is_empty() {
118            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
119        }
120        debug_assert_eq!(records.len() % record_size, 0);
121
122        let num_records = records.len() / record_size;
123        let height = num_records.next_power_of_two();
124        let mut seeker = arena.get_record_seeker::<(
125            &mut Rv32IsEqualModAdapterRecord<2, NUM_LANES, LANE_SIZE, TOTAL_LIMBS>,
126            &mut ModularIsEqualRecord<TOTAL_LIMBS>,
127        ), EmptyAdapterCoreLayout<
128            F,
129            Rv32IsEqualModAdapterExecutor<2, NUM_LANES, LANE_SIZE, TOTAL_LIMBS>,
130        >>();
131        let mut matrix_arena = MatrixRecordArena::<F>::with_capacity(height, trace_width);
132        seeker.transfer_to_matrix_arena(&mut matrix_arena, EmptyAdapterCoreLayout::new());
133        let cpu_ctx = Chip::<_, CpuBackend<SC>>::generate_proving_ctx(&self.cpu, matrix_arena);
134        cpu_proving_ctx_to_gpu(cpu_ctx, &self.device_ctx)
135    }
136}
137
138#[derive(Clone, Copy, Default)]
139pub struct AlgebraHybridProverExt;
140
141impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, ModularExtension>
142    for AlgebraHybridProverExt
143{
144    fn extend_prover(
145        &self,
146        extension: &ModularExtension,
147        inventory: &mut ChipInventory<SC, DenseRecordArena, GpuBackend>,
148    ) -> Result<(), ChipInventoryError> {
149        let range_checker_gpu = get_inventory_range_checker(inventory);
150        let timestamp_max_bits = inventory.timestamp_max_bits();
151        let pointer_max_bits = inventory.airs().pointer_max_bits();
152        let range_checker = range_checker_gpu.cpu_chip.clone().unwrap();
153        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
154        let bitwise_lu_gpu = get_or_create_bitwise_op_lookup(inventory)?;
155        let bitwise_lu = bitwise_lu_gpu.cpu_chip.clone().unwrap();
156        let device_ctx = range_checker_gpu.device_ctx.clone();
157
158        for (i, modulus) in extension.supported_moduli.iter().enumerate() {
159            // determine the number of bytes needed to represent a prime field element
160            let bytes = modulus.bits().div_ceil(8) as usize;
161            let start_offset =
162                Rv32ModularArithmeticOpcode::CLASS_OFFSET + i * Rv32ModularArithmeticOpcode::COUNT;
163
164            let modulus_limbs = big_uint_to_limbs(modulus, 8);
165
166            if bytes <= NUM_LIMBS_32 {
167                let config = ExprBuilderConfig {
168                    modulus: modulus.clone(),
169                    num_limbs: NUM_LIMBS_32,
170                    limb_bits: 8,
171                };
172
173                inventory.next_air::<ModularAir<MODULAR_BLOCKS_32, DEFAULT_BLOCK_SIZE>>()?;
174                let addsub = get_modular_addsub_chip::<F, MODULAR_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
175                    config.clone(),
176                    mem_helper.clone(),
177                    range_checker.clone(),
178                    bitwise_lu.clone(),
179                    pointer_max_bits,
180                );
181                inventory.add_executor_chip(HybridModularChip::new(addsub, device_ctx.clone()));
182
183                inventory.next_air::<ModularAir<MODULAR_BLOCKS_32, DEFAULT_BLOCK_SIZE>>()?;
184                let muldiv = get_modular_muldiv_chip::<F, MODULAR_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
185                    config,
186                    mem_helper.clone(),
187                    range_checker.clone(),
188                    bitwise_lu.clone(),
189                    pointer_max_bits,
190                );
191                inventory.add_executor_chip(HybridModularChip::new(muldiv, device_ctx.clone()));
192
193                let modulus_limbs = std::array::from_fn(|i| {
194                    if i < modulus_limbs.len() {
195                        modulus_limbs[i] as u8
196                    } else {
197                        0
198                    }
199                });
200                inventory.next_air::<ModularIsEqualAir<MODULAR_BLOCKS_32, DEFAULT_BLOCK_SIZE, NUM_LIMBS_32>>()?;
201                let is_eq = ModularIsEqualChip::<
202                    F,
203                    MODULAR_BLOCKS_32,
204                    DEFAULT_BLOCK_SIZE,
205                    NUM_LIMBS_32,
206                >::new(
207                    ModularIsEqualFiller::new(
208                        Rv32IsEqualModAdapterFiller::new(pointer_max_bits, bitwise_lu.clone()),
209                        start_offset,
210                        modulus_limbs,
211                        bitwise_lu.clone(),
212                    ),
213                    mem_helper.clone(),
214                );
215                inventory
216                    .add_executor_chip(HybridModularIsEqualChip::new(is_eq, device_ctx.clone()));
217            } else if bytes <= NUM_LIMBS_48 {
218                let config = ExprBuilderConfig {
219                    modulus: modulus.clone(),
220                    num_limbs: NUM_LIMBS_48,
221                    limb_bits: 8,
222                };
223
224                inventory.next_air::<ModularAir<MODULAR_BLOCKS_48, DEFAULT_BLOCK_SIZE>>()?;
225                let addsub = get_modular_addsub_chip::<F, MODULAR_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
226                    config.clone(),
227                    mem_helper.clone(),
228                    range_checker.clone(),
229                    bitwise_lu.clone(),
230                    pointer_max_bits,
231                );
232                inventory.add_executor_chip(HybridModularChip::new(addsub, device_ctx.clone()));
233
234                inventory.next_air::<ModularAir<MODULAR_BLOCKS_48, DEFAULT_BLOCK_SIZE>>()?;
235                let muldiv = get_modular_muldiv_chip::<F, MODULAR_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
236                    config,
237                    mem_helper.clone(),
238                    range_checker.clone(),
239                    bitwise_lu.clone(),
240                    pointer_max_bits,
241                );
242                inventory.add_executor_chip(HybridModularChip::new(muldiv, device_ctx.clone()));
243
244                let modulus_limbs = std::array::from_fn(|i| {
245                    if i < modulus_limbs.len() {
246                        modulus_limbs[i] as u8
247                    } else {
248                        0
249                    }
250                });
251                inventory.next_air::<ModularIsEqualAir<MODULAR_BLOCKS_48, DEFAULT_BLOCK_SIZE, NUM_LIMBS_48>>()?;
252                let is_eq = ModularIsEqualChip::<
253                    F,
254                    MODULAR_BLOCKS_48,
255                    DEFAULT_BLOCK_SIZE,
256                    NUM_LIMBS_48,
257                >::new(
258                    ModularIsEqualFiller::new(
259                        Rv32IsEqualModAdapterFiller::new(pointer_max_bits, bitwise_lu.clone()),
260                        start_offset,
261                        modulus_limbs,
262                        bitwise_lu.clone(),
263                    ),
264                    mem_helper.clone(),
265                );
266                inventory
267                    .add_executor_chip(HybridModularIsEqualChip::new(is_eq, device_ctx.clone()));
268            } else {
269                panic!("Modulus too large");
270            }
271        }
272
273        Ok(())
274    }
275}
276
277#[derive(derive_new::new)]
278pub struct HybridFp2Chip<F, const BLOCKS: usize, const BLOCK_SIZE: usize> {
279    cpu: Fp2Chip<F, BLOCKS, BLOCK_SIZE>,
280    device_ctx: GpuDeviceCtx,
281}
282
283impl<const BLOCKS: usize, const BLOCK_SIZE: usize> Chip<DenseRecordArena, GpuBackend>
284    for HybridFp2Chip<F, BLOCKS, BLOCK_SIZE>
285{
286    fn generate_proving_ctx(&self, mut arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
287        let total_input_limbs =
288            self.cpu.inner.num_inputs() * self.cpu.inner.expr.canonical_num_limbs();
289        let layout = AdapterCoreLayout::with_metadata(FieldExpressionMetadata::<
290            F,
291            Rv32VecHeapAdapterExecutor<2, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>,
292        >::new(total_input_limbs));
293
294        let record_size = RecordSeeker::<
295            DenseRecordArena,
296            AlgebraRecord<2, BLOCKS, BLOCK_SIZE>,
297            _,
298        >::get_aligned_record_size(&layout);
299
300        let records = arena.allocated();
301        if records.is_empty() {
302            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
303        }
304        debug_assert_eq!(records.len() % record_size, 0);
305
306        let num_records = records.len() / record_size;
307        let height = num_records.next_power_of_two();
308        let mut seeker = arena
309            .get_record_seeker::<AlgebraRecord<2, BLOCKS, BLOCK_SIZE>, AdapterCoreLayout<
310                FieldExpressionMetadata<
311                    F,
312                    Rv32VecHeapAdapterExecutor<2, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>,
313                >,
314            >>();
315        let adapter_width =
316            Rv32VecHeapAdapterCols::<F, 2, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>::width();
317        let width = adapter_width + BaseAir::<F>::width(&self.cpu.inner.expr);
318        let mut matrix_arena = MatrixRecordArena::<F>::with_capacity(height, width);
319        seeker.transfer_to_matrix_arena(&mut matrix_arena, layout);
320        let cpu_ctx = Chip::<_, CpuBackend<SC>>::generate_proving_ctx(&self.cpu, matrix_arena);
321        cpu_proving_ctx_to_gpu(cpu_ctx, &self.device_ctx)
322    }
323}
324
325impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, Fp2Extension>
326    for AlgebraHybridProverExt
327{
328    fn extend_prover(
329        &self,
330        extension: &Fp2Extension,
331        inventory: &mut ChipInventory<SC, DenseRecordArena, GpuBackend>,
332    ) -> Result<(), ChipInventoryError> {
333        let range_checker_gpu = get_inventory_range_checker(inventory);
334        let timestamp_max_bits = inventory.timestamp_max_bits();
335        let pointer_max_bits = inventory.airs().pointer_max_bits();
336        let range_checker = range_checker_gpu.cpu_chip.clone().unwrap();
337        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
338        let bitwise_lu_gpu = get_or_create_bitwise_op_lookup(inventory)?;
339        let bitwise_lu = bitwise_lu_gpu.cpu_chip.clone().unwrap();
340        let device_ctx = range_checker_gpu.device_ctx.clone();
341
342        for (_, modulus) in extension.supported_moduli.iter() {
343            // determine the number of bytes needed to represent a prime field element
344            let bytes = modulus.bits().div_ceil(8) as usize;
345
346            if bytes <= NUM_LIMBS_32 {
347                let config = ExprBuilderConfig {
348                    modulus: modulus.clone(),
349                    num_limbs: NUM_LIMBS_32,
350                    limb_bits: 8,
351                };
352
353                inventory.next_air::<Fp2Air<FP2_BLOCKS_32, DEFAULT_BLOCK_SIZE>>()?;
354                let addsub = get_fp2_addsub_chip::<F, FP2_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
355                    config.clone(),
356                    mem_helper.clone(),
357                    range_checker.clone(),
358                    bitwise_lu.clone(),
359                    pointer_max_bits,
360                );
361                inventory.add_executor_chip(HybridFp2Chip::new(addsub, device_ctx.clone()));
362
363                inventory.next_air::<Fp2Air<FP2_BLOCKS_32, DEFAULT_BLOCK_SIZE>>()?;
364                let muldiv = get_fp2_muldiv_chip::<F, FP2_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
365                    config,
366                    mem_helper.clone(),
367                    range_checker.clone(),
368                    bitwise_lu.clone(),
369                    pointer_max_bits,
370                );
371                inventory.add_executor_chip(HybridFp2Chip::new(muldiv, device_ctx.clone()));
372            } else if bytes <= NUM_LIMBS_48 {
373                let config = ExprBuilderConfig {
374                    modulus: modulus.clone(),
375                    num_limbs: NUM_LIMBS_48,
376                    limb_bits: 8,
377                };
378
379                inventory.next_air::<Fp2Air<FP2_BLOCKS_48, DEFAULT_BLOCK_SIZE>>()?;
380                let addsub = get_fp2_addsub_chip::<F, FP2_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
381                    config.clone(),
382                    mem_helper.clone(),
383                    range_checker.clone(),
384                    bitwise_lu.clone(),
385                    pointer_max_bits,
386                );
387                inventory.add_executor_chip(HybridFp2Chip::new(addsub, device_ctx.clone()));
388
389                inventory.next_air::<Fp2Air<FP2_BLOCKS_48, DEFAULT_BLOCK_SIZE>>()?;
390                let muldiv = get_fp2_muldiv_chip::<F, FP2_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
391                    config,
392                    mem_helper.clone(),
393                    range_checker.clone(),
394                    bitwise_lu.clone(),
395                    pointer_max_bits,
396                );
397                inventory.add_executor_chip(HybridFp2Chip::new(muldiv, device_ctx.clone()));
398            } else {
399                panic!("Modulus too large");
400            }
401        }
402
403        Ok(())
404    }
405}
406
407/// This builder will do tracegen for the RV32IM extensions on GPU but the modular extensions on
408/// CPU.
409#[derive(Clone)]
410pub struct Rv32ModularHybridBuilder;
411
412type E = GpuBabyBearPoseidon2Engine;
413
414impl VmBuilder<E> for Rv32ModularHybridBuilder {
415    type VmConfig = Rv32ModularConfig;
416    type SystemChipInventory = SystemChipInventoryGPU;
417    type RecordArena = DenseRecordArena;
418
419    fn create_chip_complex(
420        &self,
421        config: &Rv32ModularConfig,
422        circuit: AirInventory<SC>,
423        device_ctx: &openvm_stark_backend::EngineDeviceCtx<E>,
424    ) -> Result<
425        VmChipComplex<SC, Self::RecordArena, GpuBackend, Self::SystemChipInventory>,
426        ChipInventoryError,
427    > {
428        let mut chip_complex = VmBuilder::<E>::create_chip_complex(
429            &SystemGpuBuilder,
430            &config.system,
431            circuit,
432            device_ctx,
433        )?;
434        let inventory = &mut chip_complex.inventory;
435        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImGpuProverExt, &config.base, inventory)?;
436        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImGpuProverExt, &config.mul, inventory)?;
437        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImGpuProverExt, &config.io, inventory)?;
438        VmProverExtension::<E, _, _>::extend_prover(
439            &AlgebraHybridProverExt,
440            &config.modular,
441            inventory,
442        )?;
443        Ok(chip_complex)
444    }
445}
446
447/// This builder will do tracegen for the RV32IM extensions on GPU but the modular and complex
448/// extensions on CPU.
449#[derive(Clone)]
450pub struct Rv32ModularWithFp2HybridBuilder;
451
452impl VmBuilder<E> for Rv32ModularWithFp2HybridBuilder {
453    type VmConfig = Rv32ModularWithFp2Config;
454    type SystemChipInventory = SystemChipInventoryGPU;
455    type RecordArena = DenseRecordArena;
456
457    fn create_chip_complex(
458        &self,
459        config: &Rv32ModularWithFp2Config,
460        circuit: AirInventory<SC>,
461        device_ctx: &openvm_stark_backend::EngineDeviceCtx<E>,
462    ) -> Result<
463        VmChipComplex<SC, Self::RecordArena, GpuBackend, Self::SystemChipInventory>,
464        ChipInventoryError,
465    > {
466        let mut chip_complex = VmBuilder::<E>::create_chip_complex(
467            &Rv32ModularHybridBuilder,
468            &config.modular,
469            circuit,
470            device_ctx,
471        )?;
472        let inventory = &mut chip_complex.inventory;
473        VmProverExtension::<E, _, _>::extend_prover(
474            &AlgebraHybridProverExt,
475            &config.fp2,
476            inventory,
477        )?;
478        Ok(chip_complex)
479    }
480}
openvm_algebra_circuit/extension/hybrid.rs

openvm_algebra_circuit/extension/
hybrid.rs
