tenferro_gpu/cubecl/
memory.rs

1//! Host-to-device and device-to-host transfers via CubeCL-managed allocations.
2
3use cubecl::client::ComputeClient;
4use cubecl::prelude::CubeElement;
5use cubecl_cuda::CudaRuntime as CubeclCudaRuntime;
6use num_complex::{Complex32, Complex64};
7
8use super::dispatch;
9use crate::cubecl::runtime::CudaRuntime;
10use crate::types::{
11    CubeclBuffer, DeviceId, DeviceKind, GpuBackendKind, MemoryKind, Placement, StorageBuffer,
12    Tensor, TensorScalar, TypedTensor,
13};
14
15/// Upload a host tensor into a CubeCL-managed GPU allocation.
16///
17/// # Examples
18///
19/// ```
20/// use tenferro_gpu::{cuda::upload_tensor, cuda::CudaRuntime};
21/// use tenferro_tensor::{Result, Tensor};
22///
23/// let _upload: fn(&CudaRuntime, &Tensor) -> Result<Tensor> = upload_tensor;
24/// ```
25///
26/// # Errors
27///
28/// Returns [`crate::Error::RuntimeState`] when the source is backend-resident
29/// or belongs to another placement, [`crate::Error::Unsupported`] for a dtype
30/// unavailable in CubeCL, or [`crate::Error::BackendSource`] on allocation.
31pub fn upload_tensor(rt: &CudaRuntime, tensor: &Tensor) -> crate::Result<Tensor> {
32    let client = rt.client();
33    match tensor {
34        Tensor::F64(t) => upload_typed::<f64>(rt, client, t).map(Tensor::F64),
35        Tensor::F32(t) => upload_typed::<f32>(rt, client, t).map(Tensor::F32),
36        Tensor::I32(t) => upload_typed::<i32>(rt, client, t).map(Tensor::I32),
37        Tensor::I64(t) => upload_typed::<i64>(rt, client, t).map(Tensor::I64),
38        Tensor::Bool(t) => upload_bool(rt, client, t).map(Tensor::Bool),
39        Tensor::C64(t) => upload_typed::<Complex64>(rt, client, t).map(Tensor::C64),
40        Tensor::C32(t) => upload_typed::<Complex32>(rt, client, t).map(Tensor::C32),
41    }
42}
43
44/// Download a CubeCL-managed GPU tensor back to host memory.
45///
46/// # Examples
47///
48/// ```
49/// use tenferro_gpu::{cuda::download_tensor, cuda::CudaRuntime};
50/// use tenferro_tensor::{Result, Tensor};
51///
52/// let _download: fn(&CudaRuntime, &Tensor) -> Result<Tensor> = download_tensor;
53/// ```
54///
55/// # Errors
56///
57/// Returns [`crate::Error::RuntimeState`] for a host-backed or foreign tensor,
58/// [`crate::Error::BackendSource`] when synchronization/readback fails, or a
59/// typed validation error when device data cannot be decoded.
60pub fn download_tensor(rt: &CudaRuntime, tensor: &Tensor) -> crate::Result<Tensor> {
61    ensure_tensor_resident_on_runtime(rt, tensor, "download")?;
62    match tensor {
63        Tensor::F64(t) => download_typed::<f64>(rt, t).map(Tensor::F64),
64        Tensor::F32(t) => download_typed::<f32>(rt, t).map(Tensor::F32),
65        Tensor::I32(t) => download_typed::<i32>(rt, t).map(Tensor::I32),
66        Tensor::I64(t) => download_typed::<i64>(rt, t).map(Tensor::I64),
67        Tensor::Bool(t) => download_bool(rt, t).map(Tensor::Bool),
68        Tensor::C64(t) => download_typed::<Complex64>(rt, t).map(Tensor::C64),
69        Tensor::C32(t) => download_typed::<Complex32>(rt, t).map(Tensor::C32),
70    }
71}
72
73fn upload_typed<T: CubeElement + TensorScalar + Clone + Send + Sync + 'static>(
74    rt: &CudaRuntime,
75    client: &ComputeClient<CubeclCudaRuntime>,
76    typed: &TypedTensor<T>,
77) -> crate::Result<TypedTensor<T>> {
78    let host_data = match typed.buffer() {
79        StorageBuffer::Host(data) => data,
80        StorageBuffer::Backend(buffer) => {
81            return Err(crate::Error::runtime_state(
82                "upload",
83                format!(
84                    "expected host buffer, got `{}` backend buffer",
85                    buffer.backend_family()
86                ),
87            ));
88        }
89    };
90
91    let handle = client.create_from_slice(T::as_bytes(host_data));
92    let byte_len = T::as_bytes(host_data).len();
93    TypedTensor::from_buffer_col_major(
94        typed.shape().to_vec(),
95        StorageBuffer::Backend(Box::new(CubeclBuffer::new(
96            handle,
97            byte_len,
98            rt.device_ordinal(),
99            rt.allocation_domain_id(),
100        ))),
101        Placement {
102            memory_kind: MemoryKind::Device,
103            device: Some(DeviceId {
104                kind: DeviceKind::Gpu(GpuBackendKind::Cuda),
105                ordinal: rt.device_ordinal(),
106            }),
107            cpu_affinity: None,
108        },
109    )
110}
111
112fn download_typed<T: CubeElement + TensorScalar + Clone + 'static>(
113    rt: &CudaRuntime,
114    typed: &TypedTensor<T>,
115) -> crate::Result<TypedTensor<T>> {
116    let handle = match typed.buffer() {
117        StorageBuffer::Host(_) => {
118            return Err(crate::Error::runtime_state(
119                "download",
120                "expected CubeCL buffer",
121            ));
122        }
123        StorageBuffer::Backend(buffer) => cubecl_handle_from_backend(buffer.as_ref(), "download")?,
124    };
125
126    if typed.n_elements() == 0 {
127        return TypedTensor::from_buffer_col_major(
128            typed.shape().to_vec(),
129            StorageBuffer::Host(Vec::new()),
130            Placement::default(),
131        );
132    }
133
134    rt.synchronize()?;
135    let bytes = rt
136        .client()
137        .read_one(handle)
138        .map_err(|err| crate::Error::backend_source("download", err))?;
139    let data = T::from_bytes(&bytes).to_vec();
140    TypedTensor::from_buffer_col_major(
141        typed.shape().to_vec(),
142        StorageBuffer::Host(data),
143        Placement::default(),
144    )
145}
146
147fn upload_bool(
148    rt: &CudaRuntime,
149    client: &ComputeClient<CubeclCudaRuntime>,
150    typed: &TypedTensor<bool>,
151) -> crate::Result<TypedTensor<bool>> {
152    let host_data = match typed.buffer() {
153        StorageBuffer::Host(data) => data,
154        StorageBuffer::Backend(buffer) => {
155            return Err(crate::Error::runtime_state(
156                "upload",
157                format!(
158                    "expected host buffer, got `{}` backend buffer",
159                    buffer.backend_family()
160                ),
161            ));
162        }
163    };
164
165    let bytes: Vec<u8> = host_data.iter().map(|&value| u8::from(value)).collect();
166    let handle = client.create_from_slice(&bytes);
167    TypedTensor::from_buffer_col_major(
168        typed.shape().to_vec(),
169        StorageBuffer::Backend(Box::new(CubeclBuffer::new(
170            handle,
171            bytes.len(),
172            rt.device_ordinal(),
173            rt.allocation_domain_id(),
174        ))),
175        Placement {
176            memory_kind: MemoryKind::Device,
177            device: Some(DeviceId {
178                kind: DeviceKind::Gpu(GpuBackendKind::Cuda),
179                ordinal: rt.device_ordinal(),
180            }),
181            cpu_affinity: None,
182        },
183    )
184}
185
186fn download_bool(rt: &CudaRuntime, typed: &TypedTensor<bool>) -> crate::Result<TypedTensor<bool>> {
187    let handle = match typed.buffer() {
188        StorageBuffer::Host(_) => {
189            return Err(crate::Error::runtime_state(
190                "download",
191                "expected CubeCL buffer",
192            ));
193        }
194        StorageBuffer::Backend(buffer) => cubecl_handle_from_backend(buffer.as_ref(), "download")?,
195    };
196
197    if typed.n_elements() == 0 {
198        return TypedTensor::from_vec_col_major(typed.shape().to_vec(), Vec::new());
199    }
200
201    rt.synchronize()?;
202    let bytes = rt
203        .client()
204        .read_one(handle)
205        .map_err(|err| crate::Error::backend_source("download", err))?;
206    let data = bytes.iter().map(|&byte| byte != 0).collect();
207    TypedTensor::from_vec_col_major(typed.shape().to_vec(), data)
208}
209
210fn ensure_tensor_resident_on_runtime(
211    rt: &CudaRuntime,
212    tensor: &Tensor,
213    op: &'static str,
214) -> crate::Result<()> {
215    match tensor {
216        Tensor::F64(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
217        Tensor::F32(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
218        Tensor::I32(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
219        Tensor::I64(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
220        Tensor::Bool(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
221        Tensor::C64(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
222        Tensor::C32(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
223    }
224}
225
226fn cubecl_handle_from_backend<T: 'static>(
227    buffer: &dyn crate::BackendStorage<T>,
228    op: &'static str,
229) -> crate::Result<cubecl_runtime::server::Handle> {
230    buffer
231        .as_any()
232        .downcast_ref::<CubeclBuffer>()
233        .map(|buffer| buffer.handle().clone())
234        .ok_or_else(|| {
235            crate::Error::runtime_state(
236                op,
237                format!(
238                    "expected CubeCL buffer, got `{}` backend buffer",
239                    buffer.backend_family()
240                ),
241            )
242        })
243}
tenferro_gpu/cubecl/memory.rs

tenferro_gpu/cubecl/
memory.rs
