#[allow(
non_camel_case_types,
non_snake_case,
non_upper_case_globals,
dead_code,
missing_docs,
unsafe_op_in_unsafe_fn,
clippy::all
)]
mod ffi;
mod native;
mod report;
mod source;
use crate::{
Error, Result,
bundle::{self, Lock},
};
pub use report::{CompileReport, EntryChange, EntryResources};
use serde::{Deserialize, Serialize};
pub use source::{CxxSource, CxxStandard, Source};
use std::{
collections::{BTreeMap, HashMap, HashSet},
fs,
num::NonZeroUsize,
path::{Path, PathBuf},
sync::{Arc, Mutex},
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Severity {
Note,
Warning,
Error,
}
impl From<u32> for Severity {
fn from(native: u32) -> Self {
match native {
0 => Self::Note,
1 => Self::Warning,
_ => Self::Error,
}
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Diagnostic {
pub severity: Severity,
pub code: String,
pub message: String,
#[serde(default)]
pub source: String,
pub line: u32,
pub column: u32,
}
#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProcessorMode {
#[default]
Default,
ComputeUnit,
WorkgroupProcessor,
}
#[derive(Clone, Debug)]
pub struct CompilerOptions {
pub workers: NonZeroUsize,
pub target: crate::Target,
pub module_cache_capacity: usize,
pub processor_mode: ProcessorMode,
}
impl Default for CompilerOptions {
fn default() -> Self {
Self {
target: crate::Target::default(),
module_cache_capacity: 64,
processor_mode: ProcessorMode::Default,
workers: NonZeroUsize::new(
std::thread::available_parallelism()
.map_or(1, NonZeroUsize::get)
.min(4),
)
.unwrap(),
}
}
}
impl CompilerOptions {
#[must_use]
pub fn for_target(target: &crate::Target) -> Self {
Self {
target: target.clone(),
..Self::default()
}
}
#[must_use]
pub fn for_stream(stream: &crate::Stream) -> Self {
Self::for_target(stream.target())
}
}
struct Inner {
modules: Mutex<HashMap<String, Arc<ModuleData>>>,
native: native::Prepared,
target: crate::Target,
workers: NonZeroUsize,
module_cache_capacity: usize,
processor_mode: ProcessorMode,
path: PathBuf,
identity: String,
}
#[derive(Hash, PartialEq, Eq)]
struct SharedCompilerKey {
path: PathBuf,
target: String,
workers: NonZeroUsize,
module_cache_capacity: usize,
processor_mode: ProcessorMode,
}
#[derive(Clone)]
pub struct Compiler(Arc<Inner>);
impl std::fmt::Debug for Compiler {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Compiler")
.field("library", &self.0.path)
.field("identity", &self.0.identity)
.finish_non_exhaustive()
}
}
impl Compiler {
pub fn resolve(library: Option<&Path>) -> Result<Self> {
Self::with_options(library, CompilerOptions::default())
}
pub fn for_stream(library: Option<&Path>, stream: &crate::Stream) -> Result<Self> {
Self::shared(library, CompilerOptions::for_stream(stream))
}
pub fn for_target(library: Option<&Path>, target: &crate::Target) -> Result<Self> {
Self::shared(library, CompilerOptions::for_target(target))
}
pub fn shared(library: Option<&Path>, options: CompilerOptions) -> Result<Self> {
static RESOLVED: Mutex<Option<HashMap<SharedCompilerKey, Compiler>>> = Mutex::new(None);
let key = SharedCompilerKey {
path: Self::resolved_library(library)?,
target: options.target.as_str().to_owned(),
workers: options.workers,
module_cache_capacity: options.module_cache_capacity,
processor_mode: options.processor_mode,
};
let mut cache = RESOLVED
.lock()
.map_err(|_| Error::Message("compiler cache poisoned".into()))?;
let cache = cache.get_or_insert_with(HashMap::new);
if let Some(compiler) = cache.get(&key) {
return Ok(compiler.clone());
}
let compiler = Self::with_options(Some(&key.path), options)?;
cache.insert(key, compiler.clone());
Ok(compiler)
}
fn resolved_library(library: Option<&Path>) -> Result<PathBuf> {
match library
.map(Path::to_path_buf)
.or_else(|| std::env::var_os("HRX_LOOM_LIBRARY").map(PathBuf::from))
{
Some(p) => fs::canonicalize(&p)
.map_err(|e| Error::from(e).context(format!("compiler library {}", p.display()))),
None => Ok(fs::canonicalize(bundle::resolve()?.join("libloomc.so"))?),
}
}
pub fn with_options(library: Option<&Path>, options: CompilerOptions) -> Result<Self> {
let path = Self::resolved_library(library)?;
let identity = bundle::file_digest(&path)?;
let native = native::Prepared::open(
&path,
&identity,
options.workers.get(),
&options.target,
options.processor_mode,
)?;
if bundle::file_digest(&path)? != identity {
return Err(Error::Message(
"compiler library changed while loading".into(),
));
}
Ok(Self(Arc::new(Inner {
modules: Mutex::new(HashMap::new()),
native,
target: options.target,
workers: options.workers,
module_cache_capacity: options.module_cache_capacity,
processor_mode: options.processor_mode,
path,
identity,
})))
}
pub fn workers(&self) -> NonZeroUsize {
self.0.workers
}
pub fn compile_all(&self, requests: &[(&Module, &Specialization)]) -> Vec<Result<Artifact>> {
let limit = self.0.workers.get().min(requests.len());
if limit <= 1 {
return requests
.iter()
.map(|(module, spec)| module.compile(spec))
.collect();
}
let next = std::sync::atomic::AtomicUsize::new(0);
let slots: Vec<Mutex<Option<Result<Artifact>>>> =
requests.iter().map(|_| Mutex::new(None)).collect();
std::thread::scope(|scope| {
for _ in 0..limit {
scope.spawn(|| {
loop {
let index = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let Some((module, spec)) = requests.get(index) else {
return;
};
let outcome = module.compile(spec);
*slots[index].lock().unwrap_or_else(|e| e.into_inner()) = Some(outcome);
}
});
}
});
slots
.into_iter()
.map(|slot| {
slot.into_inner()
.unwrap_or_else(|e| e.into_inner())
.expect("every request slot is filled before the scope ends")
})
.collect()
}
pub fn target(&self) -> &crate::Target {
&self.0.target
}
pub fn identity(&self) -> &str {
&self.0.identity
}
pub fn library_path(&self) -> &Path {
&self.0.path
}
pub fn module(&self, source: &str) -> Module {
self.source_module(vec![Source::loom("kernel.loom", source)])
}
pub fn sources(&self, mut sources: Vec<Source>) -> Result<Module> {
if sources.is_empty() {
return Err(Error::Message(
"a compilation module requires source inputs".into(),
));
}
let mut identifiers = std::collections::BTreeSet::new();
for source in &mut sources {
source.validate()?;
let identifier = match source {
Source::Cxx(unit) => {
unit.normalize()?;
&unit.identifier
}
Source::Loom { identifier, .. } => identifier,
};
if !identifiers.insert(identifier.clone()) {
return Err(Error::Message("duplicate source identifier".into()));
}
}
Ok(self.source_module(sources))
}
pub fn import_cxx(&self, source: CxxSource) -> Result<Module> {
self.sources(vec![Source::Cxx(source)])
}
fn source_module(&self, sources: Vec<Source>) -> Module {
let digest = bundle::digest(
&serde_json::to_vec(&sources).expect("source inputs contain only JSON-safe values"),
);
let mut modules = self.0.modules.lock().unwrap_or_else(|e| e.into_inner());
let data = if let Some(data) = modules.get(&digest) {
data.clone()
} else {
let data = Arc::new(ModuleData {
sources,
digest: digest.clone(),
index: Mutex::new(None),
});
if self.0.module_cache_capacity != 0 {
if modules.len() >= self.0.module_cache_capacity {
let victim = modules.keys().next().cloned().unwrap();
modules.remove(&victim);
}
modules.insert(digest, data.clone());
}
data
};
Module {
compiler: self.clone(),
data,
}
}
pub fn trim(&self) {
self.0
.modules
.lock()
.unwrap_or_else(|e| e.into_inner())
.clear();
self.0.native.trim();
}
}
fn summarize(diagnostics: &[Diagnostic]) -> String {
let errors = || diagnostics.iter().filter(|d| d.severity >= Severity::Error);
let Some(first) = errors().next() else {
return diagnostics
.iter()
.map(|d| d.message.as_str())
.collect::<Vec<_>>()
.join("\n");
};
let mut message = if first.line == 0 {
first.message.clone()
} else {
format!("{}:{}: {}", first.line, first.column, first.message)
};
if let Some(hint) = hint(&first.message) {
message.push_str("\nhint: ");
message.push_str(hint);
}
let suppressed = errors().count() - 1;
if suppressed > 0 {
message.push_str(&format!(
"\n({suppressed} later error(s) suppressed as cascade; see Error::Compile diagnostics)"
));
}
message
}
fn hint(message: &str) -> Option<&'static str> {
if message.contains("Low representation contract") && message.contains(".generic.") {
return Some(
"a generic target has no low-asm contract; name a bare architecture \
such as gfx1151 in amdgpu.target<...> to use hand-written asm",
);
}
None
}
#[cfg(test)]
mod diagnostic_tests {
use super::{Diagnostic, Severity, summarize};
fn diagnostic(severity: u32, line: u32, message: &str) -> Diagnostic {
Diagnostic {
severity: Severity::from(severity),
code: String::new(),
message: message.into(),
source: String::new(),
line,
column: 1,
}
}
#[test]
fn the_first_error_leads_and_its_cascade_is_counted_not_printed() {
let mut all = vec![
diagnostic(1, 3, "unused binding"),
diagnostic(
2,
7,
"unknown Low representation contract amdgpu.gfx11.generic.core",
),
];
all.extend((0..12).map(|i| diagnostic(2, 8 + i, "expected expression")));
let rendered = summarize(&all);
assert!(rendered.starts_with("7:1: unknown Low representation contract"));
assert!(rendered.contains("hint: a generic target has no low-asm contract"));
assert!(rendered.contains("12 later error(s) suppressed"));
assert!(!rendered.contains("expected expression"));
let single = summarize(&[diagnostic(2, 4, "type mismatch")]);
assert_eq!(single, "4:1: type mismatch");
}
}
struct ModuleData {
sources: Vec<Source>,
digest: String,
index: Mutex<Option<(Arc<native::Index>, Vec<Diagnostic>)>>,
}
#[derive(Clone)]
pub struct Module {
data: Arc<ModuleData>,
compiler: Compiler,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReportMode {
#[default]
None,
Summary,
Details,
}
#[derive(Clone, Debug, Default)]
pub struct Specialization {
symbol: String,
config: BTreeMap<String, String>,
report: ReportMode,
}
impl Specialization {
pub fn new(symbol: impl Into<String>) -> Self {
Self {
symbol: symbol.into(),
..Self::default()
}
}
#[must_use]
pub fn with_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.config.insert(key.into(), value.into());
self
}
pub fn set_config(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
self.config.insert(key.into(), value.into());
self
}
pub fn set_symbol(&mut self, symbol: impl Into<String>) -> &mut Self {
self.symbol = symbol.into();
self
}
pub fn replace_config(&mut self, config: BTreeMap<String, String>) -> &mut Self {
self.config = config;
self
}
#[must_use]
pub fn with_report(mut self, report: ReportMode) -> Self {
self.report = report;
self
}
pub fn set_report(&mut self, report: ReportMode) -> &mut Self {
self.report = report;
self
}
#[must_use]
pub fn symbol(&self) -> &str {
&self.symbol
}
#[must_use]
pub fn configuration(&self) -> &BTreeMap<String, String> {
&self.config
}
#[must_use]
pub fn report_mode(&self) -> ReportMode {
self.report
}
fn validate(&self) -> Result<()> {
let valid = |s: &str| {
!s.is_empty()
&& s.bytes()
.all(|b| b.is_ascii_alphanumeric() || b"._-".contains(&b))
};
if !valid(&self.symbol) || self.config.keys().any(|k| !valid(k)) {
return Err(Error::Message(
"invalid Loom export or configuration key".into(),
));
}
Ok(())
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
struct Record {
key: String,
compiler: String,
target: String,
#[serde(default)]
processor_mode: ProcessorMode,
symbol: String,
sha256: String,
diagnostics: Vec<Diagnostic>,
report: Option<CompileReport>,
}
#[derive(Clone, Debug)]
pub struct Artifact {
bytes: Arc<[u8]>,
path: PathBuf,
record: Record,
}
impl Artifact {
pub fn bytes(&self) -> &[u8] {
&self.bytes
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn compiler_identity(&self) -> &str {
&self.record.compiler
}
pub fn target(&self) -> &str {
&self.record.target
}
pub fn processor_mode(&self) -> ProcessorMode {
self.record.processor_mode
}
pub fn symbol(&self) -> &str {
&self.record.symbol
}
pub fn diagnostics(&self) -> &[Diagnostic] {
&self.record.diagnostics
}
pub fn report(&self) -> Option<&CompileReport> {
self.record.report.as_ref()
}
}
struct Compiled {
bytes: Arc<[u8]>,
diagnostics: Vec<Diagnostic>,
report: Option<serde_json::Value>,
}
impl Module {
pub fn identity(&self) -> &str {
&self.data.digest
}
pub fn key(&self, spec: &Specialization) -> Result<String> {
spec.validate()?;
let base = bundle::digest(&serde_json::to_vec(&(
"loomc-v2",
self.compiler.identity(),
self.identity(),
&spec.symbol,
self.compiler.0.target.as_str(),
&spec.config,
spec.report,
))?);
if self.compiler.0.processor_mode == ProcessorMode::Default {
Ok(base)
} else {
Ok(bundle::digest(&serde_json::to_vec(&(
"loomc-processor-mode-v1",
base,
self.compiler.0.processor_mode,
))?))
}
}
pub fn compile(&self, spec: &Specialization) -> Result<Artifact> {
let key = self.key(spec)?;
let cache = &bundle::kernel_cache()?;
let dir = cache.join(&key);
if let Some(a) = cached(&dir, &key) {
return Ok(a);
}
bundle::create_cache_dir(cache)?;
let _lock = Lock::acquire(&cache.join(format!("{key}.lock")))?;
if let Some(a) = cached(&dir, &key) {
return Ok(a);
}
let (index, mut diagnostics) = {
let mut slot = self.data.index.lock().unwrap_or_else(|e| e.into_inner());
if slot.is_none() {
let (index, diagnostics) = self.compiler.0.native.index(&self.data.sources)?;
*slot = Some((Arc::new(index), diagnostics));
}
slot.as_ref().unwrap().clone()
};
let compiled = self.compiler.0.native.compile(&index, spec).map_err(|e| {
e.context(format!(
"compiling {} with {}",
spec.symbol,
self.compiler.library_path().display()
))
})?;
diagnostics.extend(compiled.diagnostics);
let record = Record {
key,
compiler: self.compiler.identity().into(),
target: self.compiler.0.target.as_str().into(),
processor_mode: self.compiler.0.processor_mode,
symbol: spec.symbol.clone(),
sha256: bundle::digest(&compiled.bytes),
diagnostics,
report: compiled
.report
.map(|json| {
CompileReport::new(
self.compiler.identity().into(),
self.compiler.0.target.as_str().into(),
self.compiler.0.processor_mode,
json,
)
})
.transpose()?,
};
let staging = tempfile::tempdir_in(cache)?;
let filename = self.compiler.target().artifact_filename();
let output = staging.path().join(filename);
fs::write(&output, &compiled.bytes)?;
let metadata = serde_json::to_vec(&record)?;
fs::write(staging.path().join("artifact.json"), &metadata)?;
fs::write(
staging.path().join("artifact.sha256"),
bundle::digest(&metadata),
)?;
for file in [filename, "artifact.json", "artifact.sha256"] {
fs::File::open(staging.path().join(file))?.sync_all()?;
}
fs::File::open(staging.path())?.sync_all()?;
if dir.exists() {
fs::remove_dir_all(&dir)?;
}
bundle::publish(staging, &dir)?;
fs::File::open(cache)?.sync_all()?;
Ok(Artifact {
bytes: compiled.bytes,
path: dir.join(filename),
record,
})
}
}
fn cached(dir: &Path, key: &str) -> Option<Artifact> {
let metadata = fs::read(dir.join("artifact.json")).ok()?;
if fs::read_to_string(dir.join("artifact.sha256")).ok()? != bundle::digest(&metadata) {
return None;
}
let record: Record = serde_json::from_slice(&metadata).ok()?;
if record.key != key {
return None;
}
if let Ok(file) = fs::File::open(dir.join("artifact.json")) {
let _ = file.set_times(fs::FileTimes::new().set_accessed(std::time::SystemTime::now()));
}
let target = crate::Target::new(&record.target).ok()?;
let path = dir.join(target.artifact_filename());
let bytes = fs::read(&path).ok()?;
if bytes.is_empty() || bundle::digest(&bytes) != record.sha256 {
return None;
}
Some(Artifact {
bytes: bytes.into(),
path,
record,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[ignore = "requires libloomc.so, no GPU"]
fn shared_compilers_respect_all_limits_and_reuse_matching_options() -> Result<()> {
let options = CompilerOptions {
workers: NonZeroUsize::new(2).unwrap(),
module_cache_capacity: 2,
..Default::default()
};
let first = Compiler::shared(None, options.clone())?;
let same = Compiler::shared(None, options.clone())?;
assert!(Arc::ptr_eq(&first.0, &same.0));
let serial_options = CompilerOptions {
workers: NonZeroUsize::new(1).unwrap(),
..options
};
let serial = Compiler::shared(None, serial_options.clone())?;
assert_eq!(serial.workers().get(), 1);
assert!(!Arc::ptr_eq(&first.0, &serial.0));
let uncached = Compiler::shared(
None,
CompilerOptions {
module_cache_capacity: 0,
..serial_options
},
)?;
assert!(!Arc::ptr_eq(&serial.0, &uncached.0));
serial.module("// retained");
uncached.module("// not retained");
assert!(!serial.0.modules.lock().unwrap().is_empty());
assert!(uncached.0.modules.lock().unwrap().is_empty());
Ok(())
}
#[test]
#[ignore = "requires libloomc.so, no GPU"]
fn processor_modes_are_isolated_or_explicitly_rejected() -> Result<()> {
let default = Compiler::shared(None, CompilerOptions::default())?;
let options = CompilerOptions {
processor_mode: ProcessorMode::ComputeUnit,
..Default::default()
};
let cu = match Compiler::shared(None, options.clone()) {
Ok(compiler) => compiler,
Err(error)
if std::env::var_os("HRX_LOOM_LIBRARY").is_some()
|| std::env::var_os("HRX_RUNTIME_DIR").is_some() =>
{
assert!(
error
.to_string()
.contains("AMDGPU profile option extensions are not supported"),
"{error}"
);
return Ok(());
}
Err(error) => return Err(error),
};
let same = Compiler::shared(None, options)?;
assert!(Arc::ptr_eq(&cu.0, &same.0));
assert!(!Arc::ptr_eq(&default.0, &cu.0));
let source = include_str!("../tests/kernels/euler.loom");
let spec = euler_spec();
let wgp = default.module(source).compile(&spec)?;
let unit = cu.module(source).compile(&spec)?;
assert_ne!(wgp.path(), unit.path());
assert_eq!(unit.processor_mode(), ProcessorMode::ComputeUnit);
let cached = cu.module(source).compile(&spec)?;
assert_eq!(cached.processor_mode(), ProcessorMode::ComputeUnit);
assert_eq!(unit.bytes(), cached.bytes());
let differences: Vec<_> = wgp
.bytes()
.iter()
.zip(unit.bytes())
.filter(|(a, b)| a != b)
.collect();
assert_eq!(wgp.bytes().len(), unit.bytes().len());
assert_eq!(differences.len(), 1);
assert_eq!(*differences[0].0 ^ *differences[0].1, 0x20);
Ok(())
}
fn euler_spec() -> Specialization {
let mut spec = Specialization::new("krea2_euler");
spec.set_config("krea2.euler.grid_x", "1");
spec.set_config("krea2.euler.grid_y", "1");
spec
}
#[test]
#[ignore = "requires libloomc.so, no GPU"]
fn reporting_preserves_shared_requests_and_can_be_replaced_from_a_callback() -> Result<()> {
use std::sync::atomic::{AtomicUsize, Ordering};
let kernels = Kernels::new(Compiler::resolve(None)?);
let clone = kernels.clone();
let pending =
kernels.request(include_str!("../tests/kernels/euler.loom"), &euler_spec())?;
let calls = Arc::new(AtomicUsize::new(0));
let reported = calls.clone();
let weak = Arc::downgrade(&kernels.0);
let kernels = kernels.reporting(move |_| {
reported.fetch_add(1, Ordering::SeqCst);
let next = reported.clone();
Kernels(weak.upgrade().unwrap()).reporting(move |_| {
next.fetch_add(10, Ordering::SeqCst);
});
});
assert!(Arc::ptr_eq(&kernels.0, &clone.0));
assert!(Arc::ptr_eq(&kernels.0, &pending.cache));
for cache in [&clone.0, &pending.cache] {
let mut attempts = 0;
let outcome = cache.build_with(1, |artifact| {
attempts += 1;
assert_eq!(artifact.record.key, pending.key);
Err(Error::Message("injected load failure".into()))
});
assert_eq!(outcome.unwrap_err().to_string(), "injected load failure");
assert_eq!(attempts, 1);
}
assert_eq!(calls.load(Ordering::SeqCst), 11);
assert_eq!(kernels.0.locked()?.waiting.len(), 1);
Ok(())
}
#[test]
#[ignore = "requires libloomc.so, no GPU"]
fn requests_during_build_are_deduplicated_and_new_requests_survive() -> Result<()> {
let kernels = Kernels::new(Compiler::resolve(None)?);
let source = include_str!("../tests/kernels/euler.loom");
let spec = euler_spec();
let pending = kernels.request(source, &spec)?;
let mut other_spec = spec.clone();
other_spec.set_config("krea2.euler.grid_x", "2");
let other_key = kernels.compiler().module(source).key(&other_spec)?;
let outcome = kernels.0.build_with(1, |_| {
std::thread::scope(|scope| {
scope
.spawn(|| {
for _ in 0..2 {
let duplicate = kernels.request(source, &spec).unwrap();
assert_eq!(duplicate.key, pending.key);
kernels.request(source, &other_spec).unwrap();
}
})
.join()
.unwrap();
});
assert_eq!(kernels.0.locked().unwrap().waiting.len(), 1);
Err(Error::Message("injected load failure".into()))
});
assert_eq!(outcome.unwrap_err().to_string(), "injected load failure");
{
let state = kernels.0.locked()?;
let waiting: Vec<_> = state.waiting.iter().map(|r| r.0.as_str()).collect();
assert_eq!(waiting, vec![pending.key.as_str(), other_key.as_str()]);
assert!(state.active.is_empty());
}
let mut attempts = Vec::new();
let outcome = kernels.0.build_with(1, |artifact| {
attempts.push(artifact.record.key.clone());
Err(Error::Message("injected retry failure".into()))
});
assert!(outcome.is_err());
assert_eq!(attempts, vec![pending.key, other_key]);
assert_eq!(kernels.0.locked()?.waiting.len(), 2);
Ok(())
}
#[test]
#[ignore = "requires libloomc.so, no GPU"]
fn compile_and_load_failures_remain_requested_for_retry() -> Result<()> {
let kernels = Kernels::new(Compiler::resolve(None)?);
let source = include_str!("../tests/kernels/euler.loom");
let bad = kernels.request(source, &Specialization::new("missing_export"))?;
let good = kernels.request(source, &euler_spec())?;
for _ in 0..2 {
let mut attempted = Vec::new();
let result = kernels.0.build_with(1, |artifact| {
attempted.push(artifact.record.key.clone());
Err(Error::Message("injected transient load failure".into()))
});
let error = result.unwrap_err().to_string();
assert!(
!error.contains("injected transient load failure"),
"the first compile failure must be returned: {error}"
);
assert_eq!(
attempted,
vec![good.key.clone()],
"the later request was attempted"
);
let state = kernels.0.locked()?;
let waiting: Vec<_> = state.waiting.iter().map(|r| r.0.as_str()).collect();
assert_eq!(waiting, vec![bad.key.as_str(), good.key.as_str()]);
}
assert!(good.built().is_none());
Ok(())
}
#[test]
#[ignore = "requires libloomc.so, no GPU"]
fn concurrent_build_waits_for_the_active_batch_and_retries_its_failure() -> Result<()> {
use std::sync::mpsc;
use std::time::Duration;
let kernels = Kernels::new(Compiler::resolve(None)?);
let pending =
kernels.request(include_str!("../tests/kernels/euler.loom"), &euler_spec())?;
let (entered_tx, entered_rx) = mpsc::channel();
let (release_tx, release_rx) = mpsc::channel();
let (started_tx, started_rx) = mpsc::channel();
let (done_tx, done_rx) = mpsc::channel();
std::thread::scope(|scope| {
let first_cache = &kernels.0;
let first = scope.spawn(move || {
first_cache.build_with(1, |_| {
entered_tx.send(()).unwrap();
release_rx.recv_timeout(Duration::from_secs(10)).unwrap();
Err(Error::Message("first load failed".into()))
})
});
entered_rx.recv_timeout(Duration::from_secs(10)).unwrap();
assert!(
kernels.0.locked().unwrap().waiting.is_empty(),
"batch has drained"
);
let second = scope.spawn(|| {
started_tx.send(()).unwrap();
let result = kernels
.0
.build_with(1, |_| Err(Error::Message("retry failed".into())));
done_tx.send(result.map_err(|e| e.to_string())).unwrap();
});
started_rx.recv_timeout(Duration::from_secs(10)).unwrap();
let early = done_rx.recv_timeout(Duration::from_millis(100));
release_tx.send(()).unwrap();
assert!(first.join().unwrap().is_err());
second.join().unwrap();
assert!(
matches!(early, Err(mpsc::RecvTimeoutError::Timeout)),
"build returned before the active batch finished"
);
assert_eq!(done_rx.recv().unwrap(), Err("retry failed".into()));
});
assert!(pending.built().is_none());
assert_eq!(kernels.0.locked()?.waiting.len(), 1);
Ok(())
}
#[test]
#[ignore = "requires libloomc.so, no GPU"]
fn module_cache_eviction_preserves_live_modules() -> Result<()> {
let compiler = Compiler::with_options(
None,
CompilerOptions {
module_cache_capacity: 2,
..Default::default()
},
)?;
let first = compiler.module(include_str!("../tests/kernels/euler.loom"));
let first_id = first.identity().to_owned();
for i in 0..8 {
compiler.module(&format!("// source {i}"));
assert!(compiler.0.modules.lock().unwrap().len() <= 2);
}
compiler.trim();
assert!(compiler.0.modules.lock().unwrap().is_empty());
assert_eq!(first.identity(), first_id);
let mut spec = Specialization::new("krea2_euler");
spec.set_config("krea2.euler.grid_x", "1");
spec.set_config("krea2.euler.grid_y", "1");
assert!(!first.compile(&spec)?.bytes().is_empty());
Ok(())
}
}
#[derive(Clone)]
pub struct Kernels(Arc<Cache>);
struct Cache {
compiler: Compiler,
state: Mutex<Loaded>,
build: Mutex<()>,
report: Mutex<Option<Reporter>>,
}
type Reporter = Arc<dyn Fn(&Artifact) + Send + Sync>;
#[derive(Default)]
struct Loaded {
device: Option<usize>,
ready: HashMap<String, crate::Kernel>,
waiting: Vec<(String, String, Specialization)>,
active: HashSet<String>,
}
pub struct Pending {
key: String,
cache: Arc<Cache>,
cell: std::sync::OnceLock<crate::Kernel>,
}
impl Clone for Pending {
fn clone(&self) -> Self {
Self {
key: self.key.clone(),
cache: self.cache.clone(),
cell: self.cell.clone(),
}
}
}
impl Pending {
pub fn built(&self) -> Option<&crate::Kernel> {
if self.cell.get().is_none() {
let found = self.cache.locked().ok()?.ready.get(&self.key).cloned()?;
let _ = self.cell.set(found);
}
self.cell.get()
}
pub unsafe fn resolve(&self, stream: &crate::Stream) -> Result<&crate::Kernel> {
if self.built().is_some() {
return Ok(self.cell.get().expect("just built"));
}
let outcome = unsafe { self.cache.build(stream) };
match self.built() {
Some(_) => Ok(self.cell.get().expect("just built")),
None => Err(outcome
.err()
.unwrap_or_else(|| Error::Message("kernel was never requested".into()))),
}
}
}
impl Kernels {
pub fn keyed<K>(self) -> KeyedKernels<K> {
KeyedKernels {
kernels: self,
keys: Mutex::new(HashMap::new()),
}
}
pub fn new(compiler: Compiler) -> Self {
Self(Arc::new(Cache {
compiler,
state: Mutex::new(Loaded::default()),
build: Mutex::new(()),
report: Mutex::new(None),
}))
}
pub fn reporting(self, report: impl Fn(&Artifact) + Send + Sync + 'static) -> Self {
let previous = self
.0
.report
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.replace(Arc::new(report));
drop(previous);
self
}
pub fn compiler(&self) -> &Compiler {
&self.0.compiler
}
pub fn len(&self) -> usize {
self.0.locked().map_or(0, |state| state.ready.len())
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub unsafe fn get(
&self,
stream: &crate::Stream,
source: &str,
spec: &Specialization,
) -> Result<crate::Kernel> {
let module = self.0.compiler.module(source);
let key = module.key(spec)?;
unsafe { self.get_module(stream, &module, key, spec) }
}
unsafe fn get_module(
&self,
stream: &crate::Stream,
module: &Module,
key: String,
spec: &Specialization,
) -> Result<crate::Kernel> {
{
let mut state = self.0.locked()?;
state.claim(stream)?;
if let Some(kernel) = state.ready.get(&key) {
return Ok(kernel.clone());
}
}
let _build = self
.0
.build
.lock()
.map_err(|_| Error::Message("kernel build poisoned".into()))?;
if let Some(kernel) = self.0.locked()?.ready.get(&key) {
return Ok(kernel.clone());
}
let artifact = module.compile(spec)?;
self.0.reported(&artifact);
let kernel = unsafe { stream.load_artifact(&artifact) }?;
let mut state = self.0.locked()?;
state.waiting.retain(|(waiting, ..)| *waiting != key);
state.ready.insert(key, kernel.clone());
Ok(kernel)
}
pub fn request(&self, source: &str, spec: &Specialization) -> Result<Pending> {
let key = self.0.compiler.module(source).key(spec)?;
let mut state = self.0.locked()?;
let known = state.ready.contains_key(&key)
|| state.active.contains(&key)
|| state.waiting.iter().any(|(waiting, ..)| *waiting == key);
if !known {
state
.waiting
.push((key.clone(), source.to_owned(), spec.clone()));
}
drop(state);
Ok(Pending {
key,
cache: self.0.clone(),
cell: std::sync::OnceLock::new(),
})
}
pub unsafe fn build(&self, stream: &crate::Stream) -> Result<()> {
unsafe { self.0.build(stream) }
}
}
impl Cache {
fn locked(&self) -> Result<std::sync::MutexGuard<'_, Loaded>> {
self.state
.lock()
.map_err(|_| Error::Message("kernel cache poisoned".into()))
}
fn reported(&self, artifact: &Artifact) {
let report = self
.report
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clone();
if let Some(report) = report {
report(artifact);
}
}
unsafe fn build(&self, stream: &crate::Stream) -> Result<()> {
self.build_with(stream.device_id(), |artifact| {
unsafe { stream.load_artifact(artifact) }
})
}
fn build_with(
&self,
device: usize,
mut load: impl FnMut(&Artifact) -> Result<crate::Kernel>,
) -> Result<()> {
let _build = self
.build
.lock()
.map_err(|_| Error::Message("kernel build poisoned".into()))?;
let waiting = {
let mut state = self.locked()?;
state.claim_device(device)?;
let waiting = std::mem::take(&mut state.waiting);
state
.active
.extend(waiting.iter().map(|(key, ..)| key.clone()));
waiting
};
if waiting.is_empty() {
return Ok(());
}
let modules: Vec<Module> = waiting
.iter()
.map(|(_, source, _)| self.compiler.module(source))
.collect();
let requests: Vec<(&Module, &Specialization)> = modules
.iter()
.zip(waiting.iter().map(|(_, _, spec)| spec))
.collect();
let built = self.compiler.compile_all(&requests);
let mut failure = None;
let mut unbuilt = Vec::new();
for (request, outcome) in waiting.into_iter().zip(built) {
let loaded = outcome.and_then(|artifact| {
self.reported(&artifact);
load(&artifact)
});
match loaded {
Ok(kernel) => {
self.locked()?.ready.insert(request.0.clone(), kernel);
}
Err(error) => {
unbuilt.push(request);
failure = failure.or(Some(error));
}
}
}
{
let mut state = self.locked()?;
unbuilt.append(&mut state.waiting);
state.waiting = unbuilt;
state.active.clear();
}
match failure {
Some(error) => Err(error),
None => Ok(()),
}
}
}
pub struct KeyedKernels<K> {
kernels: Kernels,
keys: Mutex<HashMap<K, String>>,
}
impl<K: Eq + std::hash::Hash> KeyedKernels<K> {
pub fn kernels(&self) -> &Kernels {
&self.kernels
}
pub unsafe fn request_or_insert_with<S: AsRef<str>>(
&self,
stream: &crate::Stream,
key: K,
request: impl FnOnce(&K) -> Result<(S, Specialization)>,
) -> Result<Pending> {
let mut keys = self
.keys
.lock()
.map_err(|_| Error::Message("kernel key index poisoned".into()))?;
self.kernels.0.locked()?.claim(stream)?;
if let Some(artifact) = keys.get(&key) {
return Ok(Pending {
key: artifact.clone(),
cache: self.kernels.0.clone(),
cell: std::sync::OnceLock::new(),
});
}
let (source, spec) = request(&key)?;
let pending = self.kernels.request(source.as_ref(), &spec)?;
keys.insert(key, pending.key.clone());
Ok(pending)
}
pub unsafe fn get_or_insert_with<'s>(
&self,
stream: &crate::Stream,
key: K,
request: impl FnOnce(&K) -> Result<(&'s str, Specialization)>,
) -> Result<crate::Kernel> {
let mut keys = self
.keys
.lock()
.map_err(|_| Error::Message("kernel key index poisoned".into()))?;
{
let mut state = self.kernels.0.locked()?;
state.claim(stream)?;
if let Some(artifact) = keys.get(&key) {
if let Some(kernel) = state.ready.get(artifact) {
return Ok(kernel.clone());
}
let pending = Pending {
key: artifact.clone(),
cache: self.kernels.0.clone(),
cell: std::sync::OnceLock::new(),
};
drop(state);
drop(keys);
return unsafe { pending.resolve(stream) }.cloned();
}
}
let (source, spec) = request(&key)?;
let module = self.kernels.compiler().module(source);
let artifact = module.key(&spec)?;
let kernel = unsafe {
self.kernels
.get_module(stream, &module, artifact.clone(), &spec)
}?;
keys.insert(key, artifact);
Ok(kernel)
}
}
impl Loaded {
fn claim(&mut self, stream: &crate::Stream) -> Result<()> {
self.claim_device(stream.device_id())
}
fn claim_device(&mut self, device: usize) -> Result<()> {
match self.device {
Some(owner) if owner != device => Err(Error::Message(
"a kernel cache serves one device: an executable loaded on another is not \
dispatchable here"
.into(),
)),
Some(_) => Ok(()),
None => {
self.device = Some(device);
Ok(())
}
}
}
}
impl std::fmt::Debug for Kernels {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Kernels")
.field("loaded", &self.len())
.finish_non_exhaustive()
}
}