openvm_circuit_primitives/range_tuple/
cuda.rs

1use std::sync::{atomic::Ordering, Arc};
2
3use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
4use openvm_cuda_common::{copy::MemCopyH2D as _, d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
5use openvm_stark_backend::prover::AirProvingContext;
6
7use crate::{cuda_abi::range_tuple::tracegen, range_tuple::RangeTupleCheckerChip, Chip};
8
9pub struct RangeTupleCheckerChipGPU<const N: usize> {
10    pub device_ctx: GpuDeviceCtx,
11    pub count: Arc<DeviceBuffer<F>>,
12    pub cpu_chip: Option<Arc<RangeTupleCheckerChip<N>>>,
13    pub sizes: [u32; N],
14}
15
16impl<const N: usize> RangeTupleCheckerChipGPU<N> {
17    pub fn new(sizes: [u32; N], device_ctx: GpuDeviceCtx) -> Self {
18        assert!(N > 1, "RangeTupleChecker requires at least 2 dimensions");
19        let range_max = sizes.iter().product::<u32>() as usize;
20        let count = Arc::new(DeviceBuffer::<F>::with_capacity_on(range_max, &device_ctx));
21        count.fill_zero_on(&device_ctx).unwrap();
22        Self {
23            device_ctx,
24            count,
25            cpu_chip: None,
26            sizes,
27        }
28    }
29
30    pub fn hybrid(cpu_chip: Arc<RangeTupleCheckerChip<N>>, device_ctx: GpuDeviceCtx) -> Self {
31        let count = Arc::new(DeviceBuffer::<F>::with_capacity_on(
32            cpu_chip.count.len(),
33            &device_ctx,
34        ));
35        count.fill_zero_on(&device_ctx).unwrap();
36        let sizes = *cpu_chip.sizes();
37        Self {
38            device_ctx,
39            count,
40            cpu_chip: Some(cpu_chip),
41            sizes,
42        }
43    }
44}
45
46impl<RA, const N: usize> Chip<RA, GpuBackend> for RangeTupleCheckerChipGPU<N> {
47    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<GpuBackend> {
48        let cpu_count = self.cpu_chip.as_ref().map(|cpu_chip| {
49            cpu_chip
50                .count
51                .iter()
52                .map(|c| c.swap(0, Ordering::Relaxed))
53                .collect::<Vec<_>>()
54                .to_device_on(&self.device_ctx)
55                .unwrap()
56        });
57        // ATTENTION: we create a new buffer to copy `count` into because this chip is stateful and
58        // `count` will be reused.
59        let trace = DeviceMatrix::<F>::with_capacity_on(self.count.len(), N + 1, &self.device_ctx);
60        // Zero padding rows so stale pool data doesn't cause constraint violations.
61        trace.buffer().fill_zero_on(&self.device_ctx).unwrap();
62        let d_sizes = self.sizes.to_device_on(&self.device_ctx).unwrap();
63        unsafe {
64            tracegen(
65                &self.count,
66                &cpu_count,
67                trace.buffer(),
68                &d_sizes,
69                self.device_ctx.stream.as_raw(),
70            )
71            .unwrap();
72        }
73        // Zero the internal count buffer because this chip is stateful and may be used again.
74        self.count.fill_zero_on(&self.device_ctx).unwrap();
75        AirProvingContext::simple_no_pis(trace)
76    }
77
78    fn constant_trace_height(&self) -> Option<usize> {
79        Some(self.count.len())
80    }
81}
openvm_circuit_primitives/range_tuple/cuda.rs

openvm_circuit_primitives/range_tuple/
cuda.rs
