cubecl_runtime/
compiler.rs

1use crate::kernel::{CompiledKernel, KernelDefinition, KernelMetadata};
2use alloc::{string::String, vec::Vec};
3use cubecl_common::ExecutionMode;
4use cubecl_ir::ElemType;
5
6/// Kernel trait with the ComputeShader that will be compiled and cached based on the
7/// provided id.
8pub trait CubeTask<C: Compiler>: KernelMetadata + Send + Sync {
9    /// Compile a kernel and return the compiled form with an optional non-text representation
10    fn compile(
11        &self,
12        compiler: &mut C,
13        compilation_options: &C::CompilationOptions,
14        mode: ExecutionMode,
15    ) -> Result<CompiledKernel<C>, CompilationError>;
16}
17
18/// JIT compilation error.
19#[derive(Debug, PartialEq, Eq, Clone, Hash)]
20#[cfg_attr(std_io, derive(serde::Serialize, serde::Deserialize))]
21pub enum CompilationError {
22    /// An instruction isn't supported.
23    UnsupportedInstruction {
24        /// The error context.
25        context: String,
26    },
27
28    /// When multiple compilation errors are detected.
29    Multiple {
30        /// The error context.
31        context: String,
32        /// The errors.
33        errors: Vec<Self>,
34    },
35
36    /// A generic compilation error.
37    Generic {
38        /// The error context.
39        context: String,
40    },
41}
42
43impl core::fmt::Display for CompilationError {
44    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
45        // TODO Make it better
46        f.write_fmt(format_args!("{:?}", self))
47    }
48}
49
50/// Compiles the representation into its own representation that can be formatted into tokens.
51pub trait Compiler: Sync + Send + 'static + Clone + core::fmt::Debug {
52    /// The representation for the compiled code.
53    type Representation: core::fmt::Display;
54    /// The compilation options used to configure the compiler
55    type CompilationOptions: Send + Default + core::fmt::Debug;
56
57    /// Compiles the [kernel definition](KernelDefinition) into the compiler's representation.
58    fn compile(
59        &mut self,
60        kernel: KernelDefinition,
61        compilation_options: &Self::CompilationOptions,
62        mode: ExecutionMode,
63    ) -> Result<Self::Representation, CompilationError>;
64
65    /// The size of the given element in bytes.
66    fn elem_size(&self, elem: ElemType) -> usize;
67
68    /// The default extension for the runtime's kernel/shader code.
69    /// Might change based on which compiler is used.
70    fn extension(&self) -> &'static str;
71}
cubecl_runtime/compiler.rs

cubecl_runtime/
compiler.rs
