cubecl_runtime/compiler.rs
1use crate::kernel::KernelDefinition;
2use crate::poison::DevicePoison;
3use alloc::string::{String, ToString};
4use cubecl_environment::backtrace::BackTrace;
5use thiserror::Error;
6
7/// JIT compilation error.
8#[derive(Error, Clone)]
9#[cfg_attr(serializable, derive(serde::Serialize, serde::Deserialize))]
10pub enum CompilationError {
11 /// An instruction isn't supported.
12 #[error(
13 "An unsupported instruction caused the compilation to fail\nCaused by:\n {reason}\nBacktrace:\n{backtrace}"
14 )]
15 UnsupportedInstruction {
16 /// The caused of the error.
17 reason: String,
18 /// The backtrace for this error.
19 #[cfg_attr(serializable, serde(skip))]
20 backtrace: BackTrace,
21 },
22
23 /// A generic compilation error.
24 #[error(
25 "An error caused the compilation to fail\nCaused by:\n {reason}\nBacktrace:\n{backtrace}"
26 )]
27 Generic {
28 /// The error context.
29 reason: String,
30 /// The backtrace for this error.
31 #[cfg_attr(serializable, serde(skip))]
32 backtrace: BackTrace,
33 },
34 /// The device was poisoned before the compiled kernel could be loaded onto
35 /// it. See [`ServerError::DevicePoisoned`](crate::server::ServerError::DevicePoisoned).
36 #[error("The device is poisoned, the kernel could not be loaded\nCaused by:\n {0}")]
37 DevicePoisoned(#[from] DevicePoison),
38
39 /// A generic compilation error.
40 #[error(
41 "A validation error caused the compilation to fail\nCaused by:\n {reason}\nBacktrace:\n{backtrace}"
42 )]
43 Validation {
44 /// The error context.
45 reason: String,
46 /// The backtrace for this error.
47 #[cfg_attr(serializable, serde(skip))]
48 backtrace: BackTrace,
49 },
50 /// The kernel's own code panicked while it expanded, an assertion in it included: a defect
51 /// in the kernel, where every other variant is a compiler turning it down.
52 #[error(
53 "Expanding the kernel `{kernel}` panicked\nCaused by:\n {message}\nBacktrace:\n{backtrace}"
54 )]
55 ExpansionPanicked {
56 /// The kernel that panicked.
57 kernel: String,
58 /// The panic's message.
59 message: String,
60 /// The backtrace for this error.
61 #[cfg_attr(serializable, serde(skip))]
62 backtrace: BackTrace,
63 },
64}
65
66impl CompilationError {
67 /// Whether the device that emitted the error is poisoned.
68 pub fn is_device_poisoned(&self) -> bool {
69 matches!(self, Self::DevicePoisoned(_))
70 }
71
72 /// Whether this is the kernel being turned down, a candidate to drop or a case to skip,
73 /// rather than a defect in it to report or something going wrong while building it.
74 ///
75 /// Every variant is, except two:
76 /// - a kernel that panicked while it expanded: a kernel declines a configuration by
77 /// returning an error, so a panic is a bug that skipping would hide;
78 /// - a device that died before the module could load.
79 pub fn is_refusal(&self) -> bool {
80 !matches!(self, Self::ExpansionPanicked { .. }) && !self.is_device_poisoned()
81 }
82}
83
84impl core::fmt::Debug for CompilationError {
85 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
86 write!(f, "{self}")
87 }
88}
89
90impl From<pliron::result::Error> for CompilationError {
91 fn from(value: pliron::result::Error) -> Self {
92 CompilationError::Validation {
93 reason: value.to_string(),
94 backtrace: BackTrace::capture(),
95 }
96 }
97}
98
99/// Compiles the representation into its own representation that can be formatted into tokens.
100pub trait Compiler: Sync + Send + 'static + Clone + core::fmt::Debug {
101 /// The representation for the compiled code.
102 type Representation: core::fmt::Display;
103 /// The compilation options used to configure the compiler
104 type CompilationOptions: Send + Default + core::fmt::Debug;
105
106 /// Compiles the [kernel definition](KernelDefinition) into the compiler's representation.
107 fn compile(
108 &mut self,
109 kernel: KernelDefinition,
110 compilation_options: &Self::CompilationOptions,
111 ) -> Result<Self::Representation, CompilationError>;
112
113 /// What the compiled kernel does with each buffer binding, by buffer
114 /// position — the visibility analysis's answer, when the representation
115 /// kept it (see [`BufferIOAttr`](crate::kernel::BufferIOAttr)).
116 ///
117 /// `None` reads as every buffer both read and written, the conservative
118 /// direction. A compiler overriding this must answer from the IR
119 /// attributes the annotate pass stamped, never from what its shader
120 /// language kept — wgpu's shader visibility, for one, is deliberately
121 /// forced wider than the kernel's own behavior.
122 fn buffer_io(
123 _repr: &Self::Representation,
124 ) -> Option<alloc::vec::Vec<crate::kernel::BufferIOAttr>> {
125 None
126 }
127
128 /// The default extension for the runtime's kernel/shader code.
129 /// Might change based on which compiler is used.
130 fn extension(&self) -> &'static str;
131
132 /// Short identifier of the language this compiler produces, such as
133 /// `"wgsl"` or `"cuda"`.
134 ///
135 /// What a [`PrecompiledSource`](crate::kernel::PrecompiledSource) has to
136 /// name to be accepted by this compiler.
137 fn lang_tag(&self) -> &'static str;
138}