openvm_sha2_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,
10        ExecutorInventoryBuilder, ExecutorInventoryError, InitFileGenerator, MatrixRecordArena,
11        RowMajorMatrixArena, SystemConfig, VmBuilder, VmChipComplex, VmCircuitExtension,
12        VmExecutionExtension, VmField, VmProverExtension,
13    },
14    system::{memory::SharedMemoryHelper, SystemChipInventory, SystemCpuBuilder, SystemExecutor},
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_cpu_backend::{CpuBackend, CpuDevice};
22use openvm_instructions::LocalOpcode;
23use openvm_rv32im_circuit::{
24    Rv32I, Rv32IExecutor, Rv32ImCpuProverExt, Rv32Io, Rv32IoExecutor, Rv32M, Rv32MExecutor,
25};
26use openvm_sha2_air::{Sha256Config, Sha512Config};
27use openvm_sha2_transpiler::Rv32Sha2Opcode;
28use openvm_stark_backend::{StarkEngine, StarkProtocolConfig, Val};
29use serde::{Deserialize, Serialize};
30
31use crate::{Sha2BlockHasherChip, Sha2BlockHasherVmAir, Sha2MainAir, Sha2MainChip, Sha2VmExecutor};
32
33cfg_if::cfg_if! {
34    if #[cfg(feature = "cuda")] {
35        mod cuda;
36        pub use self::cuda::*;
37        pub use self::cuda::Sha2GpuProverExt as Sha2ProverExt;
38        pub use self::cuda::Sha2Rv32GpuBuilder as Sha2Rv32Builder;
39    } else {
40        pub use self::Sha2CpuProverExt as Sha2ProverExt;
41        pub use self::Sha2Rv32CpuBuilder as Sha2Rv32Builder;
42    }
43}
44
45#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)]
46pub struct Sha2Rv32Config {
47    #[config(executor = "SystemExecutor<F>")]
48    pub system: SystemConfig,
49    #[extension]
50    pub rv32i: Rv32I,
51    #[extension]
52    pub rv32m: Rv32M,
53    #[extension]
54    pub io: Rv32Io,
55    #[extension]
56    pub sha2: Sha2,
57}
58
59impl Default for Sha2Rv32Config {
60    fn default() -> Self {
61        Self {
62            system: SystemConfig::default(),
63            rv32i: Rv32I,
64            rv32m: Rv32M::default(),
65            io: Rv32Io,
66            sha2: Sha2,
67        }
68    }
69}
70
71// Default implementation uses no init file
72impl InitFileGenerator for Sha2Rv32Config {}
73
74#[derive(Clone)]
75pub struct Sha2Rv32CpuBuilder;
76
77impl<E, SC> VmBuilder<E> for Sha2Rv32CpuBuilder
78where
79    SC: StarkProtocolConfig,
80    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
81    Val<SC>: VmField,
82    SC::EF: Ord,
83{
84    type VmConfig = Sha2Rv32Config;
85    type SystemChipInventory = SystemChipInventory<SC>;
86    type RecordArena = MatrixRecordArena<Val<SC>>;
87
88    fn create_chip_complex(
89        &self,
90        config: &Sha2Rv32Config,
91        circuit: AirInventory<SC>,
92        device_ctx: &openvm_stark_backend::EngineDeviceCtx<E>,
93    ) -> Result<
94        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
95        ChipInventoryError,
96    > {
97        let mut chip_complex = VmBuilder::<E>::create_chip_complex(
98            &SystemCpuBuilder,
99            &config.system,
100            circuit,
101            device_ctx,
102        )?;
103        let inventory = &mut chip_complex.inventory;
104        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32i, inventory)?;
105        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32m, inventory)?;
106        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.io, inventory)?;
107        VmProverExtension::<E, _, _>::extend_prover(&Sha2CpuProverExt, &config.sha2, inventory)?;
108        Ok(chip_complex)
109    }
110}
111
112// =================================== VM Extension Implementation =================================
113#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize)]
114pub struct Sha2;
115
116#[derive(Clone, From, AnyEnum, Executor, MeteredExecutor, PreflightExecutor)]
117#[cfg_attr(
118    feature = "aot",
119    derive(
120        openvm_circuit_derive::AotExecutor,
121        openvm_circuit_derive::AotMeteredExecutor
122    )
123)]
124pub enum Sha2Executor {
125    Sha256(Sha2VmExecutor<Sha256Config>),
126    Sha512(Sha2VmExecutor<Sha512Config>),
127}
128
129impl<F> VmExecutionExtension<F> for Sha2 {
130    type Executor = Sha2Executor;
131
132    fn extend_execution(
133        &self,
134        inventory: &mut ExecutorInventoryBuilder<F, Sha2Executor>,
135    ) -> Result<(), ExecutorInventoryError> {
136        let pointer_max_bits = inventory.pointer_max_bits();
137
138        let sha256_executor =
139            Sha2VmExecutor::<Sha256Config>::new(Rv32Sha2Opcode::CLASS_OFFSET, pointer_max_bits);
140        inventory.add_executor(sha256_executor, [Rv32Sha2Opcode::SHA256.global_opcode()])?;
141
142        let sha512_executor =
143            Sha2VmExecutor::<Sha512Config>::new(Rv32Sha2Opcode::CLASS_OFFSET, pointer_max_bits);
144        inventory.add_executor(sha512_executor, [Rv32Sha2Opcode::SHA512.global_opcode()])?;
145
146        Ok(())
147    }
148}
149
150impl<SC: StarkProtocolConfig> VmCircuitExtension<SC> for Sha2 {
151    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
152        let bitwise_lu = {
153            let existing_air = inventory.find_air::<BitwiseOperationLookupAir<8>>().next();
154            if let Some(air) = existing_air {
155                air.bus
156            } else {
157                let bus = BitwiseOperationLookupBus::new(inventory.new_bus_idx());
158                let air = BitwiseOperationLookupAir::<8>::new(bus);
159                inventory.add_air(air);
160                air.bus
161            }
162        };
163
164        // this bus will be used for communication between the block hasher chip and the main chip
165        let sha2_bus_index = inventory.new_bus_idx();
166        // the sha2 subair needs its own bus for self-interactions
167        let subair_bus_index = inventory.new_bus_idx();
168
169        // SHA-256
170        let sha256_block_hasher_air =
171            Sha2BlockHasherVmAir::<Sha256Config>::new(bitwise_lu, subair_bus_index, sha2_bus_index);
172        inventory.add_air(sha256_block_hasher_air);
173
174        let sha256_main_air = Sha2MainAir::<Sha256Config>::new(
175            inventory.system().port(),
176            bitwise_lu,
177            inventory.pointer_max_bits(),
178            sha2_bus_index,
179            Rv32Sha2Opcode::CLASS_OFFSET,
180        );
181        inventory.add_air(sha256_main_air);
182
183        // SHA-512
184        let sha512_block_hasher_air =
185            Sha2BlockHasherVmAir::<Sha512Config>::new(bitwise_lu, subair_bus_index, sha2_bus_index);
186        inventory.add_air(sha512_block_hasher_air);
187
188        let sha512_main_air = Sha2MainAir::<Sha512Config>::new(
189            inventory.system().port(),
190            bitwise_lu,
191            inventory.pointer_max_bits(),
192            sha2_bus_index,
193            Rv32Sha2Opcode::CLASS_OFFSET,
194        );
195        inventory.add_air(sha512_main_air);
196
197        Ok(())
198    }
199}
200
201pub struct Sha2CpuProverExt;
202// This implementation is specific to CpuBackend because the lookup chips (VariableRangeChecker,
203// BitwiseOperationLookupChip) are specific to CpuBackend.
204impl<E, SC, RA> VmProverExtension<E, RA, Sha2> for Sha2CpuProverExt
205where
206    SC: StarkProtocolConfig,
207    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
208    RA: RowMajorMatrixArena<Val<SC>> + Send + Sync + 'static,
209    Val<SC>: VmField,
210    SC::EF: Ord,
211{
212    fn extend_prover(
213        &self,
214        _: &Sha2,
215        inventory: &mut ChipInventory<SC, RA, CpuBackend<SC>>,
216    ) -> Result<(), ChipInventoryError> {
217        let range_checker = inventory.range_checker()?.clone();
218        let timestamp_max_bits = inventory.timestamp_max_bits();
219        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
220        let pointer_max_bits = inventory.airs().pointer_max_bits();
221
222        let bitwise_lu = {
223            let existing_chip = inventory
224                .find_chip::<SharedBitwiseOperationLookupChip<8>>()
225                .next();
226            if let Some(chip) = existing_chip {
227                chip.clone()
228            } else {
229                let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
230                let chip = Arc::new(BitwiseOperationLookupChip::new(air.bus));
231                inventory.add_periphery_chip(chip.clone());
232                chip
233            }
234        };
235
236        // We must add each block hasher chip before the main chip to ensure that main chip does its
237        // tracegen first, because the main chip will pass the records to the block hasher chip
238        // after its tracegen is done.
239
240        // SHA-256
241        inventory.next_air::<Sha2BlockHasherVmAir<Sha256Config>>()?;
242        // shared records between the main chip and the block hasher chip
243        let records = Arc::new(Mutex::new(None));
244        let sha256_block_hasher_chip = Sha2BlockHasherChip::<Val<SC>, Sha256Config>::new(
245            bitwise_lu.clone(),
246            pointer_max_bits,
247            mem_helper.clone(),
248            records.clone(),
249        );
250        inventory.add_periphery_chip(sha256_block_hasher_chip);
251
252        inventory.next_air::<Sha2MainAir<Sha256Config>>()?;
253        let sha256_main_chip = Sha2MainChip::<Val<SC>, Sha256Config>::new(
254            records,
255            bitwise_lu.clone(),
256            pointer_max_bits,
257            mem_helper.clone(),
258        );
259        inventory.add_executor_chip(sha256_main_chip);
260
261        // SHA-512
262        inventory.next_air::<Sha2BlockHasherVmAir<Sha512Config>>()?;
263        // shared records between the main chip and the block hasher chip
264        let records = Arc::new(Mutex::new(None));
265        let sha512_block_hasher_chip = Sha2BlockHasherChip::<Val<SC>, Sha512Config>::new(
266            bitwise_lu.clone(),
267            pointer_max_bits,
268            mem_helper.clone(),
269            records.clone(),
270        );
271        inventory.add_periphery_chip(sha512_block_hasher_chip);
272
273        inventory.next_air::<Sha2MainAir<Sha512Config>>()?;
274        let sha512_main_chip = Sha2MainChip::<Val<SC>, Sha512Config>::new(
275            records,
276            bitwise_lu.clone(),
277            pointer_max_bits,
278            mem_helper.clone(),
279        );
280        inventory.add_executor_chip(sha512_main_chip);
281
282        Ok(())
283    }
284}
openvm_sha2_circuit/extension/mod.rs

openvm_sha2_circuit/extension/
mod.rs
