openvm_rv32im_circuit/less_than/
cuda.rs

1use std::{mem::size_of, sync::Arc};
2
3use derive_new::new;
4use openvm_circuit::{arch::DenseRecordArena, utils::next_power_of_two_or_zero};
5use openvm_circuit_primitives::{
6    bitwise_op_lookup::BitwiseOperationLookupChipGPU, var_range::VariableRangeCheckerChipGPU,
7};
8use openvm_cuda_backend::{
9    base::DeviceMatrix, chip::get_empty_air_proving_ctx, prover_backend::GpuBackend, types::F,
10};
11use openvm_cuda_common::copy::MemCopyH2D;
12use openvm_stark_backend::{prover::types::AirProvingContext, Chip};
13
14use crate::{
15    adapters::{
16        Rv32BaseAluAdapterCols, Rv32BaseAluAdapterRecord, RV32_CELL_BITS, RV32_REGISTER_NUM_LIMBS,
17    },
18    cuda_abi::less_than_cuda::tracegen,
19    LessThanCoreCols, LessThanCoreRecord,
20};
21
22#[derive(new)]
23pub struct Rv32LessThanChipGpu {
24    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
25    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
26    pub timestamp_max_bits: usize,
27}
28
29impl Chip<DenseRecordArena, GpuBackend> for Rv32LessThanChipGpu {
30    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
31        const RECORD_SIZE: usize = size_of::<(
32            Rv32BaseAluAdapterRecord,
33            LessThanCoreRecord<RV32_REGISTER_NUM_LIMBS, RV32_CELL_BITS>,
34        )>();
35        let records = arena.allocated();
36        if records.is_empty() {
37            return get_empty_air_proving_ctx::<GpuBackend>();
38        }
39        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
40
41        let trace_width = Rv32BaseAluAdapterCols::<F>::width()
42            + LessThanCoreCols::<F, RV32_REGISTER_NUM_LIMBS, RV32_CELL_BITS>::width();
43        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
44
45        let d_records = records.to_device().unwrap();
46        let d_trace = DeviceMatrix::<F>::with_capacity(trace_height, trace_width);
47
48        unsafe {
49            tracegen(
50                d_trace.buffer(),
51                trace_height,
52                &d_records,
53                &self.range_checker.count,
54                &self.bitwise_lookup.count,
55                RV32_CELL_BITS,
56                self.timestamp_max_bits as u32,
57            )
58            .unwrap();
59        }
60        AirProvingContext::simple_no_pis(d_trace)
61    }
62}