openvm_deferral_circuit/output/
cuda.rs

1use std::{array::from_fn, mem::size_of, sync::Arc};
2
3use derive_new::new;
4use openvm_circuit::{
5    arch::{DenseRecordArena, SizedRecord},
6    utils::next_power_of_two_or_zero,
7};
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};
13use openvm_instructions::riscv::RV32_CELL_BITS;
14use openvm_stark_backend::{p3_field::PrimeCharacteristicRing, prover::AirProvingContext};
15use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
16
17use super::{
18    DeferralOutputCols, DeferralOutputLayout, DeferralOutputMetadata, DeferralOutputRecordHeader,
19    DeferralOutputRecordMut,
20};
21use crate::{
22    cuda_abi::output::{self, DeferralOutputPerCall, DeferralOutputPerRow},
23    poseidon2::{deferral_poseidon2_chip, DeferralPoseidon2SharedBuffer},
24    utils::f_commit_to_bytes,
25};
26
27#[derive(new)]
28pub struct DeferralOutputChipGpu {
29    pub range_checker: Arc<VariableRangeCheckerChipGPU>,
30    pub bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<RV32_CELL_BITS>>,
31    pub address_bits: usize,
32    pub timestamp_max_bits: usize,
33    pub count: Arc<DeviceBuffer<u32>>,
34    pub num_deferral_circuits: usize,
35    pub poseidon2: DeferralPoseidon2SharedBuffer,
36}
37
38impl Chip<DenseRecordArena, GpuBackend> for DeferralOutputChipGpu {
39    fn generate_proving_ctx(&self, mut arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
40        let records = arena.allocated_mut();
41        if records.is_empty() {
42            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
43        }
44
45        let poseidon2_chip = deferral_poseidon2_chip::<F>();
46        let mut per_call = Vec::<DeferralOutputPerCall>::new();
47        let mut per_row = Vec::<DeferralOutputPerRow>::new();
48
49        let mut offset = 0usize;
50        while offset < records.len() {
51            let header_offset = offset;
52            let header: &DeferralOutputRecordHeader = unsafe {
53                &*(records.as_ptr().add(header_offset) as *const DeferralOutputRecordHeader)
54            };
55
56            let num_rows = header.num_rows as usize;
57            let output_len = (num_rows - 1) * DIGEST_SIZE;
58            let write_bytes = unsafe {
59                std::slice::from_raw_parts(
60                    records
61                        .as_ptr()
62                        .add(header_offset + size_of::<DeferralOutputRecordHeader>()),
63                    output_len,
64                )
65            };
66
67            let mut current_poseidon2_res = [F::ZERO; DIGEST_SIZE];
68            let call_idx =
69                u32::try_from(per_call.len()).expect("deferral output call index should fit u32");
70            let header_offset_u32 =
71                u32::try_from(header_offset).expect("record byte offset should fit u32");
72
73            for section_idx in 0..num_rows {
74                let sponge_inputs = if section_idx == 0 {
75                    let mut input = [F::ZERO; DIGEST_SIZE];
76                    input[0] = F::from_u32(header.deferral_idx);
77                    input[1] = F::from_usize(output_len);
78                    input
79                } else {
80                    let base = (section_idx - 1) * DIGEST_SIZE;
81                    from_fn(|i| F::from_u8(write_bytes[base + i]))
82                };
83
84                current_poseidon2_res = poseidon2_chip.perm(
85                    &sponge_inputs,
86                    &current_poseidon2_res,
87                    section_idx + 1 == num_rows,
88                );
89
90                per_row.push(DeferralOutputPerRow {
91                    header_offset: header_offset_u32,
92                    section_idx: section_idx as u32,
93                    call_idx,
94                    poseidon2_res: current_poseidon2_res,
95                });
96            }
97
98            per_call.push(DeferralOutputPerCall {
99                output_commit: f_commit_to_bytes(&current_poseidon2_res),
100            });
101
102            let layout = DeferralOutputLayout::new(DeferralOutputMetadata { num_rows });
103            let record_size =
104                <DeferralOutputRecordMut<'_> as SizedRecord<DeferralOutputLayout>>::size(&layout);
105            let record_alignment = <DeferralOutputRecordMut<'_> as SizedRecord<
106                DeferralOutputLayout,
107            >>::alignment(&layout);
108            offset += record_size.next_multiple_of(record_alignment);
109        }
110        debug_assert_eq!(offset, records.len());
111
112        let rows_used = per_row.len();
113        let trace_height = next_power_of_two_or_zero(rows_used);
114        let trace_width = DeferralOutputCols::<F>::width();
115        let device_ctx = &self.range_checker.device_ctx;
116        let trace = DeviceMatrix::<F>::with_capacity_on(trace_height, trace_width, device_ctx);
117
118        let d_raw_records = records.to_device_on(device_ctx).unwrap();
119        let d_per_call = per_call.to_device_on(device_ctx).unwrap();
120        let d_per_row = per_row.to_device_on(device_ctx).unwrap();
121
122        unsafe {
123            output::tracegen(
124                trace.buffer(),
125                trace_height,
126                trace_width,
127                &d_raw_records,
128                &d_per_call,
129                &d_per_row,
130                rows_used,
131                &self.count,
132                self.num_deferral_circuits,
133                &self.range_checker.count,
134                self.timestamp_max_bits as u32,
135                &self.bitwise_lookup.count,
136                RV32_CELL_BITS as u32,
137                self.address_bits,
138                &self.poseidon2.records,
139                &self.poseidon2.counts,
140                &self.poseidon2.idx,
141                // Length in F elements; the CUDA side converts to record count.
142                self.poseidon2.records.len(),
143                device_ctx.stream.as_raw(),
144            )
145            .expect("Failed to generate deferral output trace");
146        }
147
148        AirProvingContext::simple_no_pis(trace)
149    }
150}
openvm_deferral_circuit/output/cuda.rs

openvm_deferral_circuit/output/
cuda.rs
