torch_sys_cuda/
cuda.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! Bindings for torch's wrappers around CUDA-related functionality.
10use std::time::Duration;
11
12use derive_more::Into;
13use monarch_types::py_global;
14use pyo3::Py;
15use pyo3::PyAny;
16use pyo3::prelude::*;
17use thiserror::Error;
18use torch_sys2::CudaDevice;
19
20// Cached imports for torch.cuda APIs
21py_global!(cuda_stream_class, "torch.cuda", "Stream");
22py_global!(cuda_event_class, "torch.cuda", "Event");
23py_global!(cuda_current_stream, "torch.cuda", "current_stream");
24py_global!(cuda_set_stream, "torch.cuda", "set_stream");
25py_global!(cuda_current_device, "torch.cuda", "current_device");
26py_global!(cuda_set_device, "torch.cuda", "set_device");
27
28/// Wrapper around a CUDA stream.
29///
30/// A CUDA stream is a linear sequence of execution that belongs to a specific
31/// device, independent from other streams.  See the documentation for
32/// `torch.cuda.Stream` for more details.
33#[derive(Debug, Into)]
34#[into(ref)]
35pub struct Stream {
36    inner: Py<PyAny>,
37}
38
39impl Clone for Stream {
40    fn clone(&self) -> Self {
41        Python::attach(|py| Self {
42            inner: self.inner.clone_ref(py),
43        })
44    }
45}
46
47impl Stream {
48    /// Create a new stream on the current device, at priority 0.
49    pub fn new() -> PyResult<Self> {
50        Python::attach(|py| {
51            let stream = cuda_stream_class(py).call0()?;
52            Ok(Self {
53                inner: stream.into(),
54            })
55        })
56    }
57
58    /// Clone this stream reference. Requires holding the GIL.
59    pub fn clone_ref(&self, py: Python<'_>) -> Self {
60        Self {
61            inner: self.inner.clone_ref(py),
62        }
63    }
64
65    /// Create a new stream on the specified device, at priority 0.
66    pub fn new_with_device(device: CudaDevice) -> PyResult<Self> {
67        Python::attach(|py| {
68            let device_idx: i8 = device.index().into();
69            let stream = cuda_stream_class(py).call1((device_idx,))?;
70            Ok(Self {
71                inner: stream.into(),
72            })
73        })
74    }
75
76    /// Get the current stream on the current device.
77    pub fn get_current_stream() -> PyResult<Self> {
78        Python::attach(|py| {
79            let stream = cuda_current_stream(py).call0()?;
80            Ok(Self {
81                inner: stream.into(),
82            })
83        })
84    }
85
86    /// Get the current stream on the specified device.
87    pub fn get_current_stream_on_device(device: CudaDevice) -> PyResult<Self> {
88        Python::attach(|py| {
89            let device_idx: i8 = device.index().into();
90            let stream = cuda_current_stream(py).call1((device_idx,))?;
91            Ok(Self {
92                inner: stream.into(),
93            })
94        })
95    }
96
97    /// Set the provided stream as the current stream. Also sets the current
98    /// device to be the same as the stream's device.
99    pub fn set_current_stream(stream: &Stream) -> PyResult<()> {
100        Python::attach(|py| {
101            let stream_obj = stream.inner.bind(py);
102
103            // Get current device and stream device
104            let current_device = cuda_current_device(py).call0()?.extract::<i64>()?;
105
106            let stream_device = stream_obj.getattr("device_index")?.extract::<i64>()?;
107
108            // Set device if different
109            if current_device != stream_device {
110                cuda_set_device(py).call1((stream_device,))?;
111            }
112
113            // Set the stream
114            cuda_set_stream(py).call1((stream_obj,))?;
115
116            Ok(())
117        })
118    }
119
120    /// Make all future work submitted to this stream wait for an event.
121    pub fn wait_event(&self, event: &mut Event) -> PyResult<()> {
122        event.wait(Some(self))
123    }
124
125    /// Synchronize with another stream.
126    ///
127    /// All future work submitted to this stream will wait until all kernels
128    /// submitted to a given stream at the time of call entry complete.
129    pub fn wait_stream(&self, stream: &Stream) -> PyResult<()> {
130        self.wait_event(&mut stream.record_event(None)?)
131    }
132
133    /// Record an event on this stream. If no event is provided one will be
134    /// created.
135    pub fn record_event(&self, event: Option<Event>) -> PyResult<Event> {
136        let mut event = match event {
137            Some(event) => event,
138            None => Event::new()?,
139        };
140        event.record(Some(self))?;
141        Ok(event)
142    }
143
144    /// Check if all work submitted to this stream has completed.
145    pub fn query(&self) -> PyResult<bool> {
146        Python::attach(|py| {
147            let stream_obj = self.inner.bind(py);
148            stream_obj.call_method0("query")?.extract::<bool>()
149        })
150    }
151
152    /// Wait for all kernels in this stream to complete.
153    pub fn synchronize(&self) -> PyResult<()> {
154        Python::attach(|py| {
155            let stream_obj = self.inner.bind(py);
156            stream_obj.call_method0("synchronize")?;
157            Ok(())
158        })
159    }
160
161    /// Get the raw CUDA stream pointer. Only available with the `cuda` feature.
162    #[cfg(feature = "cuda")]
163    pub fn stream(&self) -> PyResult<nccl_sys::cudaStream_t> {
164        Python::attach(|py| {
165            let stream_obj = self.inner.bind(py);
166            let cuda_stream = stream_obj.getattr("cuda_stream")?;
167
168            // Extract the raw pointer
169            let ptr = cuda_stream.extract::<usize>()?;
170
171            Ok(ptr as nccl_sys::cudaStream_t)
172        })
173    }
174}
175
176impl PartialEq for Stream {
177    fn eq(&self, other: &Self) -> bool {
178        #[cfg(feature = "cuda")]
179        {
180            self.stream()
181                .expect("cuda_stream attribute should be readable")
182                == other
183                    .stream()
184                    .expect("cuda_stream attribute should be readable")
185        }
186
187        #[cfg(not(feature = "cuda"))]
188        {
189            self.inner.as_ptr() == other.inner.as_ptr()
190        }
191    }
192}
193
194/// Wrapper around a CUDA event.
195///
196/// CUDA events are synchronization markers that can be used to monitor the
197/// device's progress, to accurately measure timing, and to synchronize CUDA
198/// streams.
199///
200/// The underlying CUDA events are lazily initialized when the event is first
201/// recorded or exported to another process. After creation, only streams on the
202/// same device may record the event. However, streams on any device can wait on
203/// the event.
204///
205/// See the docs of `torch.cuda.Event` for more details.
206#[derive(Debug)]
207pub struct Event {
208    inner: Py<PyAny>,
209}
210
211impl Clone for Event {
212    fn clone(&self) -> Self {
213        Python::attach(|py| Self {
214            inner: self.inner.clone_ref(py),
215        })
216    }
217}
218
219impl Event {
220    /// Create a new event.
221    // TODO: add support for flags.
222    pub fn new() -> PyResult<Self> {
223        Python::attach(|py| {
224            let event = cuda_event_class(py).call0()?;
225            Ok(Self {
226                inner: event.into(),
227            })
228        })
229    }
230
231    /// Record the event on the current stream.
232    ///
233    /// Uses the current stream if no stream is provided.
234    pub fn record(&mut self, stream: Option<&Stream>) -> PyResult<()> {
235        Python::attach(|py| {
236            let event_obj = self.inner.bind(py);
237
238            match stream {
239                Some(stream) => {
240                    let stream_obj = stream.inner.bind(py);
241                    event_obj.call_method1("record", (stream_obj,))?;
242                }
243                None => {
244                    event_obj.call_method0("record")?;
245                }
246            }
247
248            Ok(())
249        })
250    }
251
252    /// Make all future work submitted to the given stream wait for this event.
253    ///
254    /// Uses the current stream if no stream is specified.
255    pub fn wait(&mut self, stream: Option<&Stream>) -> PyResult<()> {
256        Python::attach(|py| {
257            let event_obj = self.inner.bind(py);
258
259            match stream {
260                Some(stream) => {
261                    let stream_obj = stream.inner.bind(py);
262                    event_obj.call_method1("wait", (stream_obj,))?;
263                }
264                None => {
265                    event_obj.call_method0("wait")?;
266                }
267            }
268
269            Ok(())
270        })
271    }
272
273    /// Check if all work currently captured by event has completed.
274    pub fn query(&self) -> PyResult<bool> {
275        Python::attach(|py| {
276            let event_obj = self.inner.bind(py);
277            event_obj.call_method0("query")?.extract::<bool>()
278        })
279    }
280
281    /// Return the time elapsed.
282    ///
283    /// Time reported in after the event was recorded and before the end_event
284    /// was recorded.
285    pub fn elapsed_time(&self, end_event: &Event) -> PyResult<Duration> {
286        Python::attach(|py| {
287            let event_obj = self.inner.bind(py);
288            let end_event_obj = end_event.inner.bind(py);
289
290            let elapsed_ms = event_obj
291                .call_method1("elapsed_time", (end_event_obj,))?
292                .extract::<f64>()?;
293
294            Ok(Duration::from_millis(elapsed_ms as u64))
295        })
296    }
297
298    /// Wait for the event to complete.
299    /// Waits until the completion of all work currently captured in this event.
300    /// This prevents the CPU thread from proceeding until the event completes.
301    pub fn synchronize(&self) -> PyResult<()> {
302        Python::attach(|py| {
303            let event_obj = self.inner.bind(py);
304            event_obj.call_method0("synchronize")?;
305            Ok(())
306        })
307    }
308}
309
310/// Corresponds to the CUDA error codes.
311#[derive(Debug, Error)]
312pub enum CudaError {
313    #[error(
314        "one or more parameters passed to the API call is not within an acceptable range of values"
315    )]
316    InvalidValue,
317    #[error("the API call failed due to insufficient memory or resources")]
318    MemoryAllocation,
319    #[error("failed to initialize the CUDA driver and runtime")]
320    InitializationError,
321    #[error("CUDA Runtime API call was executed after the CUDA driver has been unloaded")]
322    CudartUnloading,
323    #[error("profiler is not initialized for this run, possibly due to an external profiling tool")]
324    ProfilerDisabled,
325    #[error("deprecated. Attempted to enable/disable profiling without initialization")]
326    ProfilerNotInitialized,
327    #[error("deprecated. Profiling is already started")]
328    ProfilerAlreadyStarted,
329    #[error("deprecated. Profiling is already stopped")]
330    ProfilerAlreadyStopped,
331    #[error("kernel launch requested resources that cannot be satisfied by the current device")]
332    InvalidConfiguration,
333    #[error("one or more of the pitch-related parameters passed to the API call is out of range")]
334    InvalidPitchValue,
335    #[error("the symbol name/identifier passed to the API call is invalid")]
336    InvalidSymbol,
337    #[error("the host pointer passed to the API call is invalid")]
338    InvalidHostPointer,
339    #[error("the device pointer passed to the API call is invalid")]
340    InvalidDevicePointer,
341    #[error("the texture passed to the API call is invalid")]
342    InvalidTexture,
343    #[error("the texture binding is invalid")]
344    InvalidTextureBinding,
345    #[error("the channel descriptor passed to the API call is invalid")]
346    InvalidChannelDescriptor,
347    #[error("the direction of the memcpy operation is invalid")]
348    InvalidMemcpyDirection,
349    #[error(
350        "attempted to take the address of a constant variable, which is forbidden before CUDA 3.1"
351    )]
352    AddressOfConstant,
353    #[error("deprecated. A texture fetch operation failed")]
354    TextureFetchFailed,
355    #[error("deprecated. The texture is not bound for access")]
356    TextureNotBound,
357    #[error("a synchronization operation failed")]
358    SynchronizationError,
359    #[error(
360        "a non-float texture was accessed with linear filtering, which is not supported by CUDA"
361    )]
362    InvalidFilterSetting,
363    #[error(
364        "attempted to read a non-float texture as a normalized float, which is not supported by CUDA"
365    )]
366    InvalidNormSetting,
367    #[error("the API call is not yet implemented")]
368    NotYetImplemented,
369    #[error("an emulated device pointer exceeded the 32-bit address range")]
370    MemoryValueTooLarge,
371    #[error("the CUDA driver is a stub library")]
372    StubLibrary,
373    #[error("the installed NVIDIA CUDA driver is older than the CUDA runtime library")]
374    InsufficientDriver,
375    #[error("the API call requires a newer CUDA driver")]
376    CallRequiresNewerDriver,
377    #[error("the surface passed to the API call is invalid")]
378    InvalidSurface,
379    #[error("multiple global or constant variables share the same string name")]
380    DuplicateVariableName,
381    #[error("multiple textures share the same string name")]
382    DuplicateTextureName,
383    #[error("multiple surfaces share the same string name")]
384    DuplicateSurfaceName,
385    #[error("all CUDA devices are currently busy or unavailable")]
386    DevicesUnavailable,
387    #[error("the current CUDA context is not compatible with the runtime")]
388    IncompatibleDriverContext,
389    #[error("the device function being invoked was not previously configured")]
390    MissingConfiguration,
391    #[error("a previous kernel launch failed")]
392    PriorLaunchFailure,
393    #[error(
394        "the depth of the child grid exceeded the maximum supported number of nested grid launches"
395    )]
396    LaunchMaxDepthExceeded,
397    #[error("a grid launch did not occur because file-scoped textures are unsupported")]
398    LaunchFileScopedTex,
399    #[error("a grid launch did not occur because file-scoped surfaces are unsupported")]
400    LaunchFileScopedSurf,
401    #[error("a call to cudaDeviceSynchronize failed due to exceeding the sync depth")]
402    SyncDepthExceeded,
403    #[error(
404        "a grid launch failed because the launch exceeded the limit of pending device runtime launches"
405    )]
406    LaunchPendingCountExceeded,
407    #[error(
408        "the requested device function does not exist or is not compiled for the proper device architecture"
409    )]
410    InvalidDeviceFunction,
411    #[error("no CUDA-capable devices were detected")]
412    NoDevice,
413    #[error("the device ordinal supplied does not correspond to a valid CUDA device")]
414    InvalidDevice,
415    #[error("the device does not have a valid Grid License")]
416    DeviceNotLicensed,
417    #[error("an internal startup failure occurred in the CUDA runtime")]
418    StartupFailure,
419    #[error("the device kernel image is invalid")]
420    InvalidKernelImage,
421    #[error("the device is not initialized")]
422    DeviceUninitialized,
423    #[error("the buffer object could not be mapped")]
424    MapBufferObjectFailed,
425    #[error("the buffer object could not be unmapped")]
426    UnmapBufferObjectFailed,
427    #[error("the specified array is currently mapped and cannot be destroyed")]
428    ArrayIsMapped,
429    #[error("the resource is already mapped")]
430    AlreadyMapped,
431    #[error("there is no kernel image available that is suitable for the device")]
432    NoKernelImageForDevice,
433    #[error("the resource has already been acquired")]
434    AlreadyAcquired,
435    #[error("the resource is not mapped")]
436    NotMapped,
437    #[error("the mapped resource is not available for access as an array")]
438    NotMappedAsArray,
439    #[error("the mapped resource is not available for access as a pointer")]
440    NotMappedAsPointer,
441    #[error("an uncorrectable ECC error was detected")]
442    ECCUncorrectable,
443    #[error("the specified cudaLimit is not supported by the device")]
444    UnsupportedLimit,
445    #[error("a call tried to access an exclusive-thread device that is already in use")]
446    DeviceAlreadyInUse,
447    #[error("P2P access is not supported across the given devices")]
448    PeerAccessUnsupported,
449    #[error("a PTX compilation failed")]
450    InvalidPtx,
451    #[error("an error occurred with the OpenGL or DirectX context")]
452    InvalidGraphicsContext,
453    #[error("an uncorrectable NVLink error was detected during execution")]
454    NvlinkUncorrectable,
455    #[error("the PTX JIT compiler library was not found")]
456    JitCompilerNotFound,
457    #[error("the provided PTX was compiled with an unsupported toolchain")]
458    UnsupportedPtxVersion,
459    #[error("JIT compilation was disabled")]
460    JitCompilationDisabled,
461    #[error("the provided execution affinity is not supported by the device")]
462    UnsupportedExecAffinity,
463    #[error("the operation is not permitted when the stream is capturing")]
464    StreamCaptureUnsupported,
465    #[error(
466        "the current capture sequence on the stream has been invalidated due to a previous error"
467    )]
468    StreamCaptureInvalidated,
469    #[error("a merge of two independent capture sequences was not allowed")]
470    StreamCaptureMerge,
471    #[error("the capture was not initiated in this stream")]
472    StreamCaptureUnmatched,
473    #[error("a stream capture sequence was passed to cudaStreamEndCapture in a different thread")]
474    StreamCaptureWrongThread,
475    #[error("the wait operation has timed out")]
476    Timeout,
477    #[error("an unknown internal error occurred")]
478    Unknown,
479    #[error("the API call returned a failure")]
480    ApiFailureBase,
481}
482
483#[cfg(feature = "cuda")]
484pub fn cuda_check(result: nccl_sys::cudaError_t) -> Result<(), CudaError> {
485    match result.0 {
486        0 => Ok(()),
487        1 => Err(CudaError::InvalidValue),
488        2 => Err(CudaError::MemoryAllocation),
489        3 => Err(CudaError::InitializationError),
490        4 => Err(CudaError::CudartUnloading),
491        5 => Err(CudaError::ProfilerDisabled),
492        6 => Err(CudaError::ProfilerNotInitialized),
493        7 => Err(CudaError::ProfilerAlreadyStarted),
494        8 => Err(CudaError::ProfilerAlreadyStopped),
495        9 => Err(CudaError::InvalidConfiguration),
496        12 => Err(CudaError::InvalidPitchValue),
497        13 => Err(CudaError::InvalidSymbol),
498        16 => Err(CudaError::InvalidHostPointer),
499        17 => Err(CudaError::InvalidDevicePointer),
500        18 => Err(CudaError::InvalidTexture),
501        19 => Err(CudaError::InvalidTextureBinding),
502        20 => Err(CudaError::InvalidChannelDescriptor),
503        21 => Err(CudaError::InvalidMemcpyDirection),
504        22 => Err(CudaError::AddressOfConstant),
505        23 => Err(CudaError::TextureFetchFailed),
506        24 => Err(CudaError::TextureNotBound),
507        25 => Err(CudaError::SynchronizationError),
508        26 => Err(CudaError::InvalidFilterSetting),
509        27 => Err(CudaError::InvalidNormSetting),
510        31 => Err(CudaError::NotYetImplemented),
511        32 => Err(CudaError::MemoryValueTooLarge),
512        34 => Err(CudaError::StubLibrary),
513        35 => Err(CudaError::InsufficientDriver),
514        36 => Err(CudaError::CallRequiresNewerDriver),
515        37 => Err(CudaError::InvalidSurface),
516        43 => Err(CudaError::DuplicateVariableName),
517        44 => Err(CudaError::DuplicateTextureName),
518        45 => Err(CudaError::DuplicateSurfaceName),
519        46 => Err(CudaError::DevicesUnavailable),
520        49 => Err(CudaError::IncompatibleDriverContext),
521        52 => Err(CudaError::MissingConfiguration),
522        53 => Err(CudaError::PriorLaunchFailure),
523        65 => Err(CudaError::LaunchMaxDepthExceeded),
524        66 => Err(CudaError::LaunchFileScopedTex),
525        67 => Err(CudaError::LaunchFileScopedSurf),
526        68 => Err(CudaError::SyncDepthExceeded),
527        69 => Err(CudaError::LaunchPendingCountExceeded),
528        98 => Err(CudaError::InvalidDeviceFunction),
529        100 => Err(CudaError::NoDevice),
530        101 => Err(CudaError::InvalidDevice),
531        102 => Err(CudaError::DeviceNotLicensed),
532        127 => Err(CudaError::StartupFailure),
533        200 => Err(CudaError::InvalidKernelImage),
534        201 => Err(CudaError::DeviceUninitialized),
535        205 => Err(CudaError::MapBufferObjectFailed),
536        206 => Err(CudaError::UnmapBufferObjectFailed),
537        207 => Err(CudaError::ArrayIsMapped),
538        208 => Err(CudaError::AlreadyMapped),
539        209 => Err(CudaError::NoKernelImageForDevice),
540        210 => Err(CudaError::AlreadyAcquired),
541        211 => Err(CudaError::NotMapped),
542        212 => Err(CudaError::NotMappedAsArray),
543        213 => Err(CudaError::NotMappedAsPointer),
544        214 => Err(CudaError::ECCUncorrectable),
545        215 => Err(CudaError::UnsupportedLimit),
546        216 => Err(CudaError::DeviceAlreadyInUse),
547        217 => Err(CudaError::PeerAccessUnsupported),
548        218 => Err(CudaError::InvalidPtx),
549        219 => Err(CudaError::InvalidGraphicsContext),
550        220 => Err(CudaError::NvlinkUncorrectable),
551        221 => Err(CudaError::JitCompilerNotFound),
552        222 => Err(CudaError::UnsupportedPtxVersion),
553        223 => Err(CudaError::JitCompilationDisabled),
554        224 => Err(CudaError::UnsupportedExecAffinity),
555        900 => Err(CudaError::StreamCaptureUnsupported),
556        901 => Err(CudaError::StreamCaptureInvalidated),
557        902 => Err(CudaError::StreamCaptureMerge),
558        903 => Err(CudaError::StreamCaptureUnmatched),
559        904 => Err(CudaError::StreamCaptureWrongThread),
560        909 => Err(CudaError::Timeout),
561        999 => Err(CudaError::Unknown),
562        _ => panic!("Unknown cudaError_t: {:?}", result.0),
563    }
564}
565
566#[cfg(feature = "cuda")]
567pub fn set_device(device: CudaDevice) -> Result<(), CudaError> {
568    let index: i8 = device.index().into();
569    // SAFETY: intended usage of this function
570    unsafe { cuda_check(nccl_sys::cudaSetDevice(index.into())) }
571}
572
573#[cfg(not(feature = "cuda"))]
574pub fn set_device(_device: CudaDevice) -> Result<(), CudaError> {
575    panic!("set_device requires the `cuda` feature (CUDA/ROCm runtime)")
576}
torch_sys_cuda/cuda.rs

torch_sys_cuda/
cuda.rs
