cubecl_server/compiler/target.rs
1//! What a backend contributes to compiling a kernel, split by where each half
2//! may run: [`ArtifactCompiler`] on any thread, so several kernels compile at
3//! once; [`CompilationTarget`] only where the server does, since it holds the
4//! compilation store and loads on the server's context.
5
6use core::fmt::Debug;
7use core::hash::Hash;
8
9use super::{CompilationError, CompilationRecording};
10use crate::id::KernelId;
11use crate::kernel::CubeKernel;
12use crate::logging::ServerLogger;
13use crate::server::LaunchError;
14
15/// Which artifact a launch needs: its kernel, and what the launch adds to it.
16#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17pub struct ArtifactId<V> {
18 /// The kernel.
19 pub kernel: KernelId,
20 /// What the launch decides about the artifact that the kernel id does not:
21 /// the buffer alignment the code is specialized for, the binding layout a
22 /// pipeline is built with. `()` when the kernel id says everything.
23 pub variant: V,
24}
25
26/// The half of a backend's compilation that may run on any thread: from a
27/// kernel to the artifact the device loads, in two steps — a lowering pass
28/// through cubecl's compiler, then whatever finishes the artifact — with one
29/// chance between them to reuse an artifact already finalized. A backend
30/// whose device lets several threads create objects at once may use it here.
31///
32/// `Sync` because it is shared by every thread compiling at once.
33pub trait ArtifactCompiler: Sync {
34 /// See [`ArtifactId::variant`].
35 type Variant: Clone + Eq + Hash + Debug + Send + Sync;
36 /// A kernel lowered by cubecl's compiler for this backend, not yet
37 /// finalized into what the device loads.
38 type Lowered: Send;
39 /// What the device loads. The compilation store keeps it, or the part of
40 /// it that outlives the process: a pipeline the driver built is rebuilt
41 /// from the kept source.
42 type Artifact: Send;
43
44 /// Validates `kernel` against the device, expands it to IR and lowers it
45 /// with cubecl's compiler for this backend, telling `recording` its IR and
46 /// source and `logger` the compilation.
47 fn lower(
48 &self,
49 kernel: &dyn CubeKernel,
50 id: &ArtifactId<Self::Variant>,
51 recording: &mut CompilationRecording,
52 logger: &ServerLogger,
53 ) -> Result<Self::Lowered, LaunchError>;
54
55 /// The text an artifact finalized from `lowered` can be reused by: the
56 /// source finalizing compiles, when that is slow enough that an artifact
57 /// another kernel id already finalized from the same text is worth
58 /// reusing — a C++ source two kernel ids expand to, or one expanded again
59 /// by a rebuild that changed the build id.
60 ///
61 /// `None`, the default, finalizes every time. It says nothing about
62 /// whether `lowered` has text: a backend whose lowering already produced
63 /// the artifact, LLVM's, has text and nothing to save by reusing it.
64 fn source(lowered: &Self::Lowered) -> Option<&str> {
65 let _ = lowered;
66 None
67 }
68
69 /// Turns `lowered`, lowered for `id`, into what the device loads:
70 /// compiles its source with the platform's compiler, or takes out the
71 /// artifact lowering produced.
72 ///
73 /// A backend whose driver compiles again once handed the artifact, and
74 /// lets several threads do so, does that here too and carries the result
75 /// in the artifact, so the driver's compilations run alongside each other
76 /// rather than one by one in [`CompilationTarget::load`].
77 fn finalize(
78 &self,
79 id: &ArtifactId<Self::Variant>,
80 lowered: Self::Lowered,
81 ) -> Result<Self::Artifact, LaunchError>;
82}
83
84/// The variant an [`ArtifactCompiler`] keys its artifacts by.
85pub type VariantOf<T> = <<T as CompilationTarget>::Compiler as ArtifactCompiler>::Variant;
86/// The artifact an [`ArtifactCompiler`] produces.
87pub type ArtifactOf<T> = <<T as CompilationTarget>::Compiler as ArtifactCompiler>::Artifact;
88
89/// The half of a backend's compilation that runs only where the server does:
90/// the compilation store, and the loads on the server's context.
91pub trait CompilationTarget {
92 /// The half that may run on any thread.
93 type Compiler: ArtifactCompiler;
94 /// What a launch dispatches. Cloned on every launch, so it should be
95 /// handles and shared pointers.
96 type Loaded: Clone;
97
98 /// The half that may run on any thread.
99 fn compiler(&self) -> &Self::Compiler;
100
101 /// Whether artifacts are kept between runs, which is what binds the
102 /// loaded kernels to the active environment: a switch replaces the store
103 /// they mirror, so the [`KernelLoader`](super::KernelLoader) drops them.
104 fn persists(&self) -> bool;
105
106 /// Takes the artifact the compilation store holds for `id` out of it, to
107 /// be loaded; `None` when there is no store or it holds nothing for `id`.
108 fn stored(&mut self, id: &ArtifactId<VariantOf<Self>>) -> Option<ArtifactOf<Self>>;
109
110 /// Takes the artifact the compilation store holds for another kernel
111 /// finalized from `source` out of it — see
112 /// [`ArtifactCompiler::source`] — for [`store`](Self::store) to
113 /// put back under the kernel asking. `None`, the default, for a backend
114 /// that keeps no artifact by source.
115 fn stored_for_source(&mut self, source: &str) -> Option<ArtifactOf<Self>> {
116 let _ = source;
117 None
118 }
119
120 /// Loads `artifact` on the device.
121 fn load(
122 &mut self,
123 id: &ArtifactId<VariantOf<Self>>,
124 artifact: &ArtifactOf<Self>,
125 ) -> Result<Self::Loaded, CompilationError>;
126
127 /// Keeps a loaded `artifact` in the compilation store under `id` and, when
128 /// given, under the `source` it was finalized from. Whether the store took
129 /// it, as [`store_compiled`](super::store_compiled) answers: `false` with
130 /// no store.
131 ///
132 /// Called after [`load`](Self::load), so an artifact the device refused is
133 /// never handed back on the next run.
134 fn store(
135 &mut self,
136 id: &ArtifactId<VariantOf<Self>>,
137 artifact: ArtifactOf<Self>,
138 source: Option<&str>,
139 ) -> bool;
140}