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#[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
86impl InitFileGenerator for Rv32IConfig {}
88
89#[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
98impl 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}