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::id::KernelId;
9
10/// Implement this trait to create a [kernel definition](KernelDefinition).
11pub trait KernelMetadata: core::any::Any + Send + Sync + 'static {
12    /// Name of the kernel for debugging.
13    fn name(&self) -> &'static str {
14        core::any::type_name::<Self>()
15    }
16
17    /// Identifier for the kernel, used for caching kernel compilation.
18    fn id(&self) -> KernelId;
19
20    /// Type of addresses in this kernel
21    fn address_type(&self) -> ElemType;
22}
23
24#[allow(missing_docs)]
25pub struct KernelDefinition {
26    pub body: Scope,
27    pub info: Info,
28    pub settings: KernelSettings,
29}
30
31#[derive(Debug, PartialEq, Eq, Hash, Clone)]
32/// Global argument of a kernel.
33pub struct KernelArg {
34    /// The index of the arg.
35    pub id: usize,
36    /// The value the argument is bound to.
37    pub value: Value,
38    /// Whether the argument has metadata.
39    pub has_extended_meta: bool,
40}
41
42#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
43#[allow(missing_docs)]
44pub struct ScalarKernelArg {
45    pub ty: ElemType,
46    pub count: usize,
47}
48
49#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, Hash)]
50#[allow(missing_docs)]
51#[format]
52pub enum Visibility {
53    Uniform,
54    Read,
55    ReadWrite,
56}
57
58/// What a compiled kernel does with each buffer binding, by buffer position.
59///
60/// The IR owns this concept — the visibility analysis stamps it on the entry
61/// function's arguments — so the launch path reads that enum rather than a
62/// copy of it that could drift. Re-exported here because a backend reaching
63/// for it is holding a `CompiledKernel` (in `cubecl-server`), not an IR
64/// context.
65pub use cubecl_ir::attributes::BufferIOAttr;
66
67/// A hand-written kernel's own compiled text, standing in for what the
68/// compiler would have produced from a [`KernelDefinition`].
69///
70/// The text goes to the backend as is, so two things the compiler would have
71/// settled are the kernel's to settle:
72///
73/// - `lang` names the language the text is written in, and must equal the
74///   [`lang_tag`](crate::compiler::Compiler::lang_tag) of the compiler the
75///   client runs. `CompiledKernel::compile` in `cubecl-server` refuses a
76///   mismatch, so CUDA C++ handed to a wgpu client is a
77///   [`CompilationError`](crate::compiler::CompilationError), not a naga
78///   parse error at first launch.
79/// - The kernel's [`id`](crate::kernel::KernelMetadata::id) must cover the
80///   text, for instance through [`KernelId::info`](crate::id::KernelId::info)
81///   with a hash of it. Every compilation cache, in memory and on disk, is
82///   keyed by that id and never sees the source, so two kernels with the
83///   same id and different text would share one compiled artifact.
84///
85/// There is no representation to read a dynamic shared memory size from, so
86/// a precompiled kernel is launched with none: what it needs, it declares
87/// statically in the text.
88pub struct PrecompiledSource {
89    /// The compiled source, in the target language.
90    pub source: String,
91    /// The name of the entrypoint within `source`.
92    pub entrypoint_name: String,
93    /// The language `source` is written in, as the target compiler tags it.
94    pub lang: &'static str,
95}
96
97/// Kernel that can be defined
98pub trait CubeKernel: KernelMetadata {
99    /// Define the kernel for compilation
100    fn define(&self) -> KernelDefinition;
101
102    /// The kernel's own compiled source, for a hand-written kernel that
103    /// carries target-language text rather than IR to compile.
104    ///
105    /// `None`, the default, compiles what [`define`](Self::define) returns.
106    fn source(&self) -> Option<PrecompiledSource> {
107        None
108    }
109}
cubecl_runtime/kernel.rs

cubecl_runtime/
kernel.rs
