openvm_circuit/system/cuda/
phantom.rs1use 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}