cubecl_core/frontend/container/slice/
launch.rs

1use crate::{
2    frontend::container::{buffer_len::expand_buffer_length_native, slice},
3    prelude::*,
4};
5use alloc::boxed::Box;
6use cubecl_zspace::Tiling;
7use serde::{Deserialize, Serialize};
8
9#[derive(Clone, PartialEq, Eq, Hash, Debug, Serialize, Deserialize)]
10pub struct BufferCompilationArg {
11    pub inplace: Option<usize>,
12}
13
14/// Buffer representation with a reference to the [server handle](cubecl_runtime::server::Handle).
15pub struct BufferBinding {
16    pub handle: cubecl_runtime::server::BufferBinding,
17    pub(crate) length: [usize; 1],
18}
19
20pub enum BufferArg {
21    /// The buffer is passed with a buffer handle.
22    Handle {
23        /// The buffer handle.
24        handle: BufferBinding,
25    },
26    /// The buffer is aliasing another input buffer.
27    Alias {
28        /// The position of the input buffer.
29        input_pos: usize,
30        /// The length of the underlying handle
31        length: [usize; 1],
32    },
33}
34
35impl BufferArg {
36    /// Create a new buffer argument.
37    ///
38    /// # Safety
39    ///
40    /// Specifying the wrong length may lead to out-of-bounds reads and writes.
41    pub unsafe fn from_raw_parts(handle: cubecl_runtime::server::Handle, length: usize) -> Self {
42        unsafe {
43            BufferArg::Handle {
44                handle: BufferBinding::from_raw_parts(handle, length),
45            }
46        }
47    }
48    /// Create a new buffer argument from a binding.
49    ///
50    /// # Safety
51    ///
52    /// Specifying the wrong length may lead to out-of-bounds reads and writes.
53    pub unsafe fn from_raw_parts_binding(
54        binding: cubecl_runtime::server::BufferBinding,
55        length: usize,
56    ) -> Self {
57        unsafe {
58            BufferArg::Handle {
59                handle: BufferBinding::from_raw_parts_binding(binding, length),
60            }
61        }
62    }
63
64    /// The buffer's length in elements, as it was declared.
65    pub fn len(&self) -> usize {
66        match self {
67            BufferArg::Handle { handle } => handle.length[0],
68            BufferArg::Alias { length, .. } => length[0],
69        }
70    }
71
72    /// Whether the buffer was declared empty.
73    pub fn is_empty(&self) -> bool {
74        self.len() == 0
75    }
76
77    pub fn alias(input_pos: usize, length: usize) -> Self {
78        Self::Alias {
79            input_pos,
80            length: [length],
81        }
82    }
83
84    pub fn size(&self) -> usize {
85        match self {
86            BufferArg::Handle { handle } => handle.length[0],
87            BufferArg::Alias { length, .. } => length[0],
88        }
89    }
90
91    pub fn shape(&self) -> &[usize] {
92        match self {
93            BufferArg::Handle { handle } => &handle.length,
94            BufferArg::Alias { length, .. } => length,
95        }
96    }
97}
98
99impl BufferBinding {
100    /// Create a new buffer handle reference.
101    ///
102    /// # Safety
103    ///
104    /// Specifying the wrong length may lead to out-of-bounds reads and writes.
105    pub unsafe fn from_raw_parts(handle: cubecl_runtime::server::Handle, length: usize) -> Self {
106        unsafe { Self::from_raw_parts_binding(handle.binding(), length) }
107    }
108
109    /// Create a new buffer handle reference.
110    ///
111    /// # Safety
112    ///
113    /// Specifying the wrong length or size, may lead to out-of-bounds reads and writes.
114    pub unsafe fn from_raw_parts_binding(
115        handle: cubecl_runtime::server::BufferBinding,
116        length: usize,
117    ) -> Self {
118        Self {
119            handle,
120            length: [length],
121        }
122    }
123
124    /// Return the handle as a tensor instead of a buffer.
125    pub fn into_tensor(self) -> TensorBinding {
126        let shape = self.length.into();
127
128        TensorBinding {
129            handle: self.handle,
130            strides: [1].into(),
131            shape,
132            tiling: Tiling::UNTILED,
133        }
134    }
135}
136
137impl<C: CubePrimitive> LaunchArg for Box<[C]> {
138    type RuntimeArg = BufferArg;
139    type CompilationArg = BufferCompilationArg;
140
141    fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
142        <[C]>::register(arg, launcher)
143    }
144
145    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> NativeExpand<Box<[C]>> {
146        <[C]>::expand(arg, builder).expand.into()
147    }
148}
149
150impl<C: CubePrimitive> LaunchArg for [C] {
151    type RuntimeArg = BufferArg;
152    type CompilationArg = BufferCompilationArg;
153
154    fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
155        let elem_size = launcher.with_scope(|scope| C::__expand_size(scope));
156        let inplace = match &arg {
157            BufferArg::Handle { .. } => None,
158            BufferArg::Alias { input_pos, .. } => Some(*input_pos),
159        };
160        launcher.register_buffer(arg, elem_size);
161
162        BufferCompilationArg { inplace }
163    }
164
165    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> NativeExpand<[C]> {
166        let buffer = match arg.inplace {
167            Some(id) => builder.inplace(id),
168            None => builder.buffer(C::__expand_as_type(&builder.scope)),
169        };
170        let scope = &builder.scope;
171        let len = expand_buffer_length_native(scope, buffer);
172        let slice_var =
173            slice::from_raw_parts::<C>(scope, buffer, 0usize.into_expand(scope), len.into());
174        slice_var.expand.into()
175    }
176}
cubecl_core/frontend/container/slice/launch.rs

cubecl_core/frontend/container/slice/
launch.rs
