openvm_keccak256_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, ExecutionBridge,
7        ExecutorInventoryBuilder, ExecutorInventoryError, InitFileGenerator, MatrixRecordArena,
8        RowMajorMatrixArena, SystemConfig, VmBuilder, VmChipComplex, VmCircuitExtension,
9        VmExecutionExtension, VmProverExtension,
10    },
11    system::{
12        memory::SharedMemoryHelper, SystemChipInventory, SystemCpuBuilder, SystemExecutor,
13        SystemPort,
14    },
15};
16use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor, VmConfig};
17use openvm_circuit_primitives::bitwise_op_lookup::{
18    BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip,
19    SharedBitwiseOperationLookupChip,
20};
21use openvm_instructions::*;
22use openvm_keccak256_transpiler::Rv32KeccakOpcode;
23use openvm_rv32im_circuit::{
24    Rv32I, Rv32IExecutor, Rv32ImCpuProverExt, Rv32Io, Rv32IoExecutor, Rv32M, Rv32MExecutor,
25};
26use openvm_stark_backend::{
27    config::{StarkGenericConfig, Val},
28    p3_field::PrimeField32,
29    prover::cpu::{CpuBackend, CpuDevice},
30};
31use openvm_stark_sdk::engine::StarkEngine;
32use serde::{Deserialize, Serialize};
33use strum::IntoEnumIterator;
34
35use crate::{KeccakVmAir, KeccakVmChip, KeccakVmExecutor, KeccakVmFiller};
36
37cfg_if::cfg_if! {
38    if #[cfg(feature = "cuda")] {
39        mod cuda;
40        pub use cuda::*;
41        pub use cuda::{
42            Keccak256GpuProverExt as Keccak256ProverExt,
43            Keccak256Rv32GpuBuilder as Keccak256Rv32Builder,
44        };
45    } else {
46        pub use self::{
47            Keccak256CpuProverExt as Keccak256ProverExt,
48            Keccak256Rv32CpuBuilder as Keccak256Rv32Builder,
49        };
50    }
51}
52
53#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)]
54pub struct Keccak256Rv32Config {
55    #[config(executor = "SystemExecutor<F>")]
56    pub system: SystemConfig,
57    #[extension]
58    pub rv32i: Rv32I,
59    #[extension]
60    pub rv32m: Rv32M,
61    #[extension]
62    pub io: Rv32Io,
63    #[extension]
64    pub keccak: Keccak256,
65}
66
67impl Default for Keccak256Rv32Config {
68    fn default() -> Self {
69        Self {
70            system: SystemConfig::default(),
71            rv32i: Rv32I,
72            rv32m: Rv32M::default(),
73            io: Rv32Io,
74            keccak: Keccak256,
75        }
76    }
77}
78
79// Default implementation uses no init file
80impl InitFileGenerator for Keccak256Rv32Config {}
81
82#[derive(Clone)]
83pub struct Keccak256Rv32CpuBuilder;
84
85impl<E, SC> VmBuilder<E> for Keccak256Rv32CpuBuilder
86where
87    SC: StarkGenericConfig,
88    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
89    Val<SC>: PrimeField32,
90{
91    type VmConfig = Keccak256Rv32Config;
92    type SystemChipInventory = SystemChipInventory<SC>;
93    type RecordArena = MatrixRecordArena<Val<SC>>;
94
95    fn create_chip_complex(
96        &self,
97        config: &Keccak256Rv32Config,
98        circuit: AirInventory<SC>,
99    ) -> Result<
100        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
101        ChipInventoryError,
102    > {
103        let mut chip_complex =
104            VmBuilder::<E>::create_chip_complex(&SystemCpuBuilder, &config.system, circuit)?;
105        let inventory = &mut chip_complex.inventory;
106        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32i, inventory)?;
107        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32m, inventory)?;
108        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.io, inventory)?;
109        VmProverExtension::<E, _, _>::extend_prover(
110            &Keccak256CpuProverExt,
111            &config.keccak,
112            inventory,
113        )?;
114        Ok(chip_complex)
115    }
116}
117
118// =================================== VM Extension Implementation =================================
119#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
120pub struct Keccak256;
121
122#[derive(Clone, Copy, From, AnyEnum, Executor, MeteredExecutor, PreflightExecutor)]
123#[cfg_attr(
124    feature = "aot",
125    derive(
126        openvm_circuit_derive::AotExecutor,
127        openvm_circuit_derive::AotMeteredExecutor
128    )
129)]
130pub enum Keccak256Executor {
131    Keccak256(KeccakVmExecutor),
132}
133
134impl<F> VmExecutionExtension<F> for Keccak256 {
135    type Executor = Keccak256Executor;
136
137    fn extend_execution(
138        &self,
139        inventory: &mut ExecutorInventoryBuilder<F, Keccak256Executor>,
140    ) -> Result<(), ExecutorInventoryError> {
141        let pointer_max_bits = inventory.pointer_max_bits();
142        let keccak_step = KeccakVmExecutor::new(Rv32KeccakOpcode::CLASS_OFFSET, pointer_max_bits);
143        inventory.add_executor(
144            keccak_step,
145            Rv32KeccakOpcode::iter().map(|x| x.global_opcode()),
146        )?;
147
148        Ok(())
149    }
150}
151
152impl<SC: StarkGenericConfig> VmCircuitExtension<SC> for Keccak256 {
153    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
154        let SystemPort {
155            execution_bus,
156            program_bus,
157            memory_bridge,
158        } = inventory.system().port();
159
160        let exec_bridge = ExecutionBridge::new(execution_bus, program_bus);
161        let pointer_max_bits = inventory.pointer_max_bits();
162
163        let bitwise_lu = {
164            let existing_air = inventory.find_air::<BitwiseOperationLookupAir<8>>().next();
165            if let Some(air) = existing_air {
166                air.bus
167            } else {
168                let bus = BitwiseOperationLookupBus::new(inventory.new_bus_idx());
169                let air = BitwiseOperationLookupAir::<8>::new(bus);
170                inventory.add_air(air);
171                air.bus
172            }
173        };
174
175        let keccak = KeccakVmAir::new(
176            exec_bridge,
177            memory_bridge,
178            bitwise_lu,
179            pointer_max_bits,
180            Rv32KeccakOpcode::CLASS_OFFSET,
181        );
182        inventory.add_air(keccak);
183
184        Ok(())
185    }
186}
187
188pub struct Keccak256CpuProverExt;
189// This implementation is specific to CpuBackend because the lookup chips (VariableRangeChecker,
190// BitwiseOperationLookupChip) are specific to CpuBackend.
191impl<E, SC, RA> VmProverExtension<E, RA, Keccak256> for Keccak256CpuProverExt
192where
193    SC: StarkGenericConfig,
194    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
195    RA: RowMajorMatrixArena<Val<SC>>,
196    Val<SC>: PrimeField32,
197{
198    fn extend_prover(
199        &self,
200        _: &Keccak256,
201        inventory: &mut ChipInventory<SC, RA, CpuBackend<SC>>,
202    ) -> Result<(), ChipInventoryError> {
203        let range_checker = inventory.range_checker()?.clone();
204        let timestamp_max_bits = inventory.timestamp_max_bits();
205        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
206        let pointer_max_bits = inventory.airs().pointer_max_bits();
207
208        let bitwise_lu = {
209            let existing_chip = inventory
210                .find_chip::<SharedBitwiseOperationLookupChip<8>>()
211                .next();
212            if let Some(chip) = existing_chip {
213                chip.clone()
214            } else {
215                let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
216                let chip = Arc::new(BitwiseOperationLookupChip::new(air.bus));
217                inventory.add_periphery_chip(chip.clone());
218                chip
219            }
220        };
221
222        inventory.next_air::<KeccakVmAir>()?;
223        let keccak = KeccakVmChip::new(
224            KeccakVmFiller::new(bitwise_lu, pointer_max_bits),
225            mem_helper,
226        );
227        inventory.add_executor_chip(keccak);
228
229        Ok(())
230    }
231}
openvm_keccak256_circuit/extension/mod.rs

openvm_keccak256_circuit/extension/
mod.rs
