Skip to main content

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}