openvm_cuda_backend/
sponge.rs

1//! GPU-accelerated duplex sponge with host/device state synchronization.
2//!
3//! This module provides [`DuplexSpongeGpu`], a transcript implementation that maintains
4//! state on both host and device with explicit synchronization methods.
5
6use std::ffi::c_void;
7
8use openvm_cuda_common::{
9    copy::cuda_memcpy_on,
10    d_buffer::DeviceBuffer,
11    error::{CudaError, MemCopyError},
12    stream::GpuDeviceCtx,
13};
14use openvm_stark_backend::{
15    p3_challenger::{CanObserve, CanSample},
16    FiatShamirTranscript, StarkProtocolConfig,
17};
18use openvm_stark_sdk::config::baby_bear_poseidon2::poseidon2_perm;
19use p3_baby_bear::default_babybear_poseidon2_16;
20use p3_field::{PrimeCharacteristicRing, PrimeField32};
21use p3_symmetric::Permutation;
22
23use crate::types::{Challenger, Digest, CHUNK, F, SC, WIDTH};
24
25pub(crate) fn validate_gpu_grind_bits(bits: usize) -> Result<(), GrindError> {
26    if bits >= u32::BITS as usize || (1u64 << bits) >= u64::from(F::ORDER_U32) {
27        return Err(CudaError::new(1).into());
28    }
29    Ok(())
30}
31
32/// Device-side sponge state, matching the CUDA `DeviceSpongeState` struct.
33///
34/// This struct is `#[repr(C)]` to ensure ABI compatibility with the CUDA kernel.
35/// The state layout matches the Poseidon2 duplex sponge with overwrite mode.
36///
37/// This struct implements the same logic as `DuplexSponge` from openvm_stark_backend,
38/// but with public fields so we can sync state to/from GPU.
39#[repr(C)]
40#[derive(Clone, Debug)]
41pub struct DeviceSpongeState {
42    /// Full Poseidon2 state (WIDTH = 16 elements)
43    pub state: [F; WIDTH],
44    /// Current absorb position (0 <= absorb_idx < CHUNK)
45    pub absorb_idx: u32,
46    /// Current sample position (0 <= sample_idx <= CHUNK)
47    pub sample_idx: u32,
48}
49
50impl Default for DeviceSpongeState {
51    fn default() -> Self {
52        Self {
53            state: [F::default(); WIDTH],
54            absorb_idx: 0,
55            sample_idx: 0,
56        }
57    }
58}
59
60impl DeviceSpongeState {
61    /// Observe a value into the sponge (absorb phase).
62    ///
63    /// This matches the behavior of `DuplexSponge::observe`.
64    #[inline]
65    pub fn observe(&mut self, value: F) {
66        self.state[self.absorb_idx as usize] = value;
67        self.absorb_idx += 1;
68        if self.absorb_idx == CHUNK as u32 {
69            poseidon2_perm().permute_mut(&mut self.state);
70            self.absorb_idx = 0;
71            self.sample_idx = CHUNK as u32;
72        }
73    }
74
75    /// Sample a value from the sponge (squeeze phase).
76    ///
77    /// This matches the behavior of `DuplexSponge::sample`.
78    #[inline]
79    pub fn sample(&mut self) -> F {
80        if self.absorb_idx != 0 || self.sample_idx == 0 {
81            poseidon2_perm().permute_mut(&mut self.state);
82            self.absorb_idx = 0;
83            self.sample_idx = CHUNK as u32;
84        }
85        self.sample_idx -= 1;
86        self.state[self.sample_idx as usize]
87    }
88}
89
90impl FiatShamirTranscript<SC> for DeviceSpongeState {
91    #[inline]
92    fn observe(&mut self, value: F) {
93        DeviceSpongeState::observe(self, value);
94    }
95
96    #[inline]
97    fn sample(&mut self) -> F {
98        DeviceSpongeState::sample(self)
99    }
100
101    #[inline]
102    fn observe_commit(&mut self, digest: Digest) {
103        for x in digest {
104            self.observe(x);
105        }
106    }
107}
108
109/// GPU-accelerated duplex sponge that maintains state on both host and device.
110///
111/// The host-side state uses [`DeviceSpongeState`] which matches the behavior of
112/// `DuplexSponge` from openvm_stark_backend (and `p3_challenger::DuplexChallenger`).
113/// The device-side state is stored in GPU memory for CUDA kernel operations.
114///
115/// # State Synchronization
116///
117/// The host and device states are **independent** and must be explicitly synchronized:
118/// - Use [`sync_h2d`](Self::sync_h2d) to copy host state to device before GPU operations
119/// - Use [`sync_d2h`](Self::sync_d2h) to copy device state back to host after GPU operations
120///
121/// # Usage Example
122///
123/// ```ignore
124/// let mut sponge = DuplexSpongeGpu::default();
125///
126/// // Do some host operations
127/// sponge.observe(some_value);
128/// let challenge = sponge.sample();
129///
130/// // Sync to device before GPU grinding
131/// sponge.sync_h2d()?;
132///
133/// // ... GPU grinding kernel runs ...
134///
135/// // Sync back after GPU modifies state
136/// sponge.sync_d2h()?;
137///
138/// // Continue with host operations
139/// let next_challenge = sponge.sample();
140/// ```
141#[derive(Debug)]
142pub struct DuplexSpongeGpu {
143    /// Host-side structure
144    host: Challenger,
145    /// Device-side state buffer (allocated lazily on first sync)
146    device: DeviceBuffer<DeviceSpongeState>,
147}
148
149impl Default for DuplexSpongeGpu {
150    fn default() -> Self {
151        Self::new()
152    }
153}
154
155impl Clone for DuplexSpongeGpu {
156    fn clone(&self) -> Self {
157        // Device buffer is not cloned — caller must explicitly sync_h2d with a
158        // GpuDeviceCtx after cloning if device state is needed.
159        Self {
160            host: self.host.clone(),
161            device: DeviceBuffer::new(),
162        }
163    }
164}
165
166impl DuplexSpongeGpu {
167    /// Create a new GPU-accelerated duplex sponge with default (zeroed) state.
168    pub fn new() -> Self {
169        Self {
170            host: Challenger::new(default_babybear_poseidon2_16()),
171            device: DeviceBuffer::new(),
172        }
173    }
174
175    /// Returns true if the device buffer has been allocated.
176    pub fn is_device_allocated(&self) -> bool {
177        !self.device.is_empty()
178    }
179
180    /// Ensure the device buffer is allocated.
181    fn ensure_device_allocated(&mut self, device_ctx: &GpuDeviceCtx) {
182        if self.device.is_empty() {
183            self.device = DeviceBuffer::with_capacity_on(1, device_ctx);
184        }
185    }
186
187    /// Synchronize state from host to device (H2D memcpy).
188    ///
189    /// Call this before running GPU kernels that read/modify the sponge state.
190    ///
191    /// This converts from `DuplexChallenger`'s representation (with buffered input/output)
192    /// to `DeviceSpongeState`'s representation (with indices pointing into state).
193    pub fn sync_h2d(&mut self, device_ctx: &GpuDeviceCtx) -> Result<(), MemCopyError> {
194        self.ensure_device_allocated(device_ctx);
195
196        // Convert DuplexChallenger state to DeviceSpongeState format:
197        // - DuplexChallenger buffers input values before writing to sponge_state
198        // - DeviceSpongeState writes directly to state[absorb_idx]
199        // We need to overlay the input_buffer onto state[0..len]
200
201        let mut device_state = DeviceSpongeState {
202            state: self.host.sponge_state,
203            absorb_idx: self.host.input_buffer.len() as u32,
204            sample_idx: self.host.output_buffer.len() as u32,
205        };
206
207        // Overlay pending input_buffer values onto the beginning of state
208        // (DuplexChallenger writes to state[0..N] during duplexing, so pending
209        // values that haven't been duplexed yet need to be placed there)
210        for (i, &val) in self.host.input_buffer.iter().enumerate() {
211            device_state.state[i] = val;
212        }
213
214        // SAFETY: Copying a single DeviceSpongeState from host to device
215        // - Both pointers are valid and properly aligned
216        // - The size matches the struct size
217        unsafe {
218            cuda_memcpy_on::<false, true>(
219                self.device.as_mut_ptr() as *mut c_void,
220                &device_state as *const DeviceSpongeState as *const c_void,
221                std::mem::size_of::<DeviceSpongeState>(),
222                device_ctx,
223            )
224        }
225    }
226
227    /// Get a pointer to the device state buffer.
228    ///
229    /// Returns `None` if the device buffer hasn't been allocated yet.
230    /// Call [`sync_h2d`](Self::sync_h2d) to allocate and initialize the device buffer.
231    pub fn device_ptr(&self) -> Option<*const DeviceSpongeState> {
232        if self.device.is_empty() {
233            None
234        } else {
235            Some(self.device.as_ptr())
236        }
237    }
238
239    /// Get a mutable pointer to the device state buffer.
240    ///
241    /// Returns `None` if the device buffer hasn't been allocated yet.
242    /// Call [`sync_h2d`](Self::sync_h2d) to allocate and initialize the device buffer.
243    pub fn device_ptr_mut(&mut self) -> Option<*mut DeviceSpongeState> {
244        if self.device.is_empty() {
245            None
246        } else {
247            Some(self.device.as_mut_ptr())
248        }
249    }
250
251    /// Perform GPU-accelerated grinding to find a proof-of-work witness.
252    ///
253    /// This syncs state to device, runs the grinding kernel, and updates
254    /// the host state with the witness.
255    ///
256    /// # Arguments
257    /// * `bits` - Number of bits that must be zero in the sampled value
258    ///
259    /// # Returns
260    ///
261    /// The PoW witness value that satisfies `sample_bits(bits) == 0` after observing it.
262    ///
263    /// # Note
264    ///
265    /// After this call, the host state will have observed the witness and sampled,
266    /// matching the state after calling `check_witness(bits, witness)`.
267    pub fn grind_gpu(&mut self, bits: usize, device_ctx: &GpuDeviceCtx) -> Result<F, GrindError> {
268        validate_gpu_grind_bits(bits)?;
269        // Trivial case: 0 bits mean no PoW is required and any witness is valid.
270        if bits == 0 {
271            return Ok(F::ZERO);
272        }
273        // 1. Sync host state to device
274        self.sync_h2d(device_ctx)?;
275
276        // 2. Launch grinding kernel
277        let witness_u32 = unsafe {
278            crate::cuda::sponge::sponge_grind(
279                self.device.as_ptr(),
280                bits as u32,
281                F::ORDER_U32 - 1,
282                device_ctx,
283            )?
284        };
285
286        let witness = F::from_u32(witness_u32);
287
288        // 3. Update host state to match (observe the witness + sample)
289        // This is cheaper than syncing the full state back from device
290        debug_assert!(self.clone().check_witness(bits, witness));
291        self.host.observe(witness);
292        let _: F = self.host.sample(); // Consume the sample to advance state
293
294        Ok(witness)
295    }
296}
297
298/// Error type for GPU grinding operations.
299#[derive(Debug, thiserror::Error)]
300pub enum GrindError {
301    #[error("Memory copy error: {0}")]
302    MemCopy(#[from] MemCopyError),
303
304    #[error("CUDA error: {0}")]
305    Cuda(#[from] CudaError),
306
307    #[error("Failed to find PoW witness within search space")]
308    WitnessNotFound,
309}
310
311impl FiatShamirTranscript<SC> for DuplexSpongeGpu {
312    #[inline]
313    fn observe(&mut self, value: F) {
314        self.host.observe(value);
315    }
316
317    #[inline]
318    fn sample(&mut self) -> F {
319        self.host.sample()
320    }
321
322    #[inline]
323    fn observe_commit(&mut self, digest: Digest) {
324        for x in digest {
325            self.observe(x);
326        }
327    }
328}
329
330/// Marker trait for GPU-compatible Fiat-Shamir transcripts.
331///
332/// Extends [`FiatShamirTranscript`] with a GPU-accelerated proof-of-work grind method.
333/// Implementers maintain a device-side state buffer and can off-load the grinding loop
334/// to a CUDA kernel.
335pub trait GpuFiatShamirTranscript<Config: StarkProtocolConfig>:
336    FiatShamirTranscript<Config>
337{
338    /// GPU-accelerated proof-of-work grinding.
339    ///
340    /// Finds a witness value `w` of type `Config::F` such that after observing `w` and
341    /// sampling, the result has `bits` trailing zero bits.  Implementations should sync
342    /// host state to device, launch a CUDA grinding kernel, then update the host state
343    /// to match (observe witness + consume one sample).
344    fn grind_gpu(
345        &mut self,
346        bits: usize,
347        device_ctx: &GpuDeviceCtx,
348    ) -> Result<Config::F, GrindError>;
349}
350
351impl GpuFiatShamirTranscript<SC> for DuplexSpongeGpu {
352    fn grind_gpu(&mut self, bits: usize, device_ctx: &GpuDeviceCtx) -> Result<F, GrindError> {
353        DuplexSpongeGpu::grind_gpu(self, bits, device_ctx)
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use std::time::Instant;
360
361    use openvm_cuda_common::{
362        common::get_device,
363        stream::{CudaStream, GpuDeviceCtx, StreamGuard},
364    };
365    use openvm_stark_sdk::config::baby_bear_poseidon2::default_duplex_sponge;
366    use p3_field::PrimeCharacteristicRing;
367
368    use super::*;
369    use crate::prelude::SC;
370
371    fn test_ctx() -> GpuDeviceCtx {
372        GpuDeviceCtx {
373            device_id: get_device().unwrap() as u32,
374            stream: StreamGuard::new(CudaStream::new_non_blocking().unwrap()),
375        }
376    }
377
378    #[test]
379    fn test_device_sponge_state_size() {
380        // Verify the struct size is what we expect for FFI
381        let expected_size = std::mem::size_of::<[F; WIDTH]>() // state
382            + std::mem::size_of::<u32>() // absorb_idx
383            + std::mem::size_of::<u32>(); // sample_idx
384
385        assert_eq!(
386            std::mem::size_of::<DeviceSpongeState>(),
387            expected_size,
388            "DeviceSpongeState size mismatch - check repr(C) and padding"
389        );
390    }
391
392    #[test]
393    fn test_device_sponge_state_alignment() {
394        // Verify alignment for CUDA compatibility
395        assert!(
396            std::mem::align_of::<DeviceSpongeState>() >= 4,
397            "DeviceSpongeState should be at least 4-byte aligned"
398        );
399    }
400
401    #[test]
402    fn test_default_state() {
403        let state = DeviceSpongeState::default();
404        assert_eq!(state.absorb_idx, 0);
405        assert_eq!(state.sample_idx, 0);
406        for elem in state.state.iter() {
407            assert_eq!(*elem, F::default());
408        }
409    }
410
411    #[test]
412    fn test_sponge_gpu_new() {
413        let sponge = DuplexSpongeGpu::new();
414        assert!(!sponge.is_device_allocated());
415    }
416
417    #[test]
418    fn test_device_sponge_state_matches_duplex_sponge() {
419        // Verify our implementation matches DuplexSponge exactly
420        let mut device_state = DeviceSpongeState::default();
421        let mut duplex_sponge = default_duplex_sponge();
422
423        // Test observe/sample sequence
424        for i in 0..20 {
425            let val = F::from_u32(i * 42 + 17);
426            device_state.observe(val);
427            FiatShamirTranscript::<SC>::observe(&mut duplex_sponge, val);
428        }
429
430        for _ in 0..10 {
431            let device_sample = device_state.sample();
432            let duplex_sample = FiatShamirTranscript::<SC>::sample(&mut duplex_sponge);
433            assert_eq!(device_sample, duplex_sample);
434        }
435
436        // Interleaved observe/sample
437        for i in 0..5 {
438            let val = F::from_u32(i * 100);
439            device_state.observe(val);
440            FiatShamirTranscript::<SC>::observe(&mut duplex_sponge, val);
441
442            let device_sample = device_state.sample();
443            let duplex_sample = FiatShamirTranscript::<SC>::sample(&mut duplex_sponge);
444            assert_eq!(device_sample, duplex_sample);
445        }
446
447        // Many samples in a row
448        for _ in 0..15 {
449            let device_sample = device_state.sample();
450            let duplex_sample = FiatShamirTranscript::<SC>::sample(&mut duplex_sponge);
451            assert_eq!(device_sample, duplex_sample);
452        }
453    }
454
455    #[test]
456    fn test_sponge_gpu_uses_host_transcript() {
457        let mut gpu_sponge = DuplexSpongeGpu::default();
458        let mut cpu_sponge = default_duplex_sponge();
459
460        // Test that host operations match DuplexSponge
461        for i in 0..10 {
462            let val = F::from_u32(i * 42 + 17);
463            gpu_sponge.observe(val);
464            FiatShamirTranscript::<SC>::observe(&mut cpu_sponge, val);
465        }
466
467        for _ in 0..5 {
468            let gpu_sample = gpu_sponge.sample();
469            let cpu_sample = FiatShamirTranscript::<SC>::sample(&mut cpu_sponge);
470            assert_eq!(gpu_sample, cpu_sample);
471        }
472    }
473
474    /// Benchmark test comparing CPU vs GPU grinding performance.
475    ///
476    /// Run with: `cargo test -p openvm-cuda-backend test_grind_cpu_vs_gpu -- --nocapture`
477    ///
478    /// Note: GPU has ~20ms fixed overhead for kernel launch + sync. CPU wins for
479    /// small search spaces (low bit counts). GPU wins for larger search spaces
480    /// where parallelism amortizes the overhead.
481    #[test]
482    fn test_grind_cpu_vs_gpu() {
483        let device_ctx = test_ctx();
484        // Warmup: run one GPU grind to initialize CUDA context
485        {
486            let mut warmup = DuplexSpongeGpu::default();
487            let _ = warmup.grind_gpu(8, &device_ctx);
488        }
489
490        // Test multiple bit counts to see scaling
491        let bit_counts = [8, 12, 16, 18, 20]
492            .iter()
493            .flat_map(|x| std::iter::repeat_n(*x, 5))
494            .collect::<Vec<_>>();
495
496        eprintln!("\n{}", "=".repeat(60));
497        eprintln!("Grinding Performance: CPU vs GPU");
498        eprintln!("{}", "=".repeat(60));
499        eprintln!(
500            "{:>6} {:>12} {:>12} {:>10}",
501            "bits", "CPU (ms)", "GPU (ms)", "speedup"
502        );
503        eprintln!("{:->6} {:->12} {:->12} {:->10}", "", "", "", "");
504
505        let mut seed = 265;
506        for bits in bit_counts {
507            let mut cpu_sponge = default_duplex_sponge();
508            let mut gpu_sponge = DuplexSpongeGpu::default();
509
510            // Add some initial state
511            for _ in 0..5 {
512                let val = F::from_u32(seed);
513                seed += 228;
514                FiatShamirTranscript::<SC>::observe(&mut cpu_sponge, val);
515                gpu_sponge.observe(val);
516            }
517
518            // Time CPU grinding
519            let cpu_start = Instant::now();
520            let cpu_witness = FiatShamirTranscript::<SC>::grind(&mut cpu_sponge, bits);
521            let cpu_time = cpu_start.elapsed();
522
523            // Time GPU grinding
524            let gpu_start = Instant::now();
525            let gpu_witness = gpu_sponge
526                .grind_gpu(bits, &device_ctx)
527                .expect("GPU grinding failed");
528            let gpu_time = gpu_start.elapsed();
529
530            // Verify both found valid witnesses (witnesses may differ but both should be valid)
531            // We already validated inside grind_gpu with debug_assert
532
533            let speedup = cpu_time.as_secs_f64() / gpu_time.as_secs_f64();
534
535            eprintln!(
536                "{:>6} {:>12.2} {:>12.2} {:>10.2}x",
537                bits,
538                cpu_time.as_secs_f64() * 1000.0,
539                gpu_time.as_secs_f64() * 1000.0,
540                speedup
541            );
542
543            // Verify the witnesses are valid by checking with a fresh sponge
544            // (grind() and grind_gpu() already do this internally via check_witness)
545            let _ = (cpu_witness, gpu_witness); // suppress unused warnings
546        }
547
548        eprintln!("{}\n", "=".repeat(60));
549    }
550}
openvm_cuda_backend/sponge.rs

openvm_cuda_backend/
sponge.rs
