metrale_core/
compute.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: A build-time kernel compilation target: the [`Vendor`] parsed from `HARDWARE.toml`, the [`ComputeTarget`] trait, and its nvcc implementation.
4//!
5//! No code outside this file uses the module. The kernel build script declares
6//! its own `ComputeTarget` trait and vendor implementations in
7//! `crates/kernels/build_target.rs`. The runtime side is the `GpuBackend` trait
8//! in metrale-gpu-runtime.
9//!
10//! Owner: core.
11//! Invariants: none beyond the types.
12
13use std::path::{Path, PathBuf};
14
15/// 2026-09-25: Vendor identifier parsed from a `HARDWARE.toml` `vendor` string.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17pub enum Vendor {
18    /// 2026-09-25: `nvidia` or `cuda`; [`NvidiaTarget`] compiles for it.
19    Nvidia,
20    /// 2026-09-25: `amd`, `rocm` or `hip`; no target in this module.
21    Amd,
22    /// 2026-09-25: `apple` or `metal`; no target in this module.
23    Apple,
24    /// 2026-09-25: `intel`, `oneapi` or `sycl`; no target in this module.
25    Intel,
26}
27
28impl Vendor {
29    pub fn parse(s: &str) -> Option<Self> {
30        match s.to_lowercase().as_str() {
31            "nvidia" | "cuda" => Some(Self::Nvidia),
32            "amd" | "rocm" | "hip" => Some(Self::Amd),
33            "apple" | "metal" => Some(Self::Apple),
34            "intel" | "oneapi" | "sycl" => Some(Self::Intel),
35            _ => None,
36        }
37    }
38}
39
40impl std::fmt::Display for Vendor {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        match self {
43            Self::Nvidia => write!(f, "nvidia"),
44            Self::Amd => write!(f, "amd"),
45            Self::Apple => write!(f, "apple"),
46            Self::Intel => write!(f, "intel"),
47        }
48    }
49}
50
51/// 2026-09-25: How one vendor compiles kernel source files into loadable
52/// modules. [`NvidiaTarget`] is the only implementation.
53///
54/// The diagram shows the flow of the kernel build, which uses its own trait of
55/// the same name in `crates/kernels/build_target.rs`, not this one.
56///
57/// ```text
58/// [build.rs]                                [runtime]
59///
60/// .cu / .metal / .cl                        GpuBackend::new(modules)
61///        │                                         │
62///        ▼                                         ▼
63///  ComputeTarget::compile()              GpuBackend::kernel(name, fn)
64///        │                                         │
65///        ▼                                         ▼
66///   .ptx / .metallib / .spv              GpuBackend::launch(handle, ...)
67///        │
68///        ▼
69///  include_str!() / include_bytes!()
70///        │
71///        ▼
72///  Embedded in binary as &str / &[u8]
73/// ```
74pub trait ComputeTarget {
75    /// 2026-09-25: File extension of kernel source files, without the dot
76    /// (`"cu"` for [`NvidiaTarget`]).
77    fn source_extension(&self) -> &str;
78
79    /// 2026-09-25: File extension of compiled modules, without the dot
80    /// (`"ptx"` for [`NvidiaTarget`]).
81    fn output_extension(&self) -> &str;
82
83    /// 2026-09-25: Whether the compiled output is text (`include_str!` can
84    /// embed it) or binary (`include_bytes!` only).
85    fn output_is_text(&self) -> bool;
86
87    /// 2026-09-25: The compiler executable, or `None` if it is not installed.
88    fn find_compiler(&self) -> Option<PathBuf>;
89
90    /// 2026-09-25: Compile one kernel source file.
91    ///
92    /// - `source`: the kernel source file;
93    /// - `output`: where to write the compiled module;
94    /// - `arch`: the target architecture string (e.g. `"sm_121f"`);
95    /// - `extra_flags`: appended to the compiler arguments.
96    ///
97    /// Returns `Err` with a message if the compiler cannot be started or
98    /// fails; for [`NvidiaTarget`] the message carries nvcc's stderr, and a
99    /// path that is not UTF-8 panics.
100    fn compile(
101        &self,
102        source: &Path,
103        output: &Path,
104        arch: &str,
105        extra_flags: &[String],
106    ) -> Result<(), String>;
107
108    fn vendor(&self) -> Vendor;
109}
110
111/// 2026-09-25: NVIDIA compilation target: `.cu` to PTX with
112/// `nvcc --ptx -arch=<arch> -O3`.
113pub struct NvidiaTarget {
114    nvcc_path: PathBuf,
115}
116
117impl NvidiaTarget {
118    /// 2026-09-25: Locate nvcc (see `find_nvcc`); `None` if it is not found.
119    pub fn new() -> Option<Self> {
120        let nvcc = find_nvcc()?;
121        Some(Self { nvcc_path: nvcc })
122    }
123
124    pub fn with_compiler(nvcc_path: PathBuf) -> Self {
125        Self { nvcc_path }
126    }
127}
128
129impl ComputeTarget for NvidiaTarget {
130    fn source_extension(&self) -> &str {
131        "cu"
132    }
133
134    fn output_extension(&self) -> &str {
135        "ptx"
136    }
137
138    fn output_is_text(&self) -> bool {
139        true
140    }
141
142    fn find_compiler(&self) -> Option<PathBuf> {
143        Some(self.nvcc_path.clone())
144    }
145
146    fn compile(
147        &self,
148        source: &Path,
149        output: &Path,
150        arch: &str,
151        extra_flags: &[String],
152    ) -> Result<(), String> {
153        let arch_flag = format!("-arch={arch}");
154        let mut args = vec!["--ptx".to_string(), arch_flag, "-O3".to_string()];
155        args.extend(extra_flags.iter().cloned());
156        args.push(source.to_str().unwrap().to_string());
157        args.push("-o".to_string());
158        args.push(output.to_str().unwrap().to_string());
159
160        let result = std::process::Command::new(&self.nvcc_path)
161            .args(&args)
162            .output()
163            .map_err(|e| format!("Failed to run nvcc: {e}"))?;
164
165        if result.status.success() {
166            Ok(())
167        } else {
168            let stderr = String::from_utf8_lossy(&result.stderr);
169            Err(format!(
170                "nvcc --ptx failed for {}: {}",
171                source.display(),
172                stderr
173            ))
174        }
175    }
176
177    fn vendor(&self) -> Vendor {
178        Vendor::Nvidia
179    }
180}
181
182/// 2026-09-25: Locate nvcc: `bin/nvcc` under `CUDA_HOME`, `CUDA_PATH` or
183/// `CUDA_ROOT`, then four fixed install paths, then `PATH`.
184fn find_nvcc() -> Option<PathBuf> {
185    for var in ["CUDA_HOME", "CUDA_PATH", "CUDA_ROOT"] {
186        if let Ok(dir) = std::env::var(var) {
187            let nvcc = PathBuf::from(dir).join("bin/nvcc");
188            if nvcc.exists() {
189                return Some(nvcc);
190            }
191        }
192    }
193    for path in [
194        "/usr/local/cuda/bin/nvcc",
195        "/usr/local/cuda-13.0/bin/nvcc",
196        "/usr/local/cuda-12.0/bin/nvcc",
197        "/opt/cuda/bin/nvcc",
198    ] {
199        let p = PathBuf::from(path);
200        if p.exists() {
201            return Some(p);
202        }
203    }
204    which_in_path("nvcc")
205}
206
207fn which_in_path(name: &str) -> Option<PathBuf> {
208    std::env::var_os("PATH").and_then(|paths| {
209        std::env::split_paths(&paths)
210            .map(|dir| dir.join(name))
211            .find(|p| p.exists())
212    })
213}
214
215/// 2026-09-25: Resolve the [`ComputeTarget`] for a `HARDWARE.toml` vendor
216/// string. A missing or unrecognised vendor resolves to [`NvidiaTarget`].
217///
218/// Panics if nvcc is not found, or if the vendor is AMD, Apple or Intel.
219pub fn target_for_vendor(vendor: Option<&str>) -> Box<dyn ComputeTarget> {
220    match vendor.and_then(Vendor::parse) {
221        Some(Vendor::Nvidia) | None => {
222            Box::new(NvidiaTarget::new().expect("nvcc not found — install CUDA toolkit"))
223        }
224        Some(v) => {
225            panic!("Compute target '{v}' is not yet implemented. Only 'nvidia' is supported.")
226        }
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233
234    #[test]
235    fn vendor_parse_accepts_supported_names() {
236        assert_eq!(Vendor::parse("nvidia"), Some(Vendor::Nvidia));
237        assert_eq!(Vendor::parse("CUDA"), Some(Vendor::Nvidia));
238        assert_eq!(Vendor::parse("amd"), Some(Vendor::Amd));
239        assert_eq!(Vendor::parse("rocm"), Some(Vendor::Amd));
240        assert_eq!(Vendor::parse("hip"), Some(Vendor::Amd));
241        assert_eq!(Vendor::parse("apple"), Some(Vendor::Apple));
242        assert_eq!(Vendor::parse("metal"), Some(Vendor::Apple));
243        assert_eq!(Vendor::parse("intel"), Some(Vendor::Intel));
244        assert_eq!(Vendor::parse("oneapi"), Some(Vendor::Intel));
245        assert_eq!(Vendor::parse("sycl"), Some(Vendor::Intel));
246        assert_eq!(Vendor::parse("unknown"), None);
247    }
248
249    #[test]
250    fn nvidia_target_metadata() {
251        let target = NvidiaTarget::with_compiler(PathBuf::from("nvcc"));
252        assert_eq!(target.source_extension(), "cu");
253        assert_eq!(target.output_extension(), "ptx");
254        assert!(target.output_is_text());
255        assert_eq!(target.vendor(), Vendor::Nvidia);
256    }
257}
metrale_core/compute.rs

metrale_core/
compute.rs
