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#[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
89impl InitFileGenerator for Rv32IConfig {}
91
92#[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
101impl 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}