1#[allow(
3 non_camel_case_types,
4 non_snake_case,
5 non_upper_case_globals,
6 dead_code,
7 missing_docs,
8 unsafe_op_in_unsafe_fn,
9 clippy::all
10)]
11mod ffi;
12mod native;
13mod report;
14mod source;
15use crate::{
16 Error, Result,
17 bundle::{self, Lock},
18};
19pub use report::{
20 CompileReport, EntryChange, EntryResources, Expansion, ReportGuidance, ReportSuggestion,
21 WaitCounts, WaitReason, WaitReasonChange,
22};
23use serde::{Deserialize, Serialize};
24pub use source::{CxxFpModel, CxxPreamble, CxxSource, CxxStandard, Source};
25use std::{
26 collections::{BTreeMap, HashMap, HashSet},
27 fs,
28 num::NonZeroUsize,
29 path::{Path, PathBuf},
30 sync::{Arc, Mutex},
31};
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
37#[serde(rename_all = "lowercase")]
38pub enum Severity {
39 Note,
41 Warning,
43 Error,
45}
46impl From<u32> for Severity {
47 fn from(native: u32) -> Self {
50 match native {
51 0 => Self::Note,
52 1 => Self::Warning,
53 _ => Self::Error,
54 }
55 }
56}
57
58#[derive(Clone, Debug, Serialize, Deserialize)]
60pub struct Diagnostic {
61 pub severity: Severity,
63 pub code: String,
65 pub message: String,
67 pub formatted_text: String,
69 pub parameters: Vec<DiagnosticParameter>,
71 #[serde(default)]
73 pub source: String,
74 pub line: u32,
76 pub column: u32,
78 #[serde(default)]
80 pub related_locations: Vec<RelatedLocation>,
81 #[serde(default)]
83 pub related_location_omitted_count: usize,
84}
85
86#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
88pub struct DiagnosticParameter {
89 pub name: String,
91 pub value: String,
93}
94
95#[derive(Clone, Debug, Serialize, Deserialize)]
97pub struct RelatedLocation {
98 pub label: String,
100 pub source: String,
102 pub line: u32,
104 pub column: u32,
106 pub end_line: u32,
108 pub end_column: u32,
110}
111#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq, Serialize, Deserialize)]
113#[serde(rename_all = "snake_case")]
114pub enum ProcessorMode {
115 #[default]
117 Default,
118 ComputeUnit,
120 WorkgroupProcessor,
122}
123
124#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
126pub struct SanitizerChecks {
127 pub access: bool,
129 pub value: bool,
131 pub operation: bool,
133 pub race: bool,
135}
136impl SanitizerChecks {
137 pub(super) fn bits(self) -> u64 {
138 u64::from(self.access)
139 | (u64::from(self.value) << 1)
140 | (u64::from(self.operation) << 2)
141 | (u64::from(self.race) << 3)
142 }
143}
144#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
146pub enum SanitizerReporting {
147 #[default]
149 Default,
150 Trap,
152 ReportOnly,
154}
155#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
158pub struct SanitizerOptions {
159 pub checks: SanitizerChecks,
161 pub reporting: SanitizerReporting,
163}
164
165#[derive(Clone, Debug)]
167pub struct CompilerOptions {
168 pub workers: NonZeroUsize,
170 pub target: crate::Target,
172 pub module_cache_capacity: usize,
175 pub processor_mode: ProcessorMode,
177 pub sanitizer: SanitizerOptions,
179}
180impl Default for CompilerOptions {
181 fn default() -> Self {
182 Self {
183 target: crate::Target::default(),
184 module_cache_capacity: 64,
185 processor_mode: ProcessorMode::Default,
186 sanitizer: SanitizerOptions::default(),
187 workers: NonZeroUsize::new(
188 std::thread::available_parallelism()
189 .map_or(1, NonZeroUsize::get)
190 .min(4),
191 )
192 .unwrap(),
193 }
194 }
195}
196impl CompilerOptions {
197 #[must_use]
199 pub fn for_target(target: &crate::Target) -> Self {
200 Self {
201 target: target.clone(),
202 ..Self::default()
203 }
204 }
205 #[must_use]
207 pub fn for_stream(stream: &crate::Stream) -> Self {
208 Self::for_target(stream.target())
209 }
210}
211struct Inner {
212 modules: Mutex<HashMap<String, Arc<ModuleData>>>,
214 native: native::Prepared,
215 target: crate::Target,
216 workers: NonZeroUsize,
217 module_cache_capacity: usize,
218 processor_mode: ProcessorMode,
219 sanitizer: SanitizerOptions,
220 path: PathBuf,
221 identity: String,
222}
223
224#[derive(Hash, PartialEq, Eq)]
225struct SharedCompilerKey {
226 path: PathBuf,
227 target: String,
228 workers: NonZeroUsize,
229 module_cache_capacity: usize,
230 processor_mode: ProcessorMode,
231 sanitizer: SanitizerOptions,
232}
233#[derive(Clone)]
235pub struct Compiler(Arc<Inner>);
236impl std::fmt::Debug for Compiler {
237 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
238 f.debug_struct("Compiler")
239 .field("library", &self.0.path)
240 .field("identity", &self.0.identity)
241 .finish_non_exhaustive()
242 }
243}
244impl Compiler {
245 pub fn resolve(library: Option<&Path>) -> Result<Self> {
250 Self::with_options(library, CompilerOptions::default())
251 }
252 pub fn for_stream(library: Option<&Path>, stream: &crate::Stream) -> Result<Self> {
254 Self::shared(library, CompilerOptions::for_stream(stream))
255 }
256 pub fn for_target(library: Option<&Path>, target: &crate::Target) -> Result<Self> {
258 Self::shared(library, CompilerOptions::for_target(target))
259 }
260 pub fn shared(library: Option<&Path>, options: CompilerOptions) -> Result<Self> {
269 static RESOLVED: Mutex<Option<HashMap<SharedCompilerKey, Compiler>>> = Mutex::new(None);
270 let key = SharedCompilerKey {
271 path: Self::resolved_library(library)?,
272 target: options.target.as_str().to_owned(),
273 workers: options.workers,
274 module_cache_capacity: options.module_cache_capacity,
275 processor_mode: options.processor_mode,
276 sanitizer: options.sanitizer,
277 };
278 let mut cache = RESOLVED
279 .lock()
280 .map_err(|_| Error::Message("compiler cache poisoned".into()))?;
281 let cache = cache.get_or_insert_with(HashMap::new);
282 if let Some(compiler) = cache.get(&key) {
283 return Ok(compiler.clone());
284 }
285 let compiler = Self::with_options(Some(&key.path), options)?;
286 cache.insert(key, compiler.clone());
287 Ok(compiler)
288 }
289
290 fn resolved_library(library: Option<&Path>) -> Result<PathBuf> {
292 match library
293 .map(Path::to_path_buf)
294 .or_else(|| std::env::var_os("HRX_LOOM_LIBRARY").map(PathBuf::from))
295 {
296 Some(p) => fs::canonicalize(&p)
297 .map_err(|e| Error::from(e).context(format!("compiler library {}", p.display()))),
298 None => Ok(fs::canonicalize(bundle::resolve()?.join("libloomc.so"))?),
299 }
300 }
301
302 pub fn with_options(library: Option<&Path>, options: CompilerOptions) -> Result<Self> {
304 let path = Self::resolved_library(library)?;
305 let identity = bundle::file_digest(&path)?;
306 let native = native::Prepared::open(
307 &path,
308 &identity,
309 options.workers.get(),
310 &options.target,
311 options.processor_mode,
312 options.sanitizer,
313 )?;
314 if bundle::file_digest(&path)? != identity {
315 return Err(Error::Message(
316 "compiler library changed while loading".into(),
317 ));
318 }
319 Ok(Self(Arc::new(Inner {
320 modules: Mutex::new(HashMap::new()),
321 native,
322 target: options.target,
323 workers: options.workers,
324 module_cache_capacity: options.module_cache_capacity,
325 processor_mode: options.processor_mode,
326 sanitizer: options.sanitizer,
327 path,
328 identity,
329 })))
330 }
331 pub fn workers(&self) -> NonZeroUsize {
333 self.0.workers
334 }
335
336 pub fn compile_all(&self, requests: &[(&Module, &Specialization)]) -> Vec<Result<Artifact>> {
344 let limit = self.0.workers.get().min(requests.len());
345 if limit <= 1 {
346 return requests
347 .iter()
348 .map(|(module, spec)| module.compile(spec))
349 .collect();
350 }
351 let next = std::sync::atomic::AtomicUsize::new(0);
352 let slots: Vec<Mutex<Option<Result<Artifact>>>> =
353 requests.iter().map(|_| Mutex::new(None)).collect();
354 std::thread::scope(|scope| {
355 for _ in 0..limit {
356 scope.spawn(|| {
357 loop {
358 let index = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
359 let Some((module, spec)) = requests.get(index) else {
360 return;
361 };
362 let outcome = module.compile(spec);
363 *slots[index].lock().unwrap_or_else(|e| e.into_inner()) = Some(outcome);
364 }
365 });
366 }
367 });
368 slots
369 .into_iter()
370 .map(|slot| {
371 slot.into_inner()
372 .unwrap_or_else(|e| e.into_inner())
373 .expect("every request slot is filled before the scope ends")
374 })
375 .collect()
376 }
377
378 pub fn target(&self) -> &crate::Target {
380 &self.0.target
381 }
382 pub fn identity(&self) -> &str {
384 &self.0.identity
385 }
386 pub fn library_path(&self) -> &Path {
388 &self.0.path
389 }
390 pub fn module(&self, source: &str) -> Module {
393 self.source_module(vec![Source::loom("kernel.loom", source)])
394 }
395 pub fn sources(&self, mut sources: Vec<Source>) -> Result<Module> {
398 if sources.is_empty() {
399 return Err(Error::Message(
400 "a compilation module requires source inputs".into(),
401 ));
402 }
403 let mut identifiers = std::collections::BTreeSet::new();
404 for source in &mut sources {
405 source.validate()?;
406 let identifier = match source {
407 Source::Cxx(unit) => {
408 unit.normalize()?;
409 &unit.identifier
410 }
411 Source::Loom { identifier, .. } => identifier,
412 };
413 if !identifiers.insert(identifier.clone()) {
414 return Err(Error::Message("duplicate source identifier".into()));
415 }
416 }
417 Ok(self.source_module(sources))
418 }
419 pub fn import_cxx(&self, source: CxxSource) -> Result<Module> {
421 self.sources(vec![Source::Cxx(source)])
422 }
423 pub fn cxx_preamble(&self, mut source: CxxSource) -> Result<CxxPreamble> {
427 source.normalize()?;
428 self.0.native.preamble(&source)
429 }
430 fn source_module(&self, sources: Vec<Source>) -> Module {
431 let digest = bundle::digest(
432 &serde_json::to_vec(&sources).expect("source inputs contain only JSON-safe values"),
433 );
434 let mut modules = self.0.modules.lock().unwrap_or_else(|e| e.into_inner());
435 let data = if let Some(data) = modules.get(&digest) {
436 data.clone()
437 } else {
438 let data = Arc::new(ModuleData {
439 sources,
440 digest: digest.clone(),
441 index: Mutex::new(None),
442 });
443 if self.0.module_cache_capacity != 0 {
444 if modules.len() >= self.0.module_cache_capacity {
445 let victim = modules.keys().next().cloned().unwrap();
446 modules.remove(&victim);
447 }
448 modules.insert(digest, data.clone());
449 }
450 data
451 };
452 Module {
453 compiler: self.clone(),
454 data,
455 }
456 }
457 pub fn trim(&self) {
460 self.0
461 .modules
462 .lock()
463 .unwrap_or_else(|e| e.into_inner())
464 .clear();
465 self.0.native.trim();
466 }
467}
468fn summarize(diagnostics: &[Diagnostic]) -> String {
475 let errors = || diagnostics.iter().filter(|d| d.severity >= Severity::Error);
476 let Some(first) = errors().next() else {
477 return diagnostics
478 .iter()
479 .map(|d| d.message.as_str())
480 .collect::<Vec<_>>()
481 .join("\n");
482 };
483 let mut message = if first.line == 0 {
484 first.message.clone()
485 } else {
486 format!("{}:{}: {}", first.line, first.column, first.message)
487 };
488 if let Some(hint) = hint(&first.message) {
489 message.push_str("\nhint: ");
490 message.push_str(hint);
491 }
492 for related in &first.related_locations {
493 message.push_str(&format!(
494 "\n{}:{}:{}: {}",
495 related.source, related.line, related.column, related.label
496 ));
497 }
498 if first.related_location_omitted_count != 0 {
499 message.push_str(&format!(
500 "\n({} additional related location(s) omitted)",
501 first.related_location_omitted_count
502 ));
503 }
504 let suppressed = errors().count() - 1;
505 if suppressed > 0 {
506 message.push_str(&format!(
507 "\n({suppressed} later error(s) suppressed as cascade; see Error::Compile diagnostics)"
508 ));
509 }
510 message
511}
512
513fn hint(message: &str) -> Option<&'static str> {
515 if message.contains("Low representation contract") && message.contains(".generic.") {
516 return Some(
517 "a generic target has no low-asm contract; name a bare architecture \
518 such as gfx1151 in amdgpu.target<...> to use hand-written asm",
519 );
520 }
521 None
522}
523
524#[cfg(test)]
525mod diagnostic_tests {
526 use super::{Diagnostic, Severity, summarize};
527 fn diagnostic(severity: u32, line: u32, message: &str) -> Diagnostic {
528 Diagnostic {
529 severity: Severity::from(severity),
530 code: String::new(),
531 message: message.into(),
532 formatted_text: message.into(),
533 parameters: Vec::new(),
534 source: String::new(),
535 line,
536 column: 1,
537 related_locations: Vec::new(),
538 related_location_omitted_count: 0,
539 }
540 }
541 #[test]
542 fn the_first_error_leads_and_its_cascade_is_counted_not_printed() {
543 let mut all = vec![
544 diagnostic(1, 3, "unused binding"),
545 diagnostic(
546 2,
547 7,
548 "unknown Low representation contract amdgpu.gfx11.generic.core",
549 ),
550 ];
551 all.extend((0..12).map(|i| diagnostic(2, 8 + i, "expected expression")));
552 let rendered = summarize(&all);
553 assert!(rendered.starts_with("7:1: unknown Low representation contract"));
554 assert!(rendered.contains("hint: a generic target has no low-asm contract"));
555 assert!(rendered.contains("12 later error(s) suppressed"));
556 assert!(!rendered.contains("expected expression"));
557 let single = summarize(&[diagnostic(2, 4, "type mismatch")]);
559 assert_eq!(single, "4:1: type mismatch");
560 }
561}
562
563struct ModuleData {
564 sources: Vec<Source>,
565 digest: String,
566 index: Mutex<Option<(Arc<native::Index>, Vec<Diagnostic>)>>,
567}
568#[derive(Clone)]
570pub struct Module {
571 data: Arc<ModuleData>,
573 compiler: Compiler,
574}
575#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
577#[serde(rename_all = "snake_case")]
578pub enum ReportMode {
579 #[default]
581 None,
582 Summary,
584 Details,
586}
587
588#[derive(Clone, Copy, Debug, Default)]
590pub enum TraceFormat {
591 Text,
593 #[default]
595 JsonLines,
596}
597#[derive(Clone, Debug)]
599pub struct TraceOptions {
600 pub format: TraceFormat,
602 pub max_bytes: usize,
604 pub before: Vec<String>,
606 pub after: Vec<String>,
609}
610impl Default for TraceOptions {
611 fn default() -> Self {
612 Self {
613 format: TraceFormat::JsonLines,
614 max_bytes: 16 * 1024 * 1024,
615 before: Vec::new(),
616 after: Vec::new(),
617 }
618 }
619}
620
621#[derive(Clone, Debug, Default)]
623pub struct Specialization {
624 symbol: String,
626 config: BTreeMap<String, String>,
628 report: ReportMode,
630 listing: bool,
631}
632impl Specialization {
633 #[must_use]
636 pub fn with_listing(mut self, enabled: bool) -> Self {
637 self.listing = enabled;
638 self
639 }
640 pub fn new(symbol: impl Into<String>) -> Self {
642 Self {
643 symbol: symbol.into(),
644 ..Self::default()
645 }
646 }
647 #[must_use]
649 pub fn with_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
650 self.config.insert(key.into(), value.into());
651 self
652 }
653 pub fn set_config(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
655 self.config.insert(key.into(), value.into());
656 self
657 }
658 pub fn set_symbol(&mut self, symbol: impl Into<String>) -> &mut Self {
660 self.symbol = symbol.into();
661 self
662 }
663 pub fn replace_config(&mut self, config: BTreeMap<String, String>) -> &mut Self {
665 self.config = config;
666 self
667 }
668 #[must_use]
670 pub fn with_report(mut self, report: ReportMode) -> Self {
671 self.report = report;
672 self
673 }
674 pub fn set_report(&mut self, report: ReportMode) -> &mut Self {
676 self.report = report;
677 self
678 }
679 #[must_use]
681 pub fn symbol(&self) -> &str {
682 &self.symbol
683 }
684 #[must_use]
686 pub fn configuration(&self) -> &BTreeMap<String, String> {
687 &self.config
688 }
689 #[must_use]
691 pub fn report_mode(&self) -> ReportMode {
692 self.report
693 }
694 fn validate(&self) -> Result<()> {
695 let valid = |s: &str| {
696 !s.is_empty()
697 && s.bytes()
698 .all(|b| b.is_ascii_alphanumeric() || b"._-".contains(&b))
699 };
700 if !valid(&self.symbol) || self.config.keys().any(|k| !valid(k)) {
701 return Err(Error::Message(
702 "invalid Loom export or configuration key".into(),
703 ));
704 }
705 Ok(())
706 }
707}
708#[derive(Clone, Debug, Serialize, Deserialize)]
709struct Record {
710 key: String,
711 compiler: String,
712 target: String,
713 #[serde(default)]
714 processor_mode: ProcessorMode,
715 symbol: String,
716 sha256: String,
717 diagnostics: Vec<Diagnostic>,
718 report: Option<CompileReport>,
719 manifest: Option<serde_json::Value>,
720 launch_sha256: Option<String>,
721 #[serde(default)]
722 listing_sha256: Option<String>,
723}
724#[derive(Clone, Debug)]
726pub struct Artifact {
727 bytes: Arc<[u8]>,
728 launch_bytes: Option<Arc<[u8]>>,
729 listing: Option<Arc<str>>,
730 launch_loader: native::LaunchLoader,
731 path: PathBuf,
732 record: Record,
733}
734impl Artifact {
735 pub fn listing(&self) -> Option<&str> {
737 self.listing.as_deref()
738 }
739 pub fn bytes(&self) -> &[u8] {
741 &self.bytes
742 }
743 pub fn manifest(&self) -> Option<&serde_json::Value> {
745 self.record.manifest.as_ref()
746 }
747 pub fn launch_bytes(&self) -> Option<&[u8]> {
749 self.launch_bytes.as_deref()
750 }
751 pub fn launch_program(&self, export: &str) -> Result<LaunchProgram> {
754 self.launch_loader.load(
755 self.launch_bytes.as_deref().ok_or_else(|| {
756 Error::Unsupported("this target has no host launch program".into())
757 })?,
758 export,
759 )
760 }
761 pub fn path(&self) -> &Path {
763 &self.path
764 }
765 pub fn compiler_identity(&self) -> &str {
767 &self.record.compiler
768 }
769 pub fn target(&self) -> &str {
771 &self.record.target
772 }
773 pub fn processor_mode(&self) -> ProcessorMode {
775 self.record.processor_mode
776 }
777 pub fn symbol(&self) -> &str {
779 &self.record.symbol
780 }
781 pub fn diagnostics(&self) -> &[Diagnostic] {
783 &self.record.diagnostics
784 }
785 pub fn report(&self) -> Option<&CompileReport> {
787 self.record.report.as_ref()
788 }
789}
790#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
792pub struct LaunchConfig {
793 pub workgroup_count: [u32; 3],
795 pub workgroup_size: [u32; 3],
797 pub workgroup_cluster_size: [u32; 3],
799 pub subgroup_size: u32,
801 pub workgroup_storage_bytes: u64,
803}
804
805pub struct LaunchProgram(native::LaunchProgram);
808impl LaunchProgram {
809 pub fn evaluate(&mut self, workload_argument_bits: &[u64]) -> Result<LaunchConfig> {
812 self.0.evaluate(workload_argument_bits)
813 }
814}
815
816struct Compiled {
817 bytes: Arc<[u8]>,
818 diagnostics: Vec<Diagnostic>,
819 report: Option<serde_json::Value>,
820 manifest: Option<serde_json::Value>,
821 launch_bytes: Option<Arc<[u8]>>,
822 listing: Option<Arc<str>>,
823}
824impl Module {
825 pub fn identity(&self) -> &str {
827 &self.data.digest
828 }
829 pub fn key(&self, spec: &Specialization) -> Result<String> {
831 spec.validate()?;
832 let base = bundle::digest(&serde_json::to_vec(&(
833 "loomc-v5-sanitizers",
834 self.compiler.0.sanitizer,
835 self.compiler.identity(),
836 self.identity(),
837 &spec.symbol,
838 self.compiler.0.target.as_str(),
839 &spec.config,
840 spec.report,
841 ))?);
842 let base = if spec.listing {
843 bundle::digest(&serde_json::to_vec(&("loomc-target-listing-v1", base))?)
844 } else {
845 base
846 };
847 if self.compiler.0.processor_mode == ProcessorMode::Default {
850 Ok(base)
851 } else {
852 Ok(bundle::digest(&serde_json::to_vec(&(
853 "loomc-processor-mode-v1",
854 base,
855 self.compiler.0.processor_mode,
856 ))?))
857 }
858 }
859 pub fn compile(&self, spec: &Specialization) -> Result<Artifact> {
862 self.compile_impl(spec, None)
863 }
864 pub fn compile_traced(
868 &self,
869 spec: &Specialization,
870 options: &TraceOptions,
871 writer: &mut dyn std::io::Write,
872 ) -> Result<Artifact> {
873 self.compile_impl(spec, Some((options, writer)))
874 }
875 fn compile_impl(
876 &self,
877 spec: &Specialization,
878 trace: Option<(&TraceOptions, &mut dyn std::io::Write)>,
879 ) -> Result<Artifact> {
880 let key = self.key(spec)?;
881 let cache = &bundle::kernel_cache()?;
882 let dir = cache.join(&key);
883 if trace.is_none()
884 && let Some(a) = cached(&dir, &key, &self.compiler.0.native)
885 {
886 return Ok(a);
887 }
888 bundle::create_cache_dir(cache)?;
889 let _lock = Lock::acquire(&cache.join(format!("{key}.lock")))?;
890 if trace.is_none()
891 && let Some(a) = cached(&dir, &key, &self.compiler.0.native)
892 {
893 return Ok(a);
894 }
895 let (index, mut diagnostics) = {
896 let mut slot = self.data.index.lock().unwrap_or_else(|e| e.into_inner());
897 if slot.is_none() {
898 let (index, diagnostics) = self.compiler.0.native.index(&self.data.sources)?;
899 *slot = Some((Arc::new(index), diagnostics));
900 }
901 slot.as_ref().unwrap().clone()
902 };
903 let compiled = self
904 .compiler
905 .0
906 .native
907 .compile(&index, spec, trace)
908 .map_err(|e| {
909 e.context(format!(
910 "compiling {} with {}",
911 spec.symbol,
912 self.compiler.library_path().display()
913 ))
914 })?;
915 diagnostics.extend(compiled.diagnostics);
916 let record = Record {
917 key,
918 compiler: self.compiler.identity().into(),
919 target: self.compiler.0.target.as_str().into(),
920 processor_mode: self.compiler.0.processor_mode,
921 symbol: spec.symbol.clone(),
922 sha256: bundle::digest(&compiled.bytes),
923 launch_sha256: compiled.launch_bytes.as_deref().map(bundle::digest),
924 listing_sha256: compiled
925 .listing
926 .as_deref()
927 .map(|s| bundle::digest(s.as_bytes())),
928 manifest: compiled.manifest,
929 diagnostics,
930 report: compiled
931 .report
932 .map(|json| {
933 CompileReport::new(
934 self.compiler.identity().into(),
935 self.compiler.0.target.as_str().into(),
936 self.compiler.0.processor_mode,
937 json,
938 )
939 })
940 .transpose()?,
941 };
942 let staging = tempfile::tempdir_in(cache)?;
943 let filename = self.compiler.target().artifact_filename();
944 let output = staging.path().join(filename);
945 fs::write(&output, &compiled.bytes)?;
946 if let Some(listing) = &compiled.listing {
947 fs::write(staging.path().join("target.listing"), listing.as_bytes())?;
948 fs::File::open(staging.path().join("target.listing"))?.sync_all()?;
949 }
950 if let Some(bytes) = &compiled.launch_bytes {
951 fs::write(staging.path().join("launch.loombc"), bytes)?;
952 fs::File::open(staging.path().join("launch.loombc"))?.sync_all()?;
953 }
954 let metadata = serde_json::to_vec(&record)?;
955 fs::write(staging.path().join("artifact.json"), &metadata)?;
956 fs::write(
957 staging.path().join("artifact.sha256"),
958 bundle::digest(&metadata),
959 )?;
960 for file in [filename, "artifact.json", "artifact.sha256"] {
961 fs::File::open(staging.path().join(file))?.sync_all()?;
962 }
963 fs::File::open(staging.path())?.sync_all()?;
964 if dir.exists() {
965 fs::remove_dir_all(&dir)?;
966 }
967 bundle::publish(staging, &dir)?;
968 fs::File::open(cache)?.sync_all()?;
969 Ok(Artifact {
970 bytes: compiled.bytes,
971 listing: compiled.listing,
972 launch_bytes: compiled.launch_bytes,
973 launch_loader: self.compiler.0.native.launch_loader(),
974 path: dir.join(filename),
975 record,
976 })
977 }
978}
979fn cached(dir: &Path, key: &str, native: &native::Prepared) -> Option<Artifact> {
980 let metadata = fs::read(dir.join("artifact.json")).ok()?;
981 if fs::read_to_string(dir.join("artifact.sha256")).ok()? != bundle::digest(&metadata) {
982 return None;
983 }
984 let record: Record = serde_json::from_slice(&metadata).ok()?;
985 if record.key != key {
986 return None;
987 }
988 if let Ok(file) = fs::File::open(dir.join("artifact.json")) {
993 let _ = file.set_times(fs::FileTimes::new().set_accessed(std::time::SystemTime::now()));
994 }
995 let target = crate::Target::new(&record.target).ok()?;
996 let path = dir.join(target.artifact_filename());
997 let bytes = fs::read(&path).ok()?;
998 if bytes.is_empty() || bundle::digest(&bytes) != record.sha256 {
999 return None;
1000 }
1001 let launch_bytes = if let Some(hash) = &record.launch_sha256 {
1002 let bytes = fs::read(dir.join("launch.loombc")).ok()?;
1003 if bytes.is_empty() || &bundle::digest(&bytes) != hash {
1004 return None;
1005 }
1006 Some(Arc::from(bytes))
1007 } else {
1008 None
1009 };
1010 Some(Artifact {
1011 bytes: bytes.into(),
1012 listing: match &record.listing_sha256 {
1013 Some(hash) => {
1014 let listing = fs::read_to_string(dir.join("target.listing")).ok()?;
1015 if &bundle::digest(listing.as_bytes()) != hash {
1016 return None;
1017 }
1018 Some(listing.into())
1019 }
1020 None => None,
1021 },
1022 launch_bytes,
1023 launch_loader: native.launch_loader(),
1024 path,
1025 record,
1026 })
1027}
1028
1029#[cfg(test)]
1030mod tests {
1031 use super::*;
1032
1033 #[test]
1034 #[ignore = "requires libloomc.so, no GPU"]
1035 fn shared_compilers_respect_all_limits_and_reuse_matching_options() -> Result<()> {
1036 let options = CompilerOptions {
1037 workers: NonZeroUsize::new(2).unwrap(),
1038 module_cache_capacity: 2,
1039 ..Default::default()
1040 };
1041 let first = Compiler::shared(None, options.clone())?;
1042 let same = Compiler::shared(None, options.clone())?;
1043 assert!(Arc::ptr_eq(&first.0, &same.0));
1044 let serial_options = CompilerOptions {
1045 workers: NonZeroUsize::new(1).unwrap(),
1046 ..options
1047 };
1048 let serial = Compiler::shared(None, serial_options.clone())?;
1049 assert_eq!(serial.workers().get(), 1);
1050 assert!(!Arc::ptr_eq(&first.0, &serial.0));
1051 let uncached = Compiler::shared(
1052 None,
1053 CompilerOptions {
1054 module_cache_capacity: 0,
1055 ..serial_options
1056 },
1057 )?;
1058 assert!(!Arc::ptr_eq(&serial.0, &uncached.0));
1059 serial.module("// retained");
1060 uncached.module("// not retained");
1061 assert!(!serial.0.modules.lock().unwrap().is_empty());
1062 assert!(uncached.0.modules.lock().unwrap().is_empty());
1063 Ok(())
1064 }
1065
1066 #[test]
1067 #[ignore = "requires libloomc.so, no GPU"]
1068 fn processor_modes_have_distinct_compilers_and_artifacts() -> Result<()> {
1069 let default = Compiler::shared(None, CompilerOptions::default())?;
1070 let options = CompilerOptions {
1071 processor_mode: ProcessorMode::ComputeUnit,
1072 ..Default::default()
1073 };
1074 let cu = Compiler::shared(None, options.clone())?;
1075 let same = Compiler::shared(None, options)?;
1076 assert!(Arc::ptr_eq(&cu.0, &same.0));
1077 assert!(!Arc::ptr_eq(&default.0, &cu.0));
1078 let source = include_str!("../tests/kernels/euler.loom");
1079 let spec = euler_spec();
1080 let wgp = default.module(source).compile(&spec)?;
1081 let unit = cu.module(source).compile(&spec)?;
1082 assert_ne!(wgp.path(), unit.path());
1083 assert_eq!(unit.processor_mode(), ProcessorMode::ComputeUnit);
1084 let cached = cu.module(source).compile(&spec)?;
1085 assert_eq!(cached.processor_mode(), ProcessorMode::ComputeUnit);
1086 assert_eq!(unit.bytes(), cached.bytes());
1087 let differences: Vec<_> = wgp
1088 .bytes()
1089 .iter()
1090 .zip(unit.bytes())
1091 .filter(|(a, b)| a != b)
1092 .collect();
1093 assert_eq!(wgp.bytes().len(), unit.bytes().len());
1094 assert_eq!(differences.len(), 1);
1095 assert_eq!(*differences[0].0 ^ *differences[0].1, 0x20);
1096 Ok(())
1097 }
1098
1099 fn euler_spec() -> Specialization {
1100 let mut spec = Specialization::new("krea2_euler");
1101 spec.set_config("krea2.euler.grid_x", "1");
1102 spec.set_config("krea2.euler.grid_y", "1");
1103 spec
1104 }
1105
1106 #[test]
1107 #[ignore = "requires libloomc.so, no GPU"]
1108 fn reporting_preserves_shared_requests_and_can_be_replaced_from_a_callback() -> Result<()> {
1109 use std::sync::atomic::{AtomicUsize, Ordering};
1110 let kernels = Kernels::new(Compiler::resolve(None)?);
1111 let clone = kernels.clone();
1112 let pending =
1113 kernels.request(include_str!("../tests/kernels/euler.loom"), &euler_spec())?;
1114 let calls = Arc::new(AtomicUsize::new(0));
1115 let reported = calls.clone();
1116 let weak = Arc::downgrade(&kernels.0);
1117 let kernels = kernels.reporting(move |_| {
1118 reported.fetch_add(1, Ordering::SeqCst);
1119 let next = reported.clone();
1122 Kernels(weak.upgrade().unwrap()).reporting(move |_| {
1123 next.fetch_add(10, Ordering::SeqCst);
1124 });
1125 });
1126 assert!(Arc::ptr_eq(&kernels.0, &clone.0));
1127 assert!(Arc::ptr_eq(&kernels.0, &pending.cache));
1128 for cache in [&clone.0, &pending.cache] {
1129 let mut attempts = 0;
1130 let outcome = cache.build_with(1, |artifact| {
1131 attempts += 1;
1132 assert_eq!(artifact.record.key, pending.key);
1133 Err(Error::Message("injected load failure".into()))
1134 });
1135 assert_eq!(outcome.unwrap_err().to_string(), "injected load failure");
1136 assert_eq!(attempts, 1);
1137 }
1138 assert_eq!(calls.load(Ordering::SeqCst), 11);
1139 assert_eq!(kernels.0.locked()?.waiting.len(), 1);
1140 Ok(())
1141 }
1142
1143 #[test]
1144 #[ignore = "requires libloomc.so, no GPU"]
1145 fn requests_during_build_are_deduplicated_and_new_requests_survive() -> Result<()> {
1146 let kernels = Kernels::new(Compiler::resolve(None)?);
1147 let source = include_str!("../tests/kernels/euler.loom");
1148 let spec = euler_spec();
1149 let pending = kernels.request(source, &spec)?;
1150 let mut other_spec = spec.clone();
1151 other_spec.set_config("krea2.euler.grid_x", "2");
1152 let other_key = kernels.compiler().module(source).key(&other_spec)?;
1153 let outcome = kernels.0.build_with(1, |_| {
1154 std::thread::scope(|scope| {
1157 scope
1158 .spawn(|| {
1159 for _ in 0..2 {
1160 let duplicate = kernels.request(source, &spec).unwrap();
1161 assert_eq!(duplicate.key, pending.key);
1162 kernels.request(source, &other_spec).unwrap();
1163 }
1164 })
1165 .join()
1166 .unwrap();
1167 });
1168 assert_eq!(kernels.0.locked().unwrap().waiting.len(), 1);
1169 Err(Error::Message("injected load failure".into()))
1170 });
1171 assert_eq!(outcome.unwrap_err().to_string(), "injected load failure");
1172 {
1173 let state = kernels.0.locked()?;
1174 let waiting: Vec<_> = state.waiting.iter().map(|r| r.0.as_str()).collect();
1175 assert_eq!(waiting, vec![pending.key.as_str(), other_key.as_str()]);
1176 assert!(state.active.is_empty());
1177 }
1178 let mut attempts = Vec::new();
1179 let outcome = kernels.0.build_with(1, |artifact| {
1180 attempts.push(artifact.record.key.clone());
1181 Err(Error::Message("injected retry failure".into()))
1182 });
1183 assert!(outcome.is_err());
1184 assert_eq!(attempts, vec![pending.key, other_key]);
1185 assert_eq!(kernels.0.locked()?.waiting.len(), 2);
1186 Ok(())
1187 }
1188
1189 #[test]
1190 #[ignore = "requires libloomc.so, no GPU"]
1191 fn compile_and_load_failures_remain_requested_for_retry() -> Result<()> {
1192 let kernels = Kernels::new(Compiler::resolve(None)?);
1193 let source = include_str!("../tests/kernels/euler.loom");
1194 let bad = kernels.request(source, &Specialization::new("missing_export"))?;
1195 let good = kernels.request(source, &euler_spec())?;
1196 for _ in 0..2 {
1197 let mut attempted = Vec::new();
1198 let result = kernels.0.build_with(1, |artifact| {
1199 attempted.push(artifact.record.key.clone());
1200 Err(Error::Message("injected transient load failure".into()))
1201 });
1202 let error = result.unwrap_err().to_string();
1203 assert!(
1204 !error.contains("injected transient load failure"),
1205 "the first compile failure must be returned: {error}"
1206 );
1207 assert_eq!(
1208 attempted,
1209 vec![good.key.clone()],
1210 "the later request was attempted"
1211 );
1212 let state = kernels.0.locked()?;
1213 let waiting: Vec<_> = state.waiting.iter().map(|r| r.0.as_str()).collect();
1214 assert_eq!(waiting, vec![bad.key.as_str(), good.key.as_str()]);
1215 }
1216 assert!(good.built().is_none());
1217 Ok(())
1218 }
1219
1220 #[test]
1221 #[ignore = "requires libloomc.so, no GPU"]
1222 fn concurrent_build_waits_for_the_active_batch_and_retries_its_failure() -> Result<()> {
1223 use std::sync::mpsc;
1224 use std::time::Duration;
1225 let kernels = Kernels::new(Compiler::resolve(None)?);
1226 let pending =
1227 kernels.request(include_str!("../tests/kernels/euler.loom"), &euler_spec())?;
1228 let (entered_tx, entered_rx) = mpsc::channel();
1229 let (release_tx, release_rx) = mpsc::channel();
1230 let (started_tx, started_rx) = mpsc::channel();
1231 let (done_tx, done_rx) = mpsc::channel();
1232 std::thread::scope(|scope| {
1233 let first_cache = &kernels.0;
1234 let first = scope.spawn(move || {
1235 first_cache.build_with(1, |_| {
1236 entered_tx.send(()).unwrap();
1237 release_rx.recv_timeout(Duration::from_secs(10)).unwrap();
1238 Err(Error::Message("first load failed".into()))
1239 })
1240 });
1241 entered_rx.recv_timeout(Duration::from_secs(10)).unwrap();
1242 assert!(
1243 kernels.0.locked().unwrap().waiting.is_empty(),
1244 "batch has drained"
1245 );
1246 let second = scope.spawn(|| {
1247 started_tx.send(()).unwrap();
1248 let result = kernels
1249 .0
1250 .build_with(1, |_| Err(Error::Message("retry failed".into())));
1251 done_tx.send(result.map_err(|e| e.to_string())).unwrap();
1252 });
1253 started_rx.recv_timeout(Duration::from_secs(10)).unwrap();
1254 let early = done_rx.recv_timeout(Duration::from_millis(100));
1255 release_tx.send(()).unwrap();
1256 assert!(first.join().unwrap().is_err());
1257 second.join().unwrap();
1258 assert!(
1259 matches!(early, Err(mpsc::RecvTimeoutError::Timeout)),
1260 "build returned before the active batch finished"
1261 );
1262 assert_eq!(done_rx.recv().unwrap(), Err("retry failed".into()));
1263 });
1264 assert!(pending.built().is_none());
1265 assert_eq!(kernels.0.locked()?.waiting.len(), 1);
1266 Ok(())
1267 }
1268
1269 #[test]
1270 #[ignore = "requires libloomc.so, no GPU"]
1271 fn module_cache_eviction_preserves_live_modules() -> Result<()> {
1272 let compiler = Compiler::with_options(
1273 None,
1274 CompilerOptions {
1275 module_cache_capacity: 2,
1276 ..Default::default()
1277 },
1278 )?;
1279 let first = compiler.module(include_str!("../tests/kernels/euler.loom"));
1280 let first_id = first.identity().to_owned();
1281 for i in 0..8 {
1282 compiler.module(&format!("// source {i}"));
1283 assert!(compiler.0.modules.lock().unwrap().len() <= 2);
1284 }
1285 compiler.trim();
1286 assert!(compiler.0.modules.lock().unwrap().is_empty());
1287 assert_eq!(first.identity(), first_id);
1288 let mut spec = Specialization::new("krea2_euler");
1289 spec.set_config("krea2.euler.grid_x", "1");
1290 spec.set_config("krea2.euler.grid_y", "1");
1291 assert!(!first.compile(&spec)?.bytes().is_empty());
1292 Ok(())
1293 }
1294}
1295
1296#[derive(Clone)]
1310pub struct Kernels(Arc<Cache>);
1311
1312struct Cache {
1313 compiler: Compiler,
1314 state: Mutex<Loaded>,
1315 build: Mutex<()>,
1317 report: Mutex<Option<Reporter>>,
1318}
1319
1320type Reporter = Arc<dyn Fn(&Artifact) + Send + Sync>;
1322
1323#[derive(Default)]
1324struct Loaded {
1325 device: Option<usize>,
1327 ready: HashMap<String, crate::Kernel>,
1328 waiting: Vec<(String, String, Specialization)>,
1330 active: HashSet<String>,
1332}
1333
1334pub struct Pending {
1342 key: String,
1343 cache: Arc<Cache>,
1344 cell: std::sync::OnceLock<crate::Kernel>,
1348}
1349
1350impl Clone for Pending {
1351 fn clone(&self) -> Self {
1352 Self {
1353 key: self.key.clone(),
1354 cache: self.cache.clone(),
1355 cell: self.cell.clone(),
1356 }
1357 }
1358}
1359
1360impl Pending {
1361 pub fn built(&self) -> Option<&crate::Kernel> {
1367 if self.cell.get().is_none() {
1368 let found = self.cache.locked().ok()?.ready.get(&self.key).cloned()?;
1369 let _ = self.cell.set(found);
1370 }
1371 self.cell.get()
1372 }
1373
1374 pub unsafe fn resolve(&self, stream: &crate::Stream) -> Result<&crate::Kernel> {
1380 if self.built().is_some() {
1381 return Ok(self.cell.get().expect("just built"));
1382 }
1383 let outcome = unsafe { self.cache.build(stream) };
1388 match self.built() {
1389 Some(_) => Ok(self.cell.get().expect("just built")),
1390 None => Err(outcome
1391 .err()
1392 .unwrap_or_else(|| Error::Message("kernel was never requested".into()))),
1393 }
1394 }
1395}
1396
1397impl Kernels {
1398 pub fn keyed<K>(self) -> KeyedKernels<K> {
1405 KeyedKernels {
1406 kernels: self,
1407 keys: Mutex::new(HashMap::new()),
1408 }
1409 }
1410
1411 pub fn new(compiler: Compiler) -> Self {
1413 Self(Arc::new(Cache {
1414 compiler,
1415 state: Mutex::new(Loaded::default()),
1416 build: Mutex::new(()),
1417 report: Mutex::new(None),
1418 }))
1419 }
1420
1421 pub fn reporting(self, report: impl Fn(&Artifact) + Send + Sync + 'static) -> Self {
1431 let previous = self
1432 .0
1433 .report
1434 .lock()
1435 .unwrap_or_else(|poisoned| poisoned.into_inner())
1436 .replace(Arc::new(report));
1437 drop(previous);
1439 self
1440 }
1441
1442 pub fn compiler(&self) -> &Compiler {
1444 &self.0.compiler
1445 }
1446
1447 pub fn len(&self) -> usize {
1449 self.0.locked().map_or(0, |state| state.ready.len())
1450 }
1451
1452 pub fn is_empty(&self) -> bool {
1454 self.len() == 0
1455 }
1456
1457 pub unsafe fn get(
1464 &self,
1465 stream: &crate::Stream,
1466 source: &str,
1467 spec: &Specialization,
1468 ) -> Result<crate::Kernel> {
1469 let module = self.0.compiler.module(source);
1470 let key = module.key(spec)?;
1471 unsafe { self.get_module(stream, &module, key, spec) }
1473 }
1474
1475 unsafe fn get_module(
1476 &self,
1477 stream: &crate::Stream,
1478 module: &Module,
1479 key: String,
1480 spec: &Specialization,
1481 ) -> Result<crate::Kernel> {
1482 {
1483 let mut state = self.0.locked()?;
1484 state.claim(stream)?;
1485 if let Some(kernel) = state.ready.get(&key) {
1486 return Ok(kernel.clone());
1487 }
1488 }
1489 let _build = self
1492 .0
1493 .build
1494 .lock()
1495 .map_err(|_| Error::Message("kernel build poisoned".into()))?;
1496 if let Some(kernel) = self.0.locked()?.ready.get(&key) {
1497 return Ok(kernel.clone());
1498 }
1499 let artifact = module.compile(spec)?;
1500 self.0.reported(&artifact);
1501 let kernel = unsafe { stream.load_artifact(&artifact) }?;
1503 let mut state = self.0.locked()?;
1504 state.waiting.retain(|(waiting, ..)| *waiting != key);
1505 state.ready.insert(key, kernel.clone());
1506 Ok(kernel)
1507 }
1508
1509 pub fn request(&self, source: &str, spec: &Specialization) -> Result<Pending> {
1514 let key = self.0.compiler.module(source).key(spec)?;
1515 let mut state = self.0.locked()?;
1516 let known = state.ready.contains_key(&key)
1517 || state.active.contains(&key)
1518 || state.waiting.iter().any(|(waiting, ..)| *waiting == key);
1519 if !known {
1520 state
1521 .waiting
1522 .push((key.clone(), source.to_owned(), spec.clone()));
1523 }
1524 drop(state);
1525 Ok(Pending {
1526 key,
1527 cache: self.0.clone(),
1528 cell: std::sync::OnceLock::new(),
1529 })
1530 }
1531
1532 pub unsafe fn build(&self, stream: &crate::Stream) -> Result<()> {
1538 unsafe { self.0.build(stream) }
1540 }
1541}
1542
1543impl Cache {
1544 fn locked(&self) -> Result<std::sync::MutexGuard<'_, Loaded>> {
1545 self.state
1546 .lock()
1547 .map_err(|_| Error::Message("kernel cache poisoned".into()))
1548 }
1549
1550 fn reported(&self, artifact: &Artifact) {
1551 let report = self
1552 .report
1553 .lock()
1554 .unwrap_or_else(|poisoned| poisoned.into_inner())
1555 .clone();
1556 if let Some(report) = report {
1558 report(artifact);
1559 }
1560 }
1561
1562 unsafe fn build(&self, stream: &crate::Stream) -> Result<()> {
1570 self.build_with(stream.device_id(), |artifact| {
1571 unsafe { stream.load_artifact(artifact) }
1573 })
1574 }
1575
1576 fn build_with(
1577 &self,
1578 device: usize,
1579 mut load: impl FnMut(&Artifact) -> Result<crate::Kernel>,
1580 ) -> Result<()> {
1581 let _build = self
1582 .build
1583 .lock()
1584 .map_err(|_| Error::Message("kernel build poisoned".into()))?;
1585 let waiting = {
1586 let mut state = self.locked()?;
1587 state.claim_device(device)?;
1588 let waiting = std::mem::take(&mut state.waiting);
1589 state
1590 .active
1591 .extend(waiting.iter().map(|(key, ..)| key.clone()));
1592 waiting
1593 };
1594 if waiting.is_empty() {
1595 return Ok(());
1596 }
1597 let modules: Vec<Module> = waiting
1598 .iter()
1599 .map(|(_, source, _)| self.compiler.module(source))
1600 .collect();
1601 let requests: Vec<(&Module, &Specialization)> = modules
1602 .iter()
1603 .zip(waiting.iter().map(|(_, _, spec)| spec))
1604 .collect();
1605 let built = self.compiler.compile_all(&requests);
1606
1607 let mut failure = None;
1608 let mut unbuilt = Vec::new();
1609 for (request, outcome) in waiting.into_iter().zip(built) {
1610 let loaded = outcome.and_then(|artifact| {
1611 self.reported(&artifact);
1612 load(&artifact)
1613 });
1614 match loaded {
1615 Ok(kernel) => {
1616 self.locked()?.ready.insert(request.0.clone(), kernel);
1617 }
1618 Err(error) => {
1619 unbuilt.push(request);
1620 failure = failure.or(Some(error));
1621 }
1622 }
1623 }
1624 {
1625 let mut state = self.locked()?;
1626 unbuilt.append(&mut state.waiting);
1627 state.waiting = unbuilt;
1628 state.active.clear();
1629 }
1630 match failure {
1631 Some(error) => Err(error),
1632 None => Ok(()),
1633 }
1634 }
1635}
1636
1637pub struct KeyedKernels<K> {
1644 kernels: Kernels,
1645 keys: Mutex<HashMap<K, String>>,
1646}
1647
1648impl<K: Eq + std::hash::Hash> KeyedKernels<K> {
1649 pub fn kernels(&self) -> &Kernels {
1651 &self.kernels
1652 }
1653
1654 pub unsafe fn request_or_insert_with<S: AsRef<str>>(
1662 &self,
1663 stream: &crate::Stream,
1664 key: K,
1665 request: impl FnOnce(&K) -> Result<(S, Specialization)>,
1666 ) -> Result<Pending> {
1667 let mut keys = self
1668 .keys
1669 .lock()
1670 .map_err(|_| Error::Message("kernel key index poisoned".into()))?;
1671 self.kernels.0.locked()?.claim(stream)?;
1672 if let Some(artifact) = keys.get(&key) {
1673 return Ok(Pending {
1674 key: artifact.clone(),
1675 cache: self.kernels.0.clone(),
1676 cell: std::sync::OnceLock::new(),
1677 });
1678 }
1679 let (source, spec) = request(&key)?;
1680 let pending = self.kernels.request(source.as_ref(), &spec)?;
1681 keys.insert(key, pending.key.clone());
1682 Ok(pending)
1683 }
1684
1685 pub unsafe fn get_or_insert_with<'s>(
1696 &self,
1697 stream: &crate::Stream,
1698 key: K,
1699 request: impl FnOnce(&K) -> Result<(&'s str, Specialization)>,
1700 ) -> Result<crate::Kernel> {
1701 let mut keys = self
1702 .keys
1703 .lock()
1704 .map_err(|_| Error::Message("kernel key index poisoned".into()))?;
1705 {
1706 let mut state = self.kernels.0.locked()?;
1707 state.claim(stream)?;
1708 if let Some(artifact) = keys.get(&key) {
1709 if let Some(kernel) = state.ready.get(artifact) {
1710 return Ok(kernel.clone());
1711 }
1712 let pending = Pending {
1713 key: artifact.clone(),
1714 cache: self.kernels.0.clone(),
1715 cell: std::sync::OnceLock::new(),
1716 };
1717 drop(state);
1718 drop(keys);
1719 return unsafe { pending.resolve(stream) }.cloned();
1721 }
1722 }
1723 let (source, spec) = request(&key)?;
1724 let module = self.kernels.compiler().module(source);
1725 let artifact = module.key(&spec)?;
1726 let kernel = unsafe {
1728 self.kernels
1729 .get_module(stream, &module, artifact.clone(), &spec)
1730 }?;
1731 keys.insert(key, artifact);
1732 Ok(kernel)
1733 }
1734}
1735
1736impl Loaded {
1737 fn claim(&mut self, stream: &crate::Stream) -> Result<()> {
1739 self.claim_device(stream.device_id())
1740 }
1741
1742 fn claim_device(&mut self, device: usize) -> Result<()> {
1743 match self.device {
1744 Some(owner) if owner != device => Err(Error::Message(
1745 "a kernel cache serves one device: an executable loaded on another is not \
1746 dispatchable here"
1747 .into(),
1748 )),
1749 Some(_) => Ok(()),
1750 None => {
1751 self.device = Some(device);
1752 Ok(())
1753 }
1754 }
1755 }
1756}
1757
1758impl std::fmt::Debug for Kernels {
1759 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1760 f.debug_struct("Kernels")
1761 .field("loaded", &self.len())
1762 .finish_non_exhaustive()
1763 }
1764}