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, Mutex, OnceLock};
7
8use cubecl::client::ComputeClient;
9use cubecl::stream_id::StreamId;
10use cubecl::Runtime;
11use cubecl_cuda::{CudaDevice, CudaRuntime as CubeclCudaRuntime};
12use cubecl_runtime::config::{CubeClRuntimeConfig, RuntimeConfig};
13use cudarc::cublas::sys as cublas_sys;
14use cudarc::driver::result::DriverError;
15use cudarc::driver::sys::{CUcontext, CUdevice, CUresult};
16use cudarc::runtime::{result as cuda_result, sys as cuda_sys, sys::cudaStream_t};
17use tenferro_tensor::AllocationDomainId;
18
19use super::device::{
20    cuda_devices, unavailable_device_error, CudaDeviceError, CudaDeviceId, CudaDeviceInfo,
21};
22use super::identity::GpuExtensionCapability;
23
24/// Returns `true` if a CUDA device can initialize a CubeCL runtime.
25///
26/// Use this in test helpers to skip GPU tests on machines without hardware.
27pub fn gpu_available() -> bool {
28    let library_present = std::panic::catch_unwind(|| {
29        // SAFETY: `is_culib_present` only probes candidate library names and
30        // does not call CUDA function pointers or retain a library handle.
31        unsafe { cudarc::driver::sys::is_culib_present() }
32    })
33    .unwrap_or(false);
34    if !library_present {
35        return false;
36    }
37    let Ok(devices) = cuda_devices() else {
38        return false;
39    };
40    let Some(device_id) = devices.first().map(|device| device.id()) else {
41        return false;
42    };
43    std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
44        let Ok(runtime) = CudaRuntime::new(device_id) else {
45            return false;
46        };
47        runtime.synchronize().is_ok()
48    }))
49    .unwrap_or(false)
50}
51
52/// RAII guard that attempts to restore the thread's previous CUDA device and
53/// current context when dropped.
54///
55/// Used by the `with_raw` enter/exit protocol: the guard is created after the
56/// calling thread's device/context are saved and the tenferro primary context
57/// is activated. Drop attempts best-effort restoration of the saved state on
58/// normal return, `Err`, and unwind; a restoration failure is logged to
59/// stderr (non-panicking) and never returned.
60pub(crate) struct RawContextRestore {
61    saved_device: Result<i32, cudarc::runtime::result::RuntimeError>,
62    saved_context: Result<Option<CUcontext>, cudarc::driver::result::DriverError>,
63    op: &'static str,
64}
65
66impl RawContextRestore {
67    /// Save the current device/context, then activate `device`/`context`.
68    pub(crate) fn enter(op: &'static str, device: i32, context: CUcontext) -> crate::Result<Self> {
69        let saved_device = cudarc::runtime::result::device::get();
70        let saved_context = cudarc::driver::result::ctx::get_current();
71        cudarc::runtime::result::device::set(device)
72            .map_err(|err| crate::Error::backend_source(op, err))?;
73        if let Err(err) = unsafe { cudarc::driver::result::ctx::set_current(context) } {
74            // Roll the device and context back so a partial activation failure
75            // cannot leave the caller's thread on a different device or with a
76            // different current context (setting the device can implicitly
77            // change the thread's current context to the new primary).
78            if let Ok(previous_device) = saved_device {
79                let _ = cudarc::runtime::result::device::set(previous_device);
80            }
81            match saved_context {
82                Ok(Some(previous)) => {
83                    let _ = unsafe { cudarc::driver::result::ctx::set_current(previous) };
84                }
85                Ok(None) => {
86                    let _ =
87                        unsafe { cudarc::driver::result::ctx::set_current(std::ptr::null_mut()) };
88                }
89                Err(_) => {}
90            }
91            return Err(crate::Error::backend_source(op, err));
92        }
93        Ok(Self {
94            saved_device,
95            saved_context,
96            op,
97        })
98    }
99
100    fn restore(&self) {
101        let mut stderr = std::io::stderr();
102        if let Ok(device) = self.saved_device {
103            if let Err(err) = cudarc::runtime::result::device::set(device) {
104                let _ = writeln!(
105                    stderr,
106                    "tenferro-gpu: failed to restore CUDA device during {}: {err:?}",
107                    self.op
108                );
109            }
110        }
111        match self.saved_context {
112            Ok(Some(context)) => {
113                if let Err(err) = unsafe { cudarc::driver::result::ctx::set_current(context) } {
114                    let _ = writeln!(
115                        stderr,
116                        "tenferro-gpu: failed to restore CUDA context during {}: {err:?}",
117                        self.op
118                    );
119                }
120            }
121            // The thread had no current context before the guard; restore that
122            // state instead of leaving the tenferro primary context current.
123            Ok(None) => {
124                if let Err(err) =
125                    unsafe { cudarc::driver::result::ctx::set_current(std::ptr::null_mut()) }
126                {
127                    let _ = writeln!(
128                        stderr,
129                        "tenferro-gpu: failed to clear CUDA context during {}: {err:?}",
130                        self.op
131                    );
132                }
133            }
134            // The saved-context query itself failed; nothing can be restored.
135            Err(_) => {}
136        }
137    }
138}
139
140impl Drop for RawContextRestore {
141    fn drop(&mut self) {
142        self.restore();
143    }
144}
145
146/// Opaque identity of one exact CUDA runtime instance.
147///
148/// Cloning the identity preserves the underlying executable runtime witness;
149/// constructing another runtime, even for the same device ordinal, produces a
150/// distinct identity. The cache key intentionally carries no provider or
151/// device identifier and grants no execution authority.
152#[derive(Clone, Debug)]
153pub struct CudaRuntimeIdentity {
154    marker: Arc<u8>,
155}
156
157impl CudaRuntimeIdentity {
158    fn fresh() -> Self {
159        Self {
160            marker: Arc::new(0),
161        }
162    }
163}
164
165impl PartialEq for CudaRuntimeIdentity {
166    fn eq(&self, other: &Self) -> bool {
167        Arc::ptr_eq(&self.marker, &other.marker)
168    }
169}
170
171impl Eq for CudaRuntimeIdentity {}
172
173impl Hash for CudaRuntimeIdentity {
174    fn hash<H: Hasher>(&self, state: &mut H) {
175        // INVARIANT: `marker` is retained by every clone of this identity, so
176        // its Arc allocation address is move/clone-invariant while witnessed.
177        state.write_usize(Arc::as_ptr(&self.marker) as usize);
178    }
179}
180
181/// CubeCL CUDA runtime wrapper.
182///
183/// # Examples
184///
185/// ```
186/// use tenferro_gpu::cuda::CudaRuntime;
187///
188/// let _ctor: fn(tenferro_gpu::cuda::CudaDeviceId) ->
189///     Result<CudaRuntime, tenferro_gpu::cuda::CudaDeviceError> = CudaRuntime::new;
190/// let _sync: fn(&CudaRuntime) -> tenferro_tensor::Result<()> =
191///     CudaRuntime::synchronize;
192/// ```
193#[derive(Clone)]
194pub struct CudaRuntime {
195    inner: Arc<CudaRuntimeState>,
196}
197
198struct CudaRuntimeState {
199    client: ComputeClient<CubeclCudaRuntime>,
200    device_id: CudaDeviceId,
201    device_ordinal: usize,
202    device_info: CudaDeviceInfo,
203    primary_context: CudaPrimaryContext,
204    identity: CudaRuntimeIdentity,
205    allocation_domain: AllocationDomainId,
206    // Memoized raw CUDA stream handles keyed by the bounded CubeCL stream-pool
207    // slot, not the process-global and monotonically increasing `StreamId`.
208    //
209    // INVARIANT: in pinned CubeCL rev 5939d8e, the CUDA server maps each
210    // `StreamId` to a fixed `StreamPool` slot whose `CUstream` is created once
211    // and never destroyed or replaced while the server is alive, and the
212    // server outlives the `ComputeClient` clone owned by this state. The table
213    // is owned by this runtime object (not thread-local/global), holds one
214    // entry per CubeCL stream slot, and is dropped with the runtime.
215    raw_streams: Box<[OnceLock<u64>]>,
216    // One lazily created cuBLAS handle per bounded CubeCL stream-pool slot.
217    // Each slot lock covers pointer-mode selection and the enqueue itself:
218    // distinct `StreamId`s can map to the same physical stream, and cuBLAS
219    // handle configuration is mutable.
220    cublas_handles: Box<[Mutex<Option<CublasStreamHandle>>]>,
221    // Lazily allocated pinned-host staging slot for single-scalar downloads
222    // (`cudaHostAlloc`, `PINNED_SCALAR_BYTES` bytes). Freed in `Drop` with
223    // `cudaFreeHost` while the primary context is still retained.
224    pinned_scalar: Mutex<PinnedScalarSlot>,
225}
226
227/// One cached cuBLAS handle bound to a fixed CUDA stream.
228struct CublasStreamHandle(cublas_sys::cublasHandle_t);
229
230/// Pinned-host staging slot; `ptr` is null until the first scalar download.
231struct PinnedScalarSlot {
232    ptr: *mut std::ffi::c_void,
233}
234
235/// Size of the runtime-owned pinned staging slot: large enough for the widest
236/// supported scalar (`Complex64`, 16 bytes).
237pub(crate) const PINNED_SCALAR_BYTES: usize = 16;
238
239// SAFETY: `CudaRuntimeState` owns a retained CUDA primary context and a CubeCL
240// client for one device ordinal. Methods set the context current before raw CUDA
241// calls, and backend/executor layers serialize mutating tensor execution. The
242// raw cuBLAS handles and the pinned staging pointer are plain CUDA resource
243// addresses owned by this state and released in `Drop`.
244unsafe impl Send for CudaRuntimeState {}
245// SAFETY: Shared state access exposes immutable runtime handles; synchronization
246// and stream queries use explicit CUDA/CubeCL handles and do not mutate Rust
247// aliasing-visible fields. cuBLAS handles are locked per bounded physical
248// stream slot, and the pinned staging slot is used only while its mutex is held.
249unsafe impl Sync for CudaRuntimeState {}
250
251impl fmt::Debug for CudaRuntime {
252    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
253        f.debug_struct("CudaRuntime")
254            .field("device_id", &self.inner.device_id)
255            .finish_non_exhaustive()
256    }
257}
258
259struct CudaPrimaryContext {
260    cuda_device: CUdevice,
261    cuda_context: CUcontext,
262}
263
264impl CudaPrimaryContext {
265    fn retain(cuda_device: CUdevice) -> crate::Result<Self> {
266        let cuda_context = unsafe { cudarc::driver::result::primary_ctx::retain(cuda_device) }
267            .map_err(|err| crate::Error::backend_source("cubecl_runtime_init", err))?;
268        Ok(Self {
269            cuda_device,
270            cuda_context,
271        })
272    }
273
274    fn context(&self) -> CUcontext {
275        self.cuda_context
276    }
277}
278
279impl Drop for CudaPrimaryContext {
280    fn drop(&mut self) {
281        if let Err(err) = unsafe { cudarc::driver::result::primary_ctx::release(self.cuda_device) }
282        {
283            report_cuda_primary_context_release_error(&err);
284        }
285    }
286}
287
288#[cold]
289fn report_cuda_primary_context_release_error(err: &impl fmt::Debug) {
290    eprintln!("tenferro-gpu: failed to release CUDA primary context during Drop: {err:?}");
291}
292
293#[cold]
294fn report_cuda_runtime_drop_error(err: &crate::Error) {
295    eprintln!("tenferro-gpu: failed to synchronize CUDA runtime during Drop: {err}");
296}
297
298impl CudaRuntime {
299    /// Initialize the CubeCL CUDA runtime on the caller-selected device.
300    ///
301    /// # Examples
302    ///
303    /// ```
304    /// use tenferro_gpu::{cuda::CudaDeviceError, cuda::CudaDeviceId, cuda::CudaRuntime};
305    ///
306    /// let _ctor: fn(CudaDeviceId) -> Result<CudaRuntime, CudaDeviceError> = CudaRuntime::new;
307    /// ```
308    ///
309    /// # Errors
310    ///
311    /// Returns [`CudaDeviceError::Discovery`] when fallback discovery for an
312    /// invalid selected ordinal fails, [`CudaDeviceError::Unavailable`] when
313    /// that ordinal is not available, or [`CudaDeviceError::Initialization`]
314    /// when CUDA driver, runtime, context, or CubeCL client initialization
315    /// fails.
316    pub fn new(device_id: CudaDeviceId) -> Result<Self, CudaDeviceError> {
317        let device_ordinal = usize::try_from(device_id.ordinal()).map_err(|source| {
318            cuda_initialization_error(device_id, "convert_device_ordinal", source)
319        })?;
320        let cuda_ordinal = i32::try_from(device_id.ordinal()).map_err(|source| {
321            cuda_initialization_error(device_id, "convert_cuda_ordinal", source)
322        })?;
323        cudarc::driver::result::init()
324            .map_err(|source| cuda_initialization_error(device_id, "initialize_driver", source))?;
325        let cuda_device = match cudarc::driver::result::device::get(cuda_ordinal) {
326            Ok(cuda_device) => cuda_device,
327            Err(source) if is_invalid_device_lookup(source) => {
328                return Err(unavailable_device_error(device_id, cuda_devices()?));
329            }
330            Err(source) => {
331                return Err(cuda_initialization_error(device_id, "get_device", source));
332            }
333        };
334        let primary_context = CudaPrimaryContext::retain(cuda_device).map_err(|source| {
335            cuda_initialization_error(device_id, "retain_primary_context", source)
336        })?;
337        unsafe { cudarc::driver::result::ctx::set_current(primary_context.context()) }.map_err(
338            |source| cuda_initialization_error(device_id, "set_current_context", source),
339        )?;
340        cudarc::runtime::result::device::set(cuda_ordinal)
341            .map_err(|source| cuda_initialization_error(device_id, "set_device", source))?;
342        let device = CudaDevice::new(device_ordinal);
343        let client = CubeclCudaRuntime::client(&device);
344        let discovered = cuda_devices()?;
345        let device_info = discovered
346            .iter()
347            .find(|info| info.id() == device_id)
348            .cloned()
349            .ok_or_else(|| unavailable_device_error(device_id, discovered))?;
350        Ok(Self {
351            inner: Arc::new(CudaRuntimeState {
352                client,
353                device_id,
354                device_ordinal,
355                device_info,
356                primary_context,
357                identity: CudaRuntimeIdentity::fresh(),
358                allocation_domain: AllocationDomainId::fresh(),
359                raw_streams: (0..cubecl_stream_slots())
360                    .map(|_| OnceLock::new())
361                    .collect(),
362                cublas_handles: (0..cubecl_stream_slots())
363                    .map(|_| Mutex::new(None))
364                    .collect(),
365                pinned_scalar: Mutex::new(PinnedScalarSlot {
366                    ptr: std::ptr::null_mut(),
367                }),
368            }),
369        })
370    }
371
372    pub(crate) fn client(&self) -> &ComputeClient<CubeclCudaRuntime> {
373        &self.inner.client
374    }
375
376    /// Return the caller-selected CUDA device identity that this runtime targets.
377    ///
378    /// # Examples
379    ///
380    /// ```
381    /// use tenferro_gpu::{cuda::CudaDeviceId, cuda::CudaRuntime};
382    ///
383    /// let _device_id: fn(&CudaRuntime) -> CudaDeviceId = CudaRuntime::device_id;
384    /// ```
385    pub fn device_id(&self) -> CudaDeviceId {
386        self.inner.device_id
387    }
388
389    /// Return immutable metadata for the runtime's device.
390    ///
391    /// # Examples
392    ///
393    /// ```
394    /// use tenferro_gpu::cuda::CudaRuntime;
395    ///
396    /// let _info: fn(&CudaRuntime) -> &tenferro_gpu::cuda::CudaDeviceInfo =
397    ///     CudaRuntime::device_info;
398    /// ```
399    pub fn device_info(&self) -> &CudaDeviceInfo {
400        &self.inner.device_info
401    }
402
403    /// Return the allocation ownership domain of this runtime.
404    ///
405    /// # Examples
406    ///
407    /// ```
408    /// use tenferro_gpu::cuda::CudaRuntime;
409    ///
410    /// let _domain: fn(&CudaRuntime) -> tenferro_tensor::AllocationDomainId =
411    ///     CudaRuntime::allocation_domain;
412    /// ```
413    pub fn allocation_domain(&self) -> AllocationDomainId {
414        self.inner.allocation_domain
415    }
416
417    /// Report whether this CUDA session supports a GPU extension capability.
418    ///
419    /// The CUDA provider supports the full extension vocabulary: external
420    /// CubeCL kernels, native module loading, runtime compilation (NVRTC), raw
421    /// stream borrowing, and same-device copy. `PeerCopy` is reported as a
422    /// directional query; availability is hardware-dependent and is checked
423    /// per source/destination pair rather than here.
424    ///
425    /// # Examples
426    ///
427    /// ```
428    /// use tenferro_gpu::cuda::GpuExtensionCapability;
429    /// use tenferro_gpu::cuda::CudaRuntime;
430    ///
431    /// let _supports: fn(&CudaRuntime, GpuExtensionCapability) -> bool =
432    ///     CudaRuntime::supports_extension;
433    /// ```
434    pub fn supports_extension(&self, capability: GpuExtensionCapability) -> bool {
435        capabilities_for_device(capability)
436    }
437
438    pub(crate) fn device_ordinal(&self) -> usize {
439        self.inner.device_ordinal
440    }
441
442    pub(crate) fn primary_context(&self) -> CUcontext {
443        self.inner.primary_context.context()
444    }
445
446    /// Run `f` with the tenferro primary context current on this thread.
447    ///
448    /// Saves the calling thread's current CUDA device/context, activates the
449    /// tenferro primary context for the duration of `f`, and attempts to
450    /// restore the saved state on every exit path (normal return, `Err`, and
451    /// unwind). Restoration is best-effort: a failure to restore the
452    /// caller's previous device/context is logged to stderr rather than
453    /// returned. This is the scoped context authority used by vendor-library
454    /// lifecycle paths (plan creation/retirement) that run outside a
455    /// raw-session callback.
456    ///
457    /// # Errors
458    ///
459    /// Returns [`crate::Error::BackendSource`] when the tenferro primary
460    /// context cannot be activated (device or context driver failure); a
461    /// partial activation is best-effort rolled back before the error is
462    /// returned (rollback failures are discarded).
463    ///
464    /// # Examples
465    ///
466    /// ```
467    /// use tenferro_gpu::cuda::CudaRuntime;
468    ///
469    /// let _check: fn(&CudaRuntime) -> tenferro_tensor::Result<u64> = |rt| {
470    ///     rt.with_current_context("test.context", || 7)
471    /// };
472    /// ```
473    pub fn with_current_context<R>(
474        &self,
475        op: &'static str,
476        f: impl FnOnce() -> R,
477    ) -> crate::Result<R> {
478        let device_ordinal = i32::try_from(self.device_ordinal())
479            .map_err(|source| crate::Error::backend_source(op, source))?;
480        let _guard = RawContextRestore::enter(op, device_ordinal, self.primary_context())?;
481        Ok(f())
482    }
483
484    /// Flush pending CubeCL work on the current stream.
485    ///
486    /// Used by the raw-session enter protocol so raw library calls observe
487    /// previously enqueued CubeCL work.
488    pub(crate) fn flush_cubecl(&self, op: &'static str) -> crate::Result<()> {
489        self.inner.flush_cubecl(op)
490    }
491
492    /// Return the opaque identity of this exact executable runtime instance.
493    ///
494    /// # Examples
495    ///
496    /// ```
497    /// use tenferro_gpu::cuda::CudaRuntime;
498    ///
499    /// let _identity: fn(&CudaRuntime) -> tenferro_gpu::cuda::CudaRuntimeIdentity =
500    ///     CudaRuntime::runtime_identity;
501    /// ```
502    pub fn runtime_identity(&self) -> CudaRuntimeIdentity {
503        self.inner.identity.clone()
504    }
505
506    pub(crate) fn allocation_domain_id(&self) -> AllocationDomainId {
507        self.inner.allocation_domain
508    }
509
510    pub(crate) fn set_current_cuda_context(&self, op: &'static str) -> crate::Result<()> {
511        self.inner.set_current_cuda_context(op)
512    }
513
514    pub(crate) fn raw_cuda_stream(&self) -> crate::Result<u64> {
515        self.inner.raw_cuda_stream()
516    }
517
518    pub(crate) fn synchronize_raw_stream(
519        &self,
520        stream: u64,
521        op: &'static str,
522    ) -> crate::Result<()> {
523        self.inner.synchronize_raw_stream(stream, op)
524    }
525
526    /// Run one cuBLAS enqueue with the handle for the current CubeCL stream.
527    ///
528    /// The caller must have the tenferro primary context current on this
529    /// thread (see [`CudaRuntimeState::set_current_cuda_context`]).
530    pub(crate) fn with_cublas_handle<R>(
531        &self,
532        op: &'static str,
533        pointer_mode: cublas_sys::cublasPointerMode_t,
534        cross_stream_handles: Vec<cubecl_runtime::server::Handle>,
535        execute: impl FnOnce(cublas_sys::cublasHandle_t) -> crate::Result<R>,
536    ) -> crate::Result<R> {
537        self.inner
538            .with_cublas_handle(op, pointer_mode, cross_stream_handles, execute)
539    }
540
541    pub(crate) fn finish_vendor_enqueue<R>(
542        &self,
543        op: &'static str,
544        cross_stream_handles: Vec<cubecl_runtime::server::Handle>,
545        result: crate::Result<R>,
546    ) -> crate::Result<R> {
547        self.inner
548            .finish_vendor_enqueue(op, cross_stream_handles, result)
549    }
550
551    pub(crate) fn stream_slot(&self) -> usize {
552        self.inner.stream_slot()
553    }
554
555    pub(crate) fn stream_slot_count(&self) -> usize {
556        self.inner.raw_streams.len()
557    }
558
559    pub(crate) fn is_current_stream_slot(&self, handle: &cubecl_runtime::server::Handle) -> bool {
560        self.inner.stream_slot_for(handle.stream) == self.inner.stream_slot()
561    }
562
563    /// Download up to [`PINNED_SCALAR_BYTES`] bytes from a device address
564    /// through the runtime-owned pinned staging slot.
565    ///
566    /// Enqueues an async device-to-host copy on the current thread's CubeCL
567    /// stream and synchronizes only that stream, so previously enqueued work
568    /// on the stream is observed without a device-wide barrier.
569    pub(crate) fn download_scalar_bytes(
570        &self,
571        device_addr: u64,
572        out: &mut [u8],
573        op: &'static str,
574        retained: cubecl_runtime::server::Handle,
575    ) -> crate::Result<()> {
576        self.inner
577            .download_scalar_bytes(device_addr, out, op, retained)
578    }
579
580    /// Block the current thread until work submitted to the current CUDA stream completes.
581    ///
582    /// # Examples
583    ///
584    /// ```
585    /// use tenferro_gpu::cuda::CudaRuntime;
586    ///
587    /// let _sync: fn(&CudaRuntime) -> tenferro_tensor::Result<()> =
588    ///     CudaRuntime::synchronize;
589    /// ```
590    ///
591    /// # Errors
592    ///
593    /// Returns [`crate::Error::RuntimeState`] when CubeCL cannot expose the
594    /// current stream, or [`crate::Error::BackendSource`] when CUDA context or
595    /// stream synchronization fails.
596    pub fn synchronize(&self) -> crate::Result<()> {
597        self.inner.synchronize()
598    }
599}
600
601impl CudaRuntimeState {
602    fn stream_slot(&self) -> usize {
603        self.stream_slot_for(StreamId::current())
604    }
605
606    fn stream_slot_for(&self, stream_id: StreamId) -> usize {
607        stream_id.value as usize % self.raw_streams.len()
608    }
609
610    fn set_current_cuda_context(&self, op: &'static str) -> crate::Result<()> {
611        // Fast path: the tenferro primary context is already current on this
612        // thread. `cuCtxGetCurrent` only reads driver thread state, so this
613        // skips the per-op `cudaSetDevice` + `cuCtxSetCurrent` round trips.
614        // Runtime-API calls made afterwards operate on the current driver
615        // context, so no separate runtime-API device activation is needed.
616        if let Ok(Some(current)) = cudarc::driver::result::ctx::get_current() {
617            if current == self.primary_context.context() {
618                return Ok(());
619            }
620        }
621        // INVARIANT: CUDA ordinals are device identifiers; bad ordinals are
622        // reported by CUDA instead of indexing memory in tenferro.
623        let device_ordinal = i32::try_from(self.device_id.ordinal())
624            .map_err(|source| crate::Error::backend_source(op, source))?;
625        cudarc::runtime::result::device::set(device_ordinal)
626            .map_err(|err| crate::Error::backend_source(op, err))?;
627        unsafe { cudarc::driver::result::ctx::set_current(self.primary_context.context()) }
628            .map_err(|err| crate::Error::backend_source(op, err))
629    }
630
631    fn raw_cuda_stream(&self) -> crate::Result<u64> {
632        let stream_id = StreamId::current();
633        let slot = self.stream_slot();
634        if let Some(&stream) = self.raw_streams[slot].get() {
635            return Ok(stream);
636        }
637        let stream = self
638            .client
639            .with_server(move |server| {
640                server
641                    .raw_stream(stream_id)
642                    .map(|stream| stream as u64)
643                    .map_err(|err| crate::Error::backend_source("raw_cuda_stream", err))
644            })
645            .ok_or_else(|| {
646                crate::Error::runtime_state("raw_cuda_stream", "CubeCL server is unavailable")
647            })??;
648        Ok(*self.raw_streams[slot].get_or_init(|| stream))
649    }
650
651    fn flush_cubecl(&self, op: &'static str) -> crate::Result<()> {
652        self.client
653            .flush()
654            .map_err(|err| crate::Error::backend_source(op, err))
655    }
656
657    fn synchronize(&self) -> crate::Result<()> {
658        const OP: &str = "cubecl_runtime_synchronize";
659        // A cached raw stream does not drain CubeCL's host-side launch queue.
660        self.flush_cubecl(OP)?;
661        let stream = self.raw_cuda_stream()?;
662        self.synchronize_raw_stream(stream, OP)
663    }
664
665    fn synchronize_raw_stream(&self, stream: u64, op: &'static str) -> crate::Result<()> {
666        self.set_current_cuda_context(op)?;
667        unsafe { cuda_result::stream::synchronize(stream as usize as cudaStream_t) }
668            .map_err(|err| crate::Error::backend_source(op, err))
669    }
670
671    fn retire_initialized_streams(&self) -> bool {
672        const OP: &str = "cuda_runtime_drop";
673        if let Err(error) = self.set_current_cuda_context(OP) {
674            report_cuda_runtime_drop_error(&error);
675            return false;
676        }
677        let mut retired = true;
678        for stream in &self.raw_streams {
679            let Some(&stream) = stream.get() else {
680                continue;
681            };
682            // SAFETY: every initialized entry is a CubeCL-owned stream that
683            // remains live until this runtime state and its client are dropped.
684            if let Err(source) =
685                unsafe { cuda_result::stream::synchronize(stream as usize as cudaStream_t) }
686            {
687                retired = false;
688                report_cuda_runtime_drop_error(&crate::Error::backend_source(OP, source));
689            }
690        }
691        retired
692    }
693
694    fn with_cublas_handle<R>(
695        &self,
696        op: &'static str,
697        pointer_mode: cublas_sys::cublasPointerMode_t,
698        cross_stream_handles: Vec<cubecl_runtime::server::Handle>,
699        execute: impl FnOnce(cublas_sys::cublasHandle_t) -> crate::Result<R>,
700    ) -> crate::Result<R> {
701        let poisoned = || crate::Error::runtime_state(op, "cuBLAS handle cache lock poisoned");
702        let slot = self.stream_slot();
703        let mut cached = self.cublas_handles[slot].lock().map_err(|_| poisoned())?;
704        let handle = match *cached {
705            Some(ref handle) => handle.0,
706            None => {
707                if !cublas_library_present() {
708                    return Err(crate::Error::io_source(op, CublasLibraryMissing));
709                }
710                let stream = self.raw_cuda_stream()? as usize as cublas_sys::cudaStream_t;
711                let mut raw = std::ptr::null_mut();
712                // SAFETY: the caller holds the tenferro primary context current;
713                // the handle is created on this device and bound to this stream.
714                check_cublas(op, "cublasCreate", unsafe {
715                    cublas_sys::cublasCreate_v2(&mut raw)
716                })?;
717                // SAFETY: `raw` was just created and `stream` is owned by this
718                // runtime for the lifetime of the cached handle.
719                if let Err(err) = check_cublas(op, "cublasSetStream", unsafe {
720                    cublas_sys::cublasSetStream_v2(raw, stream)
721                }) {
722                    // SAFETY: `raw` is live and not stored anywhere else.
723                    let _ = unsafe { cublas_sys::cublasDestroy_v2(raw) };
724                    return Err(err);
725                }
726                *cached = Some(CublasStreamHandle(raw));
727                raw
728            }
729        };
730        // SAFETY: the per-stream slot lock is held across configuration and
731        // enqueue, so no caller can race this mutable handle state.
732        check_cublas(op, "cublasSetPointerMode", unsafe {
733            cublas_sys::cublasSetPointerMode_v2(handle, pointer_mode)
734        })?;
735        let result = execute(handle);
736        self.finish_vendor_enqueue(op, cross_stream_handles, result)
737    }
738
739    fn finish_vendor_enqueue<R>(
740        &self,
741        op: &'static str,
742        cross_stream_handles: Vec<cubecl_runtime::server::Handle>,
743        result: crate::Result<R>,
744    ) -> crate::Result<R> {
745        if cross_stream_handles.is_empty() {
746            return result;
747        }
748        let retirement = self.synchronize();
749        match (result, retirement) {
750            (Ok(value), Ok(())) => Ok(value),
751            (Err(error), Ok(())) => Err(error),
752            (Ok(_), Err(retirement)) => {
753                // No completion barrier was proven. Retain the foreign-stream
754                // allocations so their owners cannot reclaim or mutate them.
755                std::mem::forget(cross_stream_handles);
756                Err(crate::Error::backend_source(op, retirement))
757            }
758            (Err(error), Err(_retirement)) => {
759                std::mem::forget(cross_stream_handles);
760                Err(error)
761            }
762        }
763    }
764
765    fn download_scalar_bytes(
766        &self,
767        device_addr: u64,
768        out: &mut [u8],
769        op: &'static str,
770        retained: cubecl_runtime::server::Handle,
771    ) -> crate::Result<()> {
772        if out.len() > PINNED_SCALAR_BYTES {
773            return Err(crate::Error::Internal(format!(
774                "pinned scalar staging supports at most {PINNED_SCALAR_BYTES} bytes, got {}",
775                out.len()
776            )));
777        }
778        // Reused output addresses can already be cached: pointer lookup is not
779        // a queue barrier. Submit pending kernels before the raw D2H copy.
780        self.flush_cubecl(op)?;
781        self.set_current_cuda_context(op)?;
782        let stream = self.raw_cuda_stream()? as usize as cudaStream_t;
783        // ponytail: one shared staging slot serializes concurrent scalar
784        // downloads per runtime; add per-thread slots if that lock contends.
785        let mut slot = self
786            .pinned_scalar
787            .lock()
788            .map_err(|_| crate::Error::runtime_state(op, "pinned scalar staging lock poisoned"))?;
789        if slot.ptr.is_null() {
790            let mut ptr = std::ptr::null_mut();
791            // SAFETY: the primary context is current; the allocation is freed
792            // in this state's `Drop` with `cudaFreeHost`.
793            unsafe {
794                cuda_sys::cudaHostAlloc(
795                    &mut ptr,
796                    PINNED_SCALAR_BYTES,
797                    cuda_sys::cudaHostAllocDefault,
798                )
799            }
800            .result()
801            .map_err(|err| crate::Error::backend_source(op, err))?;
802            slot.ptr = ptr;
803        }
804        let src = super::interop::cuda_device_ptr_from_addr(device_addr, op)?;
805        // SAFETY: `slot.ptr` is a live pinned allocation of PINNED_SCALAR_BYTES
806        // bytes, `out.len()` is validated above, and the mutex guard keeps the
807        // slot exclusive until the copy below is known to have completed. Every
808        // exit that cannot prove completion abandons the slot instead of
809        // returning it, so exclusivity never rests on an unproven barrier.
810        let staging = unsafe { std::slice::from_raw_parts_mut(slot.ptr.cast::<u8>(), out.len()) };
811        // Neither submitting the copy nor waiting on it proves the device is
812        // done with `staging` and `retained` once it reports an error: an async
813        // CUDA call can surface a failure from an earlier launch on the stream,
814        // so a non-success return says nothing about what is still running.
815        // Both paths therefore leak the source allocation and abandon the
816        // staging slot rather than let a later download reuse a destination the
817        // device may still write, or let `Drop` `cudaFreeHost` it. Each failure
818        // leaks one slot and one handle; the next call allocates fresh ones.
819        // SAFETY: `src` is a residency-checked device address owned by this
820        // runtime, the copy length equals the destination slice length, and
821        // `stream` is the memoized CubeCL stream the copy is enqueued on.
822        let completed = unsafe { cuda_result::memcpy_dtoh_async(staging, src, stream) }
823            .and_then(|()| unsafe { cuda_result::stream::synchronize(stream) });
824        if let Err(err) = completed {
825            std::mem::forget(retained);
826            slot.ptr = std::ptr::null_mut();
827            return Err(crate::Error::backend_source(op, err));
828        }
829        out.copy_from_slice(staging);
830        Ok(())
831    }
832
833    /// Destroy cached cuBLAS handles and free the pinned staging slot.
834    ///
835    /// Called from `Drop` after all initialized streams have retired, which
836    /// leaves the primary context current on the dropping thread.
837    fn release_cuda_library_resources(&mut self) {
838        for cached in &self.cublas_handles {
839            if let Ok(mut handle) = cached.lock() {
840                let Some(handle) = handle.take() else {
841                    continue;
842                };
843                // SAFETY: each stored handle is live and no longer reachable.
844                if let Err(err) = unsafe { cublas_sys::cublasDestroy_v2(handle.0) }.result() {
845                    report_cuda_resource_release_error("cuBLAS handle", &err);
846                }
847            }
848        }
849        if let Ok(mut slot) = self.pinned_scalar.lock() {
850            if !slot.ptr.is_null() {
851                // SAFETY: the slot owns exactly one live cudaHostAlloc allocation.
852                if let Err(err) = unsafe { cuda_sys::cudaFreeHost(slot.ptr) }.result() {
853                    report_cuda_resource_release_error("pinned scalar staging", &err);
854                }
855                slot.ptr = std::ptr::null_mut();
856            }
857        }
858    }
859}
860
861fn cubecl_stream_slots() -> usize {
862    usize::from(CubeClRuntimeConfig::get().streaming.max_streams.max(1))
863}
864
865/// Typed load failure for the dynamically loaded cuBLAS library.
866#[derive(Debug, thiserror::Error)]
867#[error(
868    "cuBLAS shared library not found; ensure `LD_LIBRARY_PATH` includes the CUDA toolkit library directory"
869)]
870struct CublasLibraryMissing;
871
872/// Report whether the cuBLAS shared library can be dynamically loaded.
873fn cublas_library_present() -> bool {
874    use std::sync::OnceLock;
875    static PRESENT: OnceLock<bool> = OnceLock::new();
876    // SAFETY: `is_culib_present` only probes candidate library names and does
877    // not call cuBLAS function pointers or retain a library handle.
878    *PRESENT.get_or_init(|| unsafe { cublas_sys::is_culib_present() })
879}
880
881/// Map a non-success cuBLAS status to a typed provider error.
882pub(super) fn check_cublas(
883    op: &'static str,
884    call: &'static str,
885    status: cublas_sys::cublasStatus_t,
886) -> crate::Result<()> {
887    if matches!(status, cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS) {
888        Ok(())
889    } else {
890        Err(super::error::provider_status(
891            op,
892            "cuBLAS",
893            call,
894            status as i32,
895        ))
896    }
897}
898
899#[cold]
900fn report_cuda_resource_release_error(what: &'static str, err: &impl fmt::Debug) {
901    eprintln!("tenferro-gpu: failed to release {what} during Drop: {err:?}");
902}
903
904fn is_invalid_device_lookup(source: DriverError) -> bool {
905    source.0 == CUresult::CUDA_ERROR_INVALID_DEVICE
906}
907
908/// CUDA provider support for the shared GPU extension vocabulary.
909///
910/// See [`GpuExtensionCapability`](super::identity::GpuExtensionCapability) for
911/// the vocabulary. `PeerCopy` is hardware/topology dependent and is therefore
912/// reported false at the provider level; the directional query in the explicit
913/// multi-GPU copy API decides availability per source/destination pair.
914pub(crate) fn capabilities_for_device(capability: GpuExtensionCapability) -> bool {
915    !matches!(capability, GpuExtensionCapability::PeerCopy)
916}
917
918fn cuda_initialization_error<E>(
919    device: CudaDeviceId,
920    operation: &'static str,
921    source: E,
922) -> CudaDeviceError
923where
924    E: std::error::Error + Send + Sync + 'static,
925{
926    CudaDeviceError::Initialization {
927        device,
928        operation,
929        source: Box::new(source),
930    }
931}
932
933impl Drop for CudaRuntimeState {
934    fn drop(&mut self) {
935        // Drop cannot surface errors, but the runtime must not release the
936        // primary context while queued kernels on any initialized slot may
937        // still reference it.
938        if self.retire_initialized_streams() {
939            // Retirement left the primary context current; release CUDA
940            // library resources before the retained primary context drops.
941            self.release_cuda_library_resources();
942        }
943        // On retirement failure, raw library resources intentionally leak:
944        // their pointer-only owners have no Drop implementation, so Rust does
945        // not reclaim resources that may still be in use asynchronously.
946    }
947}
948
949#[cfg(test)]
950mod tests;
tenferro_gpu/cubecl/runtime.rs

tenferro_gpu/cubecl/
runtime.rs
