use alloc::boxed::Box;
use alloc::string::String;
use alloc::vec::Vec;
use cubecl_common::profile::Instant;
use cubecl_environment::collections::{HashMap, HashSet};
use cubecl_runtime::execution::{KernelHandling, KernelOutcome, KernelRegistration, SettledKernel};
use super::{
ArtifactCompiler, ArtifactId, BatchOutcome, CompilationBatchRecording, CompilationOutcome,
CompilationRecording, CompilationTarget, VariantOf,
};
use crate::config::RuntimeConfig;
use crate::kernel::CubeKernel;
use crate::logging::ServerLogger;
use crate::server::LaunchError;
pub struct KernelLoader<T: CompilationTarget> {
target: T,
loaded: CompilationCache<ArtifactId<VariantOf<T>>, T::Loaded>,
queue: KernelQueue<VariantOf<T>>,
failed: HashMap<ArtifactId<VariantOf<T>>, LaunchError>,
parallelism: usize,
}
impl<T: CompilationTarget> KernelLoader<T> {
pub fn new(target: T) -> Self {
let loaded = CompilationCache::new(target.persists());
let parallelism = crate::config::CubeClRuntimeConfig::get()
.compilation
.parallelism();
Self {
target,
loaded,
queue: KernelQueue::new(),
failed: HashMap::new(),
parallelism,
}
}
#[cfg(test)]
fn get(&mut self, id: &ArtifactId<VariantOf<T>>) -> Option<&T::Loaded> {
self.loaded.get(id)
}
pub fn load(
&mut self,
kernel: &dyn CubeKernel,
id: &ArtifactId<VariantOf<T>>,
logger: &ServerLogger,
) -> Result<T::Loaded, LaunchError> {
if self.queue.is_empty()
&& let Some(loaded) = self.loaded.get(id)
{
return Ok(loaded.clone());
}
self.load_missing(kernel, id, logger)
}
#[cold]
fn load_missing(
&mut self,
kernel: &dyn CubeKernel,
id: &ArtifactId<VariantOf<T>>,
logger: &ServerLogger,
) -> Result<T::Loaded, LaunchError> {
loop {
let reported = self.failed.contains_key(id) && self.queue.is_empty();
let missing = !self.loaded.contains(id) && !reported;
let asked = missing.then(|| Request {
id: id.clone(),
kernel,
registration: None,
});
self.compile_queue(asked, logger);
if let Some(err) = self.failed.remove(id) {
return Err(err);
}
if let Some(loaded) = self.loaded.get(id) {
return Ok(loaded.clone());
}
}
}
pub fn compile_queued(&mut self, logger: &ServerLogger) {
self.compile_queue(None, logger);
}
fn compile_queue(&mut self, asked: Option<Request<'_, VariantOf<T>>>, logger: &ServerLogger) {
if self.queue.is_empty() && asked.is_none() {
return;
}
let mut queue = self.queue.take();
self.failed.clear();
let asked = asked.map(|mut asked| {
asked.registration = match queue.iter_mut().find(|queued| queued.id == asked.id) {
Some(queued) => queued.registration.take(),
None => KernelRegistration::register(),
};
asked
});
let mut requests: Vec<Request<'_, VariantOf<T>>> = queue
.iter_mut()
.filter(|queued| asked.as_ref().is_none_or(|asked| queued.id != asked.id))
.map(Queued::request)
.collect();
requests.extend(asked);
for outcome in self.compile(requests, logger) {
self.keep(outcome);
}
}
pub fn enqueue(&mut self, kernel: Box<dyn CubeKernel>, variant: VariantOf<T>) {
let id = ArtifactId {
kernel: kernel.id(),
variant,
};
if self.loaded.contains(&id) || self.queue.contains(&id) {
return;
}
self.failed.remove(&id);
self.queue.push(id, kernel, KernelRegistration::register());
}
fn keep(&mut self, outcome: JobOutcome<T>) {
match outcome.result {
Ok(loaded) => self.loaded.insert(outcome.id, loaded),
Err(err) => {
self.failed.insert(outcome.id, err);
}
}
}
fn compile(
&mut self,
requests: Vec<Request<'_, VariantOf<T>>>,
logger: &ServerLogger,
) -> Vec<JobOutcome<T>> {
let batch = CompilationBatchRecording::new();
let mut jobs: Vec<Job<'_, T>> = requests.into_iter().map(Job::new).collect();
let jobs_len = jobs.len();
for job in jobs.iter_mut() {
job.look_up(&mut self.target);
}
jobs.iter_mut().for_each(Job::settle);
let compiler = self.target.compiler();
let mut missing: Vec<&mut Job<'_, T>> =
jobs.iter_mut().filter(|job| job.is_missing()).collect();
let threads = self.parallelism.min(missing.len()).max(1);
if missing.len() > 1 {
log::info!(
target: "cubecl::compilation",
"Compiling {} kernels together ({} asked, the rest read from the store) on {} threads",
missing.len(),
jobs_len,
threads,
);
}
let reuses = self.target.persists();
for_each_at_once(&mut missing, self.parallelism, |job| {
job.lower(compiler, logger);
if job.is_lowered() && !(reuses && job.has_source()) {
job.finalize(compiler);
}
job.settle();
});
let mut finalizing = HashSet::new();
for job in jobs.iter_mut() {
job.reuse(&mut self.target, &mut finalizing);
}
let leading: Vec<(usize, String)> = jobs
.iter()
.enumerate()
.filter_map(|(index, job)| match &job.step {
Step::Lowered {
source: Some(source),
..
} => Some((index, source.clone())),
_ => None,
})
.collect();
let compiler = self.target.compiler();
jobs.iter_mut().for_each(Job::settle);
let mut lowered: Vec<&mut Job<'_, T>> =
jobs.iter_mut().filter(|job| job.is_lowered()).collect();
for_each_at_once(&mut lowered, self.parallelism, |job| {
job.finalize(compiler);
job.settle();
});
let failed: HashMap<String, LaunchError> = leading
.into_iter()
.filter_map(|(index, source)| match &jobs[index].step {
Step::Failed(err) => Some((source, err.clone())),
_ => None,
})
.collect();
for job in jobs.iter_mut() {
if let Step::Waiting { source, .. } = &job.step
&& let Some(err) = failed.get(source)
{
job.step = Step::Failed(err.clone());
}
}
let outcomes: Vec<JobOutcome<T>> = jobs
.into_iter()
.map(|job| job.load(&mut self.target))
.collect();
batch.close(BatchOutcome {
kernels: jobs_len,
threads,
stored: outcomes.iter().any(|outcome| outcome.stored),
});
outcomes
}
}
impl<T: CompilationTarget + core::fmt::Debug> core::fmt::Debug for KernelLoader<T> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("KernelLoader")
.field("target", &self.target)
.field("queued", &self.queue.len())
.field("failed", &self.failed.len())
.field("parallelism", &self.parallelism)
.finish_non_exhaustive()
}
}
struct JobOutcome<T: CompilationTarget> {
id: ArtifactId<VariantOf<T>>,
result: Result<T::Loaded, LaunchError>,
stored: bool,
}
#[derive(Debug)]
struct CompilationCache<K, V> {
entries: HashMap<K, V>,
generation: Option<u32>,
}
impl<K: Eq + core::hash::Hash, V> CompilationCache<K, V> {
fn new(bound: bool) -> Self {
Self {
entries: HashMap::new(),
generation: bound.then(cubecl_environment::environment::generation),
}
}
fn get(&mut self, key: &K) -> Option<&V> {
self.reset_if_switched();
self.entries.get(key)
}
fn contains(&mut self, key: &K) -> bool {
self.reset_if_switched();
self.entries.contains_key(key)
}
fn insert(&mut self, key: K, value: V) {
self.reset_if_switched();
self.entries.insert(key, value);
}
fn reset_if_switched(&mut self) {
let Some(generation) = self.generation else {
return;
};
let current = cubecl_environment::environment::generation();
if current == generation {
return;
}
log::debug!("Environment switched, dropping the in-memory compilation cache");
self.generation = Some(current);
self.entries.clear();
}
}
struct KernelQueue<V> {
kernels: Vec<Queued<V>>,
ids: HashSet<ArtifactId<V>>,
}
impl<V: Clone + Eq + core::hash::Hash> KernelQueue<V> {
fn new() -> Self {
Self {
kernels: Vec::new(),
ids: HashSet::new(),
}
}
fn is_empty(&self) -> bool {
self.kernels.is_empty()
}
fn len(&self) -> usize {
self.kernels.len()
}
fn contains(&self, id: &ArtifactId<V>) -> bool {
self.ids.contains(id)
}
fn push(
&mut self,
id: ArtifactId<V>,
kernel: Box<dyn CubeKernel>,
registration: Option<KernelRegistration>,
) {
self.ids.insert(id.clone());
self.kernels.push(Queued {
id,
kernel,
registration,
});
}
fn take(&mut self) -> Vec<Queued<V>> {
self.ids.clear();
core::mem::take(&mut self.kernels)
}
}
struct Queued<V> {
id: ArtifactId<V>,
kernel: Box<dyn CubeKernel>,
registration: Option<KernelRegistration>,
}
impl<V: Clone> Queued<V> {
fn request(&mut self) -> Request<'_, V> {
Request {
id: self.id.clone(),
kernel: &*self.kernel,
registration: self.registration.take(),
}
}
}
struct Request<'k, V> {
id: ArtifactId<V>,
kernel: &'k dyn CubeKernel,
registration: Option<KernelRegistration>,
}
struct Job<'k, T: CompilationTarget> {
id: ArtifactId<VariantOf<T>>,
kernel: &'k dyn CubeKernel,
recording: CompilationRecording,
step: Step<T>,
count: JobCount,
}
enum JobCount {
Uncounted,
Registered(KernelRegistration),
Settled(SettledKernel),
}
impl JobCount {
fn record(&mut self, load: KernelHandling) {
if let Self::Settled(settled) = self {
settled.record(load);
}
}
}
enum Step<T: CompilationTarget> {
Missing,
Lowered {
lowered: <T::Compiler as ArtifactCompiler>::Lowered,
source: Option<String>,
},
Waiting {
lowered: <T::Compiler as ArtifactCompiler>::Lowered,
source: String,
},
Ready {
artifact: <T::Compiler as ArtifactCompiler>::Artifact,
source: Option<String>,
outcome: CompilationOutcome,
},
Failed(LaunchError),
}
impl<T: CompilationTarget> Step<T> {
fn outcome(&self) -> Option<KernelOutcome> {
match self {
Self::Ready {
outcome: CompilationOutcome::Compiled,
..
} => Some(KernelOutcome::Compiled),
Self::Ready {
outcome: CompilationOutcome::Loaded | CompilationOutcome::Rekeyed,
..
} => Some(KernelOutcome::Loaded),
Self::Failed(_) => Some(KernelOutcome::Failed),
Self::Missing | Self::Lowered { .. } | Self::Waiting { .. } => None,
}
}
}
impl<'k, T: CompilationTarget> Job<'k, T> {
fn new(request: Request<'k, VariantOf<T>>) -> Self {
Self {
recording: CompilationRecording::new(&request.id.kernel),
id: request.id,
kernel: request.kernel,
step: Step::Missing,
count: match request.registration {
Some(registration) => JobCount::Registered(registration),
None => JobCount::Uncounted,
},
}
}
fn settle(&mut self) {
let Some(outcome) = self.step.outcome() else {
return;
};
self.count = match core::mem::replace(&mut self.count, JobCount::Uncounted) {
JobCount::Registered(registration) => JobCount::Settled(registration.settle(outcome)),
count => count,
};
}
fn is_missing(&self) -> bool {
matches!(self.step, Step::Missing)
}
fn is_lowered(&self) -> bool {
matches!(self.step, Step::Lowered { .. })
}
fn has_source(&self) -> bool {
matches!(
self.step,
Step::Lowered {
source: Some(_),
..
}
)
}
fn look_up(&mut self, target: &mut T) {
if let Some(artifact) = target.stored(&self.id) {
self.step = Step::Ready {
artifact,
source: None,
outcome: CompilationOutcome::Loaded,
};
}
}
fn lower(&mut self, compiler: &T::Compiler, logger: &ServerLogger) {
let started = Instant::now();
let (kernel, id, recording) = (self.kernel, &self.id, &mut self.recording);
self.step = match contained(|| compiler.lower(kernel, id, recording, logger)) {
Ok(lowered) => Step::Lowered {
source: T::Compiler::source(&lowered).map(String::from),
lowered,
},
Err(err) => Step::Failed(err),
};
self.recording.worked(started.elapsed());
}
fn reuse(&mut self, target: &mut T, finalizing: &mut HashSet<String>) {
let Step::Lowered {
source: Some(source),
..
} = &self.step
else {
return;
};
if !target.persists() {
return;
}
if finalizing.contains(source) {
let Step::Lowered {
lowered,
source: Some(source),
} = core::mem::replace(&mut self.step, Step::Missing)
else {
unreachable!("matched right above");
};
self.step = Step::Waiting { lowered, source };
return;
}
finalizing.insert(source.clone());
if let Some(artifact) = target.stored_for_source(source) {
self.step = Step::Ready {
artifact,
source: Some(source.clone()),
outcome: CompilationOutcome::Rekeyed,
};
}
}
fn finalize(&mut self, compiler: &T::Compiler) {
let Step::Lowered { lowered, source } = core::mem::replace(&mut self.step, Step::Missing)
else {
unreachable!("only lowered jobs are finalized");
};
let started = Instant::now();
let id = &self.id;
self.step = match contained(|| compiler.finalize(id, lowered)) {
Ok(artifact) => Step::Ready {
artifact,
source,
outcome: CompilationOutcome::Compiled,
},
Err(err) => Step::Failed(err),
};
self.recording.worked(started.elapsed());
}
fn load(mut self, target: &mut T) -> JobOutcome<T> {
if let Step::Waiting { .. } = self.step {
let Step::Waiting { lowered, source } =
core::mem::replace(&mut self.step, Step::Missing)
else {
unreachable!("matched right above");
};
self.step = match target.stored_for_source(&source) {
Some(artifact) => Step::Ready {
artifact,
source: Some(source),
outcome: CompilationOutcome::Rekeyed,
},
None => Step::Lowered {
lowered,
source: Some(source),
},
};
if self.is_lowered() {
self.finalize(target.compiler());
}
}
self.settle();
let mut stored = false;
let result = match self.step {
Step::Ready {
artifact,
source,
outcome,
} => {
let started = Instant::now();
match target.load(&self.id, &artifact) {
Ok(loaded) => {
self.recording.worked(started.elapsed());
if outcome != CompilationOutcome::Loaded {
stored = target.store(&self.id, artifact, source.as_deref());
}
if stored {
self.count.record(KernelHandling::Stored);
}
self.recording.close(outcome, stored);
Ok(loaded)
}
Err(err) => {
self.count.record(KernelHandling::Refused);
Err(err.into())
}
}
}
Step::Failed(err) => Err(err),
Step::Missing | Step::Lowered { .. } | Step::Waiting { .. } => {
unreachable!("every job is lowered, then finalized or reused, or failed")
}
};
JobOutcome {
id: self.id,
result,
stored,
}
}
}
fn contained<R>(step: impl FnOnce() -> Result<R, LaunchError>) -> Result<R, LaunchError> {
#[cfg(feature = "std")]
{
std::panic::catch_unwind(core::panic::AssertUnwindSafe(step)).unwrap_or_else(|payload| {
let message = payload
.downcast_ref::<String>()
.map(String::as_str)
.or_else(|| payload.downcast_ref::<&str>().copied())
.unwrap_or("no message");
Err(LaunchError::CompilationError(
crate::compiler::CompilationError::Generic {
reason: alloc::format!("the compiler panicked: {message}"),
backtrace: cubecl_environment::backtrace::BackTrace::capture(),
},
))
})
}
#[cfg(not(feature = "std"))]
step()
}
fn for_each_at_once<J: Send>(jobs: &mut [J], parallelism: usize, work: impl Fn(&mut J) + Sync) {
#[cfg(all(feature = "std", not(target_family = "wasm")))]
if jobs.len() > 1 && parallelism > 1 {
use core::sync::atomic::{AtomicUsize, Ordering};
let slots: Vec<std::sync::Mutex<&mut J>> =
jobs.iter_mut().map(std::sync::Mutex::new).collect();
let next = AtomicUsize::new(0);
let panic = std::thread::scope(|scope| {
let workers: Vec<_> = (0..parallelism.min(slots.len()))
.map(|_| {
scope.spawn(|| {
while let Some(slot) = slots.get(next.fetch_add(1, Ordering::Relaxed)) {
work(&mut slot.lock().unwrap_or_else(|err| err.into_inner()));
}
})
})
.collect();
workers
.into_iter()
.fold(None, |first, worker| first.or(worker.join().err()))
});
if let Some(payload) = panic {
std::panic::resume_unwind(payload);
}
return;
}
let _ = parallelism;
jobs.iter_mut().for_each(work);
}
#[cfg(test)]
mod tests {
use super::*;
use crate::compiler::CompilationError;
use crate::id::KernelId;
use crate::kernel::{KernelDefinition, KernelMetadata};
use alloc::collections::BTreeSet;
use alloc::vec;
use core::sync::atomic::{AtomicUsize, Ordering};
use cubecl_environment::backtrace::BackTrace;
use std::sync::Mutex;
#[test]
#[serial_test::serial(records)]
fn a_load_compiles_the_queue() {
let mut loader = loader(8);
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
assert_eq!(
loader.target.compiler.lowered(),
0,
"queuing compiles nothing"
);
assert_eq!(loader.load(&Numbered(3), &id(3), &logger).unwrap(), 3);
assert_eq!(
loader.target.compiler.lowered(),
4,
"the queue compiled with it"
);
for number in 0..3 {
assert_eq!(loader.get(&id(number)), Some(&number));
}
}
#[test]
#[serial_test::serial(records)]
fn a_batch_settles_every_kernel_to_the_collector_that_registered_it() {
use cubecl_runtime::execution::{ProcessMode, ProcessModeOverride, StatisticsCollector};
let mut loader = loader(2);
let logger = ServerLogger::default();
loader.enqueue(Box::new(Numbered(9)), ());
let collector = StatisticsCollector::new();
let compile = ProcessModeOverride::new(ProcessMode::CompileOnly, &collector);
for number in 0..5 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.enqueue(Box::new(Numbered(0)), ());
core::mem::drop(compile);
let queued = collector.statistics().compilation;
assert_eq!((queued.registered, queued.settled()), (5, 0));
let tune = ProcessModeOverride::new(ProcessMode::CompileAndAutotune, &collector);
loader.load(&Numbered(5), &id(5), &logger).unwrap();
let batch = collector.statistics().compilation;
assert_eq!(
(batch.registered, batch.settled(), batch.failed),
(6, 6, 0),
"the kernel queued with no override open counts nowhere"
);
assert_eq!(batch.compiled + batch.loaded, 6);
loader.load(&Numbered(5), &id(5), &logger).unwrap();
assert_eq!(
collector.statistics().compilation,
batch,
"a loaded kernel is not registered again"
);
core::mem::drop(tune);
}
fn counted(
collector: &cubecl_runtime::execution::StatisticsCollector,
) -> (usize, usize, usize, usize, usize) {
let statistics = collector.statistics().compilation;
(
statistics.registered,
statistics.compiled,
statistics.loaded,
statistics.failed,
statistics.refused,
)
}
#[test]
#[serial_test::serial(records)]
fn a_failure_settles_as_failed_and_a_retry_registers_again() {
use cubecl_runtime::execution::{ProcessMode, ProcessModeOverride, StatisticsCollector};
let collector = StatisticsCollector::new();
let tune = ProcessModeOverride::new(ProcessMode::CompileAndAutotune, &collector);
let mut loader = loader(4);
loader.target.compiler.failing = Some(1);
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(3), &id(3), &logger).unwrap();
assert_eq!(counted(&collector), (4, 3, 0, 1, 0));
assert!(loader.load(&Numbered(1), &id(1), &logger).is_err());
assert_eq!(counted(&collector), (4, 3, 0, 1, 0), "reported, not tried");
assert!(loader.load(&Numbered(1), &id(1), &logger).is_err());
assert_eq!(counted(&collector), (5, 3, 0, 2, 0), "tried again");
loader.target.compiler.failing = None;
loader.load(&Numbered(1), &id(1), &logger).unwrap();
assert_eq!(counted(&collector), (6, 4, 0, 2, 0));
core::mem::drop(tune);
}
#[test]
#[serial_test::serial(records)]
fn kernels_sharing_a_source_settle_with_it() {
use cubecl_runtime::execution::{ProcessMode, ProcessModeOverride, StatisticsCollector};
let logger = ServerLogger::default();
for (refusing, expected) in [(false, (4, 1, 3, 0, 0)), (true, (4, 0, 0, 4, 0))] {
let collector = StatisticsCollector::new();
let tune = ProcessModeOverride::new(ProcessMode::CompileAndAutotune, &collector);
let mut loader = loader(4);
loader.target.compiler.shared_source = true;
loader.target.compiler.refusing = refusing;
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
let _ = loader.load(&Numbered(3), &id(3), &logger);
assert_eq!(counted(&collector), expected, "refusing: {refusing}");
let stored = collector.statistics().compilation.stored;
assert_eq!(stored, if refusing { 0 } else { 4 }, "refusing: {refusing}");
core::mem::drop(tune);
}
}
#[test]
#[serial_test::serial(records)]
fn a_refused_kernel_is_counted_beside_its_outcome() {
use cubecl_runtime::execution::{ProcessMode, ProcessModeOverride, StatisticsCollector};
let collector = StatisticsCollector::new();
let tune = ProcessModeOverride::new(ProcessMode::CompileAndAutotune, &collector);
let mut loader = loader(4);
loader.target.rejecting = Some(1);
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(3), &id(3), &logger).unwrap();
assert_eq!(loader.get(&id(1)), None);
assert_eq!(counted(&collector), (4, 4, 0, 0, 1));
assert_eq!(
collector.statistics().compilation.stored,
0,
"a store keeps by source, and these kernels have none"
);
core::mem::drop(tune);
}
#[test]
#[serial_test::serial(records)]
fn a_panicking_batch_settles_every_kernel() {
use cubecl_runtime::execution::{ProcessMode, ProcessModeOverride, StatisticsCollector};
let collector = StatisticsCollector::new();
let tune = ProcessModeOverride::new(ProcessMode::CompileAndAutotune, &collector);
let mut loader = loader(4);
loader.target.compiler.panicking = Some(2);
let logger = ServerLogger::default();
for number in 0..4 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(4), &id(4), &logger).unwrap();
let statistics = collector.statistics().compilation;
assert_eq!(statistics.registered, 5);
assert_eq!(statistics.settled(), 5);
assert_eq!(statistics.failed, 1, "only the kernel that panicked");
assert_eq!(statistics.abandoned, 0);
core::mem::drop(tune);
}
#[test]
#[serial_test::serial(records)]
fn the_queue_compiles_when_asked() {
let mut loader = loader(8);
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.compile_queued(&logger);
assert_eq!(loader.target.compiler.lowered(), 3);
assert_eq!(loader.load(&Numbered(0), &id(0), &logger).unwrap(), 0);
assert_eq!(
loader.target.compiler.lowered(),
3,
"nothing left to compile"
);
}
#[test]
#[serial_test::serial(records)]
fn a_long_queue_compiles_in_one_phase() {
let mut loader = loader(2);
let logger = ServerLogger::default();
for number in 0..5 {
loader.enqueue(Box::new(Numbered(number)), ());
}
assert_eq!(loader.target.compiler.lowered(), 0, "the queue waits");
loader.load(&Numbered(5), &id(5), &logger).unwrap();
assert_eq!(
loader.target.compiler.lowered(),
6,
"with the kernel launched"
);
}
#[test]
#[serial_test::serial(records)]
fn a_kernel_compiles_once() {
let mut loader = loader(8);
let logger = ServerLogger::default();
loader.enqueue(Box::new(Numbered(0)), ());
loader.load(&Numbered(0), &id(0), &logger).unwrap();
loader.enqueue(Box::new(Numbered(0)), ());
loader.load(&Numbered(0), &id(0), &logger).unwrap();
assert_eq!(loader.target.compiler.lowered(), 1);
}
#[test]
#[serial_test::serial(records)]
fn a_failure_is_reported_by_its_own_launch() {
let mut loader = loader(8);
loader.target.compiler.failing = Some(1);
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(3), &id(3), &logger).unwrap();
assert_eq!(loader.get(&id(0)), Some(&0));
assert_eq!(loader.get(&id(1)), None);
assert_eq!(loader.get(&id(2)), Some(&2));
assert_eq!(loader.target.compiler.lowered(), 4);
assert!(loader.load(&Numbered(1), &id(1), &logger).is_err());
assert_eq!(
loader.target.compiler.lowered(),
4,
"reported, not compiled again"
);
assert!(loader.load(&Numbered(1), &id(1), &logger).is_err());
assert_eq!(
loader.target.compiler.lowered(),
5,
"the next launch tries again"
);
}
#[test]
#[serial_test::serial(records)]
fn a_stale_failure_compiles_again() {
let mut loader = loader(8);
loader.target.compiler.failing = Some(1);
let logger = ServerLogger::default();
loader.enqueue(Box::new(Numbered(1)), ());
loader.load(&Numbered(0), &id(0), &logger).unwrap();
loader.enqueue(Box::new(Numbered(2)), ());
loader.load(&Numbered(0), &id(0), &logger).unwrap();
assert_eq!(loader.target.compiler.lowered(), 3);
loader.target.compiler.failing = None;
assert_eq!(loader.load(&Numbered(1), &id(1), &logger).unwrap(), 1);
assert_eq!(loader.target.compiler.lowered(), 4, "compiled again");
}
#[test]
#[serial_test::serial(records)]
fn a_failure_queued_again_is_retried() {
let mut loader = loader(8);
loader.target.compiler.failing = Some(1);
let logger = ServerLogger::default();
loader.enqueue(Box::new(Numbered(1)), ());
loader.load(&Numbered(0), &id(0), &logger).unwrap();
assert_eq!(loader.get(&id(1)), None);
loader.target.compiler.failing = None;
loader.enqueue(Box::new(Numbered(1)), ());
loader.compile_queued(&logger);
assert_eq!(loader.get(&id(1)), Some(&1));
}
#[test]
#[serial_test::serial(records)]
fn a_panic_fails_only_its_own_kernel() {
let mut loader = loader(4);
loader.target.compiler.panicking = Some(2);
let logger = ServerLogger::default();
for number in 0..4 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(4), &id(4), &logger).unwrap();
for number in [0, 1, 3] {
assert_eq!(loader.get(&id(number)), Some(&number));
}
let err = loader
.load(&Numbered(2), &id(2), &logger)
.expect_err("the kernel that panicked fails");
assert!(
alloc::format!("{err}").contains("kernel 2 panicked"),
"the failure keeps the panic's message: {err}"
);
}
#[test]
#[cfg(persistence)]
#[serial_test::serial(records)]
fn a_batch_is_recorded_beside_its_kernels() {
use crate::compiler::{CompilationBatchRecord, CompilationRecord};
use crate::config::RuntimeConfig;
use cubecl_environment::persistence::Database;
use cubecl_environment::records::{self, RecordLevel, Records, RecordsConfig};
let root = tempfile::tempdir().unwrap();
let _ = crate::config::CubeClRuntimeConfig::get();
cubecl_environment::environment::set_root(root.path());
records::configure(RecordsConfig {
level: RecordLevel::Basic,
..Default::default()
});
let mut loader = loader(4);
loader.target.compiler.slow = true;
loader.target.compiler.shared_source = true;
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(3), &id(3), &logger).unwrap();
let database = Database::open_active().unwrap();
let records = Records::new(&database);
let batches = records.read::<CompilationBatchRecord>();
let kernels = records.read::<CompilationRecord>();
assert_eq!(batches.len(), 1);
assert_eq!(kernels.len(), 4);
records::configure(RecordsConfig {
level: RecordLevel::Off,
..Default::default()
});
let batch = &batches[0].record;
assert_eq!((batch.kernels, batch.threads), (4, 4));
let work: core::time::Duration = kernels.iter().map(|kernel| kernel.record.work).sum();
assert!(
batch.wall < work,
"four threads: {:?} of wall for {work:?} of work",
batch.wall
);
}
#[test]
#[serial_test::serial(records)]
fn a_failed_source_is_finalized_once() {
let mut loader = loader(8);
loader.target.compiler.shared_source = true;
loader.target.compiler.refusing = true;
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
assert!(loader.load(&Numbered(3), &id(3), &logger).is_err());
assert_eq!(loader.target.compiler.finalized.load(Ordering::Relaxed), 1);
for number in 0..3 {
assert_eq!(loader.get(&id(number)), None);
}
}
#[test]
#[serial_test::serial(records)]
fn a_kernel_without_a_source_finalizes_where_it_lowered() {
let mut loader = loader(4);
loader.target.compiler.slow = true;
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(3), &id(3), &logger).unwrap();
let lowered = loader.target.compiler.lowered_on.lock().unwrap().clone();
let finalized = loader.target.compiler.finalized_on.lock().unwrap().clone();
assert_eq!(lowered.len(), 4);
assert_eq!(lowered, finalized);
}
#[test]
#[serial_test::serial(records)]
fn a_loaded_kernel_still_compiles_the_queue() {
let mut loader = loader(8);
let logger = ServerLogger::default();
loader.load(&Numbered(0), &id(0), &logger).unwrap();
loader.enqueue(Box::new(Numbered(1)), ());
loader.load(&Numbered(0), &id(0), &logger).unwrap();
assert_eq!(loader.get(&id(1)), Some(&1));
}
#[test]
#[serial_test::serial(records)]
fn one_source_finalizes_once() {
let mut loader = loader(8);
loader.target.compiler.shared_source = true;
let logger = ServerLogger::default();
for number in 0..4 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(4), &id(4), &logger).unwrap();
assert_eq!(loader.target.compiler.finalized.load(Ordering::Relaxed), 1);
for number in 0..5 {
assert!(loader.get(&id(number)).is_some());
}
}
#[test]
#[serial_test::serial(records)]
fn a_stored_source_serves_the_whole_batch() {
let mut loader = loader(8);
loader.target.compiler.shared_source = true;
loader.target.by_source.insert("shared".into(), 99);
let logger = ServerLogger::default();
for number in 0..2 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(2), &id(2), &logger).unwrap();
assert_eq!(loader.target.compiler.finalized.load(Ordering::Relaxed), 0);
for number in 0..3 {
assert_eq!(loader.get(&id(number)), Some(&99));
}
}
#[test]
#[serial_test::serial(records)]
fn without_a_store_each_kernel_finalizes_its_own() {
let mut loader = KernelLoader::new(Fake {
storeless: true,
..Fake::default()
});
loader.parallelism = 8;
loader.target.compiler.shared_source = true;
let logger = ServerLogger::default();
for number in 0..3 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(3), &id(3), &logger).unwrap();
assert_eq!(loader.target.compiler.finalized.load(Ordering::Relaxed), 4);
}
#[test]
#[serial_test::serial(records)]
fn a_queue_compiles_on_several_threads() {
let mut loader = loader(4);
loader.target.compiler.slow = true;
let logger = ServerLogger::default();
for number in 0..8 {
loader.enqueue(Box::new(Numbered(number)), ());
}
loader.load(&Numbered(8), &id(8), &logger).unwrap();
let threads = loader.target.compiler.threads.lock().unwrap().len();
assert!(threads > 1, "lowered on {threads} thread");
}
fn loader(parallelism: usize) -> KernelLoader<Fake> {
let mut loader = KernelLoader::new(Fake::default());
loader.parallelism = parallelism;
loader
}
fn id(number: u32) -> ArtifactId<()> {
ArtifactId {
kernel: Numbered(number).id(),
variant: (),
}
}
#[derive(Debug)]
struct Numbered(u32);
impl KernelMetadata for Numbered {
fn id(&self) -> KernelId {
KernelId::new::<Self>().info(self.0)
}
fn address_type(&self) -> cubecl_ir::ElemType {
cubecl_ir::ElemType::UInt(cubecl_ir::UIntKind::U32)
}
}
impl CubeKernel for Numbered {
fn define(&self) -> KernelDefinition {
unreachable!("the fake compiler reads the number, not the definition")
}
}
#[derive(Debug, Default)]
struct Fake {
compiler: FakeCompiler,
by_source: alloc::collections::BTreeMap<String, u32>,
storeless: bool,
rejecting: Option<u32>,
}
#[derive(Debug, Default)]
struct FakeCompiler {
lowered: AtomicUsize,
failing: Option<u32>,
panicking: Option<u32>,
slow: bool,
threads: Mutex<BTreeSet<String>>,
shared_source: bool,
finalized: AtomicUsize,
refusing: bool,
lowered_on: Mutex<alloc::collections::BTreeMap<u32, String>>,
finalized_on: Mutex<alloc::collections::BTreeMap<u32, String>>,
}
impl FakeCompiler {
fn lowered(&self) -> usize {
self.lowered.load(Ordering::Relaxed)
}
}
impl ArtifactCompiler for FakeCompiler {
type Variant = ();
type Lowered = (u32, Option<String>);
type Artifact = u32;
fn lower(
&self,
kernel: &dyn CubeKernel,
id: &ArtifactId<()>,
_recording: &mut CompilationRecording,
_logger: &ServerLogger,
) -> Result<(u32, Option<String>), LaunchError> {
self.lowered.fetch_add(1, Ordering::Relaxed);
self.threads
.lock()
.unwrap()
.insert(alloc::format!("{:?}", std::thread::current().id()));
if self.slow {
std::thread::sleep(core::time::Duration::from_millis(20));
}
let number = (0..64)
.find(|number| Numbered(*number).id() == id.kernel)
.expect("a numbered kernel");
assert_eq!(kernel.id(), id.kernel);
self.lowered_on
.lock()
.unwrap()
.insert(number, alloc::format!("{:?}", std::thread::current().id()));
if self.panicking == Some(number) {
panic!("kernel {number} panicked");
}
match self.failing == Some(number) {
true => Err(LaunchError::Unknown {
reason: "refused".into(),
backtrace: BackTrace::capture(),
}),
false => Ok((number, self.shared_source.then(|| "shared".into()))),
}
}
fn source(lowered: &(u32, Option<String>)) -> Option<&str> {
lowered.1.as_deref()
}
fn finalize(
&self,
_id: &ArtifactId<()>,
lowered: (u32, Option<String>),
) -> Result<u32, LaunchError> {
self.finalized.fetch_add(1, Ordering::Relaxed);
self.finalized_on.lock().unwrap().insert(
lowered.0,
alloc::format!("{:?}", std::thread::current().id()),
);
match self.refusing {
true => Err(LaunchError::Unknown {
reason: "refused".into(),
backtrace: BackTrace::capture(),
}),
false => Ok(lowered.0),
}
}
}
impl CompilationTarget for Fake {
type Compiler = FakeCompiler;
type Loaded = u32;
fn compiler(&self) -> &FakeCompiler {
&self.compiler
}
fn persists(&self) -> bool {
!self.storeless
}
fn stored(&mut self, _id: &ArtifactId<()>) -> Option<u32> {
None
}
fn load(
&mut self,
loaded: &ArtifactId<()>,
artifact: &u32,
) -> Result<u32, CompilationError> {
match self.rejecting.is_some_and(|number| *loaded == id(number)) {
true => Err(CompilationError::Validation {
reason: "rejected".into(),
backtrace: BackTrace::capture(),
}),
false => Ok(*artifact),
}
}
fn stored_for_source(&mut self, source: &str) -> Option<u32> {
self.by_source.remove(source)
}
fn store(&mut self, _id: &ArtifactId<()>, artifact: u32, source: Option<&str>) -> bool {
if let Some(source) = source {
self.by_source.insert(source.into(), artifact);
}
source.is_some()
}
}
}