openvm_deferral_circuit/count/
cuda.rs

1use std::sync::Arc;
2
3use derive_new::new;
4use openvm_circuit::{arch::DenseRecordArena, utils::next_power_of_two_or_zero};
5use openvm_circuit_primitives::Chip;
6use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
7use openvm_cuda_common::{d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
8use openvm_stark_backend::prover::AirProvingContext;
9
10use crate::{count::DeferralCircuitCountCols, cuda_abi::count};
11
12#[derive(new)]
13pub struct DeferralCircuitCountChipGpu {
14    pub count: Arc<DeviceBuffer<u32>>,
15    pub num_deferral_circuits: usize,
16    pub device_ctx: GpuDeviceCtx,
17}
18
19impl Chip<DenseRecordArena, GpuBackend> for DeferralCircuitCountChipGpu {
20    fn constant_trace_height(&self) -> Option<usize> {
21        Some(next_power_of_two_or_zero(self.num_deferral_circuits))
22    }
23
24    fn generate_proving_ctx(&self, _: DenseRecordArena) -> AirProvingContext<GpuBackend> {
25        if self.num_deferral_circuits == 0 {
26            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
27        }
28
29        let trace_width = DeferralCircuitCountCols::<F>::width();
30        let trace_height = next_power_of_two_or_zero(self.num_deferral_circuits);
31        let trace =
32            DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, &self.device_ctx);
33
34        unsafe {
35            count::tracegen(
36                trace.buffer(),
37                trace_height,
38                &self.count,
39                self.num_deferral_circuits,
40                self.device_ctx.stream.as_raw(),
41            )
42            .expect("Failed to generate deferral count trace");
43        }
44
45        self.count
46            .fill_zero_on(&self.device_ctx)
47            .expect("Failed to reset deferral count");
48        AirProvingContext::simple_no_pis(trace)
49    }
50}
openvm_deferral_circuit/count/cuda.rs

openvm_deferral_circuit/count/
cuda.rs
