openvm_circuit_primitives/bitwise_op_lookup/
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    bitwise_op_lookup::{
9        BitwiseOperationLookupChip, BitwiseOperationLookupCols, NUM_BITWISE_OP_LOOKUP_MULT_COLS,
10    },
11    cuda_abi::bitwise_op_lookup::tracegen,
12    Chip,
13};
14
15pub struct BitwiseOperationLookupChipGPU<const NUM_BITS: usize> {
16    pub device_ctx: GpuDeviceCtx,
17    pub count: Arc<DeviceBuffer<F>>,
18    pub cpu_chip: Option<Arc<BitwiseOperationLookupChip<NUM_BITS>>>,
19}
20
21impl<const NUM_BITS: usize> BitwiseOperationLookupChipGPU<NUM_BITS> {
22    pub const fn num_rows() -> usize {
23        1 << (2 * NUM_BITS)
24    }
25
26    pub fn new(device_ctx: GpuDeviceCtx) -> Self {
27        // The first 2^(2 * NUM_BITS) indices are for range checking, the rest are for XOR
28        let count = Arc::new(DeviceBuffer::<F>::with_capacity_on(
29            NUM_BITWISE_OP_LOOKUP_MULT_COLS * Self::num_rows(),
30            &device_ctx,
31        ));
32        count.fill_zero_on(&device_ctx).unwrap();
33        Self {
34            device_ctx,
35            count,
36            cpu_chip: None,
37        }
38    }
39
40    pub fn hybrid(
41        cpu_chip: Arc<BitwiseOperationLookupChip<NUM_BITS>>,
42        device_ctx: GpuDeviceCtx,
43    ) -> Self {
44        assert_eq!(cpu_chip.count_range.len(), Self::num_rows());
45        assert_eq!(cpu_chip.count_xor.len(), Self::num_rows());
46        let count = Arc::new(DeviceBuffer::<F>::with_capacity_on(
47            NUM_BITWISE_OP_LOOKUP_MULT_COLS * Self::num_rows(),
48            &device_ctx,
49        ));
50        count.fill_zero_on(&device_ctx).unwrap();
51        Self {
52            device_ctx,
53            count,
54            cpu_chip: Some(cpu_chip),
55        }
56    }
57}
58
59impl<RA, const NUM_BITS: usize> Chip<RA, GpuBackend> for BitwiseOperationLookupChipGPU<NUM_BITS> {
60    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<GpuBackend> {
61        let num_cols = BitwiseOperationLookupCols::<F, NUM_BITS>::width();
62        debug_assert_eq!(
63            NUM_BITWISE_OP_LOOKUP_MULT_COLS * Self::num_rows(),
64            self.count.len()
65        );
66        let cpu_count = self.cpu_chip.as_ref().map(|cpu_chip| {
67            cpu_chip
68                .count_range
69                .iter()
70                .chain(cpu_chip.count_xor.iter())
71                .map(|c| c.swap(0, Ordering::Relaxed))
72                .collect::<Vec<_>>()
73                .to_device_on(&self.device_ctx)
74                .unwrap()
75        });
76        // ATTENTION: we create a new buffer to copy `count` into because this chip is stateful and
77        // `count` will be reused.
78        let trace =
79            DeviceMatrix::<F>::with_capacity_on(Self::num_rows(), num_cols, &self.device_ctx);
80        // Zero padding rows so stale pool data doesn't cause constraint violations.
81        trace.buffer().fill_zero_on(&self.device_ctx).unwrap();
82        unsafe {
83            tracegen(
84                &self.count,
85                &cpu_count,
86                trace.buffer(),
87                NUM_BITS as u32,
88                self.device_ctx.stream.as_raw(),
89            )
90            .unwrap();
91        }
92        // Zero the internal count buffer because this chip is stateful and may be used again.
93        self.count.fill_zero_on(&self.device_ctx).unwrap();
94        AirProvingContext::simple_no_pis(trace)
95    }
96
97    fn constant_trace_height(&self) -> Option<usize> {
98        Some(Self::num_rows())
99    }
100}
openvm_circuit_primitives/bitwise_op_lookup/cuda.rs

openvm_circuit_primitives/bitwise_op_lookup/
cuda.rs
