cubecl_core/frontend/container/slice/
launch.rs

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

cubecl_core/frontend/container/slice/
launch.rs
