cubecl_core/frontend/container/tensor/
launch.rs

1use core::marker::PhantomData;
2
3use cubecl_ir::{AddressType, Id};
4use cubecl_runtime::{runtime::Runtime, server::CopyDescriptor};
5use cubecl_zspace::{Shape, Strides};
6
7use crate::{
8    self as cubecl,
9    compute::{KernelBuilder, KernelLauncher},
10    frontend::container::slice,
11    prelude::*,
12};
13
14use super::Tensor;
15
16#[derive(CubeType, CubeLaunch, Clone, Copy)]
17#[expand(derive(Clone, Copy))]
18pub struct TensorMeta {
19    pub len: usize,
20    pub rank: usize,
21}
22
23/// Argument to be used for [tensors](Tensor) passed as arguments to kernels.
24#[derive(Debug)]
25pub enum TensorArg<R: Runtime> {
26    /// The tensor is passed with a tensor handle.
27    Handle {
28        /// The tensor handle.
29        handle: TensorBinding<R>,
30    },
31    /// The tensor is aliasing another input tensor.
32    Alias {
33        /// The position of the input tensor.
34        input_pos: usize,
35        strides: Strides,
36        shape: Shape,
37    },
38}
39
40/// Tensor representation with a reference to the [server handle](cubecl_runtime::server::Handle),
41/// the strides and the shape.
42pub struct TensorBinding<R: Runtime> {
43    pub handle: cubecl_runtime::server::Binding,
44    pub strides: Strides,
45    pub shape: Shape,
46    pub runtime: PhantomData<R>,
47}
48
49impl<R: Runtime> Clone for TensorBinding<R> {
50    fn clone(&self) -> Self {
51        Self {
52            handle: self.handle.clone(),
53            strides: self.strides.clone(),
54            shape: self.shape.clone(),
55            runtime: PhantomData,
56        }
57    }
58}
59
60impl<R: Runtime> TensorBinding<R> {
61    pub fn size(&self) -> usize {
62        self.shape.iter().product()
63    }
64
65    /// Address type required to fully index this tensor handle, assuming scalar access.
66    pub fn required_address_type(&self, elem_size: usize) -> AddressType {
67        AddressType::from_len(self.handle.size() as usize / elem_size)
68    }
69}
70
71impl<R: Runtime> core::fmt::Debug for TensorBinding<R> {
72    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
73        writeln!(
74            f,
75            "TensorHandleRef {{ strides: {:?}, shape: {:?} }}",
76            self.strides, self.shape
77        )
78    }
79}
80
81/// Compilation argument for a [tensor](Tensor).
82#[derive(Clone, PartialEq, Eq, Hash, Debug)]
83pub struct TensorCompilationArg {
84    pub meta: TensorMetaCompilationArg,
85    pub buffer: BufferCompilationArg,
86}
87
88impl<C: CubePrimitive> LaunchArg for Tensor<C> {
89    type RuntimeArg<R: Runtime> = TensorArg<R>;
90    type CompilationArg = TensorCompilationArg;
91
92    fn register<R: Runtime>(
93        arg: Self::RuntimeArg<R>,
94        launcher: &mut KernelLauncher<R>,
95    ) -> Self::CompilationArg {
96        let ty = launcher.with_scope(|scope| C::__expand_as_type(scope));
97        let len = arg.size() / ty.vector_size();
98        let meta_arg = TensorMetaLaunch::new(len, arg.shape().len());
99        let buffer = match &arg {
100            TensorArg::Handle { .. } => BufferCompilationArg { inplace: None },
101            TensorArg::Alias { input_pos, .. } => BufferCompilationArg {
102                inplace: Some(*input_pos as Id),
103            },
104        };
105        launcher.register_tensor(arg, ty);
106        let meta = TensorMeta::register(meta_arg, launcher);
107        TensorCompilationArg { meta, buffer }
108    }
109
110    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> TensorExpand<C> {
111        let buffer = match arg.buffer.inplace {
112            Some(id) => builder.inplace(id),
113            None => builder.tensor(C::__expand_as_type(&builder.scope)),
114        };
115        let meta = TensorMeta::expand(&arg.meta, builder);
116        let scope = &builder.scope;
117        let len = expand_buffer_length_native(scope, buffer);
118        let buffer =
119            slice::from_raw_parts::<C>(scope, buffer, 0usize.into_expand(scope), len.into());
120        TensorExpand { meta, buffer }
121    }
122}
123
124impl<C: CubePrimitive> LaunchArg for OwnedTensor<C> {
125    type RuntimeArg<R: Runtime> = TensorArg<R>;
126    type CompilationArg = TensorCompilationArg;
127
128    fn register<R: Runtime>(
129        arg: Self::RuntimeArg<R>,
130        launcher: &mut KernelLauncher<R>,
131    ) -> Self::CompilationArg {
132        Tensor::<C>::register(arg, launcher)
133    }
134
135    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> OwnedTensorExpand<C> {
136        let tensor = Tensor::<C>::expand(arg, builder);
137        OwnedTensorExpand {
138            meta: tensor.meta,
139            buffer: tensor.buffer.expand.into(),
140        }
141    }
142}
143
144impl<R: Runtime> TensorArg<R> {
145    /// Create a new tensor argument specified with its vectorization factor.
146    ///
147    /// # Safety
148    ///
149    /// If you provide wrong strides or shapes, it might create undefined behavior caused by
150    /// out-of-bound reads and writes.
151    pub unsafe fn from_raw_parts(
152        handle: cubecl_runtime::server::Handle,
153        strides: Strides,
154        shape: Shape,
155    ) -> Self {
156        unsafe { Self::from_raw_parts_binding(handle.binding(), strides, shape) }
157    }
158
159    pub(crate) unsafe fn from_raw_parts_binding(
160        handle: cubecl_runtime::server::Binding,
161        strides: Strides,
162        shape: Shape,
163    ) -> Self {
164        unsafe {
165            Self::Handle {
166                handle: TensorBinding::from_raw_parts_binding(handle, strides, shape),
167            }
168        }
169    }
170
171    /// Create an alias argument.
172    pub fn into_alias(self, position: usize) -> Self {
173        match self {
174            TensorArg::Handle { handle } => handle.into_alias(position),
175            alias @ TensorArg::Alias { .. } => alias,
176        }
177    }
178
179    pub fn size(&self) -> usize {
180        match self {
181            TensorArg::Handle { handle } => handle.size(),
182            TensorArg::Alias { shape, .. } => shape.iter().product(),
183        }
184    }
185
186    pub fn shape(&self) -> &[usize] {
187        match self {
188            TensorArg::Handle { handle } => &handle.shape,
189            TensorArg::Alias { shape, .. } => shape,
190        }
191    }
192
193    pub fn strides(&self) -> &[usize] {
194        match self {
195            TensorArg::Handle { handle } => &handle.strides,
196            TensorArg::Alias { strides, .. } => strides,
197        }
198    }
199}
200
201impl<R: Runtime> TensorArg<R> {
202    pub fn into_buffer_arg(self) -> BufferArg<R> {
203        match self {
204            TensorArg::Handle { handle } => {
205                let handle = unsafe {
206                    let size = handle.size();
207                    BufferBinding::from_raw_parts_binding(handle.handle, size)
208                };
209                BufferArg::Handle { handle }
210            }
211            TensorArg::Alias {
212                input_pos, shape, ..
213            } => BufferArg::Alias {
214                input_pos,
215                length: [shape.iter().product()],
216            },
217        }
218    }
219}
220
221impl<R: Runtime> TensorBinding<R> {
222    /// Convert the handle into a [tensor argument](TensorArg).
223    pub fn into_tensor_arg(self) -> TensorArg<R> {
224        TensorArg::Handle { handle: self }
225    }
226    /// Convert the handle into a [tensor argument](TensorArg).
227    pub fn into_alias(self, index: usize) -> TensorArg<R> {
228        TensorArg::Alias {
229            input_pos: index,
230            strides: self.strides,
231            shape: self.shape,
232        }
233    }
234    /// Convert the handle into a [tensor argument](TensorArg).
235    pub fn as_alias(&self, index: usize) -> TensorArg<R> {
236        TensorArg::Alias {
237            input_pos: index,
238            strides: self.strides.clone(),
239            shape: self.shape.clone(),
240        }
241    }
242    /// Convert the handle into a [buffer argument](BufferArg).
243    pub fn into_buffer_arg(self) -> BufferArg<R> {
244        unsafe { BufferArg::from_raw_parts_binding(self.handle, self.shape.iter().product()) }
245    }
246
247    /// Create a handle from raw parts.
248    ///
249    /// # Safety
250    ///
251    /// If you provide wrong strides or shapes, it might create undefined behavior caused by
252    /// out-of-bounds reads and writes.
253    pub unsafe fn from_raw_parts(
254        handle: cubecl_runtime::server::Handle,
255        strides: Strides,
256        shape: Shape,
257    ) -> Self {
258        unsafe { Self::from_raw_parts_binding(handle.binding(), strides, shape) }
259    }
260
261    /// Create a handle from raw parts.
262    ///
263    /// # Safety
264    ///
265    /// If you provide wrong strides or shapes, it might create undefined behavior caused by
266    /// out-of-bounds reads and writes.
267    pub unsafe fn from_raw_parts_binding(
268        handle: cubecl_runtime::server::Binding,
269        strides: Strides,
270        shape: Shape,
271    ) -> Self {
272        Self {
273            handle,
274            strides,
275            shape,
276            runtime: PhantomData,
277        }
278    }
279
280    pub fn into_copy_descriptor(self, elem_size: usize) -> CopyDescriptor {
281        CopyDescriptor {
282            handle: self.handle,
283            shape: self.shape,
284            strides: self.strides,
285            elem_size,
286        }
287    }
288}
cubecl_core/frontend/container/tensor/launch.rs

cubecl_core/frontend/container/tensor/
launch.rs
