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