openvm_sha256_circuit/
lib.rs

1#![cfg_attr(feature = "tco", allow(incomplete_features))]
2#![cfg_attr(feature = "tco", feature(explicit_tail_calls))]
3
4use std::result::Result;
5
6use openvm_circuit::{
7    arch::{
8        AirInventory, ChipInventoryError, InitFileGenerator, MatrixRecordArena, SystemConfig,
9        VmBuilder, VmChipComplex, VmProverExtension,
10    },
11    system::{SystemChipInventory, SystemCpuBuilder, SystemExecutor},
12};
13use openvm_circuit_derive::VmConfig;
14use openvm_rv32im_circuit::{
15    Rv32I, Rv32IExecutor, Rv32ImCpuProverExt, Rv32Io, Rv32IoExecutor, Rv32M, Rv32MExecutor,
16};
17use openvm_stark_backend::{
18    config::{StarkGenericConfig, Val},
19    p3_field::PrimeField32,
20    prover::cpu::{CpuBackend, CpuDevice},
21};
22use openvm_stark_sdk::engine::StarkEngine;
23use serde::{Deserialize, Serialize};
24
25mod sha256_chip;
26pub use sha256_chip::*;
27
28mod extension;
29pub use extension::*;
30
31cfg_if::cfg_if! {
32    if #[cfg(feature = "cuda")] {
33        use openvm_circuit::arch::DenseRecordArena;
34        use openvm_circuit::system::cuda::{extensions::SystemGpuBuilder, SystemChipInventoryGPU};
35        use openvm_cuda_backend::{engine::GpuBabyBearPoseidon2Engine, prover_backend::GpuBackend};
36        use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
37        use openvm_rv32im_circuit::Rv32ImGpuProverExt;
38        pub(crate) mod cuda_abi;
39        pub use Sha256Rv32GpuBuilder as Sha256Rv32Builder;
40    } else {
41        pub use Sha256Rv32CpuBuilder as Sha256Rv32Builder;
42    }
43}
44
45#[derive(Clone, Debug, VmConfig, derive_new::new, Serialize, Deserialize)]
46pub struct Sha256Rv32Config {
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 sha256: Sha256,
57}
58
59impl Default for Sha256Rv32Config {
60    fn default() -> Self {
61        Self {
62            system: SystemConfig::default(),
63            rv32i: Rv32I,
64            rv32m: Rv32M::default(),
65            io: Rv32Io,
66            sha256: Sha256,
67        }
68    }
69}
70
71// Default implementation uses no init file
72impl InitFileGenerator for Sha256Rv32Config {}
73
74#[derive(Clone)]
75pub struct Sha256Rv32CpuBuilder;
76
77impl<E, SC> VmBuilder<E> for Sha256Rv32CpuBuilder
78where
79    SC: StarkGenericConfig,
80    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
81    Val<SC>: PrimeField32,
82{
83    type VmConfig = Sha256Rv32Config;
84    type SystemChipInventory = SystemChipInventory<SC>;
85    type RecordArena = MatrixRecordArena<Val<SC>>;
86
87    fn create_chip_complex(
88        &self,
89        config: &Sha256Rv32Config,
90        circuit: AirInventory<SC>,
91    ) -> Result<
92        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
93        ChipInventoryError,
94    > {
95        let mut chip_complex =
96            VmBuilder::<E>::create_chip_complex(&SystemCpuBuilder, &config.system, circuit)?;
97        let inventory = &mut chip_complex.inventory;
98        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32i, inventory)?;
99        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32m, inventory)?;
100        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.io, inventory)?;
101        VmProverExtension::<E, _, _>::extend_prover(&Sha2CpuProverExt, &config.sha256, inventory)?;
102        Ok(chip_complex)
103    }
104}
105
106#[cfg(feature = "cuda")]
107#[derive(Clone)]
108pub struct Sha256Rv32GpuBuilder;
109
110#[cfg(feature = "cuda")]
111impl VmBuilder<GpuBabyBearPoseidon2Engine> for Sha256Rv32GpuBuilder {
112    type VmConfig = Sha256Rv32Config;
113    type SystemChipInventory = SystemChipInventoryGPU;
114    type RecordArena = DenseRecordArena;
115
116    fn create_chip_complex(
117        &self,
118        config: &Sha256Rv32Config,
119        circuit: AirInventory<BabyBearPoseidon2Config>,
120    ) -> Result<
121        VmChipComplex<
122            BabyBearPoseidon2Config,
123            Self::RecordArena,
124            GpuBackend,
125            Self::SystemChipInventory,
126        >,
127        ChipInventoryError,
128    > {
129        let mut chip_complex = VmBuilder::<GpuBabyBearPoseidon2Engine>::create_chip_complex(
130            &SystemGpuBuilder,
131            &config.system,
132            circuit,
133        )?;
134        let inventory = &mut chip_complex.inventory;
135        VmProverExtension::<GpuBabyBearPoseidon2Engine, _, _>::extend_prover(
136            &Rv32ImGpuProverExt,
137            &config.rv32i,
138            inventory,
139        )?;
140        VmProverExtension::<GpuBabyBearPoseidon2Engine, _, _>::extend_prover(
141            &Rv32ImGpuProverExt,
142            &config.rv32m,
143            inventory,
144        )?;
145        VmProverExtension::<GpuBabyBearPoseidon2Engine, _, _>::extend_prover(
146            &Rv32ImGpuProverExt,
147            &config.io,
148            inventory,
149        )?;
150        VmProverExtension::<GpuBabyBearPoseidon2Engine, _, _>::extend_prover(
151            &Sha256GpuProverExt,
152            &config.sha256,
153            inventory,
154        )?;
155        Ok(chip_complex)
156    }
157}
openvm_sha256_circuit/lib.rs

openvm_sha256_circuit/
lib.rs
