tenferro_gpu/
lib.rs

1//! GPU backend implementations for tenferro tensors.
2//!
3//! # Examples
4//!
5//! ```rust
6//! # fn main() -> tenferro_tensor::Result<()> {
7//! #[cfg(feature = "cuda")]
8//! {
9//!     use tenferro_gpu::{download_tensor, gpu_available, upload_tensor, CudaBackend};
10//!     use tenferro_tensor::{Tensor, TensorElementwise};
11//!
12//!     if gpu_available() {
13//!         let mut backend = CudaBackend::new(0).unwrap();
14//!         let a = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
15//!         let b = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0])?;
16//!         let gpu_a = upload_tensor(backend.runtime(), &a).unwrap();
17//!         let gpu_b = upload_tensor(backend.runtime(), &b).unwrap();
18//!         let gpu_sum = backend.add(&gpu_a, &gpu_b).unwrap();
19//!         let sum = download_tensor(backend.runtime(), &gpu_sum).unwrap();
20//!         assert_eq!(sum.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
21//!     }
22//! }
23//! # Ok(())
24//! # }
25//! ```
26
27#[cfg(feature = "cuda")]
28use std::any::Any;
29
30#[cfg(feature = "cuda")]
31mod cubecl;
32#[cfg(feature = "cuda")]
33mod kernels;
34#[cfg(feature = "webgpu")]
35mod webgpu;
36
37#[cfg(feature = "cuda")]
38pub use cubecl::{
39    cuda_capabilities, device_ptr, download_tensor, gpu_available, upload_tensor, CudaBackend,
40    CudaRuntime,
41};
42#[cfg(feature = "cuda")]
43#[doc(hidden)]
44pub use cubecl::{CudaExtensionCache, CudaExtensionCacheGuard};
45#[cfg(feature = "webgpu")]
46pub use webgpu::{
47    download_webgpu_tensor, upload_webgpu_tensor, webgpu_available, WebGpuBackend, WebGpuRuntime,
48};
49
50#[cfg(feature = "cuda")]
51#[doc(hidden)]
52pub mod cuda_interop {
53    pub use crate::cubecl::interop::*;
54    pub use crate::cubecl::{CudaExtensionCache, CudaExtensionCacheGuard};
55}
56
57#[cfg(any(feature = "cuda", feature = "webgpu"))]
58use tenferro_tensor::*;
59
60#[cfg(feature = "cuda")]
61pub(crate) mod backend {
62    pub use tenferro_tensor::backend::*;
63}
64
65#[cfg(feature = "cuda")]
66pub(crate) mod config {
67    pub use tenferro_tensor::config::*;
68}
69
70#[cfg(feature = "cuda")]
71pub(crate) mod types {
72    pub(crate) use crate::CubeclBuffer;
73    pub use tenferro_tensor::types::*;
74}
75
76/// CubeCL-managed GPU buffer stored behind tensor backend-buffer trait objects.
77#[cfg(feature = "cuda")]
78#[derive(Clone)]
79pub(crate) struct CubeclBuffer<T> {
80    handle: cubecl_runtime::server::Handle,
81    len: usize,
82    pub(crate) _marker: std::marker::PhantomData<T>,
83}
84
85#[cfg(feature = "cuda")]
86impl<T> std::fmt::Debug for CubeclBuffer<T> {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.debug_struct("CubeclBuffer")
89            .field("len", &self.len)
90            .finish()
91    }
92}
93
94#[cfg(feature = "cuda")]
95impl<T> CubeclBuffer<T> {
96    pub(crate) fn new(handle: cubecl_runtime::server::Handle, len: usize) -> Self {
97        Self {
98            handle,
99            len,
100            _marker: std::marker::PhantomData,
101        }
102    }
103
104    pub(crate) fn handle(&self) -> &cubecl_runtime::server::Handle {
105        &self.handle
106    }
107
108    pub(crate) fn element_len(&self) -> usize {
109        self.len
110    }
111}
112
113#[cfg(feature = "cuda")]
114impl<T: Send + Sync + 'static> BackendBuffer<T> for CubeclBuffer<T> {
115    fn backend_family(&self) -> &'static str {
116        "cubecl"
117    }
118
119    fn len(&self) -> usize {
120        self.len
121    }
122
123    fn as_any(&self) -> &dyn Any {
124        self
125    }
126}
tenferro_gpu/lib.rs

tenferro_gpu/
lib.rs
