cubecl_std/tensor/
handle.rs

1use core::marker::PhantomData;
2use cubecl_core::calculate_cube_count_elemwise;
3use cubecl_core::prelude::*;
4use cubecl_core::tensor_line_size_parallel;
5use cubecl_core::{Runtime, server};
6use cubecl_runtime::server::Handle;
7
8/// Tensor representation containing a [server handle](Handle) as well as basic tensor metadata.,
9pub struct TensorHandle<R, E>
10where
11    R: Runtime,
12    E: CubePrimitive,
13{
14    /// The buffer where the data are stored.
15    pub handle: server::Handle,
16    pub shape: Vec<usize>,
17    pub strides: Vec<usize>,
18    elem: PhantomData<E>,
19    runtime: PhantomData<R>,
20}
21
22impl<R, E> core::fmt::Debug for TensorHandle<R, E>
23where
24    R: Runtime,
25    E: CubePrimitive,
26{
27    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
28        f.write_fmt(format_args!(
29            "Tensor {{ shape: {:?}, strides: {:?}, dtype: {}}}",
30            self.shape,
31            self.strides,
32            core::any::type_name::<E>(),
33        ))
34    }
35}
36
37impl<R, E> Clone for TensorHandle<R, E>
38where
39    R: Runtime,
40    E: CubePrimitive,
41{
42    fn clone(&self) -> Self {
43        Self {
44            handle: self.handle.clone(),
45            shape: self.shape.clone(),
46            strides: self.strides.clone(),
47            elem: PhantomData,
48            runtime: PhantomData,
49        }
50    }
51}
52
53impl<R, E> TensorHandle<R, E>
54where
55    R: Runtime,
56    E: CubePrimitive,
57{
58    /// Create a new tensor.
59    pub fn new(handle: server::Handle, shape: Vec<usize>, strides: Vec<usize>) -> Self {
60        Self {
61            handle,
62            shape,
63            strides,
64            elem: PhantomData,
65            runtime: PhantomData,
66        }
67    }
68
69    pub fn empty(client: &ComputeClient<R::Server, R::Channel>, shape: Vec<usize>) -> Self {
70        let elem_size = E::size().expect("To be a native type");
71        let (handle, strides) = client.empty_tensor(&shape, elem_size);
72
73        Self::new(handle, shape, strides)
74    }
75
76    /// Create a new tensor.
77    pub fn from_ref(handle: &TensorHandleRef<'_, R>) -> Self {
78        Self {
79            handle: handle.handle.clone(),
80            shape: handle.shape.to_vec(),
81            strides: handle.strides.to_vec(),
82            elem: PhantomData,
83            runtime: PhantomData,
84        }
85    }
86
87    /// Create a new tensor with a contiguous memory layout.
88    pub fn new_contiguous(shape: Vec<usize>, handle: Handle) -> Self {
89        let strides = Self::contiguous_strides(&shape);
90
91        Self {
92            handle,
93            shape,
94            strides,
95            elem: PhantomData,
96            runtime: PhantomData,
97        }
98    }
99
100    /// Check if the tensor is safe to mutate.
101    pub fn can_mut(&self) -> bool {
102        self.handle.can_mut()
103    }
104
105    pub fn as_ref(&self) -> TensorHandleRef<'_, R> {
106        unsafe {
107            TensorHandleRef::from_raw_parts(
108                &self.handle,
109                &self.strides,
110                &self.shape,
111                size_of::<E>(),
112            )
113        }
114    }
115
116    /// Return the reference to a tensor argument.
117    pub fn as_arg<'a>(&'a self, vectorisation: u8) -> TensorArg<'a, R> {
118        let handle: TensorHandleRef<'a, R> = self.as_ref();
119
120        unsafe {
121            TensorArg::from_raw_parts::<E>(
122                handle.handle,
123                handle.strides,
124                handle.shape,
125                vectorisation,
126            )
127        }
128    }
129
130    fn contiguous_strides(shape: &[usize]) -> Vec<usize> {
131        let mut strides = Vec::with_capacity(shape.len());
132
133        let mut current = 1;
134        shape.iter().enumerate().rev().for_each(|(_, val)| {
135            strides.push(current);
136            current *= val;
137        });
138        strides.reverse();
139        strides
140    }
141}
142impl<R, E> TensorHandle<R, E>
143where
144    R: Runtime,
145    E: Numeric,
146{
147    pub fn zeros(client: &ComputeClient<R::Server, R::Channel>, shape: Vec<usize>) -> Self {
148        let num_elements: usize = shape.iter().product();
149        let rank = shape.len();
150        let output = Self::empty(client, shape);
151
152        let vectorization_factor = tensor_line_size_parallel(
153            R::supported_line_sizes().iter().cloned(),
154            &output.shape,
155            &output.strides,
156            rank - 1,
157        );
158
159        let cube_dim = CubeDim::default();
160        let cube_count =
161            calculate_cube_count_elemwise(num_elements / vectorization_factor as usize, cube_dim);
162        let array_len = output.handle.size();
163
164        unsafe {
165            init::zeros_array::launch_unchecked::<E, R>(
166                client,
167                cube_count,
168                cube_dim,
169                ArrayArg::from_raw_parts::<E>(
170                    &output.handle,
171                    array_len as usize,
172                    vectorization_factor,
173                ),
174            )
175        };
176
177        output
178    }
179}
180
181pub(crate) mod init {
182    use cubecl::prelude::*;
183    use cubecl_core as cubecl;
184
185    #[cube(launch_unchecked)]
186    pub fn zeros_array<C: Numeric>(output: &mut Array<C>) {
187        if ABSOLUTE_POS < output.len() {
188            output[ABSOLUTE_POS] = C::from_int(0);
189        }
190    }
191}
cubecl_std/tensor/handle.rs

cubecl_std/tensor/
handle.rs
