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