tenferro_gpu/webgpu/
memory.rs

1use cubecl::client::ComputeClient;
2use cubecl::prelude::CubeElement;
3use cubecl_wgpu::WgpuRuntime;
4use num_complex::{Complex32, Complex64};
5
6use super::{
7    ensure_resident_on_runtime, prepared_webgpu_tensor, typed_from_webgpu, WebGpuBuffer,
8    WebGpuRuntime,
9};
10use crate::{Tensor, TypedTensor};
11
12/// Upload a host tensor into a CubeCL-managed WebGPU allocation.
13///
14/// # Examples
15///
16/// ```
17/// use tenferro_gpu::{webgpu::upload_webgpu_tensor, webgpu::WebGpuRuntime};
18/// use tenferro_tensor::{Result, Tensor};
19///
20/// let _upload: fn(&WebGpuRuntime, &Tensor) -> Result<Tensor> = upload_webgpu_tensor;
21/// ```
22///
23/// # Errors
24///
25/// Returns [`crate::Error::RuntimeState`] when the source buffer is backend
26/// resident or belongs to another placement, [`crate::Error::Unsupported`] for
27/// a dtype unavailable in WebGPU, or [`crate::Error::BackendSource`] on
28/// allocation.
29pub fn upload_webgpu_tensor(rt: &WebGpuRuntime, tensor: &Tensor) -> crate::Result<Tensor> {
30    match tensor {
31        Tensor::F64(t) => upload_typed::<f64>(rt, t).map(Tensor::F64),
32        Tensor::F32(t) => upload_typed::<f32>(rt, t).map(Tensor::F32),
33        Tensor::I32(t) => upload_typed::<i32>(rt, t).map(Tensor::I32),
34        Tensor::I64(t) => upload_typed::<i64>(rt, t).map(Tensor::I64),
35        Tensor::Bool(t) => upload_bool(rt, t).map(Tensor::Bool),
36        Tensor::C64(t) => upload_typed::<Complex64>(rt, t).map(Tensor::C64),
37        Tensor::C32(t) => upload_typed::<Complex32>(rt, t).map(Tensor::C32),
38    }
39}
40
41/// Download a CubeCL-managed WebGPU tensor back to host memory.
42///
43/// # Examples
44///
45/// ```
46/// use tenferro_gpu::{webgpu::download_webgpu_tensor, webgpu::WebGpuRuntime};
47/// use tenferro_tensor::{Result, Tensor};
48///
49/// let _download: fn(&WebGpuRuntime, &Tensor) -> Result<Tensor> = download_webgpu_tensor;
50/// ```
51///
52/// # Errors
53///
54/// Returns [`crate::Error::RuntimeState`] for missing or foreign device state,
55/// [`crate::Error::BackendSource`] when queue synchronization/readback fails,
56/// or a typed validation error when bytes do not match the tensor shape.
57pub fn download_webgpu_tensor(rt: &WebGpuRuntime, tensor: &Tensor) -> crate::Result<Tensor> {
58    let client = rt.client();
59    match tensor {
60        Tensor::F64(t) => download_typed::<f64>(rt, client, t).map(Tensor::F64),
61        Tensor::F32(t) => download_typed::<f32>(rt, client, t).map(Tensor::F32),
62        Tensor::I32(t) => download_typed::<i32>(rt, client, t).map(Tensor::I32),
63        Tensor::I64(t) => download_typed::<i64>(rt, client, t).map(Tensor::I64),
64        Tensor::Bool(t) => download_bool(rt, client, t).map(Tensor::Bool),
65        Tensor::C64(t) => download_typed::<Complex64>(rt, client, t).map(Tensor::C64),
66        Tensor::C32(t) => download_typed::<Complex32>(rt, client, t).map(Tensor::C32),
67    }
68}
69
70pub(super) fn upload_typed<T: CubeElement + crate::TensorScalar + Clone + Send + Sync + 'static>(
71    rt: &WebGpuRuntime,
72    typed: &TypedTensor<T>,
73) -> crate::Result<TypedTensor<T>> {
74    let host_data = typed.host_data().map_err(|error| {
75        crate::Error::runtime_state("webgpu_upload", format!("expected host buffer: {error}"))
76    })?;
77
78    let byte_len = T::as_bytes(host_data).len();
79    let handle = rt.client().create_from_slice(T::as_bytes(host_data));
80    let buffer = WebGpuBuffer::new_for_runtime(rt, handle, byte_len, "webgpu_upload")?;
81    let tensor = typed_from_webgpu(typed.shape().to_vec(), buffer, rt)?;
82    rt.record_upload(byte_len);
83    Ok(tensor)
84}
85
86pub(super) fn download_typed<T: CubeElement + crate::TensorScalar + Clone + 'static>(
87    rt: &WebGpuRuntime,
88    client: &ComputeClient<WgpuRuntime>,
89    typed: &TypedTensor<T>,
90) -> crate::Result<TypedTensor<T>> {
91    ensure_resident_on_runtime(rt, typed, "webgpu_download")?;
92    let handle = prepared_webgpu_tensor(typed, "webgpu_download")?.handle;
93
94    if typed.n_elements() == 0 {
95        return TypedTensor::from_vec_col_major(typed.shape().to_vec(), Vec::new());
96    }
97
98    let bytes = client
99        .read_one(handle)
100        .map_err(|err| crate::Error::backend_source("webgpu_download", err))?;
101    let data = T::from_bytes(&bytes).to_vec();
102    rt.record_download(bytes.len());
103    TypedTensor::from_vec_col_major(typed.shape().to_vec(), data)
104}
105
106fn upload_bool(rt: &WebGpuRuntime, typed: &TypedTensor<bool>) -> crate::Result<TypedTensor<bool>> {
107    let host_data = typed.host_data().map_err(|error| {
108        crate::Error::runtime_state("webgpu_upload", format!("expected host buffer: {error}"))
109    })?;
110
111    let bytes: Vec<u8> = host_data.iter().map(|&value| u8::from(value)).collect();
112    let handle = rt.client().create_from_slice(&bytes);
113    let buffer = WebGpuBuffer::new_for_runtime(rt, handle, bytes.len(), "webgpu_upload")?;
114    rt.record_upload(bytes.len());
115    TypedTensor::from_backend_allocation(
116        typed.shape().to_vec(),
117        Box::new(buffer),
118        super::webgpu_placement(rt),
119    )
120}
121
122fn download_bool(
123    rt: &WebGpuRuntime,
124    client: &ComputeClient<WgpuRuntime>,
125    typed: &TypedTensor<bool>,
126) -> crate::Result<TypedTensor<bool>> {
127    ensure_resident_on_runtime(rt, typed, "webgpu_download")?;
128    let handle = prepared_webgpu_tensor(typed, "webgpu_download")?.handle;
129
130    if typed.n_elements() == 0 {
131        return TypedTensor::from_vec_col_major(typed.shape().to_vec(), Vec::new());
132    }
133
134    let bytes = client
135        .read_one(handle)
136        .map_err(|err| crate::Error::backend_source("webgpu_download", err))?;
137    let data = bytes.iter().map(|&byte| byte != 0).collect();
138    rt.record_download(bytes.len());
139    TypedTensor::from_vec_col_major(typed.shape().to_vec(), data)
140}
tenferro_gpu/webgpu/memory.rs

tenferro_gpu/webgpu/
memory.rs
