Skip to main content

hrx/
loom.rs

1//! In-process Loom compilation with reusable indexed modules and verified disk artifacts.
2#[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/// How seriously the compiler meant a diagnostic. Ordered, so callers can
34/// filter with a comparison rather than a magic number: a spill remark and a
35/// type error arrive in the same list, and only one of them is worth a reader.
36#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
37#[serde(rename_all = "lowercase")]
38pub enum Severity {
39    /// A remark: informational, including backend resource notes.
40    Note,
41    /// A warning: compilation continued.
42    Warning,
43    /// An error: compilation of this specialization failed.
44    Error,
45}
46impl From<u32> for Severity {
47    /// Unknown native severities are treated as errors, so a future compiler
48    /// cannot make a diagnostic disappear from a caller's error filter.
49    fn from(native: u32) -> Self {
50        match native {
51            0 => Self::Note,
52            1 => Self::Warning,
53            _ => Self::Error,
54        }
55    }
56}
57
58/// A compiler diagnostic copied out of native result storage.
59#[derive(Clone, Debug, Serialize, Deserialize)]
60pub struct Diagnostic {
61    /// How seriously the compiler meant this diagnostic.
62    pub severity: Severity,
63    /// Stable compiler diagnostic code.
64    pub code: String,
65    /// Rendered diagnostic message.
66    pub message: String,
67    /// Full compiler-rendered diagnostic, including source context when available.
68    pub formatted_text: String,
69    /// Named values in diagnostic-definition order.
70    pub parameters: Vec<DiagnosticParameter>,
71    /// Diagnostic source identifier, when supplied by the compiler.
72    #[serde(default)]
73    pub source: String,
74    /// One-based source line, or zero when unavailable.
75    pub line: u32,
76    /// One-based source column, or zero when unavailable.
77    pub column: u32,
78    /// Labeled locations providing context, in compiler order.
79    #[serde(default)]
80    pub related_locations: Vec<RelatedLocation>,
81    /// Additional locations omitted by the compiler's diagnostic limit.
82    #[serde(default)]
83    pub related_location_omitted_count: usize,
84}
85
86/// One compiler-rendered diagnostic argument.
87#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
88pub struct DiagnosticParameter {
89    /// Stable name from the compiler diagnostic definition.
90    pub name: String,
91    /// Canonical rendered value.
92    pub value: String,
93}
94
95/// A source location related to a compiler diagnostic, owned by Rust.
96#[derive(Clone, Debug, Serialize, Deserialize)]
97pub struct RelatedLocation {
98    /// Relationship to the primary diagnostic, such as a previous use.
99    pub label: String,
100    /// Source identifier, even when the compiler did not retain source text.
101    pub source: String,
102    /// One-based starting line, or zero when unavailable.
103    pub line: u32,
104    /// One-based starting column, or zero when unavailable.
105    pub column: u32,
106    /// One-based ending line, or zero when unavailable.
107    pub end_line: u32,
108    /// One-based ending column, or zero when unavailable.
109    pub end_column: u32,
110}
111/// AMDGPU workgroup scheduling domain. Explicit modes require GFX11/GFX12.
112#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq, Serialize, Deserialize)]
113#[serde(rename_all = "snake_case")]
114pub enum ProcessorMode {
115    /// Preserve the compiler's default for the selected target.
116    #[default]
117    Default,
118    /// Schedule each workgroup within a single compute unit (GFX11/GFX12).
119    ComputeUnit,
120    /// Schedule across a workgroup processor (GFX11/GFX12).
121    WorkgroupProcessor,
122}
123
124/// Assertion classes inserted at the target's semantic assertion boundary.
125#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
126pub struct SanitizerChecks {
127    /// Memory access assertions (requires the target's access-shadow runtime).
128    pub access: bool,
129    /// SSA value assertions.
130    pub value: bool,
131    /// Operation assertions.
132    pub operation: bool,
133    /// Data-race observations (requires the target's race-shadow runtime).
134    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/// Target behavior on a failed sanitizer assertion.
145#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
146pub enum SanitizerReporting {
147    /// Native target's structured report path.
148    #[default]
149    Default,
150    /// Direct fatal trap without structured reports.
151    Trap,
152    /// Structured report without a direct trap.
153    ReportOnly,
154}
155/// Native compiler instrumentation. Runtime-required globals are described by
156/// the artifact manifest; compiling does not supply a sanitizer runtime.
157#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
158pub struct SanitizerOptions {
159    /// Independent classes of assertions.
160    pub checks: SanitizerChecks,
161    /// Behavior of failed assertions.
162    pub reporting: SanitizerReporting,
163}
164
165/// Compiler target and limits for concurrent workspaces and cached modules.
166#[derive(Clone, Debug)]
167pub struct CompilerOptions {
168    /// Maximum number of simultaneous compilations through this compiler.
169    pub workers: NonZeroUsize,
170    /// Architecture used by the Loom target profile and artifact cache.
171    pub target: crate::Target,
172    /// Maximum retained source modules; zero disables module caching.
173    /// Eviction chooses an arbitrary entry, not the least recently used module.
174    pub module_cache_capacity: usize,
175    /// Workgroup scheduling policy for the selected AMDGPU target.
176    pub processor_mode: ProcessorMode,
177    /// Instrumentation included in pipeline, shared compiler, and artifact identities.
178    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    /// Default compiler limits for an explicit device target.
198    #[must_use]
199    pub fn for_target(target: &crate::Target) -> Self {
200        Self {
201            target: target.clone(),
202            ..Self::default()
203        }
204    }
205    /// Default compiler limits targeting the device behind `stream`.
206    #[must_use]
207    pub fn for_stream(stream: &crate::Stream) -> Self {
208        Self::for_target(stream.target())
209    }
210}
211struct Inner {
212    // Indexes are released before prepared compiler state when the session ends.
213    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/// A pinned compiler library and reusable native state, shared by cheap clones.
234#[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    /// Select an explicit library, HRX_LOOM_LIBRARY, or the pinned native bundle.
246    /// This loads CPU compiler code but never opens a GPU. Library mappings stay
247    /// resident until process exit; contexts and workspaces remain session-owned.
248    /// To upgrade a loaded library, select a new path or restart the process.
249    pub fn resolve(library: Option<&Path>) -> Result<Self> {
250        Self::with_options(library, CompilerOptions::default())
251    }
252    /// Resolve the shared compiler configured for a stream's actual device target.
253    pub fn for_stream(library: Option<&Path>, stream: &crate::Stream) -> Result<Self> {
254        Self::shared(library, CompilerOptions::for_stream(stream))
255    }
256    /// Resolve the shared compiler configured for an explicit device target.
257    pub fn for_target(library: Option<&Path>, target: &crate::Target) -> Result<Self> {
258        Self::shared(library, CompilerOptions::for_target(target))
259    }
260    /// The compiler for this library, target and limits, resolved once per process.
261    ///
262    /// Resolving is not cheap: it canonicalizes the library path and digests
263    /// the whole shared object, twice. A consumer that asks for a compiler per
264    /// kernel pays that per kernel, so every consumer memoized it themselves.
265    ///
266    /// The cache includes all compiler options. `HRX_LOOM_LIBRARY` is resolved
267    /// on each call rather than frozen at the first.
268    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    /// The library a given override, `HRX_LOOM_LIBRARY` or the bundle selects.
291    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    /// Resolve a compiler with a bounded number of exclusive workspaces.
303    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    /// Maximum simultaneous compilations, and the workspace pool's bound.
332    pub fn workers(&self) -> NonZeroUsize {
333        self.0.workers
334    }
335
336    /// Compile several specializations, up to [`CompilerOptions::workers`] at a
337    /// time, and return their outcomes in request order.
338    ///
339    /// [`Module::compile`] blocks, so a single-threaded caller never reaches the
340    /// workspace pool's bound. This is the entry point that does: it owns the
341    /// thread budget rather than leaving every consumer to rebuild the same pool.
342    /// Each request is independent; one failure does not cancel the others.
343    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    /// Architecture selected for every compilation by this compiler.
379    pub fn target(&self) -> &crate::Target {
380        &self.0.target
381    }
382    /// Content identity of the loaded compiler library.
383    pub fn identity(&self) -> &str {
384        &self.0.identity
385    }
386    /// Selected compiler library path.
387    pub fn library_path(&self) -> &Path {
388        &self.0.path
389    }
390    /// Retain a source module, sharing its parsed index across specializations.
391    /// Parsing is deferred until the first artifact-cache miss.
392    pub fn module(&self, source: &str) -> Module {
393        self.source_module(vec![Source::loom("kernel.loom", source)])
394    }
395    /// Link named Loom and C/C++ inputs in one compiler-owned source module.
396    /// Source bytes, headers, macro order and import options all enter its identity.
397    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    /// Import a translation unit; parsing is deferred until compilation.
420    pub fn import_cxx(&self, source: CxxSource) -> Result<Module> {
421        self.sources(vec![Source::Cxx(source)])
422    }
423    /// Parse and retain a representative source's leading includes for reuse.
424    /// Sources without a reusable leading include return an error. Import roots
425    /// and floating-point policy can vary while using the same preamble.
426    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    /// Release idle workspace blocks and cached module references. Existing
458    /// Module handles remain valid; active compilations are unaffected.
459    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}
468/// Render the diagnostic a caller should act on, not the whole cascade.
469///
470/// A single rejected construct makes the rest of its block unparseable, so the
471/// compiler reports one real error followed by a run of consequences. Leading
472/// with all of them buries the cause; every diagnostic stays available through
473/// [`Error::Compile::diagnostics`].
474fn 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
513/// Diagnostics whose cause is a configuration mistake rather than a source bug.
514fn 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        // A cascade-free failure gains no suppression line and no hint.
558        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/// Immutable source plus compiler ownership. Clones share the frozen native index.
569#[derive(Clone)]
570pub struct Module {
571    // Release the index before the last compiler/context owner.
572    data: Arc<ModuleData>,
573    compiler: Compiler,
574}
575/// Amount of compiler evidence collected alongside an executable.
576#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
577#[serde(rename_all = "snake_case")]
578pub enum ReportMode {
579    /// Emit no report and avoid detail collection.
580    #[default]
581    None,
582    /// Collect artifact facts and summary compiler analysis.
583    Summary,
584    /// Include detailed provenance, scheduling, and resource evidence.
585    Details,
586}
587
588/// Compiler-formatted pass trace representation.
589#[derive(Clone, Copy, Debug, Default)]
590pub enum TraceFormat {
591    /// Human-readable metadata followed by source.
592    Text,
593    /// One JSON object per pass event.
594    #[default]
595    JsonLines,
596}
597/// Bounded streaming diagnostics for a fresh compile invocation.
598#[derive(Clone, Debug)]
599pub struct TraceOptions {
600    /// Serialization chosen by the compiler.
601    pub format: TraceFormat,
602    /// Maximum bytes forwarded to the writer; exceeding this aborts compilation.
603    pub max_bytes: usize,
604    /// Pass keys or compiler stage names to capture before execution.
605    pub before: Vec<String>,
606    /// Pass keys or compiler stage names to capture after execution.
607    /// Both filter lists empty selects every after boundary.
608    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/// One export and its exact configuration within a module.
622#[derive(Clone, Debug, Default)]
623pub struct Specialization {
624    /// Export/root symbol to compile.
625    symbol: String,
626    /// Fully qualified configuration keys and Loom value spellings.
627    config: BTreeMap<String, String>,
628    /// Select compiler analysis collected beside the executable.
629    report: ReportMode,
630    listing: bool,
631}
632impl Specialization {
633    /// Request target assembly alongside the executable. Unsupported targets
634    /// report a compiler error instead of silently omitting the listing.
635    #[must_use]
636    pub fn with_listing(mut self, enabled: bool) -> Self {
637        self.listing = enabled;
638        self
639    }
640    /// Select an export with no configuration overrides.
641    pub fn new(symbol: impl Into<String>) -> Self {
642        Self {
643            symbol: symbol.into(),
644            ..Self::default()
645        }
646    }
647    /// Set one fully-qualified Loom configuration value.
648    #[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    /// Set one fully-qualified Loom configuration value in place.
654    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    /// Select an export/root symbol in place.
659    pub fn set_symbol(&mut self, symbol: impl Into<String>) -> &mut Self {
660        self.symbol = symbol.into();
661        self
662    }
663    /// Replace all fully-qualified Loom configuration values.
664    pub fn replace_config(&mut self, config: BTreeMap<String, String>) -> &mut Self {
665        self.config = config;
666        self
667    }
668    /// Request or suppress the compiler resource report.
669    #[must_use]
670    pub fn with_report(mut self, report: ReportMode) -> Self {
671        self.report = report;
672        self
673    }
674    /// Request or suppress the compiler resource report in place.
675    pub fn set_report(&mut self, report: ReportMode) -> &mut Self {
676        self.report = report;
677        self
678    }
679    /// Export/root symbol selected by this specialization.
680    #[must_use]
681    pub fn symbol(&self) -> &str {
682        &self.symbol
683    }
684    /// Fully-qualified compiler configuration.
685    #[must_use]
686    pub fn configuration(&self) -> &BTreeMap<String, String> {
687        &self.config
688    }
689    /// Compiler evidence requested by this specialization.
690    #[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/// Owned native executable bytes and their verified compilation metadata.
725#[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    /// Compiler-produced target assembly, when requested with `with_listing`.
736    pub fn listing(&self) -> Option<&str> {
737        self.listing.as_deref()
738    }
739    /// Native executable bytes, independent of compiler/result lifetimes.
740    pub fn bytes(&self) -> &[u8] {
741        &self.bytes
742    }
743    /// Compiler-owned description of exports, arguments, globals, and target facts.
744    pub fn manifest(&self) -> Option<&serde_json::Value> {
745        self.record.manifest.as_ref()
746    }
747    /// Serialized host launch program produced alongside these executable bytes.
748    pub fn launch_bytes(&self) -> Option<&[u8]> {
749        self.launch_bytes.as_deref()
750    }
751    /// Bind an exact executable export to an independent launch evaluator.
752    /// The matching native evaluator library stays owned by the artifact.
753    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    /// Durable verified cache artifact path.
762    pub fn path(&self) -> &Path {
763        &self.path
764    }
765    /// Compiler library content identity.
766    pub fn compiler_identity(&self) -> &str {
767        &self.record.compiler
768    }
769    /// Exact GPU target key.
770    pub fn target(&self) -> &str {
771        &self.record.target
772    }
773    /// Scheduling domain used to compile this artifact.
774    pub fn processor_mode(&self) -> ProcessorMode {
775        self.record.processor_mode
776    }
777    /// Export compiled into this artifact.
778    pub fn symbol(&self) -> &str {
779        &self.record.symbol
780    }
781    /// Diagnostics retained from successful preparation and compilation.
782    pub fn diagnostics(&self) -> &[Diagnostic] {
783        &self.record.diagnostics
784    }
785    /// Optional structured compiler evidence, in the pinned compiler's schema.
786    pub fn report(&self) -> Option<&CompileReport> {
787        self.record.report.as_ref()
788    }
789}
790/// Concrete launch properties evaluated from compiler-authored workload arguments.
791#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
792pub struct LaunchConfig {
793    /// Workgroup count along each axis.
794    pub workgroup_count: [u32; 3],
795    /// Invocations per workgroup along each axis.
796    pub workgroup_size: [u32; 3],
797    /// Cooperative cluster dimensions; [1, 1, 1] for an ordinary launch.
798    pub workgroup_cluster_size: [u32; 3],
799    /// Hardware subgroup width, or zero when the target does not fix it.
800    pub subgroup_size: u32,
801    /// Total workgroup-local storage, including statically allocated storage.
802    pub workgroup_storage_bytes: u64,
803}
804
805/// One bound launch function with exclusively owned invocation scratch.
806/// Create one per independently prepared execution slot.
807pub struct LaunchProgram(native::LaunchProgram);
808impl LaunchProgram {
809    /// Evaluate positional raw scalar bits without compiling source or looking up names.
810    /// Evaluation may grow scratch; perform it while preparing a dispatch.
811    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    /// Identity of the exact source bytes.
826    pub fn identity(&self) -> &str {
827        &self.data.digest
828    }
829    /// Cache identity including compiler, pipeline contract, source and specialization.
830    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        // Retain existing default artifacts; explicit execution modes own a
848        // separate namespace even if their output bytes happen to match.
849        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    /// Compile in process or return verified cached bytes. Publication is atomic
860    /// and serialized per key across threads and processes; failures are retryable.
861    pub fn compile(&self, spec: &Specialization) -> Result<Artifact> {
862        self.compile_impl(spec, None)
863    }
864    /// Recompile and stream compiler-owned pass traces, even on a cache hit.
865    /// A writer error, panic, or exhausted byte limit aborts compilation and
866    /// leaves any existing cached artifact intact. The writer is not retained.
867    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    // `bundle::collect` evicts on access time. Reading the artifact below updates
989    // it on a normal mount, but a `noatime` mount never would, so set it here:
990    // utimensat is unaffected by the mount option that suppresses the implicit
991    // update, and without this a busy cache would look idle to the collector.
992    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            // Replacing a reporter must not hold its lock while invoking it.
1120            // A weak reference avoids a callback/cache ownership cycle.
1121            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            // Another caller requests both the active kernel and a new one
1155            // while loading is paused here. No native GPU load is performed.
1156            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/// Loaded kernels, keyed by the artifact their specialization compiles to.
1297///
1298/// The compiler already caches *compilation* on disk. This caches the loaded
1299/// executable, which is what a dispatch loop needs: turning a cached artifact
1300/// into a [`crate::gpu::Kernel`] is a native load every time, and a model that dispatches
1301/// the same kernel per block would pay it per block. Every consumer of this
1302/// crate wrote this cache, so it lives here instead.
1303///
1304/// The key is the artifact's own identity — the same digest the disk cache uses
1305/// — rather than whatever the caller derived a request from, so two requests
1306/// that compile to the same artifact share one executable.
1307///
1308/// Kernels are device-scoped, so a cache serves one device and refuses another.
1309#[derive(Clone)]
1310pub struct Kernels(Arc<Cache>);
1311
1312struct Cache {
1313    compiler: Compiler,
1314    state: Mutex<Loaded>,
1315    /// Covers draining, compiling, publishing and requeuing one complete batch.
1316    build: Mutex<()>,
1317    report: Mutex<Option<Reporter>>,
1318}
1319
1320/// Called with each artifact a [`Kernels`] builds. See [`Kernels::reporting`].
1321type Reporter = Arc<dyn Fn(&Artifact) + Send + Sync>;
1322
1323#[derive(Default)]
1324struct Loaded {
1325    /// Set by the first request; kernels from one device are useless on another.
1326    device: Option<usize>,
1327    ready: HashMap<String, crate::Kernel>,
1328    /// Requested but not yet built, in request order.
1329    waiting: Vec<(String, String, Specialization)>,
1330    /// Keep drained requests known until their results are published or requeued.
1331    active: HashSet<String>,
1332}
1333
1334/// A kernel that has been asked for but not yet built.
1335///
1336/// Compiling is what a cold cache costs and requests are independent, so
1337/// [`Kernels::request`] records one and returns this rather than stopping to
1338/// build it. The outstanding set is then built together, on as many threads as
1339/// the compiler allows, by [`Kernels::build`] or by the first
1340/// [`Pending::resolve`] that needs any one of them.
1341pub struct Pending {
1342    key: String,
1343    cache: Arc<Cache>,
1344    /// This handle's own copy, so a built kernel can be borrowed for as long as
1345    /// the handle lives. A graph records `&'g Kernel`, and the cache's own copy
1346    /// is behind a lock with no lifetime to lend.
1347    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    /// The kernel, if its batch has already been built.
1362    ///
1363    /// For callers with no stream to build one on — recording into a graph,
1364    /// which cannot load an executable. Everywhere else wants
1365    /// [`Pending::resolve`].
1366    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    /// The kernel, building everything outstanding if this is the first call
1375    /// that needs it.
1376    ///
1377    /// # Safety
1378    /// As [`Kernels::get`], for every source passed to [`Kernels::request`].
1379    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        // A batch reports the first kernel in it that would not build, which is
1384        // not necessarily this one, so ask again before passing that failure on
1385        // as though it were ours.
1386        // Safety: the caller vouched for every requested source.
1387        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    /// Index this cache by a caller's compact request key.
1399    ///
1400    /// Source hashing and specialization construction happen only on a key
1401    /// miss. Distinct caller keys that name the same artifact still share one
1402    /// loaded executable. Include every input that changes the source or
1403    /// specialization in the key.
1404    pub fn keyed<K>(self) -> KeyedKernels<K> {
1405        KeyedKernels {
1406            kernels: self,
1407            keys: Mutex::new(HashMap::new()),
1408        }
1409    }
1410
1411    /// An empty cache over a compiler. Clone it to share one cache.
1412    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    /// Call `report` with each artifact as it is built, once, before it loads.
1422    ///
1423    /// A cache that returned only kernels would swallow what the compiler said
1424    /// about them: warnings, and the backend remarks that are the only warning
1425    /// a kernel is spilling registers. Cache hits report nothing, having built
1426    /// nothing.
1427    /// Installing a reporter updates every clone and existing pending handle.
1428    /// An invocation already in progress may finish using the previous reporter.
1429    /// The callback must not build or resolve requests on this same cache.
1430    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        // Captured values can have arbitrary destructors; drop outside the lock.
1438        drop(previous);
1439        self
1440    }
1441
1442    /// The compiler these were built with.
1443    pub fn compiler(&self) -> &Compiler {
1444        &self.0.compiler
1445    }
1446
1447    /// How many kernels are loaded.
1448    pub fn len(&self) -> usize {
1449        self.0.locked().map_or(0, |state| state.ready.len())
1450    }
1451
1452    /// Whether any kernel is loaded.
1453    pub fn is_empty(&self) -> bool {
1454        self.len() == 0
1455    }
1456
1457    /// Compile `spec` from `source` if needed and load it, or return the
1458    /// loaded kernel. Outstanding [`Kernels::request`]s are left alone.
1459    ///
1460    /// # Safety
1461    /// `source` must be trusted Loom: its compiled output is loaded as native
1462    /// code, as for [`crate::Stream::load_artifact`].
1463    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        // Safety: the caller vouched for this module's source.
1472        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        // Share the batch publication lock so direct requests cannot load a
1490        // second executable while a batch or another direct miss is in flight.
1491        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        // Safety: the caller vouched for the source this artifact came from.
1502        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    /// Ask for a kernel without building it. Nothing is compiled until
1510    /// [`Kernels::build`] or the first [`Pending::resolve`], and then every
1511    /// outstanding request is built at once across the compiler's workspaces —
1512    /// so a consumer can name its whole set first and pay for it in one batch.
1513    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    /// Build every outstanding request, in one batch.
1533    /// Concurrent builds wait for the active batch before checking for more work.
1534    ///
1535    /// # Safety
1536    /// As [`Kernels::get`], for every source passed to [`Kernels::request`].
1537    pub unsafe fn build(&self, stream: &crate::Stream) -> Result<()> {
1538        // Safety: the caller vouched for every requested source.
1539        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        // A reporter may replace itself or inspect the cache.
1557        if let Some(report) = report {
1558            report(artifact);
1559        }
1560    }
1561
1562    /// Compile off the device, then load in order.
1563    ///
1564    /// A specialization that will not build is that one's failure and no one
1565    /// else's, so the batch carries on and the rest is loaded. The failure goes
1566    /// back on the queue still wanted: dropping it would leave the handles that
1567    /// asked for it waiting on nothing, with no way to say why, while returning
1568    /// it means the next attempt reports the same failure again.
1569    unsafe fn build(&self, stream: &crate::Stream) -> Result<()> {
1570        self.build_with(stream.device_id(), |artifact| {
1571            // Safety: the caller vouched for every requested source.
1572            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
1637/// A caller-key index over HRX's artifact-keyed loaded kernels.
1638///
1639/// Create with [`Kernels::keyed`]. A hit checks device ownership and returns a
1640/// shared executable without hashing source, serializing a specialization or
1641/// invoking the request factory. The index stores artifact identities, not a
1642/// second collection of executables.
1643pub struct KeyedKernels<K> {
1644    kernels: Kernels,
1645    keys: Mutex<HashMap<K, String>>,
1646}
1647
1648impl<K: Eq + std::hash::Hash> KeyedKernels<K> {
1649    /// The underlying batch cache and its compiler.
1650    pub fn kernels(&self) -> &Kernels {
1651        &self.kernels
1652    }
1653
1654    /// Index a pending request without compiling. Hits do not call the factory
1655    /// or hash source. Failed compilation remains queued for a later retry.
1656    /// The factory must not reenter this index.
1657    ///
1658    /// # Safety
1659    /// Equal keys must describe the same source and specialization. Loading is
1660    /// still governed by [`Pending::resolve`]'s trusted-code contract.
1661    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    /// Return the kernel named by `key`, constructing its request only on a miss.
1686    ///
1687    /// Failed requests are not indexed and can be retried. The factory is run
1688    /// under the index lock; it and the reporting callback must not reenter
1689    /// this index. The underlying cache's device restriction applies to hits too.
1690    ///
1691    /// # Safety
1692    /// As [`Kernels::get`]. Equal keys must always describe the same source and
1693    /// specialization, including compiler-report settings. The factory is not
1694    /// evaluated on a hit, so changing it does not change a cached kernel.
1695    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                // Safety: equal keys identify the source vouched for by the caller.
1720                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        // Safety: the caller vouched for the source and the request key.
1727        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    /// Bind this cache to a device on first use, and hold it there.
1738    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}