openvm_circuit/system/cuda/
phantom.rs

1use std::mem::size_of;
2
3use derive_new::new;
4use openvm_circuit::{
5    arch::DenseRecordArena,
6    primitives::Chip,
7    system::phantom::{PhantomCols, PhantomRecord},
8    utils::next_power_of_two_or_zero,
9};
10use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
11use openvm_cuda_common::{copy::MemCopyH2D, stream::GpuDeviceCtx};
12use openvm_stark_backend::prover::{AirProvingContext, MatrixDimensions};
13
14use crate::cuda_abi::phantom;
15
16#[derive(new)]
17pub struct PhantomChipGPU {
18    device_ctx: GpuDeviceCtx,
19}
20
21impl PhantomChipGPU {
22    pub fn trace_height(arena: &DenseRecordArena) -> usize {
23        let record_size = size_of::<PhantomRecord>();
24        let records_len = arena.allocated().len();
25        assert_eq!(records_len % record_size, 0);
26        records_len / record_size
27    }
28
29    pub fn trace_width() -> usize {
30        PhantomCols::<F>::width()
31    }
32}
33
34impl Chip<DenseRecordArena, GpuBackend> for PhantomChipGPU {
35    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
36        let num_records = Self::trace_height(&arena);
37        if num_records == 0 {
38            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
39        }
40        let trace_height = next_power_of_two_or_zero(num_records);
41        let trace = DeviceMatrix::<F>::with_capacity_on(
42            trace_height,
43            Self::trace_width(),
44            &self.device_ctx,
45        );
46        trace.buffer().fill_zero_on(&self.device_ctx).unwrap();
47        unsafe {
48            phantom::tracegen(
49                trace.buffer(),
50                trace.height(),
51                trace.width(),
52                &arena.allocated().to_device_on(&self.device_ctx).unwrap(),
53                self.device_ctx.stream.as_raw(),
54            )
55            .expect("Failed to generate trace");
56        }
57        AirProvingContext::simple_no_pis(trace)
58    }
59}