openvm_recursion_circuit/primitives/exp_bits_len/
cuda.rs

1use std::ops::Deref;
2
3use openvm_cuda_backend::{base::DeviceMatrix, prelude::F};
4use openvm_cuda_common::{memory_manager::MemTracker, stream::GpuDeviceCtx};
5use p3_matrix::dense::RowMajorMatrix;
6
7use super::{ExpBitsLenCols, ExpBitsLenCpuTraceGenerator};
8use crate::{cuda::to_device_or_nullptr_on, primitives::cuda_abi::exp_bits_len_tracegen};
9
10pub struct ExpBitsLenGpuTraceGenerator {
11    pub cpu: ExpBitsLenCpuTraceGenerator,
12    pub device_ctx: GpuDeviceCtx,
13}
14
15impl Deref for ExpBitsLenGpuTraceGenerator {
16    type Target = ExpBitsLenCpuTraceGenerator;
17
18    fn deref(&self) -> &Self::Target {
19        &self.cpu
20    }
21}
22
23impl ExpBitsLenGpuTraceGenerator {
24    pub fn new(device_ctx: GpuDeviceCtx) -> Self {
25        Self {
26            cpu: ExpBitsLenCpuTraceGenerator::default(),
27            device_ctx,
28        }
29    }
30
31    pub fn generate_trace_row_major(
32        self,
33        required_height: Option<usize>,
34    ) -> Option<RowMajorMatrix<F>> {
35        self.cpu.generate_trace_row_major(required_height)
36    }
37
38    #[tracing::instrument(name = "generate_trace", level = "trace", skip_all)]
39    pub fn generate_trace_device(self, required_height: Option<usize>) -> Option<DeviceMatrix<F>> {
40        let mem = MemTracker::start("tracegen.exp_bits_len");
41        let records = self.cpu.requests.into_inner().unwrap();
42        let num_valid_rows = records.last().map(|record| record.end_row()).unwrap_or(0);
43        let height = if let Some(height) = required_height {
44            if height < num_valid_rows {
45                return None;
46            }
47            height
48        } else {
49            num_valid_rows.next_power_of_two()
50        };
51        let width = ExpBitsLenCols::<u8>::width();
52
53        let trace = DeviceMatrix::with_capacity_on(height, width, &self.device_ctx);
54        trace.buffer().fill_zero_on(&self.device_ctx).unwrap();
55
56        let records = to_device_or_nullptr_on(&records, &self.device_ctx).unwrap();
57        unsafe {
58            exp_bits_len_tracegen(
59                &records,
60                records.len(),
61                trace.buffer(),
62                height,
63                num_valid_rows,
64                self.device_ctx.stream.as_raw(),
65            )
66            .unwrap();
67        }
68
69        mem.emit_metrics();
70        Some(trace)
71    }
72}
openvm_recursion_circuit/primitives/exp_bits_len/cuda.rs

openvm_recursion_circuit/primitives/exp_bits_len/
cuda.rs
