cubecl_core/compute/
launcher.rs

1use alloc::{boxed::Box, vec::Vec};
2
3use crate::prelude::{BufferArg, TensorArg, TensorMapArg, TensorMapKind};
4use crate::{InfoBuilder, ScalarArgType};
5#[cfg(feature = "std")]
6use core::cell::RefCell;
7use cubecl_ir::{AddressType, ElemType, Scope, settings::KernelSettings};
8use cubecl_runtime::kernel::BufferIOAttr;
9use cubecl_runtime::server::{BufferBinding, CubeCount, KernelResource, TensorMapBinding};
10use cubecl_runtime::{client::Client, kernel::CubeKernel, server::KernelArguments};
11
12#[cfg(feature = "std")]
13std::thread_local! {
14    static INFO: RefCell<InfoBuilder> = RefCell::new(InfoBuilder::default());
15    // Only used for resolving types
16    static SCOPE: RefCell<Scope> = RefCell::new(Scope::dummy());
17}
18
19/// Prepare a kernel for [launch](KernelLauncher::launch).
20pub struct KernelLauncher {
21    resources: Vec<KernelResource>,
22    /// What the caller declared each resource is for, indexed like
23    /// `resources` — see [`declare_io`](Self::declare_io).
24    declared_io: Vec<BufferIOAttr>,
25    /// The declaration the next registered resources fall under.
26    declaring: BufferIOAttr,
27    address_type: AddressType,
28    pub settings: KernelSettings,
29    #[cfg(not(feature = "std"))]
30    info: InfoBuilder,
31    #[cfg(not(feature = "std"))]
32    pub scope: Scope,
33}
34
35impl KernelLauncher {
36    #[cfg(feature = "std")]
37    pub fn with_scope<T>(&mut self, fun: impl FnMut(&Scope) -> T) -> T {
38        SCOPE.with_borrow(fun)
39    }
40
41    #[cfg(not(feature = "std"))]
42    pub fn with_scope<T>(&mut self, mut fun: impl FnMut(&Scope) -> T) -> T {
43        fun(&self.scope)
44    }
45
46    #[cfg(feature = "std")]
47    fn with_info<T>(&mut self, fun: impl FnMut(&mut InfoBuilder) -> T) -> T {
48        INFO.with_borrow_mut(fun)
49    }
50
51    #[cfg(not(feature = "std"))]
52    fn with_info<T>(&mut self, mut fun: impl FnMut(&mut InfoBuilder) -> T) -> T {
53        fun(&mut self.info)
54    }
55
56    /// Register a scalar to be launched.
57    pub fn register_scalar<C: ScalarArgType>(&mut self, scalar: C) {
58        self.with_info(|info| info.scalars.push(scalar));
59    }
60
61    /// Register a scalar to be launched from raw data.
62    pub fn register_scalar_raw(&mut self, bytes: &[u8], dtype: ElemType) {
63        self.with_info(|info| info.scalars.push_raw(bytes, dtype));
64    }
65
66    /// Launch the kernel.
67    #[track_caller]
68    pub fn launch<K: CubeKernel>(self, cube_count: CubeCount, kernel: K, client: &Client) {
69        let bindings = self.into_bindings();
70        let kernel = Box::new(kernel);
71
72        client.launch(kernel, cube_count, bindings)
73    }
74
75    /// Drop a launcher that will never launch, releasing what it registered.
76    ///
77    /// With `std` a launcher's scalars and metadata accumulate in a
78    /// thread-local [`InfoBuilder`] that only building the bindings drains, so
79    /// a launcher built to register arguments and then dropped — what the
80    /// `create_dummy_kernel` launch variant does — would leave that state
81    /// behind for the next real launch on the same thread to pick up as extra
82    /// arguments. Discarding drains it instead.
83    pub fn discard(self) {
84        let _ = self.into_bindings();
85    }
86
87    /// We need to create the bindings in the same order they are defined in the compilation step.
88    ///
89    /// The function [`crate::KernelIntegrator::integrate`] stars by registering the input tensors followed
90    /// by the output tensors. Then the tensor metadata, and the scalars at the end. The scalars
91    /// are registered in the same order they are added. This is why we store the scalar data type
92    /// in the `scalar_order` vector, so that we can register them in the same order.
93    ///
94    /// Also returns an ordered list of constant bindings. The ordering between constants and tensors
95    /// is up to the runtime.
96    fn into_bindings(mut self) -> KernelArguments {
97        let mut bindings = KernelArguments::new();
98        let address_type = self.address_type;
99        let info = self.with_info(|info| info.finish(address_type));
100
101        bindings.resources = self.resources;
102        bindings.declared_io = self.declared_io;
103        bindings.info = info;
104
105        bindings
106    }
107}
108
109// Tensors/arrays
110impl KernelLauncher {
111    /// Declare what the kernel does with the buffers registered from here on,
112    /// until the next declaration.
113    ///
114    /// The generated launch functions call this before each argument with
115    /// what the signature proves — `&Tensor` cannot be written, `&mut Tensor`
116    /// may be read — so a launch that fails before running, a kernel that
117    /// does not compile above all, taints only the buffers the kernel could
118    /// have written. The compiled kernel's own answer still wins once it
119    /// exists; this one is the answer that survives compilation failing. A
120    /// launcher that never declares leaves every resource
121    /// [`ReadWrite`](BufferIOAttr::ReadWrite), the loud fallback.
122    pub fn declare_io(&mut self, io: BufferIOAttr) {
123        self.declaring = io;
124    }
125
126    /// An aliasing argument writes the buffer it aliases in place, however
127    /// that buffer's own argument was declared — the aliased buffer usually
128    /// arrives through a `&Tensor`, and it is the one buffer an in-place
129    /// kernel exists to produce. The alias registers no resource of its own,
130    /// so its declaration lands on the buffer at `input_pos` instead: a
131    /// declaration built from each signature position alone would call that
132    /// buffer read-only and leave the in-place output unnamed by a failure,
133    /// which is silent garbage on a read.
134    fn alias_io(&mut self, input_pos: usize) {
135        if self.declaring.is_writable()
136            && let Some(io) = self.declared_io.get_mut(input_pos)
137        {
138            *io = BufferIOAttr::ReadWrite;
139        }
140    }
141
142    /// Record a resource.
143    fn push_resource(&mut self, resource: KernelResource) {
144        let io = match &resource {
145            // A tensor map's global side is written through TMA operations no
146            // signature shows — a map registered from a `&TensorMap` can
147            // still be a store's destination — so the declaration is clamped
148            // to the same answer the visibility analysis gives it.
149            KernelResource::TensorMap(_) => BufferIOAttr::ReadWrite,
150            KernelResource::Buffer(_) => self.declaring,
151        };
152        self.declared_io.push(io);
153        self.resources.push(resource);
154    }
155
156    /// Push a new input tensor to the state.
157    pub fn register_tensor(&mut self, tensor: TensorArg, elem_size: usize) {
158        if let Some(tensor) = self.process_tensor(tensor, elem_size) {
159            self.push_resource(KernelResource::Buffer(tensor));
160        }
161    }
162
163    fn process_tensor(&mut self, tensor: TensorArg, elem_size: usize) -> Option<BufferBinding> {
164        let tensor = match tensor {
165            TensorArg::Handle { handle, .. } => handle,
166            TensorArg::Alias { input_pos, .. } => {
167                self.alias_io(input_pos);
168                return None;
169            }
170        };
171
172        let buffer_len = tensor.handle.size_in_used() / elem_size as u64;
173        let address_type = self.address_type;
174
175        self.with_info(|info| {
176            info.metadata.register_tensor(
177                buffer_len,
178                tensor.shape.clone(),
179                tensor.strides.clone(),
180                address_type,
181            )
182        });
183        Some(tensor.handle)
184    }
185
186    /// Push a new input array to the state.
187    pub fn register_buffer(&mut self, array: BufferArg, elem_size: usize) {
188        if let Some(tensor) = self.process_buffer(array, elem_size) {
189            self.push_resource(KernelResource::Buffer(tensor));
190        }
191    }
192
193    fn process_buffer(&mut self, array: BufferArg, elem_size: usize) -> Option<BufferBinding> {
194        let array = match array {
195            BufferArg::Handle { handle, .. } => handle,
196            BufferArg::Alias { input_pos, .. } => {
197                self.alias_io(input_pos);
198                return None;
199            }
200        };
201
202        let buffer_len = array.handle.size_in_used() / elem_size as u64;
203        let address_type = self.address_type;
204        self.with_info(|info| info.metadata.register_buffer(buffer_len, address_type));
205        Some(array.handle)
206    }
207
208    /// Push a new tensor to the state.
209    pub fn register_tensor_map<K: TensorMapKind>(
210        &mut self,
211        map: TensorMapArg<K>,
212        elem_size: usize,
213    ) {
214        let binding = self
215            .process_tensor(map.tensor, elem_size)
216            .expect("Can't use alias for TensorMap");
217
218        let map = map.metadata.clone();
219        self.push_resource(KernelResource::TensorMap(TensorMapBinding { binding, map }));
220    }
221}
222
223impl KernelLauncher {
224    pub fn new(settings: KernelSettings) -> Self {
225        Self {
226            address_type: settings.address_type,
227            settings,
228            resources: Vec::new(),
229            declared_io: Vec::new(),
230            declaring: BufferIOAttr::ReadWrite,
231            #[cfg(not(feature = "std"))]
232            info: InfoBuilder::default(),
233            #[cfg(not(feature = "std"))]
234            scope: Scope::dummy(),
235        }
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use cubecl_ir::settings::{Dim3, ExecutionMode};
243
244    fn settings() -> KernelSettings {
245        KernelSettings::new(Dim3::new_single(), ExecutionMode::Checked, AddressType::U32)
246    }
247
248    fn info_of(launcher: KernelLauncher) -> Vec<u64> {
249        launcher.into_bindings().info.data
250    }
251
252    /// `create_dummy_kernel` registers arguments into a launcher it never
253    /// launches. With `std` those registrations land in a thread-local that
254    /// only building the bindings drains, so the launcher has to be discarded
255    /// rather than dropped — otherwise the next real launch on the same
256    /// thread inherits them as extra arguments.
257    #[test]
258    fn a_discarded_launcher_leaves_nothing_for_the_next_launch() {
259        let empty = info_of(KernelLauncher::new(settings()));
260
261        // A registered scalar is visible in the info a launcher produces, so
262        // the equality below is a real claim about the thread-local, not a
263        // comparison of two things that could never differ.
264        let mut registered = KernelLauncher::new(settings());
265        registered.register_scalar(1u32);
266        assert_ne!(info_of(registered), empty);
267
268        let mut dummy = KernelLauncher::new(settings());
269        dummy.register_scalar(1u32);
270        dummy.discard();
271
272        assert_eq!(
273            info_of(KernelLauncher::new(settings())),
274            empty,
275            "a discarded launcher left its scalars behind for the next launch"
276        );
277    }
278}
cubecl_core/compute/launcher.rs

cubecl_core/compute/
launcher.rs
