metrale_core/
arch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Decide whether PTX built for one SM architecture runs on a device, and name the target to rebuild when it does not.
4//!
5//! A kernel build compiles for the single arch in `kernels/<hw>/HARDWARE.toml`
6//! `[hardware].arch` (one `-arch=` flag per nvcc call), so a binary and a GPU
7//! can disagree. The CUDA preflight (`arch_preflight::check_arch`) calls
8//! [`ptx_arch_runs_on_device`], and `--check-kernels` reads the resulting
9//! [`ArchMismatch`] (`serve_phases/kernel_gate.rs`).
10//!
11//! Owner: core.
12//! Invariants: every function here is pure; the module uses only `std`, so it
13//! needs no GPU and no driver.
14
15/// 2026-09-25: The suffix on an `sm_XY…` architecture string. It selects the
16/// compatibility rule in [`ptx_arch_runs_on_device`].
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum SmSuffix {
19    /// 2026-09-25: No suffix: accepted on any device with a compute capability
20    /// at or above the compiled one.
21    None,
22    /// 2026-09-25: `a`, architecture-specific: accepted only on exactly the
23    /// compiled compute capability.
24    Arch,
25    /// 2026-09-25: `f`, family-specific: accepted on devices of the same major
26    /// version at or above the compiled compute capability.
27    Family,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct SmArch {
32    /// 2026-09-25: Compute-capability major (the `12` of `sm_121f`).
33    pub major: u32,
34    /// 2026-09-25: Compute-capability minor (the `1` of `sm_121f`).
35    pub minor: u32,
36    pub suffix: SmSuffix,
37}
38
39/// 2026-09-25: Parse an NVIDIA `sm_XY[a|f]` architecture string.
40///
41/// Returns `None` for anything that is not an NVIDIA SM arch: `gfx1151`,
42/// `metal3.1`, fewer than two digits, or a non-digit. `None` is the "not an
43/// NVIDIA target" marker.
44///
45/// The last digit is the minor version and the rest is the major
46/// (`sm_90` = 9.0, `sm_100` = 10.0, `sm_121` = 12.1).
47pub fn parse_sm_arch(arch: &str) -> Option<SmArch> {
48    let rest = arch.strip_prefix("sm_")?;
49    let (digits, suffix) = match rest.as_bytes().last()? {
50        b'a' => (&rest[..rest.len() - 1], SmSuffix::Arch),
51        b'f' => (&rest[..rest.len() - 1], SmSuffix::Family),
52        _ => (rest, SmSuffix::None),
53    };
54    // 2026-09-25: Two digits minimum, one for the major and one for the minor;
55    // `sm_9` is refused rather than guessed at.
56    if digits.len() < 2 || !digits.bytes().all(|b| b.is_ascii_digit()) {
57        return None;
58    }
59    let (major, minor) = digits.split_at(digits.len() - 1);
60    Some(SmArch {
61        major: major.parse().ok()?,
62        minor: minor.parse().ok()?,
63        suffix,
64    })
65}
66
67/// 2026-09-25: Which `kernels/<hw>/` target ships for a device compute
68/// capability, so the mismatch message can say what to rebuild. `None` means
69/// no target ships for that capability.
70///
71/// The table is hand-written although each NVIDIA `kernels/<hw>/HARDWARE.toml`
72/// declares its `compute_capability`: metrale-core has no TOML parser
73/// dependency. `crates/kernels/tests/target_hints.rs` keeps the two in step: it
74/// asserts that every `vendor = "nvidia"` HARDWARE.toml's capability maps back
75/// to its own directory, and that every name returned here is a directory.
76pub fn target_hint(device_cc: (u32, u32)) -> Option<&'static str> {
77    match device_cc {
78        (9, 0) => Some("hopper"),
79        // 2026-09-25: b200 (`sm_100a`) and b300 (`sm_103a`) are separate targets:
80        // an `a` arch runs only on its own capability.
81        (10, 0) => Some("b200"),
82        (10, 3) => Some("b300"),
83        (12, 1) => Some("gb10"),
84        _ => None,
85    }
86}
87
88#[derive(Debug, Clone, PartialEq, Eq)]
89pub struct ArchMismatch {
90    /// 2026-09-25: The arch string the kernels were compiled for, verbatim.
91    pub compiled_arch: String,
92    /// 2026-09-25: The parsed form of `compiled_arch`; its suffix picks the
93    /// reason printed by `Display`.
94    pub compiled: SmArch,
95    /// 2026-09-25: `(major, minor)` compute capability of the device.
96    pub device_cc: (u32, u32),
97}
98
99impl std::fmt::Display for ArchMismatch {
100    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101        let (major, minor) = self.device_cc;
102        let arch = &self.compiled_arch;
103        let reason = match self.compiled.suffix {
104            SmSuffix::None => format!(
105                "portable PTX for {arch} needs compute capability {}.{} or later",
106                self.compiled.major, self.compiled.minor
107            ),
108            SmSuffix::Arch => format!(
109                "architecture-specific PTX ({arch}) runs only on compute capability {}.{}",
110                self.compiled.major, self.compiled.minor
111            ),
112            SmSuffix::Family => format!(
113                "family-specific PTX ({arch}) runs only on compute capability {}.{} or later \
114                 within the {}.x family",
115                self.compiled.major, self.compiled.minor, self.compiled.major
116            ),
117        };
118        let fix = match target_hint(self.device_cc) {
119            Some(hw) => format!(
120                "rebuild with METRALE_TARGET_HW={hw} (kernels/{hw}/HARDWARE.toml arch must match \
121                 this GPU) or use the image built for this GPU"
122            ),
123            None => format!(
124                "no shipped target matches compute capability {major}.{minor} \
125                 (kernels/<hw>/HARDWARE.toml arch must match this GPU) — \
126                 use the image built for this GPU"
127            ),
128        };
129        write!(
130            f,
131            "kernels compiled for {arch} cannot run on this GPU \
132             (compute capability {major}.{minor}): {reason}; fix: {fix}"
133        )
134    }
135}
136
137impl std::error::Error for ArchMismatch {}
138
139/// 2026-09-25: Can PTX compiled for `compiled_arch` run on a device of
140/// `device_cc`? Three rules:
141///
142/// 1. plain `sm_XY` runs on any device with CC >= X.Y;
143/// 2. `sm_XYa` runs only on CC == X.Y;
144/// 3. `sm_XYf` runs on the same major version with CC >= X.Y.
145///
146/// A `compiled_arch` that is not an NVIDIA SM arch (`gfx1151`, `metal3.1`)
147/// returns `Ok(())`, so AMD and Apple builds are never refused here. Callers
148/// that need to tell the cases apart use [`parse_sm_arch`], whose `None` is the
149/// marker.
150pub fn ptx_arch_runs_on_device(
151    compiled_arch: &str,
152    device_cc: (u32, u32),
153) -> Result<(), ArchMismatch> {
154    let Some(compiled) = parse_sm_arch(compiled_arch) else {
155        return Ok(());
156    };
157    let compiled_cc = (compiled.major, compiled.minor);
158    let runs = match compiled.suffix {
159        SmSuffix::None => device_cc >= compiled_cc,
160        SmSuffix::Arch => device_cc == compiled_cc,
161        SmSuffix::Family => device_cc.0 == compiled.major && device_cc >= compiled_cc,
162    };
163    if runs {
164        return Ok(());
165    }
166    Err(ArchMismatch {
167        compiled_arch: compiled_arch.to_string(),
168        compiled,
169        device_cc,
170    })
171}
172
173#[cfg(test)]
174#[path = "arch_tests.rs"]
175mod tests;
metrale_core/arch.rs

metrale_core/
arch.rs
