openvm_sha2_circuit/cuda/
mod.rs

1use std::{
2    marker::PhantomData,
3    sync::{Arc, Mutex},
4};
5
6use openvm_circuit::{
7    arch::{DenseRecordArena, RecordSeeker},
8    utils::next_power_of_two_or_zero,
9};
10use openvm_circuit_primitives::{
11    bitwise_op_lookup::BitwiseOperationLookupChipGPU, var_range::VariableRangeCheckerChipGPU, Chip,
12};
13use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
14use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer};
15use openvm_sha2_air::{Sha256Config, Sha2Variant, Sha512Config};
16use openvm_stark_backend::prover::AirProvingContext;
17
18use crate::{Sha2Config, Sha2RecordLayout, Sha2RecordMut};
19
20mod cuda_abi;
21
22pub struct Sha2SharedRecordsGpu {
23    d_records: DeviceBuffer<u8>,
24    d_record_offsets: DeviceBuffer<usize>,
25    num_records: usize,
26}
27
28pub struct Sha2MainChipGpu<C: Sha2Config> {
29    records: Arc<Mutex<Option<Sha2SharedRecordsGpu>>>,
30    range_checker: Arc<VariableRangeCheckerChipGPU>,
31    bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<8>>,
32    pointer_max_bits: u32,
33    timestamp_max_bits: u32,
34    _marker: PhantomData<C>,
35}
36
37impl<C: Sha2Config> Sha2MainChipGpu<C> {
38    pub fn new(
39        records: Arc<Mutex<Option<Sha2SharedRecordsGpu>>>,
40        range_checker: Arc<VariableRangeCheckerChipGPU>,
41        bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<8>>,
42        pointer_max_bits: u32,
43        timestamp_max_bits: u32,
44    ) -> Self {
45        Self {
46            records,
47            range_checker,
48            bitwise_lookup,
49            pointer_max_bits,
50            timestamp_max_bits,
51            _marker: PhantomData,
52        }
53    }
54}
55
56impl<C> Chip<DenseRecordArena, GpuBackend> for Sha2MainChipGpu<C>
57where
58    C: Sha2Config,
59{
60    fn generate_proving_ctx(&self, mut arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
61        let records = arena.allocated_mut();
62        if records.is_empty() {
63            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
64        }
65
66        let mut record_offsets = Vec::<usize>::new();
67        let mut offset = 0usize;
68        while offset < records.len() {
69            record_offsets.push(offset);
70            let _record =
71                RecordSeeker::<DenseRecordArena, Sha2RecordMut, Sha2RecordLayout>::get_record_at(
72                    &mut offset,
73                    records,
74                );
75        }
76
77        let num_records = record_offsets.len();
78        let trace_height = next_power_of_two_or_zero(num_records);
79        let device_ctx = &self.range_checker.device_ctx;
80        let trace =
81            DeviceMatrix::<F>::with_capacity_on(trace_height, C::MAIN_CHIP_WIDTH, device_ctx);
82
83        let d_records = records.to_device_on(device_ctx).unwrap();
84        let d_record_offsets = record_offsets.to_device_on(device_ctx).unwrap();
85
86        unsafe {
87            match C::VARIANT {
88                Sha2Variant::Sha256 => {
89                    cuda_abi::sha256::sha256_main_tracegen(
90                        trace.buffer(),
91                        trace_height,
92                        &d_records,
93                        num_records,
94                        &d_record_offsets,
95                        self.pointer_max_bits,
96                        &self.range_checker.count,
97                        &self.bitwise_lookup.count,
98                        8,
99                        self.timestamp_max_bits,
100                        device_ctx.stream.as_raw(),
101                    )
102                    .unwrap();
103                }
104                Sha2Variant::Sha512 | Sha2Variant::Sha384 => {
105                    cuda_abi::sha512::sha512_main_tracegen(
106                        trace.buffer(),
107                        trace_height,
108                        &d_records,
109                        num_records,
110                        &d_record_offsets,
111                        self.pointer_max_bits,
112                        &self.range_checker.count,
113                        &self.bitwise_lookup.count,
114                        8,
115                        self.timestamp_max_bits,
116                        device_ctx.stream.as_raw(),
117                    )
118                    .unwrap();
119                }
120            }
121        }
122
123        // Pass the records to Sha2BlockHasherChip
124        *self.records.lock().unwrap() = Some(Sha2SharedRecordsGpu {
125            d_records,
126            d_record_offsets,
127            num_records,
128        });
129
130        AirProvingContext::simple_no_pis(trace)
131    }
132}
133
134/// Generic hybrid GPU wrapper that reuses CPU block-hasher tracegen.
135pub struct Sha2BlockHasherChipGpu<C: Sha2Config> {
136    records: Arc<Mutex<Option<Sha2SharedRecordsGpu>>>,
137    bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<8>>,
138    _marker: PhantomData<C>,
139}
140
141impl<C, R> Chip<R, GpuBackend> for Sha2BlockHasherChipGpu<C>
142where
143    C: Sha2Config,
144{
145    /// We don't use the record arena associated with this chip. Instead, we will use the record
146    /// arena provided by the main chip, which will be passed to this chip after the main chip's
147    /// tracegen is done.
148    fn generate_proving_ctx(&self, _: R) -> AirProvingContext<GpuBackend> {
149        let mut records = self.records.lock().unwrap();
150        if records.is_none() {
151            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
152        }
153
154        let Sha2SharedRecordsGpu {
155            d_records,
156            d_record_offsets,
157            num_records,
158        } = records.take().unwrap();
159
160        if num_records == 0 {
161            return AirProvingContext::simple_no_pis(DeviceMatrix::dummy());
162        }
163
164        let rows_used = num_records * C::ROWS_PER_BLOCK;
165        let trace_height = next_power_of_two_or_zero(rows_used);
166        let device_ctx = &self.bitwise_lookup.device_ctx;
167        let trace =
168            DeviceMatrix::<F>::with_capacity_on(trace_height, C::BLOCK_HASHER_WIDTH, device_ctx);
169
170        // one record per block, right now
171        let num_blocks: u32 = num_records as u32;
172
173        // prev_hashes
174        unsafe {
175            match C::VARIANT {
176                Sha2Variant::Sha256 => {
177                    let d_prev_hashes = DeviceBuffer::<u32>::with_capacity_on(
178                        num_blocks as usize * C::HASH_WORDS,
179                        device_ctx,
180                    );
181                    cuda_abi::sha256::sha256_hash_computation(
182                        &d_records,
183                        num_records,
184                        &d_record_offsets,
185                        &d_prev_hashes,
186                        num_blocks,
187                        device_ctx.stream.as_raw(),
188                    )
189                    .unwrap();
190
191                    // Scratch for three-phase tracegen: state[8] + w_buf[BLOCK_WORDS] u32s per
192                    // row per block.
193                    // 17 rows * (8 + 16) * 4 bytes = 1632 bytes/block, vs
194                    // 17 * 456 * 4 = 31008 bytes/block for the trace matrix (~5.3% overhead).
195                    let scratch_words_per_block = C::ROWS_PER_BLOCK * (8 + C::BLOCK_WORDS);
196                    let d_scratch = DeviceBuffer::<u32>::with_capacity_on(
197                        num_blocks as usize * scratch_words_per_block,
198                        device_ctx,
199                    );
200
201                    cuda_abi::sha256::sha256_first_pass_tracegen(
202                        trace.buffer(),
203                        trace_height,
204                        &d_records,
205                        num_records,
206                        &d_record_offsets,
207                        num_blocks,
208                        &d_prev_hashes,
209                        &self.bitwise_lookup.count,
210                        8,
211                        &d_scratch,
212                        device_ctx.stream.as_raw(),
213                    )
214                    .unwrap();
215
216                    cuda_abi::sha256::sha256_fill_invalid_rows(
217                        trace.buffer(),
218                        trace_height,
219                        rows_used,
220                        &d_prev_hashes,
221                        device_ctx.stream.as_raw(),
222                    )
223                    .unwrap();
224                    cuda_abi::sha256::sha256_second_pass_dependencies(
225                        trace.buffer(),
226                        trace_height,
227                        rows_used,
228                        device_ctx.stream.as_raw(),
229                    )
230                    .unwrap();
231                }
232                Sha2Variant::Sha512 | Sha2Variant::Sha384 => {
233                    let d_prev_hashes = DeviceBuffer::<u64>::with_capacity_on(
234                        num_blocks as usize * C::HASH_WORDS,
235                        device_ctx,
236                    );
237                    cuda_abi::sha512::sha512_hash_computation(
238                        &d_records,
239                        num_records,
240                        &d_record_offsets,
241                        &d_prev_hashes,
242                        num_blocks,
243                        device_ctx.stream.as_raw(),
244                    )
245                    .unwrap();
246
247                    // Scratch for three-phase tracegen: state[8] + w_buf[BLOCK_WORDS] u64s per
248                    // row per block.
249                    // 21 rows * (8 + 16) * 8 bytes = 4032 bytes/block, vs
250                    // 21 * 903 * 4 = 75852 bytes/block for the trace matrix (~5.3% overhead).
251                    let scratch_words_per_block = C::ROWS_PER_BLOCK * (8 + C::BLOCK_WORDS);
252                    let d_scratch = DeviceBuffer::<u64>::with_capacity_on(
253                        num_blocks as usize * scratch_words_per_block,
254                        device_ctx,
255                    );
256
257                    cuda_abi::sha512::sha512_first_pass_tracegen(
258                        trace.buffer(),
259                        trace_height,
260                        &d_records,
261                        num_records,
262                        &d_record_offsets,
263                        num_blocks,
264                        &d_prev_hashes,
265                        &self.bitwise_lookup.count,
266                        8,
267                        &d_scratch,
268                        device_ctx.stream.as_raw(),
269                    )
270                    .unwrap();
271
272                    cuda_abi::sha512::sha512_fill_invalid_rows(
273                        trace.buffer(),
274                        trace_height,
275                        rows_used,
276                        &d_prev_hashes,
277                        device_ctx.stream.as_raw(),
278                    )
279                    .unwrap();
280                    cuda_abi::sha512::sha512_second_pass_dependencies(
281                        trace.buffer(),
282                        trace_height,
283                        rows_used,
284                        device_ctx.stream.as_raw(),
285                    )
286                    .unwrap();
287                }
288            }
289        }
290
291        AirProvingContext::simple_no_pis(trace)
292    }
293}
294
295impl<C: Sha2Config> Sha2BlockHasherChipGpu<C> {
296    pub fn new(
297        records: Arc<Mutex<Option<Sha2SharedRecordsGpu>>>,
298        bitwise_lookup: Arc<BitwiseOperationLookupChipGPU<8>>,
299    ) -> Self {
300        Self {
301            records,
302            bitwise_lookup,
303            _marker: PhantomData,
304        }
305    }
306}
307
308// Convenience aliases for the common SHA-2 variants.
309pub type Sha256VmChipGpu = Sha2MainChipGpu<Sha256Config>;
310pub type Sha256BlockHasherChipGpu = Sha2BlockHasherChipGpu<Sha256Config>;
311pub type Sha512VmChipGpu = Sha2MainChipGpu<Sha512Config>;
312pub type Sha512BlockHasherChipGpu = Sha2BlockHasherChipGpu<Sha512Config>;
openvm_sha2_circuit/cuda/mod.rs

openvm_sha2_circuit/cuda/
mod.rs
