openvm_circuit/system/cuda/
boundary.rs

1use openvm_circuit::{
2    arch::DEFAULT_BLOCK_SIZE, system::memory::persistent::PersistentBoundaryCols,
3    utils::next_power_of_two_or_zero,
4};
5use openvm_circuit_primitives::Chip;
6use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
7use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
8use openvm_stark_backend::prover::{AirProvingContext, MatrixDimensions};
9
10use super::{poseidon2::SharedBuffer, DIGEST_WIDTH};
11use crate::cuda_abi::boundary::persistent_boundary_tracegen;
12
13pub struct BoundaryChipGPU {
14    pub device_ctx: GpuDeviceCtx,
15    pub poseidon2_buffer: SharedBuffer<F>,
16    /// A `Vec` of pointers to the copied guest memory on device.
17    /// This struct cannot own the device memory, hence we take extra care not to use memory we
18    /// don't own. TODO: use `Arc<DeviceBuffer>` instead?
19    pub initial_leaves: Vec<*const std::ffi::c_void>,
20    pub records: Option<DeviceBuffer<u32>>,
21    pub num_records: Option<usize>,
22    pub trace_width: Option<usize>,
23}
24
25const BLOCKS_PER_CHUNK: usize = DIGEST_WIDTH / DEFAULT_BLOCK_SIZE;
26
27#[repr(C)]
28#[derive(Clone, Copy)]
29pub struct PersistentBoundaryRecord {
30    pub address_space: u32,
31    pub ptr: u32,
32    pub timestamps: [u32; BLOCKS_PER_CHUNK],
33    pub values: [F; DIGEST_WIDTH],
34}
35
36impl BoundaryChipGPU {
37    pub fn new(poseidon2_buffer: SharedBuffer<F>, device_ctx: GpuDeviceCtx) -> Self {
38        Self {
39            device_ctx,
40            poseidon2_buffer,
41            initial_leaves: Vec::new(),
42            records: None,
43            num_records: None,
44            trace_width: None,
45        }
46    }
47
48    pub fn finalize_records<const CHUNK: usize>(&mut self, records: Vec<PersistentBoundaryRecord>) {
49        self.num_records = Some(records.len());
50        self.trace_width = Some(PersistentBoundaryCols::<F, CHUNK>::width());
51        self.records = Some(if records.is_empty() {
52            DeviceBuffer::new()
53        } else {
54            records
55                .to_device_on(&self.device_ctx)
56                .unwrap()
57                .as_buffer::<u32>()
58        });
59    }
60
61    pub fn finalize_records_device<const CHUNK: usize>(
62        &mut self,
63        records: DeviceBuffer<u32>,
64        num_records: usize,
65    ) {
66        self.num_records = Some(num_records);
67        self.trace_width = Some(PersistentBoundaryCols::<F, CHUNK>::width());
68        self.records = Some(records);
69    }
70
71    pub fn trace_width(&self) -> usize {
72        self.trace_width.expect("Finalize records to get width")
73    }
74
75    pub fn records(&self) -> &DeviceBuffer<u32> {
76        self.records
77            .as_ref()
78            .expect("Finalize records to get buffer")
79    }
80}
81
82impl<RA> Chip<RA, GpuBackend> for BoundaryChipGPU {
83    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<GpuBackend> {
84        let num_records = self.num_records.unwrap();
85        if num_records == 0 {
86            // Boundary AIR should always be present, so return a single zero-filled
87            // padding row.
88            let trace =
89                DeviceMatrix::<F>::with_capacity_on(1, self.trace_width(), &self.device_ctx);
90            trace.buffer().fill_zero_on(&self.device_ctx).unwrap();
91            return AirProvingContext::simple_no_pis(trace);
92        }
93        let unpadded_height = 2 * num_records;
94        let trace_height = next_power_of_two_or_zero(unpadded_height);
95        let trace =
96            DeviceMatrix::<F>::with_capacity_on(trace_height, self.trace_width(), &self.device_ctx);
97        trace.buffer().fill_zero_on(&self.device_ctx).unwrap();
98        let mem_ptrs = self.initial_leaves.to_device_on(&self.device_ctx).unwrap();
99        let poseidon2_records = self.poseidon2_buffer.records();
100        unsafe {
101            persistent_boundary_tracegen(
102                trace.buffer(),
103                trace.height(),
104                trace.width(),
105                &mem_ptrs,
106                self.records.as_ref().unwrap(),
107                num_records,
108                &poseidon2_records,
109                &self.poseidon2_buffer.idx,
110                self.device_ctx.stream.as_raw(),
111            )
112            .expect("Failed to generate boundary trace");
113        }
114        AirProvingContext::simple_no_pis(trace)
115    }
116}
openvm_circuit/system/cuda/boundary.rs

openvm_circuit/system/cuda/
boundary.rs
