atlas_core/
compute.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Hardware-agnostic compute target abstraction.
4//!
5//! Defines the [`ComputeTarget`] trait: the contract that any GPU compilation
6//! and runtime target must satisfy (NVIDIA/PTX, AMD/HSACO, Apple/Metal, etc.).
7//!
8//! Atlas is designed so that:
9//! - **Build time**: kernel source files are compiled by a target-specific
10//!   compiler into a target-specific binary format (PTX, SPIR-V, metallib).
11//! - **Runtime**: the binary modules are loaded via `GpuBackend::kernel()`
12//!   and executed via `GpuBackend::launch()`.
13//!
14//! This module covers the **build-time** contract. The runtime contract is
15//! defined by `GpuBackend` in spark-runtime.
16//!
17//! # Extending to new hardware
18//!
19//! 1. Create a `HARDWARE.toml` in `kernels/<hw>/` with `vendor = "<vendor>"`.
20//! 2. Implement [`ComputeTarget`] for the new vendor.
21//! 3. Write kernel source files in the vendor's language (`.cu`, `.metal`, `.cl`).
22//! 4. Implement `GpuBackend` in spark-runtime for the vendor's runtime API.
23
24use std::path::{Path, PathBuf};
25
26/// Vendor identifier parsed from `HARDWARE.toml`.
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
28pub enum Vendor {
29    /// NVIDIA CUDA — compiles `.cu` → PTX via nvcc.
30    Nvidia,
31    /// AMD ROCm — compiles `.cu`/`.hip` → HSACO via hipcc (future).
32    Amd,
33    /// Apple Metal — compiles `.metal` → metallib via xcrun (future).
34    Apple,
35    /// Intel oneAPI — compiles `.cl`/`.sycl` → SPIR-V via icpx (future).
36    Intel,
37}
38
39impl Vendor {
40    pub fn parse(s: &str) -> Option<Self> {
41        match s.to_lowercase().as_str() {
42            "nvidia" | "cuda" => Some(Self::Nvidia),
43            "amd" | "rocm" | "hip" => Some(Self::Amd),
44            "apple" | "metal" => Some(Self::Apple),
45            "intel" | "oneapi" | "sycl" => Some(Self::Intel),
46            _ => None,
47        }
48    }
49}
50
51impl std::fmt::Display for Vendor {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        match self {
54            Self::Nvidia => write!(f, "nvidia"),
55            Self::Amd => write!(f, "amd"),
56            Self::Apple => write!(f, "apple"),
57            Self::Intel => write!(f, "intel"),
58        }
59    }
60}
61
62/// Build-time compilation target contract.
63///
64/// Each hardware vendor implements this trait to describe how kernel source
65/// files are compiled into loadable binary modules. The build script
66/// (`atlas-kernels/build.rs`) uses this trait to compile kernels without
67/// knowing the specific compiler or binary format.
68///
69/// # Lifecycle
70///
71/// ```text
72/// [build.rs]                                [runtime]
73///
74/// .cu / .metal / .cl                        GpuBackend::new(modules)
75///        │                                         │
76///        ▼                                         ▼
77///  ComputeTarget::compile()              GpuBackend::kernel(name, fn)
78///        │                                         │
79///        ▼                                         ▼
80///   .ptx / .metallib / .spv              GpuBackend::launch(handle, ...)
81///        │
82///        ▼
83///  include_str!() / include_bytes!()
84///        │
85///        ▼
86///  Embedded in binary as &str / &[u8]
87/// ```
88pub trait ComputeTarget {
89    /// File extension for kernel source files (without dot).
90    ///
91    /// Examples: `"cu"` (NVIDIA), `"metal"` (Apple), `"cl"` (OpenCL).
92    fn source_extension(&self) -> &str;
93
94    /// File extension for compiled kernel modules (without dot).
95    ///
96    /// Examples: `"ptx"` (NVIDIA), `"metallib"` (Apple), `"spv"` (SPIR-V).
97    fn output_extension(&self) -> &str;
98
99    /// Whether compiled output is text (can use `include_str!`) or binary
100    /// (must use `include_bytes!`).
101    ///
102    /// PTX is text; SPIR-V and metallib are binary.
103    fn output_is_text(&self) -> bool;
104
105    /// Find the compiler executable path.
106    ///
107    /// Returns `None` if the compiler is not installed.
108    fn find_compiler(&self) -> Option<PathBuf>;
109
110    /// Compile a single kernel source file to the target binary format.
111    ///
112    /// - `source`: path to the kernel source file (e.g., `rms_norm.cu`)
113    /// - `output`: path to write the compiled output (e.g., `rms_norm.ptx`)
114    /// - `arch`: target architecture string from HARDWARE.toml (e.g., `"sm_121f"`)
115    /// - `extra_flags`: additional compiler flags from KERNEL.toml
116    ///
117    /// Returns `Ok(())` on success, `Err` with compiler output on failure.
118    fn compile(
119        &self,
120        source: &Path,
121        output: &Path,
122        arch: &str,
123        extra_flags: &[String],
124    ) -> Result<(), String>;
125
126    /// Hardware vendor for this target.
127    fn vendor(&self) -> Vendor;
128}
129
130/// NVIDIA CUDA compilation target: `.cu` → PTX via `nvcc`.
131///
132/// This is the only concrete implementation today. Other vendors can be
133/// added by implementing [`ComputeTarget`] and wiring into `build.rs`.
134pub struct NvidiaTarget {
135    nvcc_path: PathBuf,
136}
137
138impl NvidiaTarget {
139    /// Create a new NVIDIA target, locating nvcc from CUDA_HOME or PATH.
140    pub fn new() -> Option<Self> {
141        let nvcc = find_nvcc()?;
142        Some(Self { nvcc_path: nvcc })
143    }
144
145    /// Create with an explicit nvcc path.
146    pub fn with_compiler(nvcc_path: PathBuf) -> Self {
147        Self { nvcc_path }
148    }
149}
150
151impl ComputeTarget for NvidiaTarget {
152    fn source_extension(&self) -> &str {
153        "cu"
154    }
155
156    fn output_extension(&self) -> &str {
157        "ptx"
158    }
159
160    fn output_is_text(&self) -> bool {
161        true // PTX is human-readable text
162    }
163
164    fn find_compiler(&self) -> Option<PathBuf> {
165        Some(self.nvcc_path.clone())
166    }
167
168    fn compile(
169        &self,
170        source: &Path,
171        output: &Path,
172        arch: &str,
173        extra_flags: &[String],
174    ) -> Result<(), String> {
175        let arch_flag = format!("-arch={arch}");
176        let mut args = vec!["--ptx".to_string(), arch_flag, "-O3".to_string()];
177        args.extend(extra_flags.iter().cloned());
178        args.push(source.to_str().unwrap().to_string());
179        args.push("-o".to_string());
180        args.push(output.to_str().unwrap().to_string());
181
182        let result = std::process::Command::new(&self.nvcc_path)
183            .args(&args)
184            .output()
185            .map_err(|e| format!("Failed to run nvcc: {e}"))?;
186
187        if result.status.success() {
188            Ok(())
189        } else {
190            let stderr = String::from_utf8_lossy(&result.stderr);
191            Err(format!(
192                "nvcc --ptx failed for {}: {}",
193                source.display(),
194                stderr
195            ))
196        }
197    }
198
199    fn vendor(&self) -> Vendor {
200        Vendor::Nvidia
201    }
202}
203
204/// Locate nvcc from CUDA_HOME, CUDA_PATH, or standard install locations.
205fn find_nvcc() -> Option<PathBuf> {
206    // Check CUDA_HOME / CUDA_PATH environment variables
207    for var in ["CUDA_HOME", "CUDA_PATH", "CUDA_ROOT"] {
208        if let Ok(dir) = std::env::var(var) {
209            let nvcc = PathBuf::from(dir).join("bin/nvcc");
210            if nvcc.exists() {
211                return Some(nvcc);
212            }
213        }
214    }
215    // Check standard install locations
216    for path in [
217        "/usr/local/cuda/bin/nvcc",
218        "/usr/local/cuda-13.0/bin/nvcc",
219        "/usr/local/cuda-12.0/bin/nvcc",
220        "/opt/cuda/bin/nvcc",
221    ] {
222        let p = PathBuf::from(path);
223        if p.exists() {
224            return Some(p);
225        }
226    }
227    // Check PATH
228    which_in_path("nvcc")
229}
230
231fn which_in_path(name: &str) -> Option<PathBuf> {
232    std::env::var_os("PATH").and_then(|paths| {
233        std::env::split_paths(&paths)
234            .map(|dir| dir.join(name))
235            .find(|p| p.exists())
236    })
237}
238
239/// Resolve the appropriate [`ComputeTarget`] from a HARDWARE.toml vendor field.
240///
241/// Falls back to [`NvidiaTarget`] if no vendor is specified (backward compat).
242pub fn target_for_vendor(vendor: Option<&str>) -> Box<dyn ComputeTarget> {
243    match vendor.and_then(Vendor::parse) {
244        Some(Vendor::Nvidia) | None => {
245            Box::new(NvidiaTarget::new().expect("nvcc not found — install CUDA toolkit"))
246        }
247        Some(v) => {
248            panic!("Compute target '{v}' is not yet implemented. Only 'nvidia' is supported.")
249        }
250    }
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    #[test]
258    fn vendor_parse_accepts_supported_names() {
259        assert_eq!(Vendor::parse("nvidia"), Some(Vendor::Nvidia));
260        assert_eq!(Vendor::parse("CUDA"), Some(Vendor::Nvidia));
261        assert_eq!(Vendor::parse("amd"), Some(Vendor::Amd));
262        assert_eq!(Vendor::parse("rocm"), Some(Vendor::Amd));
263        assert_eq!(Vendor::parse("hip"), Some(Vendor::Amd));
264        assert_eq!(Vendor::parse("apple"), Some(Vendor::Apple));
265        assert_eq!(Vendor::parse("metal"), Some(Vendor::Apple));
266        assert_eq!(Vendor::parse("intel"), Some(Vendor::Intel));
267        assert_eq!(Vendor::parse("oneapi"), Some(Vendor::Intel));
268        assert_eq!(Vendor::parse("sycl"), Some(Vendor::Intel));
269        assert_eq!(Vendor::parse("unknown"), None);
270    }
271
272    #[test]
273    fn nvidia_target_metadata() {
274        let target = NvidiaTarget::with_compiler(PathBuf::from("nvcc"));
275        assert_eq!(target.source_extension(), "cu");
276        assert_eq!(target.output_extension(), "ptx");
277        assert!(target.output_is_text());
278        assert_eq!(target.vendor(), Vendor::Nvidia);
279    }
280}
atlas_core/compute.rs

atlas_core/
compute.rs
