openvm_keccak256_circuit/cuda/
mod.rs

1use std::{
2    mem::size_of,
3    sync::{Arc, Mutex},
4};
5
6use derive_new::new;
7use openvm_circuit::{arch::DenseRecordArena, utils::next_power_of_two_or_zero};
8use openvm_circuit_primitives::{
9    bitwise_op_lookup::BitwiseOperationLookupChipGPU, var_range::VariableRangeCheckerChipGPU, Chip,
10};
11use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
12use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
13use openvm_instructions::riscv::RV32_CELL_BITS;
14use openvm_stark_backend::prover::AirProvingContext;
15use p3_keccak_air::NUM_ROUNDS;
16
17use crate::{
18    keccakf_op::{columns::NUM_KECCAKF_OP_COLS, trace::KeccakfRecord, NUM_OP_ROWS_PER_INS},
19    keccakf_perm::NUM_KECCAKF_PERM_COLS,
20    xorin::{columns::NUM_XORIN_VM_COLS, trace::XorinVmRecordHeader},
21};
22
23mod cuda_abi;
24
25// ========================== XorinVmChipGpu ==========================
26
27#[derive(new)]
28pub struct XorinVmChipGpu {
29    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
30    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
31    pub pointer_max_bits: usize,
32    pub timestamp_max_bits: u32,
33}
34
35impl Chip<DenseRecordArena, GpuBackend> for XorinVmChipGpu {
36    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
37        const RECORD_SIZE: usize = size_of::<XorinVmRecordHeader>();
38        let records = arena.allocated();
39        if records.is_empty() {
40            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
41        }
42        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
43
44        let trace_width = NUM_XORIN_VM_COLS;
45        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
46        let device_ctx = &self.range_checker.device_ctx;
47
48        let d_records = records.to_device_on(device_ctx).unwrap();
49        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
50
51        unsafe {
52            cuda_abi::xorin::tracegen(
53                d_trace.buffer(),
54                trace_height,
55                &d_records,
56                &self.range_checker.count,
57                &self.bitwise_lookup.count,
58                RV32_CELL_BITS,
59                self.pointer_max_bits as u32,
60                self.timestamp_max_bits,
61                device_ctx.stream.as_raw(),
62            )
63            .unwrap();
64        }
65
66        AirProvingContext::simple_no_pis(d_trace)
67    }
68}
69
70// ========================== Shared state for KeccakfOp <-> KeccakfPerm ==========================
71
72/// Shared state to pass records from KeccakfOpChipGpu to KeccakfPermChipGpu
73/// The OpChip generates first and stores the device buffer, then PermChip takes it.
74#[derive(Default)]
75pub struct SharedKeccakfRecords {
76    /// Device buffer containing records (set by OpChip, consumed by PermChip)
77    pub d_records: Option<DeviceBuffer<u8>>,
78    /// Number of records
79    pub num_records: usize,
80}
81
82pub type SharedKeccakfRecordsGpu = Arc<Mutex<SharedKeccakfRecords>>;
83
84// ========================== KeccakfOpChipGpu ==========================
85
86#[derive(new)]
87pub struct KeccakfOpChipGpu {
88    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
89    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
90    pub pointer_max_bits: usize,
91    pub timestamp_max_bits: u32,
92    pub shared_records: SharedKeccakfRecordsGpu,
93}
94
95impl Chip<DenseRecordArena, GpuBackend> for KeccakfOpChipGpu {
96    fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
97        const RECORD_SIZE: usize = size_of::<KeccakfRecord>();
98        let records = arena.allocated();
99        if records.is_empty() {
100            // Store empty state for PermChip
101            let mut shared = self.shared_records.lock().unwrap();
102            shared.d_records = None;
103            shared.num_records = 0;
104            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
105        }
106        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
107
108        let num_records = records.len() / RECORD_SIZE;
109        let trace_width = NUM_KECCAKF_OP_COLS;
110        let trace_height = next_power_of_two_or_zero(num_records * NUM_OP_ROWS_PER_INS);
111        let device_ctx = &self.range_checker.device_ctx;
112
113        // Transfer records to GPU
114        let d_records = records.to_device_on(device_ctx).unwrap();
115        let d_trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
116
117        unsafe {
118            cuda_abi::keccakf_op::tracegen(
119                d_trace.buffer(),
120                trace_height,
121                &d_records,
122                &self.range_checker.count,
123                &self.bitwise_lookup.count,
124                RV32_CELL_BITS,
125                self.pointer_max_bits as u32,
126                self.timestamp_max_bits,
127                device_ctx.stream.as_raw(),
128            )
129            .unwrap();
130        }
131
132        // Store records in shared state for PermChip
133        {
134            let mut shared = self.shared_records.lock().unwrap();
135            shared.d_records = Some(d_records);
136            shared.num_records = num_records;
137        }
138
139        AirProvingContext::simple_no_pis(d_trace)
140    }
141}
142
143// ========================== KeccakfPermChipGpu ==========================
144
145#[derive(new)]
146pub struct KeccakfPermChipGpu {
147    pub shared_records: SharedKeccakfRecordsGpu,
148    pub device_ctx: GpuDeviceCtx,
149}
150
151impl Chip<DenseRecordArena, GpuBackend> for KeccakfPermChipGpu {
152    fn generate_proving_ctx(&self, _arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
153        // Take records from shared state (set by OpChip)
154        let (d_records, num_records) = {
155            let mut shared = self.shared_records.lock().unwrap();
156            (shared.d_records.take(), shared.num_records)
157        };
158
159        let Some(d_records) = d_records else {
160            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
161        };
162
163        if num_records == 0 {
164            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
165        }
166
167        let trace_width = NUM_KECCAKF_PERM_COLS;
168        let trace_height = next_power_of_two_or_zero(num_records * NUM_ROUNDS);
169
170        let d_trace =
171            DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, &self.device_ctx);
172        // Scratch buffer for two-phase tracegen: 25 u64 lanes per round per permutation.
173        // 24 rounds * 25 lanes * 8 bytes = 4800 bytes/perm, vs 24 * 2634 * 4 = 252864 bytes/perm
174        // for the trace matrix (~1.9% overhead).
175        let blocks_to_fill = trace_height.div_ceil(NUM_ROUNDS);
176        let d_round_states = DeviceBuffer::<u64>::with_capacity_on(
177            blocks_to_fill * NUM_ROUNDS * 25,
178            &self.device_ctx,
179        );
180
181        unsafe {
182            cuda_abi::keccakf_perm::tracegen(
183                d_trace.buffer(),
184                trace_height,
185                &d_records,
186                num_records,
187                &d_round_states,
188                self.device_ctx.stream.as_raw(),
189            )
190            .unwrap();
191        }
192
193        AirProvingContext::simple_no_pis(d_trace)
194    }
195}
openvm_keccak256_circuit/cuda/mod.rs

openvm_keccak256_circuit/cuda/
mod.rs
