openvm_keccak256_circuit/extension/
cuda.rs

1use openvm_circuit::{
2    arch::DenseRecordArena,
3    system::cuda::{
4        extensions::{
5            get_inventory_range_checker, get_or_create_bitwise_op_lookup, SystemGpuBuilder,
6        },
7        SystemChipInventoryGPU,
8    },
9};
10use openvm_cuda_backend::{engine::GpuBabyBearPoseidon2Engine, prover_backend::GpuBackend};
11use openvm_rv32im_circuit::Rv32ImGpuProverExt;
12use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
13
14use super::*;
15use crate::{air::KeccakVmAir, cuda::Keccak256ChipGpu};
16
17pub struct Keccak256GpuProverExt;
18
19// This implementation is specific to GpuBackend because the lookup chips
20// (VariableRangeCheckerChipGPU, BitwiseOperationLookupChipGPU) are specific to GpuBackend.
21impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, Keccak256>
22    for Keccak256GpuProverExt
23{
24    fn extend_prover(
25        &self,
26        _: &Keccak256,
27        inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
28    ) -> Result<(), ChipInventoryError> {
29        let pointer_max_bits = inventory.airs().pointer_max_bits();
30        let timestamp_max_bits = inventory.timestamp_max_bits();
31
32        let range_checker = get_inventory_range_checker(inventory);
33        let bitwise_lu = get_or_create_bitwise_op_lookup(inventory)?;
34
35        // These calls to next_air are not strictly necessary to construct the chips, but provide a
36        // safeguard to ensure that chip construction matches the circuit definition
37        inventory.next_air::<KeccakVmAir>()?;
38        let keccak = Keccak256ChipGpu::new(
39            range_checker.clone(),
40            bitwise_lu.clone(),
41            pointer_max_bits as u32,
42            timestamp_max_bits as u32,
43        );
44        inventory.add_executor_chip(keccak);
45
46        Ok(())
47    }
48}
49
50pub struct Keccak256Rv32GpuBuilder;
51
52type E = GpuBabyBearPoseidon2Engine;
53
54impl VmBuilder<E> for Keccak256Rv32GpuBuilder {
55    type VmConfig = Keccak256Rv32Config;
56    type SystemChipInventory = SystemChipInventoryGPU;
57    type RecordArena = DenseRecordArena;
58
59    fn create_chip_complex(
60        &self,
61        config: &Keccak256Rv32Config,
62        circuit: AirInventory<<E as StarkEngine>::SC>,
63    ) -> Result<
64        VmChipComplex<
65            <E as StarkEngine>::SC,
66            Self::RecordArena,
67            <E as StarkEngine>::PB,
68            Self::SystemChipInventory,
69        >,
70        ChipInventoryError,
71    > {
72        let mut chip_complex =
73            VmBuilder::<E>::create_chip_complex(&SystemGpuBuilder, &config.system, circuit)?;
74        let inventory = &mut chip_complex.inventory;
75        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImGpuProverExt, &config.rv32i, inventory)?;
76        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImGpuProverExt, &config.rv32m, inventory)?;
77        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImGpuProverExt, &config.io, inventory)?;
78        VmProverExtension::<E, _, _>::extend_prover(
79            &Keccak256GpuProverExt,
80            &config.keccak,
81            inventory,
82        )?;
83        Ok(chip_complex)
84    }
85}
openvm_keccak256_circuit/extension/cuda.rs

openvm_keccak256_circuit/extension/
cuda.rs
