openvm_keccak256_circuit/extension/
mod.rs

1use std::{
2    result::Result,
3    sync::{Arc, Mutex},
4};
5
6use derive_more::derive::From;
7use openvm_circuit::{
8    arch::{
9        AirInventory, AirInventoryError, ChipInventory, ChipInventoryError, ExecutionBridge,
10        ExecutorInventoryBuilder, ExecutorInventoryError, InitFileGenerator, MatrixRecordArena,
11        RowMajorMatrixArena, SystemConfig, VmBuilder, VmChipComplex, VmCircuitExtension,
12        VmExecutionExtension, VmField, VmProverExtension,
13    },
14    system::{
15        memory::SharedMemoryHelper, SystemChipInventory, SystemCpuBuilder, SystemExecutor,
16        SystemPort,
17    },
18};
19use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor, VmConfig};
20use openvm_circuit_primitives::bitwise_op_lookup::{
21    BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip,
22    SharedBitwiseOperationLookupChip,
23};
24use openvm_cpu_backend::{CpuBackend, CpuDevice};
25use openvm_instructions::*;
26use openvm_keccak256_transpiler::{KeccakfOpcode, XorinOpcode};
27use openvm_rv32im_circuit::{
28    Rv32I, Rv32IExecutor, Rv32ImCpuProverExt, Rv32Io, Rv32IoExecutor, Rv32M, Rv32MExecutor,
29};
30use openvm_stark_backend::{
31    interaction::PermutationCheckBus, p3_field::PrimeField32, StarkEngine, StarkProtocolConfig, Val,
32};
33use serde::{Deserialize, Serialize};
34use strum::IntoEnumIterator;
35
36use crate::{
37    keccakf_op::{KeccakfExecutor, KeccakfOpAir, KeccakfOpChip},
38    keccakf_perm::{KeccakfPermAir, KeccakfPermChip},
39    xorin::{air::XorinVmAir, XorinVmChip, XorinVmExecutor, XorinVmFiller},
40};
41
42#[cfg(feature = "cuda")]
43mod cuda;
44#[cfg(feature = "cuda")]
45pub use cuda::*;
46
47#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)]
48pub struct Keccak256Rv32Config {
49    #[config(executor = "SystemExecutor<F>")]
50    pub system: SystemConfig,
51    #[extension]
52    pub rv32i: Rv32I,
53    #[extension]
54    pub rv32m: Rv32M,
55    #[extension]
56    pub io: Rv32Io,
57    #[extension]
58    pub keccak: Keccak256,
59}
60
61impl Default for Keccak256Rv32Config {
62    fn default() -> Self {
63        Self {
64            system: SystemConfig::default(),
65            rv32i: Rv32I,
66            rv32m: Rv32M::default(),
67            io: Rv32Io,
68            keccak: Keccak256,
69        }
70    }
71}
72
73// Default implementation uses no init file
74impl InitFileGenerator for Keccak256Rv32Config {}
75
76#[derive(Clone)]
77pub struct Keccak256Rv32CpuBuilder;
78
79impl<SC, E> VmBuilder<E> for Keccak256Rv32CpuBuilder
80where
81    SC: StarkProtocolConfig,
82    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
83    Val<SC>: VmField,
84    SC::EF: Ord,
85{
86    type VmConfig = Keccak256Rv32Config;
87    type SystemChipInventory = SystemChipInventory<SC>;
88    type RecordArena = MatrixRecordArena<Val<SC>>;
89
90    fn create_chip_complex(
91        &self,
92        config: &Keccak256Rv32Config,
93        circuit: AirInventory<SC>,
94        device_ctx: &openvm_stark_backend::EngineDeviceCtx<E>,
95    ) -> Result<
96        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
97        ChipInventoryError,
98    > {
99        let mut chip_complex = VmBuilder::<E>::create_chip_complex(
100            &SystemCpuBuilder,
101            &config.system,
102            circuit,
103            device_ctx,
104        )?;
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    Keccakf(KeccakfExecutor),
132    Xorin(XorinVmExecutor),
133}
134
135impl<F> VmExecutionExtension<F> for Keccak256 {
136    type Executor = Keccak256Executor;
137
138    fn extend_execution(
139        &self,
140        inventory: &mut ExecutorInventoryBuilder<F, Keccak256Executor>,
141    ) -> Result<(), ExecutorInventoryError> {
142        let pointer_max_bits = inventory.pointer_max_bits();
143
144        let xorin_executor = XorinVmExecutor::new(XorinOpcode::CLASS_OFFSET, pointer_max_bits);
145        inventory.add_executor(
146            xorin_executor,
147            XorinOpcode::iter().map(|x| x.global_opcode()),
148        )?;
149
150        let keccak_executor = KeccakfExecutor::new(KeccakfOpcode::CLASS_OFFSET, pointer_max_bits);
151        inventory.add_executor(
152            keccak_executor,
153            KeccakfOpcode::iter().map(|x| x.global_opcode()),
154        )?;
155
156        Ok(())
157    }
158}
159
160impl<SC: StarkProtocolConfig> VmCircuitExtension<SC> for Keccak256 {
161    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
162        let SystemPort {
163            execution_bus,
164            program_bus,
165            memory_bridge,
166        } = inventory.system().port();
167
168        let exec_bridge = ExecutionBridge::new(execution_bus, program_bus);
169        let pointer_max_bits = inventory.pointer_max_bits();
170
171        let bitwise_lu = {
172            let existing_air = inventory.find_air::<BitwiseOperationLookupAir<8>>().next();
173            if let Some(air) = existing_air {
174                air.bus
175            } else {
176                let bus = BitwiseOperationLookupBus::new(inventory.new_bus_idx());
177                let air = BitwiseOperationLookupAir::<8>::new(bus);
178                inventory.add_air(air);
179                air.bus
180            }
181        };
182
183        let xorin_air = XorinVmAir::new(
184            exec_bridge,
185            memory_bridge,
186            bitwise_lu,
187            pointer_max_bits,
188            XorinOpcode::CLASS_OFFSET,
189        );
190        inventory.add_air(xorin_air);
191
192        let keccakf_state_bus = PermutationCheckBus::new(inventory.new_bus_idx());
193        let periphery_air = KeccakfPermAir::new(keccakf_state_bus);
194        inventory.add_air(periphery_air);
195
196        let op_air = KeccakfOpAir::new(
197            exec_bridge,
198            memory_bridge,
199            bitwise_lu,
200            keccakf_state_bus,
201            pointer_max_bits,
202            KeccakfOpcode::CLASS_OFFSET,
203        );
204        inventory.add_air(op_air);
205
206        Ok(())
207    }
208}
209
210pub struct Keccak256CpuProverExt;
211// This implementation is specific to CpuBackend because the lookup chips (VariableRangeChecker,
212// BitwiseOperationLookupChip) are specific to CpuBackend.
213impl<SC, E, RA> VmProverExtension<E, RA, Keccak256> for Keccak256CpuProverExt
214where
215    SC: StarkProtocolConfig,
216    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
217    RA: RowMajorMatrixArena<Val<SC>>,
218    Val<SC>: PrimeField32,
219    SC::EF: Ord,
220{
221    fn extend_prover(
222        &self,
223        _: &Keccak256,
224        inventory: &mut ChipInventory<SC, RA, CpuBackend<SC>>,
225    ) -> Result<(), ChipInventoryError> {
226        let range_checker = inventory.range_checker()?.clone();
227        let timestamp_max_bits = inventory.timestamp_max_bits();
228        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
229        let pointer_max_bits = inventory.airs().pointer_max_bits();
230
231        let bitwise_lu = {
232            let existing_chip = inventory
233                .find_chip::<SharedBitwiseOperationLookupChip<8>>()
234                .next();
235
236            if let Some(chip) = existing_chip {
237                chip.clone()
238            } else {
239                let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
240                let chip = Arc::new(BitwiseOperationLookupChip::new(air.bus));
241                inventory.add_periphery_chip(chip.clone());
242                chip
243            }
244        };
245
246        inventory.next_air::<XorinVmAir>()?;
247        let xorin_chip = XorinVmChip::new(
248            XorinVmFiller::new(bitwise_lu.clone(), pointer_max_bits),
249            mem_helper.clone(),
250        );
251        inventory.add_executor_chip(xorin_chip);
252
253        inventory.next_air::<KeccakfPermAir>()?;
254        let shared_records = Arc::new(Mutex::new(Vec::new()));
255        let periphery_chip = KeccakfPermChip::new(shared_records.clone());
256        // Clone the Arc Mutex and pass to the OpChip.
257        // WARNING: the OpChip must be added _after_ the periphery chip so that its tracegen is done
258        // _first_. After OpChip tracegen, the shared_record is set to the execution records,
259        // effectively passing the records to the periphery chip.
260        inventory.add_periphery_chip(periphery_chip);
261
262        inventory.next_air::<KeccakfOpAir>()?;
263        let op_chip = KeccakfOpChip::new(
264            bitwise_lu,
265            pointer_max_bits,
266            mem_helper.clone(),
267            shared_records,
268        );
269        inventory.add_executor_chip(op_chip);
270
271        Ok(())
272    }
273}
openvm_keccak256_circuit/extension/mod.rs

openvm_keccak256_circuit/extension/
mod.rs
