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