openvm_rv32im_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))]
5use openvm_circuit::{
6    arch::{
7        AirInventory, ChipInventoryError, InitFileGenerator, MatrixRecordArena, SystemConfig,
8        VmBuilder, VmChipComplex, VmField, VmProverExtension,
9    },
10    system::{SystemChipInventory, SystemCpuBuilder, SystemExecutor},
11};
12use openvm_circuit_derive::{Executor, MeteredExecutor, PreflightExecutor, VmConfig};
13use openvm_stark_backend::{
14    config::{StarkGenericConfig, Val},
15    engine::StarkEngine,
16    prover::cpu::{CpuBackend, CpuDevice},
17};
18use serde::{Deserialize, Serialize};
19
20pub mod adapters;
21mod auipc;
22mod base_alu;
23mod branch_eq;
24mod branch_lt;
25pub mod common;
26mod divrem;
27mod hintstore;
28mod jal_lui;
29mod jalr;
30mod less_than;
31mod load_sign_extend;
32mod loadstore;
33mod mul;
34mod mulh;
35mod shift;
36
37pub use auipc::*;
38pub use base_alu::*;
39pub use branch_eq::*;
40pub use branch_lt::*;
41pub use divrem::*;
42pub use hintstore::*;
43pub use jal_lui::*;
44pub use jalr::*;
45pub use less_than::*;
46pub use load_sign_extend::*;
47pub use loadstore::*;
48pub use mul::*;
49pub use mulh::*;
50pub use shift::*;
51
52mod extension;
53pub use extension::*;
54
55cfg_if::cfg_if! {
56    if #[cfg(feature = "cuda")] {
57        use openvm_circuit::arch::DenseRecordArena;
58        use openvm_circuit::system::cuda::{extensions::SystemGpuBuilder, SystemChipInventoryGPU};
59        use openvm_cuda_backend::{engine::GpuBabyBearPoseidon2Engine, prover_backend::GpuBackend};
60        use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Config;
61        pub(crate) mod cuda_abi;
62        pub use self::{
63            Rv32IGpuBuilder as Rv32IBuilder,
64            Rv32ImGpuBuilder as Rv32ImBuilder,
65        };
66    } else {
67        pub use self::{
68            Rv32ICpuBuilder as Rv32IBuilder,
69            Rv32ImCpuBuilder as Rv32ImBuilder,
70        };
71    }
72}
73
74#[cfg(any(test, feature = "test-utils"))]
75mod test_utils;
76
77// Config for a VM with base extension and IO extension
78#[derive(Clone, Debug, derive_new::new, VmConfig, Serialize, Deserialize)]
79pub struct Rv32IConfig {
80    #[config(executor = "SystemExecutor<F>")]
81    pub system: SystemConfig,
82    #[extension]
83    pub base: Rv32I,
84    #[extension]
85    pub io: Rv32Io,
86}
87
88// Default implementation uses no init file
89impl InitFileGenerator for Rv32IConfig {}
90
91/// Config for a VM with base extension, IO extension, and multiplication extension
92#[derive(Clone, Debug, Default, VmConfig, derive_new::new, Serialize, Deserialize)]
93pub struct Rv32ImConfig {
94    #[config]
95    pub rv32i: Rv32IConfig,
96    #[extension]
97    pub mul: Rv32M,
98}
99
100// Default implementation uses no init file
101impl InitFileGenerator for Rv32ImConfig {}
102
103impl Default for Rv32IConfig {
104    fn default() -> Self {
105        let system = SystemConfig::default();
106        Self {
107            system,
108            base: Default::default(),
109            io: Default::default(),
110        }
111    }
112}
113
114impl Rv32IConfig {
115    pub fn with_public_values(public_values: usize) -> Self {
116        let system = SystemConfig::default().with_public_values(public_values);
117        Self {
118            system,
119            base: Default::default(),
120            io: Default::default(),
121        }
122    }
123
124    pub fn with_public_values_and_segment_len(public_values: usize, segment_len: usize) -> Self {
125        let system = SystemConfig::default()
126            .with_public_values(public_values)
127            .with_max_segment_len(segment_len);
128        Self {
129            system,
130            base: Default::default(),
131            io: Default::default(),
132        }
133    }
134}
135
136impl Rv32ImConfig {
137    pub fn with_public_values(public_values: usize) -> Self {
138        Self {
139            rv32i: Rv32IConfig::with_public_values(public_values),
140            mul: Default::default(),
141        }
142    }
143
144    pub fn with_public_values_and_segment_len(public_values: usize, segment_len: usize) -> Self {
145        Self {
146            rv32i: Rv32IConfig::with_public_values_and_segment_len(public_values, segment_len),
147            mul: Default::default(),
148        }
149    }
150}
151
152#[derive(Clone)]
153pub struct Rv32ICpuBuilder;
154
155impl<E, SC> VmBuilder<E> for Rv32ICpuBuilder
156where
157    SC: StarkGenericConfig,
158    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
159    Val<SC>: VmField,
160{
161    type VmConfig = Rv32IConfig;
162    type SystemChipInventory = SystemChipInventory<SC>;
163    type RecordArena = MatrixRecordArena<Val<SC>>;
164
165    fn create_chip_complex(
166        &self,
167        config: &Rv32IConfig,
168        circuit: AirInventory<SC>,
169    ) -> Result<
170        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
171        ChipInventoryError,
172    > {
173        let mut chip_complex =
174            VmBuilder::<E>::create_chip_complex(&SystemCpuBuilder, &config.system, circuit)?;
175        let inventory = &mut chip_complex.inventory;
176        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.base, inventory)?;
177        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.io, inventory)?;
178        Ok(chip_complex)
179    }
180}
181
182#[derive(Clone)]
183pub struct Rv32ImCpuBuilder;
184
185impl<E, SC> VmBuilder<E> for Rv32ImCpuBuilder
186where
187    SC: StarkGenericConfig,
188    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
189    Val<SC>: VmField,
190{
191    type VmConfig = Rv32ImConfig;
192    type SystemChipInventory = SystemChipInventory<SC>;
193    type RecordArena = MatrixRecordArena<Val<SC>>;
194
195    fn create_chip_complex(
196        &self,
197        config: &Self::VmConfig,
198        circuit: AirInventory<SC>,
199    ) -> Result<
200        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
201        ChipInventoryError,
202    > {
203        let mut chip_complex =
204            VmBuilder::<E>::create_chip_complex(&Rv32ICpuBuilder, &config.rv32i, circuit)?;
205        let inventory = &mut chip_complex.inventory;
206        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.mul, inventory)?;
207        Ok(chip_complex)
208    }
209}
210
211#[cfg(feature = "cuda")]
212#[derive(Clone)]
213pub struct Rv32IGpuBuilder;
214
215#[cfg(feature = "cuda")]
216impl VmBuilder<GpuBabyBearPoseidon2Engine> for Rv32IGpuBuilder {
217    type VmConfig = Rv32IConfig;
218    type SystemChipInventory = SystemChipInventoryGPU;
219    type RecordArena = DenseRecordArena;
220
221    fn create_chip_complex(
222        &self,
223        config: &Rv32IConfig,
224        circuit: AirInventory<BabyBearPoseidon2Config>,
225    ) -> Result<
226        VmChipComplex<
227            BabyBearPoseidon2Config,
228            Self::RecordArena,
229            GpuBackend,
230            Self::SystemChipInventory,
231        >,
232        ChipInventoryError,
233    > {
234        let mut chip_complex = VmBuilder::<GpuBabyBearPoseidon2Engine>::create_chip_complex(
235            &SystemGpuBuilder,
236            &config.system,
237            circuit,
238        )?;
239        let inventory = &mut chip_complex.inventory;
240        VmProverExtension::<GpuBabyBearPoseidon2Engine, _, _>::extend_prover(
241            &Rv32ImGpuProverExt,
242            &config.base,
243            inventory,
244        )?;
245        VmProverExtension::<GpuBabyBearPoseidon2Engine, _, _>::extend_prover(
246            &Rv32ImGpuProverExt,
247            &config.io,
248            inventory,
249        )?;
250        Ok(chip_complex)
251    }
252}
253
254#[cfg(feature = "cuda")]
255#[derive(Clone)]
256pub struct Rv32ImGpuBuilder;
257
258#[cfg(feature = "cuda")]
259impl VmBuilder<GpuBabyBearPoseidon2Engine> for Rv32ImGpuBuilder {
260    type VmConfig = Rv32ImConfig;
261    type SystemChipInventory = SystemChipInventoryGPU;
262    type RecordArena = DenseRecordArena;
263
264    fn create_chip_complex(
265        &self,
266        config: &Self::VmConfig,
267        circuit: AirInventory<BabyBearPoseidon2Config>,
268    ) -> Result<
269        VmChipComplex<
270            BabyBearPoseidon2Config,
271            Self::RecordArena,
272            GpuBackend,
273            Self::SystemChipInventory,
274        >,
275        ChipInventoryError,
276    > {
277        let mut chip_complex = VmBuilder::<GpuBabyBearPoseidon2Engine>::create_chip_complex(
278            &Rv32IGpuBuilder,
279            &config.rv32i,
280            circuit,
281        )?;
282        let inventory = &mut chip_complex.inventory;
283        VmProverExtension::<GpuBabyBearPoseidon2Engine, _, _>::extend_prover(
284            &Rv32ImGpuProverExt,
285            &config.mul,
286            inventory,
287        )?;
288        Ok(chip_complex)
289    }
290}