openvm_ecc_circuit/extension/
hybrid.rs

1//! Prover extension for the GPU backend which still does trace generation on CPU.
2
3use openvm_algebra_circuit::Rv32ModularHybridBuilder;
4use openvm_circuit::{
5    arch::{DEFAULT_BLOCK_SIZE, *},
6    system::{
7        cuda::{
8            extensions::{get_inventory_range_checker, get_or_create_bitwise_op_lookup},
9            SystemChipInventoryGPU,
10        },
11        memory::SharedMemoryHelper,
12    },
13};
14use openvm_circuit_primitives::{hybrid_chip::cpu_proving_ctx_to_gpu, Chip};
15use openvm_cpu_backend::CpuBackend;
16use openvm_cuda_backend::{
17    base::DeviceMatrix,
18    prelude::{F, SC},
19    BabyBearPoseidon2GpuEngine as GpuBabyBearPoseidon2Engine, GpuBackend,
20};
21use openvm_cuda_common::stream::GpuDeviceCtx;
22use openvm_mod_circuit_builder::{ExprBuilderConfig, FieldExpressionMetadata};
23use openvm_rv32_adapters::{Rv32VecHeapAdapterCols, Rv32VecHeapAdapterExecutor};
24use openvm_stark_backend::{p3_air::BaseAir, prover::AirProvingContext};
25
26use crate::{
27    get_ec_addne_chip, get_ec_double_chip, EccRecord, Rv32WeierstrassConfig, WeierstrassAir,
28    WeierstrassChip, WeierstrassExtension, ECC_BLOCKS_32, ECC_BLOCKS_48, NUM_LIMBS_32,
29    NUM_LIMBS_48,
30};
31
32#[derive(derive_new::new)]
33pub struct HybridWeierstrassChip<
34    F,
35    const NUM_READS: usize,
36    const BLOCKS: usize,
37    const BLOCK_SIZE: usize,
38> {
39    cpu: WeierstrassChip<F, NUM_READS, BLOCKS, BLOCK_SIZE>,
40    device_ctx: GpuDeviceCtx,
41}
42
43// Auto-implementation of Chip for GpuBackend for a Cpu Chip by doing conversion
44// of Dense->Matrix Record Arena, cpu tracegen, and then H2D transfer of the trace matrix.
45impl<const NUM_READS: usize, const BLOCKS: usize, const BLOCK_SIZE: usize>
46    Chip<DenseRecordArena, GpuBackend> for HybridWeierstrassChip<F, NUM_READS, BLOCKS, BLOCK_SIZE>
47{
48    fn generate_proving_ctx(&self, mut arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
49        let total_input_limbs =
50            self.cpu.inner.num_inputs() * self.cpu.inner.expr.canonical_num_limbs();
51        let layout = AdapterCoreLayout::with_metadata(FieldExpressionMetadata::<
52            F,
53            Rv32VecHeapAdapterExecutor<NUM_READS, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>,
54        >::new(total_input_limbs));
55
56        let record_size = RecordSeeker::<
57            DenseRecordArena,
58            EccRecord<NUM_READS, BLOCKS, BLOCK_SIZE>,
59            _,
60        >::get_aligned_record_size(&layout);
61
62        let records = arena.allocated();
63        if records.is_empty() {
64            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
65        }
66        debug_assert_eq!(records.len() % record_size, 0);
67
68        let num_records = records.len() / record_size;
69        let height = num_records.next_power_of_two();
70        let mut seeker = arena
71            .get_record_seeker::<EccRecord<NUM_READS, BLOCKS, BLOCK_SIZE>, AdapterCoreLayout<
72                FieldExpressionMetadata<
73                    F,
74                    Rv32VecHeapAdapterExecutor<NUM_READS, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>,
75                >,
76            >>();
77        let adapter_width =
78            Rv32VecHeapAdapterCols::<F, NUM_READS, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>::width();
79        let width = adapter_width + BaseAir::<F>::width(&self.cpu.inner.expr);
80        let mut matrix_arena = MatrixRecordArena::<F>::with_capacity(height, width);
81        seeker.transfer_to_matrix_arena(&mut matrix_arena, layout);
82        let cpu_ctx = Chip::<_, CpuBackend<SC>>::generate_proving_ctx(&self.cpu, matrix_arena);
83        cpu_proving_ctx_to_gpu(cpu_ctx, &self.device_ctx)
84    }
85}
86
87#[derive(Clone, Copy, Default)]
88pub struct EccHybridProverExt;
89
90impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, WeierstrassExtension>
91    for EccHybridProverExt
92{
93    fn extend_prover(
94        &self,
95        extension: &WeierstrassExtension,
96        inventory: &mut ChipInventory<SC, DenseRecordArena, GpuBackend>,
97    ) -> Result<(), ChipInventoryError> {
98        let range_checker_gpu = get_inventory_range_checker(inventory);
99        let timestamp_max_bits = inventory.timestamp_max_bits();
100        let pointer_max_bits = inventory.airs().pointer_max_bits();
101        let range_checker = range_checker_gpu.cpu_chip.clone().unwrap();
102        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
103
104        let bitwise_lu_gpu = get_or_create_bitwise_op_lookup(inventory)?;
105        let bitwise_lu = bitwise_lu_gpu.cpu_chip.clone().unwrap();
106        let device_ctx = range_checker_gpu.device_ctx.clone();
107
108        for curve in extension.supported_curves.iter() {
109            let bytes = curve.modulus.bits().div_ceil(8) as usize;
110
111            if bytes <= NUM_LIMBS_32 {
112                let config = ExprBuilderConfig {
113                    modulus: curve.modulus.clone(),
114                    num_limbs: NUM_LIMBS_32,
115                    limb_bits: 8,
116                };
117
118                inventory.next_air::<WeierstrassAir<2, ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>>()?;
119                let addne = get_ec_addne_chip::<F, ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
120                    config.clone(),
121                    mem_helper.clone(),
122                    range_checker.clone(),
123                    bitwise_lu.clone(),
124                    pointer_max_bits,
125                );
126                inventory.add_executor_chip(HybridWeierstrassChip::new(addne, device_ctx.clone()));
127
128                inventory.next_air::<WeierstrassAir<1, ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>>()?;
129                let double = get_ec_double_chip::<F, ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
130                    config,
131                    mem_helper.clone(),
132                    range_checker.clone(),
133                    bitwise_lu.clone(),
134                    pointer_max_bits,
135                    curve.a.clone(),
136                );
137                inventory.add_executor_chip(HybridWeierstrassChip::new(double, device_ctx.clone()));
138            } else if bytes <= NUM_LIMBS_48 {
139                let config = ExprBuilderConfig {
140                    modulus: curve.modulus.clone(),
141                    num_limbs: NUM_LIMBS_48,
142                    limb_bits: 8,
143                };
144
145                inventory.next_air::<WeierstrassAir<2, ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>>()?;
146                let addne = get_ec_addne_chip::<F, ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
147                    config.clone(),
148                    mem_helper.clone(),
149                    range_checker.clone(),
150                    bitwise_lu.clone(),
151                    pointer_max_bits,
152                );
153                inventory.add_executor_chip(HybridWeierstrassChip::new(addne, device_ctx.clone()));
154
155                inventory.next_air::<WeierstrassAir<1, ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>>()?;
156                let double = get_ec_double_chip::<F, ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
157                    config,
158                    mem_helper.clone(),
159                    range_checker.clone(),
160                    bitwise_lu.clone(),
161                    pointer_max_bits,
162                    curve.a.clone(),
163                );
164                inventory.add_executor_chip(HybridWeierstrassChip::new(double, device_ctx.clone()));
165            } else {
166                panic!("Modulus too large");
167            }
168        }
169
170        Ok(())
171    }
172}
173
174/// This builder will do tracegen for the RV32IM extensions on GPU but the modular and ecc
175/// extensions on CPU.
176#[derive(Clone)]
177pub struct Rv32WeierstrassHybridBuilder;
178
179type E = GpuBabyBearPoseidon2Engine;
180
181impl VmBuilder<E> for Rv32WeierstrassHybridBuilder {
182    type VmConfig = Rv32WeierstrassConfig;
183    type SystemChipInventory = SystemChipInventoryGPU;
184    type RecordArena = DenseRecordArena;
185
186    fn create_chip_complex(
187        &self,
188        config: &Rv32WeierstrassConfig,
189        circuit: AirInventory<SC>,
190        device_ctx: &openvm_stark_backend::EngineDeviceCtx<E>,
191    ) -> Result<
192        VmChipComplex<SC, Self::RecordArena, GpuBackend, Self::SystemChipInventory>,
193        ChipInventoryError,
194    > {
195        let mut chip_complex = VmBuilder::<E>::create_chip_complex(
196            &Rv32ModularHybridBuilder,
197            &config.modular,
198            circuit,
199            device_ctx,
200        )?;
201        let inventory = &mut chip_complex.inventory;
202        VmProverExtension::<E, _, _>::extend_prover(
203            &EccHybridProverExt,
204            &config.weierstrass,
205            inventory,
206        )?;
207
208        Ok(chip_complex)
209    }
210}
openvm_ecc_circuit/extension/hybrid.rs

openvm_ecc_circuit/extension/
hybrid.rs
