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)]
123pub enum Keccak256Executor {
124    Keccak256(KeccakVmExecutor),
125}
126
127impl<F> VmExecutionExtension<F> for Keccak256 {
128    type Executor = Keccak256Executor;
129
130    fn extend_execution(
131        &self,
132        inventory: &mut ExecutorInventoryBuilder<F, Keccak256Executor>,
133    ) -> Result<(), ExecutorInventoryError> {
134        let pointer_max_bits = inventory.pointer_max_bits();
135        let keccak_step = KeccakVmExecutor::new(Rv32KeccakOpcode::CLASS_OFFSET, pointer_max_bits);
136        inventory.add_executor(
137            keccak_step,
138            Rv32KeccakOpcode::iter().map(|x| x.global_opcode()),
139        )?;
140
141        Ok(())
142    }
143}
144
145impl<SC: StarkGenericConfig> VmCircuitExtension<SC> for Keccak256 {
146    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
147        let SystemPort {
148            execution_bus,
149            program_bus,
150            memory_bridge,
151        } = inventory.system().port();
152
153        let exec_bridge = ExecutionBridge::new(execution_bus, program_bus);
154        let pointer_max_bits = inventory.pointer_max_bits();
155
156        let bitwise_lu = {
157            let existing_air = inventory.find_air::<BitwiseOperationLookupAir<8>>().next();
158            if let Some(air) = existing_air {
159                air.bus
160            } else {
161                let bus = BitwiseOperationLookupBus::new(inventory.new_bus_idx());
162                let air = BitwiseOperationLookupAir::<8>::new(bus);
163                inventory.add_air(air);
164                air.bus
165            }
166        };
167
168        let keccak = KeccakVmAir::new(
169            exec_bridge,
170            memory_bridge,
171            bitwise_lu,
172            pointer_max_bits,
173            Rv32KeccakOpcode::CLASS_OFFSET,
174        );
175        inventory.add_air(keccak);
176
177        Ok(())
178    }
179}
180
181pub struct Keccak256CpuProverExt;
182// This implementation is specific to CpuBackend because the lookup chips (VariableRangeChecker,
183// BitwiseOperationLookupChip) are specific to CpuBackend.
184impl<E, SC, RA> VmProverExtension<E, RA, Keccak256> for Keccak256CpuProverExt
185where
186    SC: StarkGenericConfig,
187    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
188    RA: RowMajorMatrixArena<Val<SC>>,
189    Val<SC>: PrimeField32,
190{
191    fn extend_prover(
192        &self,
193        _: &Keccak256,
194        inventory: &mut ChipInventory<SC, RA, CpuBackend<SC>>,
195    ) -> Result<(), ChipInventoryError> {
196        let range_checker = inventory.range_checker()?.clone();
197        let timestamp_max_bits = inventory.timestamp_max_bits();
198        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
199        let pointer_max_bits = inventory.airs().pointer_max_bits();
200
201        let bitwise_lu = {
202            let existing_chip = inventory
203                .find_chip::<SharedBitwiseOperationLookupChip<8>>()
204                .next();
205            if let Some(chip) = existing_chip {
206                chip.clone()
207            } else {
208                let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
209                let chip = Arc::new(BitwiseOperationLookupChip::new(air.bus));
210                inventory.add_periphery_chip(chip.clone());
211                chip
212            }
213        };
214
215        inventory.next_air::<KeccakVmAir>()?;
216        let keccak = KeccakVmChip::new(
217            KeccakVmFiller::new(bitwise_lu, pointer_max_bits),
218            mem_helper,
219        );
220        inventory.add_executor_chip(keccak);
221
222        Ok(())
223    }
224}
openvm_keccak256_circuit/extension/mod.rs

openvm_keccak256_circuit/extension/
mod.rs
