xlog_cuda/provider/
kernel_paths.rs

1use std::env;
2use std::path::{Path, PathBuf};
3
4/// Locates staged CUDA kernel artifacts in build, package, or override dirs.
5#[derive(Debug, Clone, Default)]
6pub struct KernelArtifactLocator {
7    cubin_dir: Option<PathBuf>,
8    package_kernels_dir: Option<PathBuf>,
9    additional_kernel_dir: Option<PathBuf>,
10}
11
12impl KernelArtifactLocator {
13    pub fn new(
14        cubin_dir: Option<PathBuf>,
15        package_kernels_dir: Option<PathBuf>,
16        additional_kernel_dir: Option<PathBuf>,
17    ) -> Self {
18        Self {
19            cubin_dir,
20            package_kernels_dir,
21            additional_kernel_dir,
22        }
23    }
24
25    /// Build a locator from the current process environment.
26    ///
27    /// Precedence matches runtime loading:
28    /// 1. `XLOG_CUBIN_DIR`
29    /// 2. binary-adjacent `kernels/`
30    ///
31    /// Cargo's build output directory is intentionally excluded because embedding
32    /// its absolute path makes otherwise identical release artifacts differ by
33    /// build location. Source-tree consumers can use `XLOG_CUBIN_DIR`; packaged
34    /// consumers discover their staged kernels through the package bootstrap.
35    pub fn from_env() -> Self {
36        let cubin_dir = env::var_os("XLOG_CUBIN_DIR").map(PathBuf::from);
37        let package_kernels_dir = env::current_exe()
38            .ok()
39            .and_then(|exe| exe.parent().map(|dir| dir.join("kernels")));
40        Self::new(cubin_dir, package_kernels_dir, None)
41    }
42
43    /// Resolve a module to a staged cubin or portable PTX.
44    pub fn resolve_module_path(&self, name: &str, cc: u32) -> Option<(PathBuf, bool)> {
45        self.resolve_module_paths(name, cc).into_iter().next()
46    }
47
48    /// Resolve all viable module artifacts in runtime precedence order.
49    ///
50    /// A matching cubin remains the first choice, but portable PTX is retained
51    /// as a fallback when a target driver rejects a cubin image.
52    pub fn resolve_module_paths(&self, name: &str, cc: u32) -> Vec<(PathBuf, bool)> {
53        let cubin_name = format!("{name}.sm_{cc}.cubin");
54        let ptx_name = format!("{name}.portable.ptx");
55        let mut found_paths = Vec::new();
56
57        for dir in [
58            self.cubin_dir.as_ref(),
59            self.package_kernels_dir.as_ref(),
60            self.additional_kernel_dir.as_ref(),
61        ] {
62            let dir = match dir {
63                Some(dir) => dir,
64                None => continue,
65            };
66
67            if let Some(found) = Self::resolve_in_dir(dir, &cubin_name, true) {
68                found_paths.push(found);
69            }
70            if let Some(found) = Self::resolve_in_dir(dir, &ptx_name, false) {
71                found_paths.push(found);
72            }
73        }
74
75        found_paths
76    }
77
78    fn resolve_in_dir(dir: &Path, file_name: &str, is_cubin: bool) -> Option<(PathBuf, bool)> {
79        let path = dir.join(file_name);
80        if path.exists() {
81            Some((path, is_cubin))
82        } else {
83            None
84        }
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::KernelArtifactLocator;
91    use std::fs;
92
93    #[test]
94    fn runtime_locator_ignores_ephemeral_build_output_directory() {
95        let module_name = format!("runtime_locator_ephemeral_output_{}", std::process::id());
96        let generated_ptx =
97            std::path::PathBuf::from(env!("OUT_DIR")).join(format!("{module_name}.portable.ptx"));
98        fs::write(&generated_ptx, ".version 8.0\n").expect("write generated PTX sentinel");
99
100        let locator = KernelArtifactLocator::from_env();
101        let resolved = locator.resolve_module_path(&module_name, 999);
102        fs::remove_file(&generated_ptx).expect("remove generated PTX sentinel");
103
104        assert!(
105            resolved.is_none(),
106            "runtime kernel lookup must ignore Cargo's ephemeral OUT_DIR"
107        );
108    }
109
110    #[test]
111    fn resolves_in_precedence_order() {
112        let root = std::env::temp_dir().join(format!(
113            "xlog-kernel-paths-{}-{}",
114            std::process::id(),
115            std::time::SystemTime::now()
116                .duration_since(std::time::UNIX_EPOCH)
117                .expect("system clock before UNIX_EPOCH")
118                .as_nanos()
119        ));
120        let cubin_dir = root.join("cubin");
121        let package_dir = root.join("bin").join("kernels");
122        let out_dir = root.join("out");
123        fs::create_dir_all(&cubin_dir).expect("create cubin dir");
124        fs::create_dir_all(&package_dir).expect("create package kernels dir");
125        fs::create_dir_all(&out_dir).expect("create out dir");
126
127        let name = "xlog_join";
128        let cc = 75;
129        let cubin_path = cubin_dir.join(format!("{name}.sm_{cc}.cubin"));
130        let package_path = package_dir.join(format!("{name}.sm_{cc}.cubin"));
131        let out_path = out_dir.join(format!("{name}.sm_{cc}.cubin"));
132        fs::write(&cubin_path, b"cubin").expect("write cubin file");
133        fs::write(&package_path, b"package").expect("write package file");
134        fs::write(&out_path, b"out").expect("write out file");
135
136        let locator = KernelArtifactLocator::new(
137            Some(cubin_dir.clone()),
138            Some(package_dir.clone()),
139            Some(out_dir.clone()),
140        );
141
142        let (path, is_cubin) = locator
143            .resolve_module_path(name, cc)
144            .expect("expected a kernel artifact");
145        assert_eq!(path, cubin_path);
146        assert!(is_cubin);
147
148        fs::remove_file(&cubin_path).expect("remove cubin file");
149        let (path, is_cubin) = locator
150            .resolve_module_path(name, cc)
151            .expect("expected package kernel artifact");
152        assert_eq!(path, package_path);
153        assert!(is_cubin);
154
155        fs::remove_file(&package_path).expect("remove package file");
156        let (path, is_cubin) = locator
157            .resolve_module_path(name, cc)
158            .expect("expected out dir kernel artifact");
159        assert_eq!(path, out_path);
160        assert!(is_cubin);
161
162        let _ = fs::remove_dir_all(&root);
163    }
164}
xlog_cuda/provider/kernel_paths.rs

xlog_cuda/provider/
kernel_paths.rs
