tenferro_gpu/
lib.rs

1//! GPU backend implementations for tenferro tensors.
2//!
3//! # Examples
4//!
5//! ```rust
6//! #[cfg(feature = "cuda")]
7//! use tenferro_gpu::{cuda::cuda_devices, cuda::CudaBackend, cuda::CudaDeviceError};
8//!
9//! #[cfg(feature = "cuda")]
10//! fn first_cuda_backend() -> Result<Option<CudaBackend>, CudaDeviceError> {
11//!     let devices = cuda_devices()?;
12//!     let Some(device) = devices.first() else {
13//!         return Ok(None);
14//!     };
15//!     Ok(Some(CudaBackend::new(device.id())?))
16//! }
17//!
18//! // This ordinary doctest checks the discovery-based selection API without
19//! // requiring CUDA hardware at test time.
20//! #[cfg(feature = "cuda")]
21//! let _example: fn() -> Result<Option<CudaBackend>, CudaDeviceError> = first_cuda_backend;
22//! ```
23
24#[cfg(feature = "cuda")]
25use std::any::Any;
26
27#[cfg(feature = "cuda")]
28mod cubecl;
29#[cfg(any(feature = "cuda", feature = "webgpu"))]
30mod event_domain_admission;
31#[cfg(any(feature = "cuda", feature = "webgpu"))]
32mod event_retirement;
33#[cfg(any(feature = "cuda", feature = "webgpu"))]
34mod kernels;
35#[cfg(any(feature = "cuda", feature = "webgpu"))]
36mod native_permutation;
37#[cfg(feature = "webgpu")]
38pub mod webgpu;
39
40/// CUDA provider namespace.
41#[cfg(feature = "cuda")]
42pub mod cuda {
43    pub use super::cubecl::{
44        cuda_capabilities, cuda_devices, cuda_runtime_engine_registration,
45        cuda_runtime_hardware_class, download_tensor, gpu_available, upload_tensor,
46        with_cuda_exec_session, CudaBackend, CudaComputeCapability, CudaDeviceError, CudaDeviceId,
47        CudaDeviceInfo, CudaDeviceUuid, CudaExecSession, CudaExtensionCache,
48        CudaExtensionCacheGuard, CudaRuntime, CudaRuntimeIdentity, GpuExtensionCapability,
49    };
50
51    /// Public tenferro-wide CubeCL session (issue #1597).
52    ///
53    /// Exposes a narrow prelude of the CubeCL types needed to write and launch
54    /// `#[cube]` kernels against tenferro's GPU runtime. This module does not
55    /// re-export the whole of `cubecl`; downstream crates declare the framework
56    /// `t4a-cubecl` package explicitly.
57    pub mod cubecl {
58        pub use super::super::cubecl::session_cubecl::Session;
59        // Narrow prelude: only the types needed to describe a CubeCL launch.
60        pub use ::cubecl::prelude::{ArrayArg, CubeCount, CubeDim, TensorBinding};
61    }
62
63    /// Type-safe raw CUDA extension session (issue #1597).
64    pub mod raw {
65        pub use super::super::cubecl::raw::{
66            CudaResourceGuard, DeviceBytes, Function, KernelArg, LaunchConfig, Module,
67            NvrtcOptions, Session, StreamRef, TensorMut, TensorRef,
68        };
69    }
70}
71
72/// Apple shared-allocation provider namespace.
73#[cfg(feature = "webgpu")]
74pub mod apple {
75    pub use super::webgpu::{AppleContext, AppleTransferStats};
76}
77
78#[cfg(any(feature = "cuda", feature = "webgpu"))]
79use tenferro_tensor::*;
80
81#[cfg(feature = "cuda")]
82pub(crate) mod backend {
83    pub use tenferro_tensor::backend::*;
84}
85
86#[cfg(feature = "cuda")]
87pub(crate) mod config {
88    pub use tenferro_tensor::config::*;
89}
90
91#[cfg(feature = "cuda")]
92pub(crate) mod types {
93    pub(crate) use crate::CubeclBuffer;
94    pub use tenferro_tensor::types::*;
95}
96
97/// Scalar-independent CubeCL allocation stored behind tensor backend-buffer
98/// trait objects; dtype is carried by the borrowed tensor descriptor.
99#[cfg(feature = "cuda")]
100pub(crate) struct CubeclBuffer {
101    handle: cubecl_runtime::server::Handle,
102    byte_len: usize,
103    device_ordinal: usize,
104    allocation_domain: AllocationDomainId,
105    allocation_id: AllocationId,
106    // Memoized device address resolved by the first raw-FFI access.
107    //
108    // INVARIANT: in pinned CubeCL rev 5939d8e, a retained handle's memory
109    // slice keeps its storage offset (pool coalescing merges only free
110    // slices) and its backing storage is never deallocated while any of its
111    // slices is live, so the resolved address is stable for this buffer's
112    // lifetime. Raw-FFI callers must still route cross-stream accesses
113    // through `get_resource` for CubeCL's stream alignment.
114    device_addr: std::sync::OnceLock<u64>,
115}
116
117#[cfg(feature = "cuda")]
118static NEXT_CUDA_ALLOCATION_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
119
120#[cfg(feature = "cuda")]
121impl std::fmt::Debug for CubeclBuffer {
122    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123        f.debug_struct("CubeclBuffer")
124            .field("byte_len", &self.byte_len)
125            .field("device_ordinal", &self.device_ordinal)
126            .field("allocation_domain", &self.allocation_domain)
127            .field("allocation_id", &self.allocation_id)
128            .finish()
129    }
130}
131
132#[cfg(feature = "cuda")]
133impl CubeclBuffer {
134    pub(crate) fn new(
135        handle: cubecl_runtime::server::Handle,
136        byte_len: usize,
137        device_ordinal: usize,
138        allocation_domain: AllocationDomainId,
139    ) -> Self {
140        Self {
141            handle,
142            byte_len,
143            device_ordinal,
144            allocation_domain,
145            allocation_id: AllocationId::from_backend_id(
146                NEXT_CUDA_ALLOCATION_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
147            ),
148            device_addr: std::sync::OnceLock::new(),
149        }
150    }
151
152    pub(crate) fn handle(&self) -> &cubecl_runtime::server::Handle {
153        &self.handle
154    }
155
156    /// Return the memoized device address, if one was resolved.
157    pub(crate) fn cached_device_addr(&self) -> Option<u64> {
158        self.device_addr.get().copied()
159    }
160
161    /// Memoize the device address resolved through `get_resource` for this
162    /// buffer's handle; later calls keep the first stored value.
163    pub(crate) fn memoize_device_addr(&self, addr: u64) {
164        let _ = self.device_addr.set(addr);
165    }
166
167    pub(crate) fn element_len<T: 'static>(&self) -> usize {
168        let element_size = std::mem::size_of::<T>();
169        debug_assert!(element_size != 0 && self.byte_len.is_multiple_of(element_size));
170        self.byte_len / element_size
171    }
172
173    pub(crate) fn device_ordinal(&self) -> usize {
174        self.device_ordinal
175    }
176
177    pub(crate) fn allocation_domain(&self) -> AllocationDomainId {
178        self.allocation_domain
179    }
180}
181
182#[cfg(feature = "cuda")]
183impl<T: Send + Sync + 'static> BackendStorage<T> for CubeclBuffer {
184    fn backend_family(&self) -> &'static str {
185        "cubecl"
186    }
187
188    fn len(&self) -> usize {
189        self.element_len::<T>()
190    }
191
192    fn allocation_domain(&self) -> Option<AllocationDomainId> {
193        Some(self.allocation_domain)
194    }
195
196    fn allocation_id(&self) -> Option<AllocationId> {
197        Some(self.allocation_id)
198    }
199
200    fn prepare_device_access(
201        &self,
202        request: DeviceAccessRequest<'_>,
203    ) -> std::result::Result<Box<dyn PreparedDeviceAccess>, DeviceAccessError> {
204        Ok(Box::new(crate::cubecl::dispatch::prepare_cubecl_access(
205            self, request,
206        )?))
207    }
208
209    fn as_any(&self) -> &dyn Any {
210        self
211    }
212}
tenferro_gpu/lib.rs

tenferro_gpu/
lib.rs
