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}
107
108#[cfg(feature = "cuda")]
109static NEXT_CUDA_ALLOCATION_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
110
111#[cfg(feature = "cuda")]
112impl std::fmt::Debug for CubeclBuffer {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        f.debug_struct("CubeclBuffer")
115            .field("byte_len", &self.byte_len)
116            .field("device_ordinal", &self.device_ordinal)
117            .field("allocation_domain", &self.allocation_domain)
118            .field("allocation_id", &self.allocation_id)
119            .finish()
120    }
121}
122
123#[cfg(feature = "cuda")]
124impl CubeclBuffer {
125    pub(crate) fn new(
126        handle: cubecl_runtime::server::Handle,
127        byte_len: usize,
128        device_ordinal: usize,
129        allocation_domain: AllocationDomainId,
130    ) -> Self {
131        Self {
132            handle,
133            byte_len,
134            device_ordinal,
135            allocation_domain,
136            allocation_id: AllocationId::from_backend_id(
137                NEXT_CUDA_ALLOCATION_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
138            ),
139        }
140    }
141
142    pub(crate) fn handle(&self) -> &cubecl_runtime::server::Handle {
143        &self.handle
144    }
145
146    pub(crate) fn element_len<T: 'static>(&self) -> usize {
147        let element_size = std::mem::size_of::<T>();
148        debug_assert!(element_size != 0 && self.byte_len.is_multiple_of(element_size));
149        self.byte_len / element_size
150    }
151
152    pub(crate) fn device_ordinal(&self) -> usize {
153        self.device_ordinal
154    }
155
156    pub(crate) fn allocation_domain(&self) -> AllocationDomainId {
157        self.allocation_domain
158    }
159}
160
161#[cfg(feature = "cuda")]
162impl<T: Send + Sync + 'static> BackendStorage<T> for CubeclBuffer {
163    fn backend_family(&self) -> &'static str {
164        "cubecl"
165    }
166
167    fn len(&self) -> usize {
168        self.element_len::<T>()
169    }
170
171    fn allocation_domain(&self) -> Option<AllocationDomainId> {
172        Some(self.allocation_domain)
173    }
174
175    fn allocation_id(&self) -> Option<AllocationId> {
176        Some(self.allocation_id)
177    }
178
179    fn prepare_device_access(
180        &self,
181        request: DeviceAccessRequest<'_>,
182    ) -> std::result::Result<Box<dyn PreparedDeviceAccess>, DeviceAccessError> {
183        Ok(Box::new(crate::cubecl::dispatch::prepare_cubecl_access(
184            self, request,
185        )?))
186    }
187
188    fn as_any(&self) -> &dyn Any {
189        self
190    }
191}
tenferro_gpu/lib.rs

tenferro_gpu/
lib.rs
