cubecl_runtime/kernel.rs
1use alloc::string::String;
2use core::hash::Hash;
3use pliron::derive::format;
4
5use cubecl_ir::{ElemType, Scope, metadata::Info, pliron::value::Value, settings::KernelSettings};
6use serde::{Deserialize, Serialize};
7
8use crate::compiler::CompilationError;
9use crate::id::KernelId;
10#[cfg(feature = "std")]
11use cubecl_environment::backtrace::BackTrace;
12
13/// Implement this trait to create a [kernel definition](KernelDefinition).
14pub trait KernelMetadata: core::any::Any + Send + Sync + 'static {
15 /// Name of the kernel for debugging.
16 fn name(&self) -> &'static str {
17 core::any::type_name::<Self>()
18 }
19
20 /// Identifier for the kernel, used for caching kernel compilation.
21 fn id(&self) -> KernelId;
22
23 /// Type of addresses in this kernel
24 fn address_type(&self) -> ElemType;
25}
26
27#[allow(missing_docs)]
28pub struct KernelDefinition {
29 pub body: Scope,
30 pub info: Info,
31 pub settings: KernelSettings,
32}
33
34#[derive(Debug, PartialEq, Eq, Hash, Clone)]
35/// Global argument of a kernel.
36pub struct KernelArg {
37 /// The index of the arg.
38 pub id: usize,
39 /// The value the argument is bound to.
40 pub value: Value,
41 /// Whether the argument has metadata.
42 pub has_extended_meta: bool,
43}
44
45#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
46#[allow(missing_docs)]
47pub struct ScalarKernelArg {
48 pub ty: ElemType,
49 pub count: usize,
50}
51
52#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Hash)]
53#[allow(missing_docs)]
54#[format]
55pub enum Visibility {
56 Uniform,
57 Read,
58 ReadWrite,
59}
60
61/// What a compiled kernel does with each buffer binding, by buffer position.
62///
63/// The IR owns this concept — the visibility analysis stamps it on the entry
64/// function's arguments — so the launch path reads that enum rather than a
65/// copy of it that could drift. Re-exported here because a backend reaching
66/// for it is holding a `CompiledKernel` (in `cubecl-server`), not an IR
67/// context.
68pub use cubecl_ir::attributes::BufferIOAttr;
69
70/// A hand-written kernel's own compiled text, standing in for what the
71/// compiler would have produced from a [`KernelDefinition`].
72///
73/// The text goes to the backend as is, so two things the compiler would have
74/// settled are the kernel's to settle:
75///
76/// - `lang` names the language the text is written in, and must equal the
77/// [`lang_tag`](crate::compiler::Compiler::lang_tag) of the compiler the
78/// client runs. `CompiledKernel::compile` in `cubecl-server` refuses a
79/// mismatch, so CUDA C++ handed to a wgpu client is a
80/// [`CompilationError`](crate::compiler::CompilationError), not a naga
81/// parse error at first launch.
82/// - The kernel's [`id`](crate::kernel::KernelMetadata::id) must cover the
83/// text, for instance through [`KernelId::info`](crate::id::KernelId::info)
84/// with a hash of it. Every compilation cache, in memory and on disk, is
85/// keyed by that id and never sees the source, so two kernels with the
86/// same id and different text would share one compiled artifact.
87///
88/// There is no representation to read a dynamic shared memory size from, so
89/// a precompiled kernel is launched with none: what it needs, it declares
90/// statically in the text.
91pub struct PrecompiledSource {
92 /// The compiled source, in the target language.
93 pub source: String,
94 /// The name of the entrypoint within `source`.
95 pub entrypoint_name: String,
96 /// The language `source` is written in, as the target compiler tags it.
97 pub lang: &'static str,
98}
99
100/// Kernel that can be defined
101pub trait CubeKernel: KernelMetadata {
102 /// Define the kernel for compilation
103 ///
104 /// A backend calls [`define_kernel`] rather than this, so a panic while expanding fails the
105 /// launch instead of leaving its outputs unwritten.
106 fn define(&self) -> KernelDefinition;
107
108 /// The kernel's own compiled source, for a hand-written kernel that
109 /// carries target-language text rather than IR to compile.
110 ///
111 /// `None`, the default, compiles what [`define`](Self::define) returns.
112 fn source(&self) -> Option<PrecompiledSource> {
113 None
114 }
115}
116
117/// `kernel`'s [definition](CubeKernel::define), or the panic its expansion raised, as a
118/// [`CompilationError`].
119///
120/// Expansion runs the kernel's own code, every assertion in it included, and a backend runs it
121/// on the device thread inside a launch nobody waits on. A panic there is caught by the device
122/// channel and only logged: the launch returns, its outputs are never written, and nothing marks
123/// them, so a later read hands back whatever the memory held before. As a [`CompilationError`]
124/// the panic takes the path a compiler's refusal takes, and the outputs carry the failure to
125/// whoever reads them; as [`ExpansionPanicked`](CompilationError::ExpansionPanicked) it is not
126/// a refusal, so a harness that skips refused kernels reports it instead.
127pub fn define_kernel<K: CubeKernel + ?Sized>(
128 kernel: &K,
129) -> Result<KernelDefinition, CompilationError> {
130 #[cfg(feature = "std")]
131 {
132 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.define())).map_err(
133 |payload| {
134 let message = payload
135 .downcast_ref::<&str>()
136 .map(|message| String::from(*message))
137 .or_else(|| payload.downcast_ref::<String>().cloned())
138 .unwrap_or_else(|| String::from("a panic with no message"));
139 CompilationError::ExpansionPanicked {
140 kernel: kernel.name().into(),
141 message,
142 backtrace: BackTrace::capture(),
143 }
144 },
145 )
146 }
147 #[cfg(not(feature = "std"))]
148 {
149 Ok(kernel.define())
150 }
151}