openvm_cuda_common/
stream.rs

1use std::{
2    borrow::Cow,
3    ffi::c_void,
4    ops::Deref,
5    sync::{Arc, Mutex},
6};
7
8use crate::error::{check, CudaError};
9
10#[link(name = "cudart")]
11extern "C" {
12    fn cudaDeviceSynchronize() -> i32;
13    fn cudaStreamCreateWithFlags(stream: *mut cudaStream_t, flags: u32) -> i32;
14    fn cudaStreamDestroy(stream: cudaStream_t) -> i32;
15    fn cudaStreamSynchronize(stream: cudaStream_t) -> i32;
16    fn cudaStreamWaitEvent(stream: cudaStream_t, event: cudaEvent_t, flags: u32) -> i32;
17    fn cudaEventCreate(event: *mut cudaEvent_t) -> i32;
18    fn cudaEventRecord(event: cudaEvent_t, stream: cudaStream_t) -> i32;
19    fn cudaEventSynchronize(event: cudaEvent_t) -> i32;
20    fn cudaEventQuery(event: cudaEvent_t) -> i32;
21    fn cudaEventDestroy(event: cudaEvent_t) -> i32;
22    fn cudaEventElapsedTime(ms: *mut f32, start: cudaEvent_t, end: cudaEvent_t) -> i32;
23}
24
25pub fn device_synchronize() -> Result<(), CudaError> {
26    check(unsafe { cudaDeviceSynchronize() })
27}
28
29#[allow(non_camel_case_types)]
30pub type cudaStream_t = *mut c_void;
31
32pub struct CudaStream {
33    stream: cudaStream_t,
34    host_event: Mutex<CudaEvent>,
35}
36
37impl std::fmt::Debug for CudaStream {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        f.debug_struct("CudaStream")
40            .field("stream", &self.stream)
41            .finish()
42    }
43}
44
45unsafe impl Send for CudaStream {}
46unsafe impl Sync for CudaStream {}
47
48/// `cudaStreamNonBlocking` flag: no implicit synchronization with the legacy
49/// default stream (stream 0).
50const CUDA_STREAM_NON_BLOCKING: u32 = 0x1;
51
52impl CudaStream {
53    /// Creates a new non-blocking CUDA stream using `cudaStreamCreateWithFlags`
54    /// with `cudaStreamNonBlocking`. Non-blocking streams have no implicit
55    /// synchronization with the legacy default stream (stream 0).
56    pub fn new_non_blocking() -> Result<Self, CudaError> {
57        let mut stream: cudaStream_t = std::ptr::null_mut();
58        check(unsafe { cudaStreamCreateWithFlags(&mut stream, CUDA_STREAM_NON_BLOCKING) })?;
59        let host_event = Mutex::new(CudaEvent::new()?);
60        Ok(Self { stream, host_event })
61    }
62
63    /// Get the raw CUDA stream handle.
64    #[inline]
65    pub fn as_raw(&self) -> cudaStream_t {
66        self.stream
67    }
68
69    /// Synchronize this stream.
70    pub fn synchronize(&self) -> Result<(), CudaError> {
71        check(unsafe { cudaStreamSynchronize(self.stream) })
72    }
73
74    /// Wait for the given event.
75    pub fn wait(&self, event: &CudaEvent) -> Result<(), CudaError> {
76        check(unsafe { cudaStreamWaitEvent(self.stream, event.event, 0) })
77    }
78
79    /// Record a per-stream event and synchronize to complete all pending D2H copies.
80    /// Uses event-based sync rather than cudaStreamSynchronize because this waits
81    /// only for work up to the event point, allowing future selective sync patterns
82    /// (e.g., wait for a specific copy without draining the entire stream).
83    pub fn to_host_sync(&self) -> Result<(), CudaError> {
84        let event = self.host_event.lock().unwrap();
85        unsafe { event.record(self.stream) }?;
86        event.synchronize()
87    }
88}
89
90impl Drop for CudaStream {
91    fn drop(&mut self) {
92        if !self.stream.is_null() {
93            // Non-blocking: CUDA defers destruction until the stream is idle
94            let err = unsafe { cudaStreamDestroy(self.stream) };
95            debug_assert_eq!(err, 0, "cudaStreamDestroy failed with error code: {err}");
96            self.stream = std::ptr::null_mut();
97        }
98    }
99}
100
101// ---------------------------------------------------------------------------
102// StreamGuard — keeps a CudaStream alive for allocation records
103// ---------------------------------------------------------------------------
104
105/// Keeps a `CudaStream` alive for the lifetime of an allocation record.
106#[derive(Clone, Debug)]
107pub struct StreamGuard(Arc<CudaStream>);
108
109impl StreamGuard {
110    pub fn new(stream: CudaStream) -> Self {
111        Self(Arc::new(stream))
112    }
113}
114
115impl PartialEq for StreamGuard {
116    fn eq(&self, other: &Self) -> bool {
117        Arc::ptr_eq(&self.0, &other.0)
118    }
119}
120
121impl Eq for StreamGuard {}
122
123impl Deref for StreamGuard {
124    type Target = CudaStream;
125    fn deref(&self) -> &CudaStream {
126        &self.0
127    }
128}
129
130// ---------------------------------------------------------------------------
131// GpuDeviceCtx — bundles device ID with a stream
132// ---------------------------------------------------------------------------
133
134/// Thin context for all GPU operations in the explicit-stream path.
135#[derive(Clone, Debug)]
136pub struct GpuDeviceCtx {
137    pub device_id: u32,
138    pub stream: StreamGuard,
139}
140
141impl GpuDeviceCtx {
142    /// Creates a new `GpuDeviceCtx` for the given device.
143    ///
144    /// NOTE: This calls `set_device_by_id` as a side effect, changing the
145    /// current CUDA device for the calling thread.
146    pub fn for_device(device_id: u32) -> Result<Self, CudaError> {
147        crate::common::set_device_by_id(device_id as i32)?;
148        Ok(Self {
149            device_id,
150            stream: StreamGuard::new(CudaStream::new_non_blocking()?),
151        })
152    }
153
154    pub fn for_current_device() -> Result<Self, CudaError> {
155        let device_id = crate::common::get_device()? as u32;
156        Self::for_device(device_id)
157    }
158}
159
160/// Synchronize the given explicit CUDA stream, blocking until all previously
161/// enqueued work on `stream` has completed.
162///
163/// # Safety
164/// The caller must ensure that `stream` is a valid CUDA stream handle.
165pub unsafe fn sync_stream(stream: cudaStream_t) -> Result<(), CudaError> {
166    check(cudaStreamSynchronize(stream))
167}
168
169// ---------------------------------------------------------------------------
170// CudaEvent
171// ---------------------------------------------------------------------------
172
173#[allow(non_camel_case_types)]
174pub type cudaEvent_t = *mut c_void;
175
176#[derive(Debug)]
177pub enum CudaEventStatus {
178    Completed,
179    NotReady,
180    Error(CudaError),
181}
182
183impl PartialEq for CudaEventStatus {
184    fn eq(&self, other: &Self) -> bool {
185        use CudaEventStatus::*;
186        matches!((self, other), (Completed, Completed) | (NotReady, NotReady))
187    }
188}
189
190impl Eq for CudaEventStatus {}
191
192impl PartialOrd for CudaEventStatus {
193    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
194        Some(self.cmp(other))
195    }
196}
197
198// Completed < NotReady < Error
199impl Ord for CudaEventStatus {
200    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
201        use std::cmp::Ordering;
202
203        use CudaEventStatus::*;
204
205        match (self, other) {
206            (Completed, Completed) => Ordering::Equal,
207            (Completed, _) => Ordering::Less,
208            (_, Completed) => Ordering::Greater,
209            (NotReady, NotReady) => Ordering::Equal,
210            (NotReady, Error(_)) => Ordering::Less,
211            (Error(_), NotReady) => Ordering::Greater,
212            (Error(_), Error(_)) => Ordering::Equal,
213        }
214    }
215}
216
217#[derive(Debug)]
218pub struct CudaEvent {
219    event: cudaEvent_t,
220}
221
222unsafe impl Send for CudaEvent {}
223unsafe impl Sync for CudaEvent {}
224
225impl CudaEvent {
226    pub fn new() -> Result<Self, CudaError> {
227        let mut event: cudaEvent_t = std::ptr::null_mut();
228        check(unsafe { cudaEventCreate(&mut event) })?;
229        Ok(Self { event })
230    }
231
232    /// # Safety
233    /// The caller must ensure that `stream` is a valid stream.
234    pub unsafe fn record(&self, stream: cudaStream_t) -> Result<(), CudaError> {
235        check(cudaEventRecord(self.event, stream))
236    }
237
238    /// Record this event on the given `CudaStream` (safe wrapper).
239    pub fn record_on(&self, stream: &CudaStream) -> Result<(), CudaError> {
240        check(unsafe { cudaEventRecord(self.event, stream.as_raw()) })
241    }
242
243    pub fn synchronize(&self) -> Result<(), CudaError> {
244        check(unsafe { cudaEventSynchronize(self.event) })
245    }
246
247    /// # Safety
248    /// The caller must ensure that `stream` is a valid stream.
249    pub unsafe fn record_and_wait(&self, stream: cudaStream_t) -> Result<(), CudaError> {
250        self.record(stream)?;
251        check(cudaEventSynchronize(self.event))
252    }
253
254    pub fn status(&self) -> CudaEventStatus {
255        let status = unsafe { cudaEventQuery(self.event) };
256        match status {
257            0 => CudaEventStatus::Completed,  // CUDA_SUCCESS
258            600 => CudaEventStatus::NotReady, // CUDA_ERROR_NOT_READY
259            _ => CudaEventStatus::Error(CudaError::new(status)),
260        }
261    }
262
263    pub fn completed(&self) -> bool {
264        self.status() == CudaEventStatus::Completed
265    }
266}
267
268impl Drop for CudaEvent {
269    fn drop(&mut self) {
270        let err = unsafe { cudaEventDestroy(self.event) };
271        debug_assert_eq!(err, 0, "cudaEventDestroy failed with error code: {err}");
272    }
273}
274
275// ---------------------------------------------------------------------------
276// GPU metrics spans
277// ---------------------------------------------------------------------------
278
279/// A GPU-aware span that collects a gauge metric using CUDA events on an explicit stream.
280pub fn gpu_metrics_span_on<R, F: FnOnce() -> R>(
281    name: impl Into<Cow<'static, str>>,
282    stream: &CudaStream,
283    f: F,
284) -> Result<R, CudaError> {
285    let start = CudaEvent::new()?;
286    let stop = CudaEvent::new()?;
287    start.record_on(stream)?;
288    let res = f();
289    stop.record_on(stream)?;
290    stop.synchronize()?;
291
292    let mut elapsed_ms = 0f32;
293    unsafe {
294        check(cudaEventElapsedTime(
295            &mut elapsed_ms,
296            start.event,
297            stop.event,
298        ))?
299    };
300
301    metrics::gauge!(name.into()).set(elapsed_ms as f64);
302    Ok(res)
303}
openvm_cuda_common/stream.rs

openvm_cuda_common/
stream.rs
