cubecl_utils_rs/
tensor.rs

1//! GPU-resident tensor.
2
3use cubecl::prelude::*;
4use cubecl::server::Handle;
5use cubecl::zspace::striding::row_major_contiguous_strides;
6use cubecl::zspace::{Shape, Strides};
7use std::marker::PhantomData;
8
9use crate::errors::CubeclUtilsErrors;
10use crate::limits::{fits_binding, GpuLimits};
11
12///////////////
13// GpuTensor //
14///////////////
15
16/// GPU-resident tensor for use with CubeCL kernels.
17pub struct GpuTensor<R: Runtime, F: CubeElement + Numeric> {
18    /// Handle to the GPU buffer containing tensor data
19    data: Handle,
20    /// Dimensions of the tensor (e.g. `[n_rows, n_cols]`)
21    shape: Vec<usize>,
22    /// Memory strides for each dimension in row-major order
23    strides: Vec<usize>,
24    /// Phantom marker for the runtime type
25    _r: PhantomData<R>,
26    /// Phantom marker for the element type
27    _f: PhantomData<F>,
28}
29
30impl<R: Runtime, F: CubeElement + Numeric> Clone for GpuTensor<R, F> {
31    fn clone(&self) -> Self {
32        Self {
33            data: self.data.clone(),
34            shape: self.shape.clone(),
35            strides: self.strides.clone(),
36            _r: PhantomData,
37            _f: PhantomData,
38        }
39    }
40}
41
42impl<R: Runtime, F: Numeric + CubeElement> GpuTensor<R, F> {
43    /// Byte size of a tensor with the given shape.
44    ///
45    /// ### Params
46    ///
47    /// * `shape` - Dimensions of the tensor
48    ///
49    /// ### Returns
50    ///
51    /// Element count multiplied by the element size.
52    fn byte_size(shape: &[usize]) -> u64 {
53        (shape.iter().product::<usize>() * core::mem::size_of::<F>()) as u64
54    }
55
56    /// Create a tensor from CPU data.
57    ///
58    /// ### Params
59    ///
60    /// * `data` - Slice of values to upload
61    /// * `shape` - Dimensions of the tensor
62    /// * `client` - GPU compute client for memory allocation
63    ///
64    /// ### Returns
65    ///
66    /// A new tensor with the data copied to GPU memory, or `BindingTooLarge`
67    /// when the allocation exceeds what this device binds in one go.
68    pub fn from_slice(
69        data: &[F],
70        shape: Vec<usize>,
71        client: &ComputeClient<R>,
72    ) -> Result<Self, CubeclUtilsErrors> {
73        fits_binding(Self::byte_size(&shape), &GpuLimits::from_client(client))?;
74
75        let handle = client.create_from_slice(F::as_bytes(data));
76        let strides = row_major_contiguous_strides(&shape).to_vec();
77        Ok(Self {
78            data: handle,
79            shape,
80            strides,
81            _r: PhantomData,
82            _f: PhantomData,
83        })
84    }
85
86    /// Create an uninitialised tensor.
87    ///
88    /// ### Params
89    ///
90    /// * `shape` - Dimensions of the tensor
91    /// * `client` - GPU compute client for memory allocation
92    ///
93    /// ### Returns
94    ///
95    /// A new tensor with allocated but uninitialised GPU memory, or
96    /// `BindingTooLarge` when the allocation exceeds what this device binds in
97    /// one go.
98    ///
99    /// ### Note
100    ///
101    /// The allocation returns quickly but its pages are not backed until
102    /// something writes them. On a large buffer the first kernel write pays
103    /// that fault and can cost more than the kernel itself, so prefer reusing
104    /// one scratch tensor over allocating per call. See [`Self::reshaped_view`].
105    pub fn empty(shape: Vec<usize>, client: &ComputeClient<R>) -> Result<Self, CubeclUtilsErrors> {
106        let size = Self::byte_size(&shape);
107        fits_binding(size, &GpuLimits::from_client(client))?;
108
109        let handle = client.empty(size as usize);
110        let strides = row_major_contiguous_strides(&shape).to_vec();
111        Ok(Self {
112            data: handle,
113            shape,
114            strides,
115            _r: PhantomData,
116            _f: PhantomData,
117        })
118    }
119
120    /// Convert to a `TensorArg` for kernel launches.
121    ///
122    /// ### Returns
123    ///
124    /// A `TensorArg` suitable for passing to CubeCL kernels. Vectorisation
125    /// width is not set per tensor; it is passed once at launch as the argument
126    /// for the kernel's `N: Size` generic.
127    pub fn into_tensor_arg(&self) -> TensorArg<R> {
128        unsafe {
129            TensorArg::from_raw_parts(
130                self.data.clone(),
131                Strides::new(&self.strides),
132                Shape::from(self.shape.clone()),
133            )
134        }
135    }
136
137    /// Read tensor data back to CPU.
138    ///
139    /// Consumes the tensor and transfers data from GPU to CPU memory.
140    ///
141    /// ### Params
142    ///
143    /// * `client` - GPU compute client for memory transfer
144    ///
145    /// ### Returns
146    ///
147    /// Vector of exactly [`Self::len`] elements.
148    ///
149    /// ### Note
150    ///
151    /// The read is truncated to the shape. That only matters for a tensor
152    /// produced by [`Self::reshaped_view`], where the underlying allocation is
153    /// larger than the view: the runtime hands back the whole binding, and
154    /// returning that would silently include another view's data.
155    pub fn read(self, client: &ComputeClient<R>) -> Result<Vec<F>, CubeclUtilsErrors> {
156        let len = self.len();
157        let bytes = client.read_one(self.data)?;
158        let mut values = F::from_bytes(&bytes).to_vec();
159        values.truncate(len);
160        Ok(values)
161    }
162
163    /// Dimensions of the tensor.
164    ///
165    /// ### Returns
166    ///
167    /// Slice of the per-dimension extents, outermost first.
168    pub fn shape(&self) -> &[usize] {
169        &self.shape
170    }
171
172    /// Number of elements the underlying allocation holds.
173    ///
174    /// ### Returns
175    ///
176    /// Product of the shape dimensions.
177    pub fn len(&self) -> usize {
178        self.shape.iter().product()
179    }
180
181    /// Whether the tensor holds no elements.
182    ///
183    /// ### Returns
184    ///
185    /// True if the element count is zero.
186    pub fn is_empty(&self) -> bool {
187        self.len() == 0
188    }
189
190    /// Reinterpret an existing allocation under a smaller shape.
191    ///
192    /// Shares the underlying buffer rather than allocating, so callers can keep
193    /// one scratch tensor alive across several differently shaped uses. The
194    /// first kernel write to a fresh allocation faults its pages in, which for
195    /// a large buffer costs more than the kernel itself, so reuse is worth the
196    /// sharp edge.
197    ///
198    /// ### Params
199    ///
200    /// * `shape` - New shape; its element count must not exceed the current one
201    ///
202    /// ### Returns
203    ///
204    /// A tensor sharing this one's buffer, with row-major strides for `shape`.
205    ///
206    /// ### Note
207    ///
208    /// The returned tensor aliases `self`. Writing through both concurrently is
209    /// a data race the type system does not prevent here.
210    pub fn reshaped_view(&self, shape: Vec<usize>) -> Self {
211        debug_assert!(
212            shape.iter().product::<usize>() <= self.len(),
213            "reshaped_view would exceed the allocation"
214        );
215        let strides = row_major_contiguous_strides(&shape).to_vec();
216        Self {
217            data: self.data.clone(),
218            shape,
219            strides,
220            _r: PhantomData,
221            _f: PhantomData,
222        }
223    }
224
225    /// Size of the tensor on the GPU, in bytes.
226    ///
227    /// ### Returns
228    ///
229    /// Element count multiplied by the element size.
230    pub fn vram_bytes(&self) -> usize {
231        self.shape.iter().product::<usize>() * std::mem::size_of::<F>()
232    }
233
234    /// Return the handle of the tensor.
235    ///
236    /// Escape hatch for crates that need to hand the raw buffer to their own
237    /// kernels or to a library matmul without going through
238    /// [`GpuTensor::into_tensor_arg`].
239    ///
240    /// ### Returns
241    ///
242    /// A reference to the underlying `Handle`.
243    pub fn handle(&self) -> &Handle {
244        &self.data
245    }
246}
247
248///////////
249// Tests //
250///////////
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use cubecl::cpu::{CpuDevice, CpuRuntime};
256
257    #[test]
258    fn test_tensor_from_slice_and_read() {
259        let client = CpuRuntime::client(&CpuDevice);
260
261        let data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
262        let tensor = GpuTensor::<CpuRuntime, f32>::from_slice(&data, vec![2, 3], &client).unwrap();
263
264        assert_eq!(tensor.read(&client).unwrap(), data);
265    }
266
267    #[test]
268    fn test_tensor_empty() {
269        let client = CpuRuntime::client(&CpuDevice);
270
271        let tensor = GpuTensor::<CpuRuntime, f32>::empty(vec![3, 4], &client).unwrap();
272
273        assert_eq!(tensor.shape(), &[3, 4]);
274        assert_eq!(tensor.len(), 12);
275        assert!(!tensor.is_empty());
276        assert_eq!(tensor.vram_bytes(), 48);
277    }
278
279    #[test]
280    fn test_tensor_reshaped_view_shares_the_buffer() {
281        let client = CpuRuntime::client(&CpuDevice);
282
283        let data: Vec<f32> = (0..12).map(|i| i as f32).collect();
284        let tensor = GpuTensor::<CpuRuntime, f32>::from_slice(&data, vec![3, 4], &client).unwrap();
285        let view = tensor.reshaped_view(vec![2, 3]);
286
287        assert_eq!(view.shape(), &[2, 3]);
288        // The runtime hands back the whole 12-element binding; `read` truncates
289        // to the view's own shape.
290        assert_eq!(view.read(&client).unwrap(), &data[..6]);
291    }
292
293    #[test]
294    fn test_tensor_rejects_an_oversized_binding() {
295        let client = CpuRuntime::client(&CpuDevice);
296        let limit = GpuLimits::from_client(&client).max_binding_bytes;
297
298        // One element past what a single binding takes.
299        let too_many = (limit / 4) as usize + 1;
300        assert!(matches!(
301            GpuTensor::<CpuRuntime, f32>::empty(vec![too_many], &client),
302            Err(CubeclUtilsErrors::BindingTooLarge { .. })
303        ));
304    }
305}
cubecl_utils_rs/tensor.rs

cubecl_utils_rs/
tensor.rs
