openvm_rv32im_circuit/
lib.rs

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