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