openvm_sha256_circuit/extension/
mod.rs

1use std::{result::Result, sync::Arc};
2
3use derive_more::derive::From;
4use openvm_circuit::{
5    arch::{
6        AirInventory, AirInventoryError, ChipInventory, ChipInventoryError,
7        ExecutorInventoryBuilder, ExecutorInventoryError, RowMajorMatrixArena, VmCircuitExtension,
8        VmExecutionExtension, VmProverExtension,
9    },
10    system::memory::SharedMemoryHelper,
11};
12use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor};
13use openvm_circuit_primitives::bitwise_op_lookup::{
14    BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip,
15    SharedBitwiseOperationLookupChip,
16};
17use openvm_instructions::*;
18use openvm_sha256_transpiler::Rv32Sha256Opcode;
19use openvm_stark_backend::{
20    config::{StarkGenericConfig, Val},
21    p3_field::PrimeField32,
22    prover::cpu::{CpuBackend, CpuDevice},
23};
24use openvm_stark_sdk::engine::StarkEngine;
25use serde::{Deserialize, Serialize};
26use strum::IntoEnumIterator;
27
28use crate::*;
29
30cfg_if::cfg_if! {
31    if #[cfg(feature = "cuda")] {
32        mod cuda;
33        pub use self::cuda::*;
34        pub use self::cuda::Sha256GpuProverExt as Sha256ProverExt;
35    } else {
36        pub use self::Sha2CpuProverExt as Sha256ProverExt;
37    }
38}
39
40// =================================== VM Extension Implementation =================================
41#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
42pub struct Sha256;
43
44#[derive(Clone, From, AnyEnum, Executor, MeteredExecutor, PreflightExecutor)]
45#[cfg_attr(
46    feature = "aot",
47    derive(
48        openvm_circuit_derive::AotExecutor,
49        openvm_circuit_derive::AotMeteredExecutor
50    )
51)]
52pub enum Sha256Executor {
53    Sha256(Sha256VmExecutor),
54}
55
56impl<F> VmExecutionExtension<F> for Sha256 {
57    type Executor = Sha256Executor;
58
59    fn extend_execution(
60        &self,
61        inventory: &mut ExecutorInventoryBuilder<F, Sha256Executor>,
62    ) -> Result<(), ExecutorInventoryError> {
63        let pointer_max_bits = inventory.pointer_max_bits();
64        let sha256_step = Sha256VmExecutor::new(Rv32Sha256Opcode::CLASS_OFFSET, pointer_max_bits);
65        inventory.add_executor(
66            sha256_step,
67            Rv32Sha256Opcode::iter().map(|x| x.global_opcode()),
68        )?;
69
70        Ok(())
71    }
72}
73
74impl<SC: StarkGenericConfig> VmCircuitExtension<SC> for Sha256 {
75    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
76        let pointer_max_bits = inventory.pointer_max_bits();
77
78        let bitwise_lu = {
79            let existing_air = inventory.find_air::<BitwiseOperationLookupAir<8>>().next();
80            if let Some(air) = existing_air {
81                air.bus
82            } else {
83                let bus = BitwiseOperationLookupBus::new(inventory.new_bus_idx());
84                let air = BitwiseOperationLookupAir::<8>::new(bus);
85                inventory.add_air(air);
86                air.bus
87            }
88        };
89
90        let sha256 = Sha256VmAir::new(
91            inventory.system().port(),
92            bitwise_lu,
93            pointer_max_bits,
94            inventory.new_bus_idx(),
95        );
96        inventory.add_air(sha256);
97
98        Ok(())
99    }
100}
101
102pub struct Sha2CpuProverExt;
103// This implementation is specific to CpuBackend because the lookup chips (VariableRangeChecker,
104// BitwiseOperationLookupChip) are specific to CpuBackend.
105impl<E, SC, RA> VmProverExtension<E, RA, Sha256> for Sha2CpuProverExt
106where
107    SC: StarkGenericConfig,
108    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
109    RA: RowMajorMatrixArena<Val<SC>>,
110    Val<SC>: PrimeField32,
111{
112    fn extend_prover(
113        &self,
114        _: &Sha256,
115        inventory: &mut ChipInventory<SC, RA, CpuBackend<SC>>,
116    ) -> Result<(), ChipInventoryError> {
117        let range_checker = inventory.range_checker()?.clone();
118        let timestamp_max_bits = inventory.timestamp_max_bits();
119        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
120        let pointer_max_bits = inventory.airs().pointer_max_bits();
121
122        let bitwise_lu = {
123            let existing_chip = inventory
124                .find_chip::<SharedBitwiseOperationLookupChip<8>>()
125                .next();
126            if let Some(chip) = existing_chip {
127                chip.clone()
128            } else {
129                let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
130                let chip = Arc::new(BitwiseOperationLookupChip::new(air.bus));
131                inventory.add_periphery_chip(chip.clone());
132                chip
133            }
134        };
135
136        inventory.next_air::<Sha256VmAir>()?;
137        let sha256 = Sha256VmChip::new(
138            Sha256VmFiller::new(bitwise_lu, pointer_max_bits),
139            mem_helper,
140        );
141        inventory.add_executor_chip(sha256);
142
143        Ok(())
144    }
145}
openvm_sha256_circuit/extension/mod.rs

openvm_sha256_circuit/extension/
mod.rs
