vyre_driver/
aot.rs

1//! Backend-neutral AOT emission and launcher registries.
2
3use std::collections::BTreeMap;
4use std::path::PathBuf;
5
6use crate::BackendError;
7
8/// Stable validated AOT target identity.
9pub use vyre_foundation::operation::TargetId as AotTargetId;
10
11/// One dependency entry required by a generated launcher crate.
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct LauncherDependency {
14    /// Dependency name in the emitted `Cargo.toml`.
15    pub name: &'static str,
16    /// Inline dependency spec, for example `{ version = "1", features = ["derive"] }`.
17    pub spec: &'static str,
18}
19
20/// Backend-neutral launcher emission request.
21#[derive(Debug)]
22pub struct AotLauncherRequest<'a> {
23    /// Stable target id matching the selected launcher emitter.
24    pub target: AotTargetId,
25    /// Generated launcher crate name.
26    pub crate_name: &'a str,
27    /// Whether to include target-owned collective/multi-rank support.
28    pub include_collectives: bool,
29    /// Whether to include a built-in eval-time training loop.
30    pub include_ttt_loop: bool,
31}
32
33/// Source files and manifest additions produced by a target-owned launcher emitter.
34#[derive(Debug, Clone, Default)]
35pub struct AotLauncherFiles {
36    /// Additional dependencies required by target-specific launcher files.
37    pub dependencies: Vec<LauncherDependency>,
38    /// Source files keyed by launcher-crate-relative path.
39    pub files: BTreeMap<PathBuf, String>,
40}
41
42impl AotLauncherFiles {
43    /// Build launcher files from a fixed backend emission list.
44    ///
45    /// Backends should emit files in a deterministic order and delegate the
46    /// final path-keyed container construction here instead of open-coding
47    /// per-backend map assembly.
48    #[must_use]
49    pub fn from_entries(
50        dependencies: Vec<LauncherDependency>,
51        entries: impl IntoIterator<Item = (PathBuf, String)>,
52    ) -> Self {
53        Self {
54            dependencies,
55            files: entries.into_iter().collect(),
56        }
57    }
58}
59
60/// One backend-owned launcher source emitter.
61pub struct AotLauncherEmitter {
62    /// Stable target identifier.
63    pub target: AotTargetId,
64    /// Emit target-owned launcher files for `request`.
65    pub emit: fn(&AotLauncherRequest<'_>) -> Result<AotLauncherFiles, String>,
66}
67
68inventory::collect!(AotLauncherEmitter);
69
70/// Return every linked launcher emitter.
71#[must_use]
72pub fn registered_aot_launcher_emitters() -> Vec<&'static AotLauncherEmitter> {
73    let emitter_count = inventory::iter::<AotLauncherEmitter>.into_iter().count();
74    let mut emitters = Vec::new();
75    let _ = emitters.try_reserve_exact(emitter_count);
76    emitters.extend(inventory::iter::<AotLauncherEmitter>);
77    emitters
78}
79
80/// Emit target-owned launcher files through the linked emitter matching `target`.
81///
82/// # Errors
83///
84/// Returns [`BackendError::UnsupportedFeature`] when no linked backend owns
85/// launcher generation for `target`, or [`BackendError::KernelCompileFailed`]
86/// when the concrete launcher emitter rejects the request.
87pub fn emit_aot_launcher_target(
88    target: &AotTargetId,
89    request: &AotLauncherRequest<'_>,
90) -> Result<AotLauncherFiles, BackendError> {
91    let Some(emitter) = inventory::iter::<AotLauncherEmitter>
92        .into_iter()
93        .find(|emitter| &emitter.target == target)
94    else {
95        return Err(BackendError::UnsupportedFeature {
96            name: format!("aot launcher target `{target}`"),
97            backend: "vyre-driver".to_string(),
98        });
99    };
100    (emitter.emit)(request).map_err(|compiler_message| BackendError::KernelCompileFailed {
101        backend: target.to_string(),
102        compiler_message,
103    })
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn launcher_files_constructor_centralizes_path_keyed_container_assembly() {
112        let files = AotLauncherFiles::from_entries(
113            vec![LauncherDependency {
114                name: "libc",
115                spec: "\"0.2\"",
116            }],
117            [
118                (PathBuf::from("src/main.rs"), String::from("fn main() {}")),
119                (PathBuf::from("src/cuda_ffi.rs"), String::from("mod ffi {}")),
120            ],
121        );
122
123        assert_eq!(files.dependencies.len(), 1);
124        assert_eq!(files.files.len(), 2);
125        assert_eq!(
126            files.files[&PathBuf::from("src/main.rs")],
127            "fn main() {}",
128            "Fix: launcher file construction must preserve emitted file contents while centralizing the map-shaped public API."
129        );
130    }
131}
vyre_driver/aot.rs

vyre_driver/
aot.rs
