cubecl_core/frontend/container/slice/
launch.rs

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

cubecl_core/frontend/container/slice/
launch.rs
