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

openvm_sha256_circuit/
lib.rs
