openvm_sha256_circuit/
lib.rs

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

openvm_sha256_circuit/
lib.rs
