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#[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
87impl InitFileGenerator for Rv32IConfig {}
89
90#[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
99impl 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}