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, PINNED_SCALAR_BYTES};
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
112/// Read a compact tensor of at most `PINNED_SCALAR_BYTES` through the runtime's
113/// pinned staging slot.
114fn download_scalar<T: CubeElement + TensorScalar + Clone + 'static>(
115    rt: &CudaRuntime,
116    typed: &TypedTensor<T>,
117    byte_len: usize,
118) -> crate::Result<Vec<T>> {
119    let ptr = super::gemm::typed_device_ptr(rt, typed, "download")?;
120    let retained = dispatch::cubecl_buffer(typed, "download")?.handle().clone();
121    let mut bytes = vec![0_u8; byte_len];
122    rt.download_scalar_bytes(ptr as u64, &mut bytes, "download", retained)?;
123    Ok(T::from_bytes(&bytes).to_vec())
124}
125
126fn download_typed<T: CubeElement + TensorScalar + Clone + 'static>(
127    rt: &CudaRuntime,
128    typed: &TypedTensor<T>,
129) -> crate::Result<TypedTensor<T>> {
130    let handle = match typed.buffer() {
131        StorageBuffer::Host(_) => {
132            return Err(crate::Error::runtime_state(
133                "download",
134                "expected CubeCL buffer",
135            ));
136        }
137        StorageBuffer::Backend(buffer) => cubecl_handle_from_backend(buffer.as_ref(), "download")?,
138    };
139
140    if typed.n_elements() == 0 {
141        return TypedTensor::from_buffer_col_major(
142            typed.shape().to_vec(),
143            StorageBuffer::Host(Vec::new()),
144            Placement::default(),
145        );
146    }
147
148    // Small-payload fast path. A Krylov loop reads back one reduction result
149    // per iteration. CubeCL's `read_one` already stages through pinned memory,
150    // so this is not a pageable-to-pinned conversion; what it avoids is
151    // CubeCL's staging-reservation machinery and one of two barriers. The
152    // general path below synchronizes the stream and then lets `read_one` wait
153    // on its own fence after the copy, while this issues the copy and waits
154    // once. Neither path synchronizes the device. The gate is a byte length, so
155    // a short vector takes it too, not only a scalar.
156    let byte_len = typed
157        .n_elements()
158        .checked_mul(size_of::<T>())
159        .ok_or_else(|| {
160            crate::Error::invalid_argument("download", "shape", "byte length overflows")
161        })?;
162    if byte_len <= PINNED_SCALAR_BYTES {
163        let data = download_scalar(rt, typed, byte_len)?;
164        return TypedTensor::from_buffer_col_major(
165            typed.shape().to_vec(),
166            StorageBuffer::Host(data),
167            Placement::default(),
168        );
169    }
170
171    rt.synchronize()?;
172    let bytes = rt
173        .client()
174        .read_one(handle)
175        .map_err(|err| crate::Error::backend_source("download", err))?;
176    let data = T::from_bytes(&bytes).to_vec();
177    TypedTensor::from_buffer_col_major(
178        typed.shape().to_vec(),
179        StorageBuffer::Host(data),
180        Placement::default(),
181    )
182}
183
184fn upload_bool(
185    rt: &CudaRuntime,
186    client: &ComputeClient<CubeclCudaRuntime>,
187    typed: &TypedTensor<bool>,
188) -> crate::Result<TypedTensor<bool>> {
189    let host_data = match typed.buffer() {
190        StorageBuffer::Host(data) => data,
191        StorageBuffer::Backend(buffer) => {
192            return Err(crate::Error::runtime_state(
193                "upload",
194                format!(
195                    "expected host buffer, got `{}` backend buffer",
196                    buffer.backend_family()
197                ),
198            ));
199        }
200    };
201
202    let bytes: Vec<u8> = host_data.iter().map(|&value| u8::from(value)).collect();
203    let handle = client.create_from_slice(&bytes);
204    TypedTensor::from_buffer_col_major(
205        typed.shape().to_vec(),
206        StorageBuffer::Backend(Box::new(CubeclBuffer::new(
207            handle,
208            bytes.len(),
209            rt.device_ordinal(),
210            rt.allocation_domain_id(),
211        ))),
212        Placement {
213            memory_kind: MemoryKind::Device,
214            device: Some(DeviceId {
215                kind: DeviceKind::Gpu(GpuBackendKind::Cuda),
216                ordinal: rt.device_ordinal(),
217            }),
218            cpu_affinity: None,
219        },
220    )
221}
222
223fn download_bool(rt: &CudaRuntime, typed: &TypedTensor<bool>) -> crate::Result<TypedTensor<bool>> {
224    let handle = match typed.buffer() {
225        StorageBuffer::Host(_) => {
226            return Err(crate::Error::runtime_state(
227                "download",
228                "expected CubeCL buffer",
229            ));
230        }
231        StorageBuffer::Backend(buffer) => cubecl_handle_from_backend(buffer.as_ref(), "download")?,
232    };
233
234    if typed.n_elements() == 0 {
235        return TypedTensor::from_vec_col_major(typed.shape().to_vec(), Vec::new());
236    }
237
238    rt.synchronize()?;
239    let bytes = rt
240        .client()
241        .read_one(handle)
242        .map_err(|err| crate::Error::backend_source("download", err))?;
243    let data = bytes.iter().map(|&byte| byte != 0).collect();
244    TypedTensor::from_vec_col_major(typed.shape().to_vec(), data)
245}
246
247fn ensure_tensor_resident_on_runtime(
248    rt: &CudaRuntime,
249    tensor: &Tensor,
250    op: &'static str,
251) -> crate::Result<()> {
252    match tensor {
253        Tensor::F64(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
254        Tensor::F32(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
255        Tensor::I32(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
256        Tensor::I64(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
257        Tensor::Bool(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
258        Tensor::C64(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
259        Tensor::C32(tensor) => dispatch::ensure_resident_on_runtime(rt, tensor, op),
260    }
261}
262
263fn cubecl_handle_from_backend<T: 'static>(
264    buffer: &dyn crate::BackendStorage<T>,
265    op: &'static str,
266) -> crate::Result<cubecl_runtime::server::Handle> {
267    buffer
268        .as_any()
269        .downcast_ref::<CubeclBuffer>()
270        .map(|buffer| buffer.handle().clone())
271        .ok_or_else(|| {
272            crate::Error::runtime_state(
273                op,
274                format!(
275                    "expected CubeCL buffer, got `{}` backend buffer",
276                    buffer.backend_family()
277                ),
278            )
279        })
280}
tenferro_gpu/cubecl/memory.rs

tenferro_gpu/cubecl/
memory.rs
