tenferro_gpu/cubecl/
runtime.rs

1//! CubeCL CUDA runtime initialization and synchronization.
2
3use std::fmt;
4use std::hash::{Hash, Hasher};
5use std::io::Write;
6use std::sync::Arc;
7
8use cubecl::client::ComputeClient;
9use cubecl::stream_id::StreamId;
10use cubecl::Runtime;
11use cubecl_cuda::{CudaDevice, CudaRuntime as CubeclCudaRuntime};
12use cudarc::driver::result::DriverError;
13use cudarc::driver::sys::{CUcontext, CUdevice, CUresult};
14use cudarc::runtime::{result as cuda_result, sys::cudaStream_t};
15use tenferro_tensor::AllocationDomainId;
16
17use super::device::{
18    cuda_devices, unavailable_device_error, CudaDeviceError, CudaDeviceId, CudaDeviceInfo,
19};
20use super::identity::GpuExtensionCapability;
21
22/// Returns `true` if a CUDA device can initialize a CubeCL runtime.
23///
24/// Use this in test helpers to skip GPU tests on machines without hardware.
25pub fn gpu_available() -> bool {
26    let library_present = std::panic::catch_unwind(|| {
27        // SAFETY: `is_culib_present` only probes candidate library names and
28        // does not call CUDA function pointers or retain a library handle.
29        unsafe { cudarc::driver::sys::is_culib_present() }
30    })
31    .unwrap_or(false);
32    if !library_present {
33        return false;
34    }
35    let Ok(devices) = cuda_devices() else {
36        return false;
37    };
38    let Some(device_id) = devices.first().map(|device| device.id()) else {
39        return false;
40    };
41    std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
42        let Ok(runtime) = CudaRuntime::new(device_id) else {
43            return false;
44        };
45        runtime.synchronize().is_ok()
46    }))
47    .unwrap_or(false)
48}
49
50/// RAII guard that attempts to restore the thread's previous CUDA device and
51/// current context when dropped.
52///
53/// Used by the `with_raw` enter/exit protocol: the guard is created after the
54/// calling thread's device/context are saved and the tenferro primary context
55/// is activated. Drop attempts best-effort restoration of the saved state on
56/// normal return, `Err`, and unwind; a restoration failure is logged to
57/// stderr (non-panicking) and never returned.
58pub(crate) struct RawContextRestore {
59    saved_device: Result<i32, cudarc::runtime::result::RuntimeError>,
60    saved_context: Result<Option<CUcontext>, cudarc::driver::result::DriverError>,
61    op: &'static str,
62}
63
64impl RawContextRestore {
65    /// Save the current device/context, then activate `device`/`context`.
66    pub(crate) fn enter(op: &'static str, device: i32, context: CUcontext) -> crate::Result<Self> {
67        let saved_device = cudarc::runtime::result::device::get();
68        let saved_context = cudarc::driver::result::ctx::get_current();
69        cudarc::runtime::result::device::set(device)
70            .map_err(|err| crate::Error::backend_source(op, err))?;
71        if let Err(err) = unsafe { cudarc::driver::result::ctx::set_current(context) } {
72            // Roll the device and context back so a partial activation failure
73            // cannot leave the caller's thread on a different device or with a
74            // different current context (setting the device can implicitly
75            // change the thread's current context to the new primary).
76            if let Ok(previous_device) = saved_device {
77                let _ = cudarc::runtime::result::device::set(previous_device);
78            }
79            match saved_context {
80                Ok(Some(previous)) => {
81                    let _ = unsafe { cudarc::driver::result::ctx::set_current(previous) };
82                }
83                Ok(None) => {
84                    let _ =
85                        unsafe { cudarc::driver::result::ctx::set_current(std::ptr::null_mut()) };
86                }
87                Err(_) => {}
88            }
89            return Err(crate::Error::backend_source(op, err));
90        }
91        Ok(Self {
92            saved_device,
93            saved_context,
94            op,
95        })
96    }
97
98    fn restore(&self) {
99        let mut stderr = std::io::stderr();
100        if let Ok(device) = self.saved_device {
101            if let Err(err) = cudarc::runtime::result::device::set(device) {
102                let _ = writeln!(
103                    stderr,
104                    "tenferro-gpu: failed to restore CUDA device during {}: {err:?}",
105                    self.op
106                );
107            }
108        }
109        match self.saved_context {
110            Ok(Some(context)) => {
111                if let Err(err) = unsafe { cudarc::driver::result::ctx::set_current(context) } {
112                    let _ = writeln!(
113                        stderr,
114                        "tenferro-gpu: failed to restore CUDA context during {}: {err:?}",
115                        self.op
116                    );
117                }
118            }
119            // The thread had no current context before the guard; restore that
120            // state instead of leaving the tenferro primary context current.
121            Ok(None) => {
122                if let Err(err) =
123                    unsafe { cudarc::driver::result::ctx::set_current(std::ptr::null_mut()) }
124                {
125                    let _ = writeln!(
126                        stderr,
127                        "tenferro-gpu: failed to clear CUDA context during {}: {err:?}",
128                        self.op
129                    );
130                }
131            }
132            // The saved-context query itself failed; nothing can be restored.
133            Err(_) => {}
134        }
135    }
136}
137
138impl Drop for RawContextRestore {
139    fn drop(&mut self) {
140        self.restore();
141    }
142}
143
144/// Opaque identity of one exact CUDA runtime instance.
145///
146/// Cloning the identity preserves the underlying executable runtime witness;
147/// constructing another runtime, even for the same device ordinal, produces a
148/// distinct identity. The cache key intentionally carries no provider or
149/// device identifier and grants no execution authority.
150#[derive(Clone, Debug)]
151pub struct CudaRuntimeIdentity {
152    marker: Arc<u8>,
153}
154
155impl CudaRuntimeIdentity {
156    fn fresh() -> Self {
157        Self {
158            marker: Arc::new(0),
159        }
160    }
161}
162
163impl PartialEq for CudaRuntimeIdentity {
164    fn eq(&self, other: &Self) -> bool {
165        Arc::ptr_eq(&self.marker, &other.marker)
166    }
167}
168
169impl Eq for CudaRuntimeIdentity {}
170
171impl Hash for CudaRuntimeIdentity {
172    fn hash<H: Hasher>(&self, state: &mut H) {
173        // INVARIANT: `marker` is retained by every clone of this identity, so
174        // its Arc allocation address is move/clone-invariant while witnessed.
175        state.write_usize(Arc::as_ptr(&self.marker) as usize);
176    }
177}
178
179/// CubeCL CUDA runtime wrapper.
180///
181/// # Examples
182///
183/// ```
184/// use tenferro_gpu::cuda::CudaRuntime;
185///
186/// let _ctor: fn(tenferro_gpu::cuda::CudaDeviceId) ->
187///     Result<CudaRuntime, tenferro_gpu::cuda::CudaDeviceError> = CudaRuntime::new;
188/// let _sync: fn(&CudaRuntime) -> tenferro_tensor::Result<()> =
189///     CudaRuntime::synchronize;
190/// ```
191#[derive(Clone)]
192pub struct CudaRuntime {
193    inner: Arc<CudaRuntimeState>,
194}
195
196struct CudaRuntimeState {
197    client: ComputeClient<CubeclCudaRuntime>,
198    device_id: CudaDeviceId,
199    device_ordinal: usize,
200    device_info: CudaDeviceInfo,
201    primary_context: CudaPrimaryContext,
202    identity: CudaRuntimeIdentity,
203    allocation_domain: AllocationDomainId,
204}
205
206// SAFETY: `CudaRuntimeState` owns a retained CUDA primary context and a CubeCL
207// client for one device ordinal. Methods set the context current before raw CUDA
208// calls, and backend/executor layers serialize mutating tensor execution.
209unsafe impl Send for CudaRuntimeState {}
210// SAFETY: Shared state access exposes immutable runtime handles; synchronization
211// and stream queries use explicit CUDA/CubeCL handles and do not mutate Rust
212// aliasing-visible fields.
213unsafe impl Sync for CudaRuntimeState {}
214
215impl fmt::Debug for CudaRuntime {
216    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
217        f.debug_struct("CudaRuntime")
218            .field("device_id", &self.inner.device_id)
219            .finish_non_exhaustive()
220    }
221}
222
223struct CudaPrimaryContext {
224    cuda_device: CUdevice,
225    cuda_context: CUcontext,
226}
227
228impl CudaPrimaryContext {
229    fn retain(cuda_device: CUdevice) -> crate::Result<Self> {
230        let cuda_context = unsafe { cudarc::driver::result::primary_ctx::retain(cuda_device) }
231            .map_err(|err| crate::Error::backend_source("cubecl_runtime_init", err))?;
232        Ok(Self {
233            cuda_device,
234            cuda_context,
235        })
236    }
237
238    fn context(&self) -> CUcontext {
239        self.cuda_context
240    }
241}
242
243impl Drop for CudaPrimaryContext {
244    fn drop(&mut self) {
245        if let Err(err) = unsafe { cudarc::driver::result::primary_ctx::release(self.cuda_device) }
246        {
247            report_cuda_primary_context_release_error(&err);
248        }
249    }
250}
251
252#[cold]
253fn report_cuda_primary_context_release_error(err: &impl fmt::Debug) {
254    eprintln!("tenferro-gpu: failed to release CUDA primary context during Drop: {err:?}");
255}
256
257#[cold]
258fn report_cuda_runtime_drop_error(err: &crate::Error) {
259    eprintln!("tenferro-gpu: failed to synchronize CUDA runtime during Drop: {err}");
260}
261
262impl CudaRuntime {
263    /// Initialize the CubeCL CUDA runtime on the caller-selected device.
264    ///
265    /// # Examples
266    ///
267    /// ```
268    /// use tenferro_gpu::{cuda::CudaDeviceError, cuda::CudaDeviceId, cuda::CudaRuntime};
269    ///
270    /// let _ctor: fn(CudaDeviceId) -> Result<CudaRuntime, CudaDeviceError> = CudaRuntime::new;
271    /// ```
272    ///
273    /// # Errors
274    ///
275    /// Returns [`CudaDeviceError::Discovery`] when fallback discovery for an
276    /// invalid selected ordinal fails, [`CudaDeviceError::Unavailable`] when
277    /// that ordinal is not available, or [`CudaDeviceError::Initialization`]
278    /// when CUDA driver, runtime, context, or CubeCL client initialization
279    /// fails.
280    pub fn new(device_id: CudaDeviceId) -> Result<Self, CudaDeviceError> {
281        let device_ordinal = usize::try_from(device_id.ordinal()).map_err(|source| {
282            cuda_initialization_error(device_id, "convert_device_ordinal", source)
283        })?;
284        let cuda_ordinal = i32::try_from(device_id.ordinal()).map_err(|source| {
285            cuda_initialization_error(device_id, "convert_cuda_ordinal", source)
286        })?;
287        cudarc::driver::result::init()
288            .map_err(|source| cuda_initialization_error(device_id, "initialize_driver", source))?;
289        let cuda_device = match cudarc::driver::result::device::get(cuda_ordinal) {
290            Ok(cuda_device) => cuda_device,
291            Err(source) if is_invalid_device_lookup(source) => {
292                return Err(unavailable_device_error(device_id, cuda_devices()?));
293            }
294            Err(source) => {
295                return Err(cuda_initialization_error(device_id, "get_device", source));
296            }
297        };
298        let primary_context = CudaPrimaryContext::retain(cuda_device).map_err(|source| {
299            cuda_initialization_error(device_id, "retain_primary_context", source)
300        })?;
301        unsafe { cudarc::driver::result::ctx::set_current(primary_context.context()) }.map_err(
302            |source| cuda_initialization_error(device_id, "set_current_context", source),
303        )?;
304        cudarc::runtime::result::device::set(cuda_ordinal)
305            .map_err(|source| cuda_initialization_error(device_id, "set_device", source))?;
306        let device = CudaDevice::new(device_ordinal);
307        let client = CubeclCudaRuntime::client(&device);
308        let discovered = cuda_devices()?;
309        let device_info = discovered
310            .iter()
311            .find(|info| info.id() == device_id)
312            .cloned()
313            .ok_or_else(|| unavailable_device_error(device_id, discovered))?;
314        Ok(Self {
315            inner: Arc::new(CudaRuntimeState {
316                client,
317                device_id,
318                device_ordinal,
319                device_info,
320                primary_context,
321                identity: CudaRuntimeIdentity::fresh(),
322                allocation_domain: AllocationDomainId::fresh(),
323            }),
324        })
325    }
326
327    pub(crate) fn client(&self) -> &ComputeClient<CubeclCudaRuntime> {
328        &self.inner.client
329    }
330
331    /// Return the caller-selected CUDA device identity that this runtime targets.
332    ///
333    /// # Examples
334    ///
335    /// ```
336    /// use tenferro_gpu::{cuda::CudaDeviceId, cuda::CudaRuntime};
337    ///
338    /// let _device_id: fn(&CudaRuntime) -> CudaDeviceId = CudaRuntime::device_id;
339    /// ```
340    pub fn device_id(&self) -> CudaDeviceId {
341        self.inner.device_id
342    }
343
344    /// Return immutable metadata for the runtime's device.
345    ///
346    /// # Examples
347    ///
348    /// ```
349    /// use tenferro_gpu::cuda::CudaRuntime;
350    ///
351    /// let _info: fn(&CudaRuntime) -> &tenferro_gpu::cuda::CudaDeviceInfo =
352    ///     CudaRuntime::device_info;
353    /// ```
354    pub fn device_info(&self) -> &CudaDeviceInfo {
355        &self.inner.device_info
356    }
357
358    /// Return the allocation ownership domain of this runtime.
359    ///
360    /// # Examples
361    ///
362    /// ```
363    /// use tenferro_gpu::cuda::CudaRuntime;
364    ///
365    /// let _domain: fn(&CudaRuntime) -> tenferro_tensor::AllocationDomainId =
366    ///     CudaRuntime::allocation_domain;
367    /// ```
368    pub fn allocation_domain(&self) -> AllocationDomainId {
369        self.inner.allocation_domain
370    }
371
372    /// Report whether this CUDA session supports a GPU extension capability.
373    ///
374    /// The CUDA provider supports the full extension vocabulary: external
375    /// CubeCL kernels, native module loading, runtime compilation (NVRTC), raw
376    /// stream borrowing, and same-device copy. `PeerCopy` is reported as a
377    /// directional query; availability is hardware-dependent and is checked
378    /// per source/destination pair rather than here.
379    ///
380    /// # Examples
381    ///
382    /// ```
383    /// use tenferro_gpu::cuda::GpuExtensionCapability;
384    /// use tenferro_gpu::cuda::CudaRuntime;
385    ///
386    /// let _supports: fn(&CudaRuntime, GpuExtensionCapability) -> bool =
387    ///     CudaRuntime::supports_extension;
388    /// ```
389    pub fn supports_extension(&self, capability: GpuExtensionCapability) -> bool {
390        capabilities_for_device(capability)
391    }
392
393    pub(crate) fn device_ordinal(&self) -> usize {
394        self.inner.device_ordinal
395    }
396
397    pub(crate) fn primary_context(&self) -> CUcontext {
398        self.inner.primary_context.context()
399    }
400
401    /// Run `f` with the tenferro primary context current on this thread.
402    ///
403    /// Saves the calling thread's current CUDA device/context, activates the
404    /// tenferro primary context for the duration of `f`, and attempts to
405    /// restore the saved state on every exit path (normal return, `Err`, and
406    /// unwind). Restoration is best-effort: a failure to restore the
407    /// caller's previous device/context is logged to stderr rather than
408    /// returned. This is the scoped context authority used by vendor-library
409    /// lifecycle paths (plan creation/retirement) that run outside a
410    /// raw-session callback.
411    ///
412    /// # Errors
413    ///
414    /// Returns [`crate::Error::BackendSource`] when the tenferro primary
415    /// context cannot be activated (device or context driver failure); a
416    /// partial activation is best-effort rolled back before the error is
417    /// returned (rollback failures are discarded).
418    ///
419    /// # Examples
420    ///
421    /// ```
422    /// use tenferro_gpu::cuda::CudaRuntime;
423    ///
424    /// let _check: fn(&CudaRuntime) -> tenferro_tensor::Result<u64> = |rt| {
425    ///     rt.with_current_context("test.context", || 7)
426    /// };
427    /// ```
428    pub fn with_current_context<R>(
429        &self,
430        op: &'static str,
431        f: impl FnOnce() -> R,
432    ) -> crate::Result<R> {
433        let device_ordinal = i32::try_from(self.device_ordinal())
434            .map_err(|source| crate::Error::backend_source(op, source))?;
435        let _guard = RawContextRestore::enter(op, device_ordinal, self.primary_context())?;
436        Ok(f())
437    }
438
439    /// Flush pending CubeCL work on the current stream.
440    ///
441    /// Used by the raw-session enter protocol so raw library calls observe
442    /// previously enqueued CubeCL work.
443    pub(crate) fn flush_cubecl(&self, op: &'static str) -> crate::Result<()> {
444        self.client()
445            .flush()
446            .map_err(|err| crate::Error::backend_source(op, err))
447    }
448
449    /// Return the opaque identity of this exact executable runtime instance.
450    ///
451    /// # Examples
452    ///
453    /// ```
454    /// use tenferro_gpu::cuda::CudaRuntime;
455    ///
456    /// let _identity: fn(&CudaRuntime) -> tenferro_gpu::cuda::CudaRuntimeIdentity =
457    ///     CudaRuntime::runtime_identity;
458    /// ```
459    pub fn runtime_identity(&self) -> CudaRuntimeIdentity {
460        self.inner.identity.clone()
461    }
462
463    pub(crate) fn allocation_domain_id(&self) -> AllocationDomainId {
464        self.inner.allocation_domain
465    }
466
467    pub(crate) fn set_current_cuda_context(&self, op: &'static str) -> crate::Result<()> {
468        self.inner.set_current_cuda_context(op)
469    }
470
471    pub(crate) fn raw_cuda_stream(&self) -> crate::Result<u64> {
472        self.inner.raw_cuda_stream()
473    }
474
475    /// Block the current thread until work submitted to the current CUDA stream completes.
476    ///
477    /// # Examples
478    ///
479    /// ```
480    /// use tenferro_gpu::cuda::CudaRuntime;
481    ///
482    /// let _sync: fn(&CudaRuntime) -> tenferro_tensor::Result<()> =
483    ///     CudaRuntime::synchronize;
484    /// ```
485    ///
486    /// # Errors
487    ///
488    /// Returns [`crate::Error::RuntimeState`] when CubeCL cannot expose the
489    /// current stream, or [`crate::Error::BackendSource`] when CUDA context or
490    /// stream synchronization fails.
491    pub fn synchronize(&self) -> crate::Result<()> {
492        self.inner.synchronize()
493    }
494}
495
496impl CudaRuntimeState {
497    fn set_current_cuda_context(&self, op: &'static str) -> crate::Result<()> {
498        // INVARIANT: CUDA ordinals are device identifiers; bad ordinals are
499        // reported by CUDA instead of indexing memory in tenferro.
500        let device_ordinal = i32::try_from(self.device_id.ordinal())
501            .map_err(|source| crate::Error::backend_source(op, source))?;
502        cudarc::runtime::result::device::set(device_ordinal)
503            .map_err(|err| crate::Error::backend_source(op, err))?;
504        unsafe { cudarc::driver::result::ctx::set_current(self.primary_context.context()) }
505            .map_err(|err| crate::Error::backend_source(op, err))
506    }
507
508    fn raw_cuda_stream(&self) -> crate::Result<u64> {
509        self.client
510            .with_server(|server| {
511                server
512                    .raw_stream(StreamId::current())
513                    .map(|stream| stream as u64)
514                    .map_err(|err| crate::Error::backend_source("raw_cuda_stream", err))
515            })
516            .ok_or_else(|| {
517                crate::Error::runtime_state("raw_cuda_stream", "CubeCL server is unavailable")
518            })?
519    }
520
521    fn synchronize(&self) -> crate::Result<()> {
522        const OP: &str = "cubecl_runtime_synchronize";
523        self.set_current_cuda_context(OP)?;
524        let stream = self.raw_cuda_stream()? as usize as cudaStream_t;
525        unsafe { cuda_result::stream::synchronize(stream) }
526            .map_err(|err| crate::Error::backend_source(OP, err))
527    }
528}
529
530fn is_invalid_device_lookup(source: DriverError) -> bool {
531    source.0 == CUresult::CUDA_ERROR_INVALID_DEVICE
532}
533
534/// CUDA provider support for the shared GPU extension vocabulary.
535///
536/// See [`GpuExtensionCapability`](super::identity::GpuExtensionCapability) for
537/// the vocabulary. `PeerCopy` is hardware/topology dependent and is therefore
538/// reported false at the provider level; the directional query in the explicit
539/// multi-GPU copy API decides availability per source/destination pair.
540pub(crate) fn capabilities_for_device(capability: GpuExtensionCapability) -> bool {
541    !matches!(capability, GpuExtensionCapability::PeerCopy)
542}
543
544fn cuda_initialization_error<E>(
545    device: CudaDeviceId,
546    operation: &'static str,
547    source: E,
548) -> CudaDeviceError
549where
550    E: std::error::Error + Send + Sync + 'static,
551{
552    CudaDeviceError::Initialization {
553        device,
554        operation,
555        source: Box::new(source),
556    }
557}
558
559impl Drop for CudaRuntimeState {
560    fn drop(&mut self) {
561        // Drop cannot surface errors, but the runtime must not release the
562        // primary context while queued kernels may still reference it.
563        if let Err(err) = self.synchronize() {
564            report_cuda_runtime_drop_error(&err);
565        }
566    }
567}
568
569#[cfg(test)]
570mod tests;
tenferro_gpu/cubecl/runtime.rs

tenferro_gpu/cubecl/
runtime.rs
