openvm_rv32im_circuit/auipc/
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, Chip,
7};
8use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
9use openvm_cuda_common::copy::MemCopyH2D;
10use openvm_stark_backend::prover::AirProvingContext;
11
12use crate::{
13    adapters::{Rv32RdWriteAdapterCols, Rv32RdWriteAdapterRecord, RV32_CELL_BITS},
14    cuda_abi::auipc_cuda::tracegen,
15    Rv32AuipcCoreCols, Rv32AuipcCoreRecord,
16};
17
18#[derive(new)]
19pub struct Rv32AuipcChipGpu {
20    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
21    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
22    pub timestamp_max_bits: usize,
23}
24
25impl Chip<DenseRecordArena, GpuBackend> for Rv32AuipcChipGpu {
26    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
27        const RECORD_SIZE: usize = size_of::<(Rv32RdWriteAdapterRecord, Rv32AuipcCoreRecord)>();
28        let records = arena.allocated();
29        if records.is_empty() {
30            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
31        }
32        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
33
34        let trace_width = Rv32AuipcCoreCols::<F>::width() + Rv32RdWriteAdapterCols::<F>::width();
35        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
36        let device_ctx = &self.range_checker.device_ctx;
37
38        let d_records = tracing::info_span!("trace_gen.h2d_records")
39            .in_scope(|| records.to_device_on(device_ctx))
40            .unwrap();
41        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
42
43        unsafe {
44            tracegen(
45                d_trace.buffer(),
46                trace_height,
47                &d_records,
48                &self.range_checker.count,
49                &self.bitwise_lookup.count,
50                RV32_CELL_BITS,
51                self.timestamp_max_bits as u32,
52                device_ctx.stream.as_raw(),
53            )
54            .unwrap();
55        }
56        AirProvingContext::simple_no_pis(d_trace)
57    }
58}
openvm_rv32im_circuit/auipc/cuda.rs

openvm_rv32im_circuit/auipc/
cuda.rs
