openvm_circuit_primitives/var_range/
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::{
8    cuda_abi::var_range::tracegen,
9    var_range::{VariableRangeCheckerBus, VariableRangeCheckerChip, NUM_VARIABLE_RANGE_COLS},
10    Chip,
11};
12
13pub struct VariableRangeCheckerChipGPU {
14    pub device_ctx: GpuDeviceCtx,
15    pub count: Arc<DeviceBuffer<F>>,
16    pub cpu_chip: Option<Arc<VariableRangeCheckerChip>>,
17}
18
19/// The preprocessed trace contains `value` and `bits`; the generated trace contains `count`.
20impl VariableRangeCheckerChipGPU {
21    pub fn new(bus: VariableRangeCheckerBus, device_ctx: GpuDeviceCtx) -> Self {
22        let num_rows = (1 << (bus.range_max_bits + 1)) as usize;
23        let count = Arc::new(DeviceBuffer::<F>::with_capacity_on(num_rows, &device_ctx));
24        count.fill_zero_on(&device_ctx).unwrap();
25        Self {
26            device_ctx,
27            count,
28            cpu_chip: None,
29        }
30    }
31
32    pub fn hybrid(cpu_chip: Arc<VariableRangeCheckerChip>, device_ctx: GpuDeviceCtx) -> Self {
33        let count = Arc::new(DeviceBuffer::<F>::with_capacity_on(
34            cpu_chip.count.len(),
35            &device_ctx,
36        ));
37        count.fill_zero_on(&device_ctx).unwrap();
38        Self {
39            device_ctx,
40            count,
41            cpu_chip: Some(cpu_chip),
42        }
43    }
44}
45
46impl<RA> Chip<RA, GpuBackend> for VariableRangeCheckerChipGPU {
47    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<GpuBackend> {
48        assert_eq!(size_of::<F>(), size_of::<u32>());
49        let cpu_count = self.cpu_chip.as_ref().map(|cpu_chip| {
50            cpu_chip
51                .count
52                .iter()
53                .map(|c| c.swap(0, Ordering::Relaxed))
54                .collect::<Vec<_>>()
55                .to_device_on(&self.device_ctx)
56                .unwrap()
57        });
58        // ATTENTION: we create a new buffer to copy `count` into because this chip is stateful and
59        // `count` will be reused.
60        let trace = DeviceMatrix::<F>::with_capacity_on(
61            self.count.len(),
62            NUM_VARIABLE_RANGE_COLS,
63            &self.device_ctx,
64        );
65        // Zero padding rows so stale pool data doesn't cause constraint violations.
66        trace.buffer().fill_zero_on(&self.device_ctx).unwrap();
67        unsafe {
68            tracegen(
69                &self.count,
70                &cpu_count,
71                trace.buffer(),
72                self.device_ctx.stream.as_raw(),
73            )
74            .unwrap();
75        }
76        // Zero the internal count buffer because this chip is stateful and may be used again.
77        self.count.fill_zero_on(&self.device_ctx).unwrap();
78        AirProvingContext::simple_no_pis(trace)
79    }
80
81    fn constant_trace_height(&self) -> Option<usize> {
82        Some(self.count.len())
83    }
84}
openvm_circuit_primitives/var_range/cuda.rs

openvm_circuit_primitives/var_range/
cuda.rs
