openvm_bigint_circuit/cuda/
mod.rs

1use std::{mem::size_of, sync::Arc};
2
3use derive_new::new;
4use openvm_circuit::{
5    arch::{DenseRecordArena, DEFAULT_BLOCK_SIZE},
6    utils::next_power_of_two_or_zero,
7};
8use openvm_circuit_primitives::{
9    bitwise_op_lookup::BitwiseOperationLookupChipGPU, cuda_abi::UInt2,
10    range_tuple::RangeTupleCheckerChipGPU, var_range::VariableRangeCheckerChipGPU, Chip,
11};
12use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
13use openvm_cuda_common::copy::MemCopyH2D;
14use openvm_rv32_adapters::{
15    Rv32VecHeapAdapterCols, Rv32VecHeapAdapterRecord, Rv32VecHeapBranchAdapterCols,
16    Rv32VecHeapBranchAdapterRecord,
17};
18use openvm_rv32im_circuit::{
19    adapters::{INT256_NUM_LIMBS, RV32_CELL_BITS},
20    BaseAluCoreCols, BaseAluCoreRecord, BranchEqualCoreCols, BranchEqualCoreRecord,
21    BranchLessThanCoreCols, BranchLessThanCoreRecord, LessThanCoreCols, LessThanCoreRecord,
22    MultiplicationCoreCols, MultiplicationCoreRecord, ShiftCoreCols, ShiftCoreRecord,
23};
24use openvm_stark_backend::prover::AirProvingContext;
25
26mod cuda_abi;
27
28use crate::INT256_NUM_BLOCKS;
29
30//////////////////////////////////////////////////////////////////////////////////////
31/// ALU
32//////////////////////////////////////////////////////////////////////////////////////
33pub type BaseAlu256AdapterRecord = Rv32VecHeapAdapterRecord<
34    2,
35    INT256_NUM_BLOCKS,
36    INT256_NUM_BLOCKS,
37    DEFAULT_BLOCK_SIZE,
38    DEFAULT_BLOCK_SIZE,
39>;
40pub type BaseAlu256CoreRecord = BaseAluCoreRecord<INT256_NUM_LIMBS>;
41
42#[derive(new)]
43pub struct BaseAlu256ChipGpu {
44    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
45    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
46    pub pointer_max_bits: usize,
47    pub timestamp_max_bits: usize,
48}
49
50impl Chip<DenseRecordArena, GpuBackend> for BaseAlu256ChipGpu {
51    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
52        const RECORD_SIZE: usize = size_of::<(BaseAlu256AdapterRecord, BaseAlu256CoreRecord)>();
53        let records = arena.allocated();
54        if records.is_empty() {
55            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
56        }
57        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
58
59        let trace_width = BaseAluCoreCols::<F, INT256_NUM_LIMBS, RV32_CELL_BITS>::width()
60            + Rv32VecHeapAdapterCols::<
61                F,
62                2,
63                INT256_NUM_BLOCKS,
64                INT256_NUM_BLOCKS,
65                DEFAULT_BLOCK_SIZE,
66                DEFAULT_BLOCK_SIZE,
67            >::width();
68        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
69        let device_ctx = &self.range_checker.device_ctx;
70
71        let d_records = records.to_device_on(device_ctx).unwrap();
72        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
73
74        unsafe {
75            cuda_abi::alu256::tracegen(
76                d_trace.buffer(),
77                trace_height,
78                &d_records,
79                &self.range_checker.count,
80                &self.bitwise_lookup.count,
81                RV32_CELL_BITS,
82                self.pointer_max_bits as u32,
83                self.timestamp_max_bits as u32,
84                device_ctx.stream.as_raw(),
85            )
86            .unwrap();
87        }
88
89        AirProvingContext::simple_no_pis(d_trace)
90    }
91}
92
93//////////////////////////////////////////////////////////////////////////////////////
94/// Branch Equal
95//////////////////////////////////////////////////////////////////////////////////////
96pub type BranchEqual256AdapterRecord =
97    Rv32VecHeapBranchAdapterRecord<2, INT256_NUM_BLOCKS, DEFAULT_BLOCK_SIZE>;
98pub type BranchEqual256CoreRecord = BranchEqualCoreRecord<INT256_NUM_LIMBS>;
99
100#[derive(new)]
101pub struct BranchEqual256ChipGpu {
102    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
103    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
104    pub pointer_max_bits: usize,
105    pub timestamp_max_bits: usize,
106}
107
108impl Chip<DenseRecordArena, GpuBackend> for BranchEqual256ChipGpu {
109    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
110        const RECORD_SIZE: usize =
111            size_of::<(BranchEqual256AdapterRecord, BranchEqual256CoreRecord)>();
112        let records = arena.allocated();
113        if records.is_empty() {
114            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
115        }
116        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
117
118        let trace_width = BranchEqualCoreCols::<F, INT256_NUM_LIMBS>::width()
119            + Rv32VecHeapBranchAdapterCols::<F, 2, INT256_NUM_BLOCKS, DEFAULT_BLOCK_SIZE>::width();
120        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
121        let device_ctx = &self.range_checker.device_ctx;
122
123        let d_records = records.to_device_on(device_ctx).unwrap();
124        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
125
126        unsafe {
127            cuda_abi::beq256::tracegen(
128                d_trace.buffer(),
129                trace_height,
130                &d_records,
131                &self.range_checker.count,
132                &self.bitwise_lookup.count,
133                RV32_CELL_BITS,
134                self.pointer_max_bits as u32,
135                self.timestamp_max_bits as u32,
136                device_ctx.stream.as_raw(),
137            )
138            .unwrap();
139        }
140
141        AirProvingContext::simple_no_pis(d_trace)
142    }
143}
144
145//////////////////////////////////////////////////////////////////////////////////////
146/// Less Than
147//////////////////////////////////////////////////////////////////////////////////////
148pub type LessThan256AdapterRecord = Rv32VecHeapAdapterRecord<
149    2,
150    INT256_NUM_BLOCKS,
151    INT256_NUM_BLOCKS,
152    DEFAULT_BLOCK_SIZE,
153    DEFAULT_BLOCK_SIZE,
154>;
155pub type LessThan256CoreRecord = LessThanCoreRecord<INT256_NUM_LIMBS, RV32_CELL_BITS>;
156
157#[derive(new)]
158pub struct LessThan256ChipGpu {
159    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
160    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
161    pub pointer_max_bits: usize,
162    pub timestamp_max_bits: usize,
163}
164
165impl Chip<DenseRecordArena, GpuBackend> for LessThan256ChipGpu {
166    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
167        const RECORD_SIZE: usize = size_of::<(LessThan256AdapterRecord, LessThan256CoreRecord)>();
168        let records = arena.allocated();
169        if records.is_empty() {
170            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
171        }
172        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
173
174        let trace_width = LessThanCoreCols::<F, INT256_NUM_LIMBS, RV32_CELL_BITS>::width()
175            + Rv32VecHeapAdapterCols::<
176                F,
177                2,
178                INT256_NUM_BLOCKS,
179                INT256_NUM_BLOCKS,
180                DEFAULT_BLOCK_SIZE,
181                DEFAULT_BLOCK_SIZE,
182            >::width();
183        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
184        let device_ctx = &self.range_checker.device_ctx;
185
186        let d_records = records.to_device_on(device_ctx).unwrap();
187        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
188
189        unsafe {
190            cuda_abi::lt256::tracegen(
191                d_trace.buffer(),
192                trace_height,
193                &d_records,
194                &self.range_checker.count,
195                &self.bitwise_lookup.count,
196                RV32_CELL_BITS,
197                self.pointer_max_bits as u32,
198                self.timestamp_max_bits as u32,
199                device_ctx.stream.as_raw(),
200            )
201            .unwrap();
202        }
203
204        AirProvingContext::simple_no_pis(d_trace)
205    }
206}
207
208//////////////////////////////////////////////////////////////////////////////////////
209/// Branch Less Than
210//////////////////////////////////////////////////////////////////////////////////////
211pub type BranchLessThan256AdapterRecord =
212    Rv32VecHeapBranchAdapterRecord<2, INT256_NUM_BLOCKS, DEFAULT_BLOCK_SIZE>;
213pub type BranchLessThan256CoreRecord = BranchLessThanCoreRecord<INT256_NUM_LIMBS, RV32_CELL_BITS>;
214
215#[derive(new)]
216pub struct BranchLessThan256ChipGpu {
217    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
218    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
219    pub pointer_max_bits: usize,
220    pub timestamp_max_bits: usize,
221}
222
223impl Chip<DenseRecordArena, GpuBackend> for BranchLessThan256ChipGpu {
224    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
225        const RECORD_SIZE: usize =
226            size_of::<(BranchLessThan256AdapterRecord, BranchLessThan256CoreRecord)>();
227        let records = arena.allocated();
228        if records.is_empty() {
229            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
230        }
231        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
232
233        let trace_width = BranchLessThanCoreCols::<F, INT256_NUM_LIMBS, RV32_CELL_BITS>::width()
234            + Rv32VecHeapBranchAdapterCols::<F, 2, INT256_NUM_BLOCKS, DEFAULT_BLOCK_SIZE>::width();
235        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
236        let device_ctx = &self.range_checker.device_ctx;
237
238        let d_records = records.to_device_on(device_ctx).unwrap();
239        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
240
241        unsafe {
242            cuda_abi::blt256::tracegen(
243                d_trace.buffer(),
244                trace_height,
245                &d_records,
246                &self.range_checker.count,
247                &self.bitwise_lookup.count,
248                RV32_CELL_BITS,
249                self.pointer_max_bits as u32,
250                self.timestamp_max_bits as u32,
251                device_ctx.stream.as_raw(),
252            )
253            .unwrap();
254        }
255
256        AirProvingContext::simple_no_pis(d_trace)
257    }
258}
259
260//////////////////////////////////////////////////////////////////////////////////////
261/// Shift
262//////////////////////////////////////////////////////////////////////////////////////
263pub type Shift256AdapterRecord = Rv32VecHeapAdapterRecord<
264    2,
265    INT256_NUM_BLOCKS,
266    INT256_NUM_BLOCKS,
267    DEFAULT_BLOCK_SIZE,
268    DEFAULT_BLOCK_SIZE,
269>;
270pub type Shift256CoreRecord = ShiftCoreRecord<INT256_NUM_LIMBS, RV32_CELL_BITS>;
271
272#[derive(new)]
273pub struct Shift256ChipGpu {
274    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
275    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
276    pub pointer_max_bits: usize,
277    pub timestamp_max_bits: usize,
278}
279
280impl Chip<DenseRecordArena, GpuBackend> for Shift256ChipGpu {
281    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
282        const RECORD_SIZE: usize = size_of::<(Shift256AdapterRecord, Shift256CoreRecord)>();
283        let records = arena.allocated();
284        if records.is_empty() {
285            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
286        }
287        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
288
289        let trace_width = ShiftCoreCols::<F, INT256_NUM_LIMBS, RV32_CELL_BITS>::width()
290            + Rv32VecHeapAdapterCols::<
291                F,
292                2,
293                INT256_NUM_BLOCKS,
294                INT256_NUM_BLOCKS,
295                DEFAULT_BLOCK_SIZE,
296                DEFAULT_BLOCK_SIZE,
297            >::width();
298        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
299        let device_ctx = &self.range_checker.device_ctx;
300
301        let d_records = records.to_device_on(device_ctx).unwrap();
302        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
303
304        unsafe {
305            cuda_abi::shift256::tracegen(
306                d_trace.buffer(),
307                trace_height,
308                &d_records,
309                &self.range_checker.count,
310                &self.bitwise_lookup.count,
311                RV32_CELL_BITS,
312                self.pointer_max_bits as u32,
313                self.timestamp_max_bits as u32,
314                device_ctx.stream.as_raw(),
315            )
316            .unwrap();
317        }
318
319        AirProvingContext::simple_no_pis(d_trace)
320    }
321}
322
323//////////////////////////////////////////////////////////////////////////////////////
324/// Multiplication
325//////////////////////////////////////////////////////////////////////////////////////
326pub type Multiplication256AdapterRecord = Rv32VecHeapAdapterRecord<
327    2,
328    INT256_NUM_BLOCKS,
329    INT256_NUM_BLOCKS,
330    DEFAULT_BLOCK_SIZE,
331    DEFAULT_BLOCK_SIZE,
332>;
333pub type Multiplication256CoreRecord = MultiplicationCoreRecord<INT256_NUM_LIMBS, RV32_CELL_BITS>;
334
335#[derive(new)]
336pub struct Multiplication256ChipGpu {
337    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
338    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
339    pub range_tuple_checker: Arc<RangeTupleCheckerChipGPU<2>>,
340    pub pointer_max_bits: usize,
341    pub timestamp_max_bits: usize,
342}
343
344impl Chip<DenseRecordArena, GpuBackend> for Multiplication256ChipGpu {
345    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
346        const RECORD_SIZE: usize =
347            size_of::<(Multiplication256AdapterRecord, Multiplication256CoreRecord)>();
348        let records = arena.allocated();
349        if records.is_empty() {
350            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
351        }
352        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
353
354        let trace_width = MultiplicationCoreCols::<F, INT256_NUM_LIMBS, RV32_CELL_BITS>::width()
355            + Rv32VecHeapAdapterCols::<
356                F,
357                2,
358                INT256_NUM_BLOCKS,
359                INT256_NUM_BLOCKS,
360                DEFAULT_BLOCK_SIZE,
361                DEFAULT_BLOCK_SIZE,
362            >::width();
363        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
364        let device_ctx = &self.range_checker.device_ctx;
365
366        let d_records = records.to_device_on(device_ctx).unwrap();
367        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
368
369        let sizes = self.range_tuple_checker.sizes;
370        let d_sizes = UInt2 {
371            x: sizes[0],
372            y: sizes[1],
373        };
374        unsafe {
375            cuda_abi::mul256::tracegen(
376                d_trace.buffer(),
377                trace_height,
378                &d_records,
379                &self.range_checker.count,
380                &self.bitwise_lookup.count,
381                RV32_CELL_BITS,
382                &self.range_tuple_checker.count,
383                d_sizes,
384                self.pointer_max_bits as u32,
385                self.timestamp_max_bits as u32,
386                device_ctx.stream.as_raw(),
387            )
388            .unwrap();
389        }
390
391        AirProvingContext::simple_no_pis(d_trace)
392    }
393}
openvm_bigint_circuit/cuda/mod.rs

openvm_bigint_circuit/cuda/
mod.rs
