openvm_sha256_circuit/extension/
cuda.rs

1use openvm_circuit::{
2    arch::DenseRecordArena,
3    system::cuda::extensions::{get_inventory_range_checker, get_or_create_bitwise_op_lookup},
4};
5use openvm_cuda_backend::{engine::GpuBabyBearPoseidon2Engine, prover_backend::GpuBackend};
6use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
7
8use super::*;
9
10pub struct Sha256GpuProverExt;
11
12impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, Sha256>
13    for Sha256GpuProverExt
14{
15    fn extend_prover(
16        &self,
17        _: &Sha256,
18        inventory: &mut ChipInventory<BabyBearPoseidon2Config, DenseRecordArena, GpuBackend>,
19    ) -> Result<(), ChipInventoryError> {
20        let pointer_max_bits = inventory.airs().pointer_max_bits();
21        let timestamp_max_bits = inventory.timestamp_max_bits();
22
23        let range_checker = get_inventory_range_checker(inventory);
24        let bitwise_lu = get_or_create_bitwise_op_lookup(inventory)?;
25
26        // These calls to next_air are not strictly necessary to construct the chips, but provide a
27        // safeguard to ensure that chip construction matches the circuit definition
28        inventory.next_air::<Sha256VmAir>()?;
29        let sha256 = Sha256VmChipGpu::new(
30            range_checker.clone(),
31            bitwise_lu,
32            pointer_max_bits as u32,
33            timestamp_max_bits as u32,
34        );
35        inventory.add_executor_chip(sha256);
36
37        Ok(())
38    }
39}