openvm_circuit/system/cuda/
mod.rs

1use std::sync::Arc;
2
3use connector::VmConnectorChipGPU;
4use memory::MemoryInventoryGPU;
5use openvm_circuit::{
6    arch::{DenseRecordArena, SystemConfig},
7    system::{
8        connector::VmConnectorChip, memory::online::GuestMemory, SystemChipComplex, SystemRecords,
9    },
10};
11use openvm_circuit_primitives::{var_range::VariableRangeCheckerChipGPU, Chip};
12use openvm_cuda_backend::{prelude::F, GpuBackend};
13use openvm_cuda_common::stream::GpuDeviceCtx;
14use openvm_stark_backend::prover::{AirProvingContext, CommittedTraceData};
15use poseidon2::Poseidon2PeripheryChipGPU;
16use program::ProgramChipGPU;
17
18use crate::system::memory::CHUNK;
19
20pub(crate) const DIGEST_WIDTH: usize = 8;
21
22pub mod boundary;
23pub mod connector;
24pub mod extensions;
25pub mod memory;
26pub mod merkle_tree;
27pub mod phantom;
28pub mod poseidon2;
29pub mod program;
30
31pub struct SystemChipInventoryGPU {
32    pub program: ProgramChipGPU,
33    pub connector: VmConnectorChipGPU,
34    pub memory_inventory: MemoryInventoryGPU,
35}
36
37impl SystemChipInventoryGPU {
38    pub fn new(
39        config: &SystemConfig,
40        range_checker: Arc<VariableRangeCheckerChipGPU>,
41        hasher_chip: Arc<Poseidon2PeripheryChipGPU>,
42        device_ctx: GpuDeviceCtx,
43    ) -> Self {
44        let cpu_range_checker = range_checker.cpu_chip.clone().unwrap();
45
46        // We create an empty program chip: the program should be loaded later (and can be swapped
47        // out). The execution frequencies are supplied only after execution.
48        let program_chip = ProgramChipGPU::new(device_ctx.clone());
49        let connector_chip = VmConnectorChipGPU::new(
50            VmConnectorChip::new(
51                cpu_range_checker.clone(),
52                config.memory_config.timestamp_max_bits,
53            ),
54            device_ctx.clone(),
55        );
56
57        let memory_inventory = MemoryInventoryGPU::new(
58            config.memory_config.clone(),
59            hasher_chip,
60            device_ctx.clone(),
61        );
62
63        Self {
64            program: program_chip,
65            connector: connector_chip,
66            memory_inventory,
67        }
68    }
69}
70
71impl SystemChipComplex<DenseRecordArena, GpuBackend> for SystemChipInventoryGPU {
72    fn load_program(&mut self, cached_program_trace: CommittedTraceData<GpuBackend>) {
73        self.program.cached.replace(cached_program_trace);
74    }
75
76    fn transport_init_memory_to_device(&mut self, memory: &GuestMemory) {
77        self.memory_inventory.set_initial_memory(&memory.memory);
78    }
79
80    fn generate_proving_ctx(
81        &mut self,
82        system_records: SystemRecords<F>,
83        _record_arenas: Vec<DenseRecordArena>,
84    ) -> Vec<AirProvingContext<GpuBackend>> {
85        let SystemRecords {
86            from_state,
87            to_state,
88            exit_code,
89            filtered_exec_frequencies,
90            touched_memory,
91        } = system_records;
92
93        let program_ctx = self.program.generate_proving_ctx(filtered_exec_frequencies);
94
95        self.connector.cpu_chip.begin(from_state);
96        self.connector.cpu_chip.end(to_state, exit_code);
97        let connector_ctx = self.connector.generate_proving_ctx(());
98
99        let memory_ctxs = self.memory_inventory.generate_proving_ctxs(touched_memory);
100
101        [program_ctx, connector_ctx]
102            .into_iter()
103            .chain(memory_ctxs)
104            .collect()
105    }
106
107    fn memory_top_tree(&self) -> Option<&[[F; CHUNK]]> {
108        let top_tree = &self.memory_inventory.merkle_tree.top_roots_host;
109        (!top_tree.is_empty()).then_some(top_tree.as_slice())
110    }
111}
openvm_circuit/system/cuda/mod.rs

openvm_circuit/system/cuda/
mod.rs
