1use std::sync::Arc;
2
3use derive_more::derive::From;
4use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor};
5use openvm_circuit_primitives::var_range::{
6 SharedVariableRangeCheckerChip, VariableRangeCheckerAir, VariableRangeCheckerBus,
7 VariableRangeCheckerChip,
8};
9use openvm_instructions::{
10 LocalOpcode, PhantomDiscriminant, PublishOpcode, SysPhantom, SystemOpcode,
11};
12use openvm_stark_backend::{
13 config::{StarkGenericConfig, Val},
14 engine::StarkEngine,
15 interaction::{LookupBus, PermutationCheckBus},
16 p3_field::{Field, PrimeField32},
17 prover::{
18 cpu::{CpuBackend, CpuDevice},
19 hal::{MatrixDimensions, ProverBackend},
20 types::{AirProvingContext, CommittedTraceData},
21 },
22 AirRef, Chip,
23};
24use rustc_hash::FxHashMap;
25
26use self::{connector::VmConnectorAir, program::ProgramAir, public_values::PublicValuesAir};
27use crate::{
28 arch::{
29 vm_poseidon2_config, AirInventory, AirInventoryError, BusIndexManager, ChipInventory,
30 ChipInventoryError, DenseRecordArena, ExecutionBridge, ExecutionBus, ExecutionState,
31 ExecutorInventory, ExecutorInventoryError, MatrixRecordArena, PhantomSubExecutor,
32 RowMajorMatrixArena, SystemConfig, VmAirWrapper, VmBuilder, VmChipComplex, VmChipWrapper,
33 VmCircuitConfig, VmExecutionConfig, CONNECTOR_AIR_ID, PROGRAM_AIR_ID, PUBLIC_VALUES_AIR_ID,
34 },
35 system::{
36 connector::VmConnectorChip,
37 memory::{
38 interface::MemoryInterfaceAirs,
39 offline_checker::{MemoryBridge, MemoryBus},
40 online::GuestMemory,
41 MemoryAirInventory, MemoryController, TimestampedEquipartition, CHUNK,
42 },
43 native_adapter::{NativeAdapterAir, NativeAdapterExecutor},
44 phantom::{
45 CycleEndPhantomExecutor, CycleStartPhantomExecutor, NopPhantomExecutor, PhantomAir,
46 PhantomChip, PhantomExecutor, PhantomFiller,
47 },
48 poseidon2::{
49 air::Poseidon2PeripheryAir, new_poseidon2_periphery_air, Poseidon2PeripheryChip,
50 },
51 program::{ProgramBus, ProgramChip},
52 public_values::{
53 PublicValuesChip, PublicValuesCoreAir, PublicValuesExecutor, PublicValuesFiller,
54 },
55 },
56};
57
58pub mod connector;
59#[cfg(feature = "cuda")]
60pub mod cuda;
61pub mod memory;
62pub mod native_adapter;
64pub mod phantom;
65pub mod poseidon2;
66pub mod program;
67pub mod public_values;
68
69const POSEIDON2_INSERTION_IDX: usize = 1;
71pub(crate) const PV_EXECUTOR_IDX: usize = 0;
73
74pub trait SystemChipComplex<RA, PB: ProverBackend> {
83 fn load_program(&mut self, cached_program_trace: CommittedTraceData<PB>);
85
86 fn transport_init_memory_to_device(&mut self, memory: &GuestMemory);
89
90 fn generate_proving_ctx(
93 &mut self,
94 system_records: SystemRecords<PB::Val>,
95 record_arenas: Vec<RA>,
96 ) -> Vec<AirProvingContext<PB>>;
97
98 #[cfg(feature = "metrics")]
108 fn finalize_trace_heights(&self, _heights: &mut [usize]) {}
109}
110
111pub trait SystemWithFixedTraceHeights {
113 fn override_trace_heights(&mut self, heights: &[u32]);
116}
117
118pub struct SystemRecords<F> {
119 pub from_state: ExecutionState<u32>,
120 pub to_state: ExecutionState<u32>,
121 pub exit_code: Option<u32>,
122 pub filtered_exec_frequencies: Vec<u32>,
125 pub access_adapter_records: DenseRecordArena,
128 pub touched_memory: TouchedMemory<F>,
130 pub public_values: Vec<F>,
133}
134
135pub enum TouchedMemory<F> {
136 Persistent(TimestampedEquipartition<F, CHUNK>),
137 Volatile(TimestampedEquipartition<F, 1>),
138}
139
140#[derive(Clone, AnyEnum, Executor, MeteredExecutor, PreflightExecutor, From)]
141pub enum SystemExecutor<F: Field> {
142 PublicValues(PublicValuesExecutor<F>),
143 Phantom(PhantomExecutor<F>),
144}
145
146#[derive(Clone, Copy)]
148pub struct SystemPort {
149 pub execution_bus: ExecutionBus,
150 pub program_bus: ProgramBus,
151 pub memory_bridge: MemoryBridge,
152}
153
154#[derive(Clone)]
155pub struct SystemAirInventory<SC: StarkGenericConfig> {
156 pub program: ProgramAir,
157 pub connector: VmConnectorAir,
158 pub memory: MemoryAirInventory<SC>,
159 pub public_values: Option<PublicValuesAir>,
162}
163
164impl<SC: StarkGenericConfig> SystemAirInventory<SC> {
165 pub fn new(
166 config: &SystemConfig,
167 port: SystemPort,
168 merkle_compression_buses: Option<(PermutationCheckBus, PermutationCheckBus)>,
169 ) -> Self {
170 let SystemPort {
171 execution_bus,
172 program_bus,
173 memory_bridge,
174 } = port;
175 let range_bus = memory_bridge.range_bus();
176 let program = ProgramAir::new(program_bus);
177 let connector = VmConnectorAir::new(
178 execution_bus,
179 program_bus,
180 range_bus,
181 config.memory_config.timestamp_max_bits,
182 );
183 assert_eq!(
184 config.continuation_enabled,
185 merkle_compression_buses.is_some()
186 );
187
188 let memory = MemoryAirInventory::new(
189 memory_bridge,
190 &config.memory_config,
191 range_bus,
192 merkle_compression_buses,
193 );
194
195 let public_values = if config.has_public_values_chip() {
196 let air = VmAirWrapper::new(
197 NativeAdapterAir::new(
198 ExecutionBridge::new(execution_bus, program_bus),
199 memory_bridge,
200 ),
201 PublicValuesCoreAir::new(
202 config.num_public_values,
203 config.max_constraint_degree as u32 - 1,
204 ),
205 );
206 Some(air)
207 } else {
208 None
209 };
210
211 Self {
212 program,
213 connector,
214 memory,
215 public_values,
216 }
217 }
218
219 pub fn port(&self) -> SystemPort {
220 SystemPort {
221 memory_bridge: self.memory.bridge,
222 program_bus: self.program.bus,
223 execution_bus: self.connector.execution_bus,
224 }
225 }
226
227 pub fn into_airs(self) -> Vec<AirRef<SC>> {
228 let mut airs: Vec<AirRef<SC>> = Vec::new();
229 airs.push(Arc::new(self.program));
230 airs.push(Arc::new(self.connector));
231 if let Some(public_values) = self.public_values {
232 airs.push(Arc::new(public_values));
233 }
234 airs.extend(self.memory.into_airs());
235 airs
236 }
237}
238
239impl<F: PrimeField32> VmExecutionConfig<F> for SystemConfig {
240 type Executor = SystemExecutor<F>;
241
242 fn create_executors(
246 &self,
247 ) -> Result<ExecutorInventory<Self::Executor>, ExecutorInventoryError> {
248 let mut inventory = ExecutorInventory::new(self.clone());
249 if self.has_public_values_chip() {
251 assert_eq!(inventory.executors().len(), PV_EXECUTOR_IDX);
252
253 let public_values = PublicValuesExecutor::new(NativeAdapterExecutor::default());
254 inventory.add_executor(public_values, [PublishOpcode::PUBLISH.global_opcode()])?;
255 }
256 let phantom_opcode = SystemOpcode::PHANTOM.global_opcode();
257 let mut phantom_executors: FxHashMap<PhantomDiscriminant, Arc<dyn PhantomSubExecutor<F>>> =
258 FxHashMap::default();
259 phantom_executors.insert(
261 PhantomDiscriminant(SysPhantom::DebugPanic as u16),
262 Arc::new(NopPhantomExecutor),
263 );
264 phantom_executors.insert(
265 PhantomDiscriminant(SysPhantom::Nop as u16),
266 Arc::new(NopPhantomExecutor),
267 );
268 phantom_executors.insert(
269 PhantomDiscriminant(SysPhantom::CtStart as u16),
270 Arc::new(CycleStartPhantomExecutor),
271 );
272 phantom_executors.insert(
273 PhantomDiscriminant(SysPhantom::CtEnd as u16),
274 Arc::new(CycleEndPhantomExecutor),
275 );
276 let phantom = PhantomExecutor::new(phantom_executors, phantom_opcode);
277 inventory.add_executor(phantom, [phantom_opcode])?;
278
279 Ok(inventory)
280 }
281}
282
283impl<SC: StarkGenericConfig> VmCircuitConfig<SC> for SystemConfig {
284 fn create_airs(&self) -> Result<AirInventory<SC>, AirInventoryError> {
287 let mut bus_idx_mgr = BusIndexManager::new();
288 let execution_bus = ExecutionBus::new(bus_idx_mgr.new_bus_idx());
289 let memory_bus = MemoryBus::new(bus_idx_mgr.new_bus_idx());
290 let program_bus = ProgramBus::new(bus_idx_mgr.new_bus_idx());
291 let range_bus =
292 VariableRangeCheckerBus::new(bus_idx_mgr.new_bus_idx(), self.memory_config.decomp);
293
294 let merkle_compression_buses = if self.continuation_enabled {
295 let merkle_bus = PermutationCheckBus::new(bus_idx_mgr.new_bus_idx());
296 let compression_bus = PermutationCheckBus::new(bus_idx_mgr.new_bus_idx());
297 Some((merkle_bus, compression_bus))
298 } else {
299 None
300 };
301 let memory_bridge =
302 MemoryBridge::new(memory_bus, self.memory_config.timestamp_max_bits, range_bus);
303 let system_port = SystemPort {
304 execution_bus,
305 program_bus,
306 memory_bridge,
307 };
308 let system = SystemAirInventory::new(self, system_port, merkle_compression_buses);
309
310 let mut inventory = AirInventory::new(self.clone(), system, bus_idx_mgr);
311
312 let range_checker = VariableRangeCheckerAir::new(range_bus);
313 inventory.add_air(range_checker);
315
316 if self.continuation_enabled {
317 assert_eq!(inventory.ext_airs().len(), POSEIDON2_INSERTION_IDX);
318 let (_, compression_bus) = merkle_compression_buses.unwrap();
322 let direct_bus_idx = compression_bus.index;
323 let air = new_poseidon2_periphery_air(
324 vm_poseidon2_config(),
325 LookupBus::new(direct_bus_idx),
326 self.max_constraint_degree,
327 );
328 inventory.add_air_ref(air);
329 }
330 let execution_bridge = ExecutionBridge::new(execution_bus, program_bus);
331 let phantom = PhantomAir {
332 execution_bridge,
333 phantom_opcode: SystemOpcode::PHANTOM.global_opcode(),
334 };
335 inventory.add_air(phantom);
336
337 Ok(inventory)
338 }
339}
340
341pub struct SystemChipInventory<SC: StarkGenericConfig> {
346 pub program_chip: ProgramChip<SC>,
347 pub connector_chip: VmConnectorChip<Val<SC>>,
348 pub memory_controller: MemoryController<Val<SC>>,
350 pub public_values_chip: Option<PublicValuesChip<Val<SC>>>,
351}
352
353impl<SC: StarkGenericConfig> SystemChipInventory<SC>
356where
357 Val<SC>: PrimeField32,
358{
359 pub fn new(
360 config: &SystemConfig,
361 mem_inventory: &MemoryAirInventory<SC>,
362 range_checker: SharedVariableRangeCheckerChip,
363 hasher_chip: Option<Arc<Poseidon2PeripheryChip<Val<SC>>>>,
364 ) -> Self {
365 let program_chip = ProgramChip::unloaded();
368 let connector_chip = VmConnectorChip::<Val<SC>>::new(
369 range_checker.clone(),
370 config.memory_config.timestamp_max_bits,
371 );
372 let memory_bus = mem_inventory.bridge.memory_bus();
373 let memory_controller = match &mem_inventory.interface {
374 MemoryInterfaceAirs::Persistent {
375 boundary: _,
376 merkle,
377 } => {
378 assert!(config.continuation_enabled);
379 MemoryController::<Val<SC>>::with_persistent_memory(
380 memory_bus,
381 config.memory_config.clone(),
382 range_checker.clone(),
383 merkle.merkle_bus,
384 merkle.compression_bus,
385 hasher_chip.unwrap(),
386 )
387 }
388 MemoryInterfaceAirs::Volatile { boundary: _ } => {
389 assert!(!config.continuation_enabled);
390 MemoryController::with_volatile_memory(
391 memory_bus,
392 config.memory_config.clone(),
393 range_checker.clone(),
394 )
395 }
396 };
397
398 let public_values_chip = config.has_public_values_chip().then(|| {
399 VmChipWrapper::new(
400 PublicValuesFiller::new(
401 NativeAdapterExecutor::default(),
402 config.num_public_values,
403 (config.max_constraint_degree as u32)
404 .checked_sub(1)
405 .unwrap(),
406 ),
407 memory_controller.helper(),
408 )
409 });
410
411 Self {
412 program_chip,
413 connector_chip,
414 memory_controller,
415 public_values_chip,
416 }
417 }
418}
419
420impl<RA, SC> SystemChipComplex<RA, CpuBackend<SC>> for SystemChipInventory<SC>
421where
422 RA: RowMajorMatrixArena<Val<SC>>,
423 SC: StarkGenericConfig,
424 Val<SC>: PrimeField32,
425{
426 fn load_program(&mut self, cached_program_trace: CommittedTraceData<CpuBackend<SC>>) {
427 let _ = self.program_chip.cached.replace(cached_program_trace);
428 }
429
430 fn transport_init_memory_to_device(&mut self, memory: &GuestMemory) {
431 self.memory_controller
432 .set_initial_memory(memory.memory.clone());
433 }
434
435 fn generate_proving_ctx(
436 &mut self,
437 system_records: SystemRecords<Val<SC>>,
438 mut record_arenas: Vec<RA>,
439 ) -> Vec<AirProvingContext<CpuBackend<SC>>> {
440 let SystemRecords {
441 from_state,
442 to_state,
443 exit_code,
444 filtered_exec_frequencies,
445 access_adapter_records,
446 touched_memory,
447 public_values,
448 } = system_records;
449
450 if let Some(chip) = &mut self.public_values_chip {
451 chip.inner.set_public_values(public_values);
452 }
453 self.program_chip.filtered_exec_frequencies = filtered_exec_frequencies;
454 let program_ctx = self.program_chip.generate_proving_ctx(());
455 self.connector_chip.begin(from_state);
456 self.connector_chip.end(to_state, exit_code);
457 let connector_ctx = self.connector_chip.generate_proving_ctx(());
458
459 let pv_ctx = self.public_values_chip.as_ref().map(|chip| {
460 let arena = record_arenas.remove(PUBLIC_VALUES_AIR_ID);
461 chip.generate_proving_ctx(arena)
462 });
463
464 let memory_ctxs = self
465 .memory_controller
466 .generate_proving_ctx(access_adapter_records, touched_memory);
467
468 [program_ctx, connector_ctx]
469 .into_iter()
470 .chain(pv_ctx)
471 .chain(memory_ctxs)
472 .collect()
473 }
474
475 #[cfg(feature = "metrics")]
476 fn finalize_trace_heights(&self, heights: &mut [usize]) {
477 use openvm_stark_backend::ChipUsageGetter;
478
479 use crate::system::memory::interface::MemoryInterface;
480
481 let boundary_idx = PUBLIC_VALUES_AIR_ID + usize::from(self.public_values_chip.is_some());
482 let mut access_adapter_offset = boundary_idx + 1;
483 match &self.memory_controller.interface_chip {
484 MemoryInterface::Volatile { boundary_chip } => {
485 let boundary_height = boundary_chip
486 .final_memory
487 .as_ref()
488 .map(|m| m.len())
489 .unwrap_or(0);
490 heights[boundary_idx] = boundary_height;
491 }
492 MemoryInterface::Persistent {
493 boundary_chip,
494 merkle_chip,
495 ..
496 } => {
497 let boundary_height = 2 * boundary_chip.touched_labels.len();
498 heights[boundary_idx] = boundary_height;
499 heights[boundary_idx + 1] = merkle_chip.current_height;
500 access_adapter_offset += 1;
501
502 let poseidon_chip = self.memory_controller.hasher_chip.as_ref().unwrap();
505 let poseidon_height = poseidon_chip.current_trace_height();
506 let poseidon_idx = heights.len() - 1 - POSEIDON2_INSERTION_IDX;
509 heights[poseidon_idx] = poseidon_height;
510 }
511 }
512 let access_heights = &self
513 .memory_controller
514 .access_adapter_inventory
515 .trace_heights;
516 heights[access_adapter_offset..access_adapter_offset + access_heights.len()]
517 .copy_from_slice(access_heights);
518 }
519}
520
521#[derive(Clone)]
522pub struct SystemCpuBuilder;
523
524impl<SC, E> VmBuilder<E> for SystemCpuBuilder
525where
526 SC: StarkGenericConfig,
527 E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
528 Val<SC>: PrimeField32,
529{
530 type VmConfig = SystemConfig;
531 type RecordArena = MatrixRecordArena<Val<SC>>;
532 type SystemChipInventory = SystemChipInventory<SC>;
533
534 fn create_chip_complex(
535 &self,
536 config: &SystemConfig,
537 airs: AirInventory<SC>,
538 ) -> Result<
539 VmChipComplex<SC, MatrixRecordArena<Val<SC>>, CpuBackend<SC>, SystemChipInventory<SC>>,
540 ChipInventoryError,
541 > {
542 let range_bus = airs.range_checker().bus;
543 let range_checker = Arc::new(VariableRangeCheckerChip::new(range_bus));
544
545 let mut inventory = ChipInventory::new(airs);
546 if config.has_public_values_chip() {
548 assert_eq!(
549 inventory.executor_idx_to_insertion_idx.len(),
550 PV_EXECUTOR_IDX
551 );
552 let insertion_idx = inventory
556 .airs()
557 .num_airs()
558 .checked_sub(1 + PUBLIC_VALUES_AIR_ID)
559 .unwrap();
560 inventory.executor_idx_to_insertion_idx.push(insertion_idx);
561 }
562 inventory.next_air::<VariableRangeCheckerAir>()?;
563 inventory.add_periphery_chip(range_checker.clone());
564
565 let hasher_chip = if config.continuation_enabled {
566 assert_eq!(inventory.chips().len(), POSEIDON2_INSERTION_IDX);
567 let direct_bus = if config.max_constraint_degree >= 7 {
569 inventory
570 .next_air::<Poseidon2PeripheryAir<Val<SC>, 0>>()?
571 .bus
572 } else {
573 inventory
574 .next_air::<Poseidon2PeripheryAir<Val<SC>, 1>>()?
575 .bus
576 };
577 let chip = Arc::new(Poseidon2PeripheryChip::new(
578 vm_poseidon2_config(),
579 direct_bus.index,
580 config.max_constraint_degree,
581 ));
582 inventory.add_periphery_chip(chip.clone());
583 Some(chip)
584 } else {
585 None
586 };
587 let system = SystemChipInventory::new(
588 config,
589 &inventory.airs().system().memory,
590 range_checker,
591 hasher_chip,
592 );
593
594 let phantom_chip = PhantomChip::new(PhantomFiller, system.memory_controller.helper());
595 inventory.add_executor_chip(phantom_chip);
596
597 Ok(VmChipComplex { system, inventory })
598 }
599}
600
601impl<SC: StarkGenericConfig> SystemWithFixedTraceHeights for SystemChipInventory<SC>
602where
603 Val<SC>: PrimeField32,
604{
605 fn override_trace_heights(&mut self, heights: &[u32]) {
608 assert_eq!(
609 heights[PROGRAM_AIR_ID] as usize,
610 self.program_chip
611 .cached
612 .as_ref()
613 .expect("program not loaded")
614 .trace
615 .height()
616 );
617 assert_eq!(heights[CONNECTOR_AIR_ID], 2);
618 let mut memory_start_idx = PUBLIC_VALUES_AIR_ID;
619 if self.public_values_chip.is_some() {
620 memory_start_idx += 1;
621 }
622 self.memory_controller
623 .set_override_trace_heights(&heights[memory_start_idx..]);
624 }
625}