nuvec_gpu/
info.rs

1//! Device facts.
2
3/// What the GPU is and what it can do.
4#[derive(Debug, Clone, PartialEq, Eq)]
5#[non_exhaustive]
6pub struct DeviceInfo {
7    /// Marketing name, e.g. `Apple M4 Pro`.
8    pub name: String,
9    /// Metal's architecture string, e.g. `applegpu_g16s`.
10    pub architecture: String,
11    /// GPU generation parsed from the architecture string (16 for `applegpu_g16s`), if any.
12    /// Neural-accelerator kernels need generation 17 or later.
13    pub generation: Option<u32>,
14    /// Class letter after the generation (`s` for `applegpu_g16s`), if any. Kept verbatim:
15    /// the letter-to-tier mapping is not documented by Apple.
16    pub class: Option<char>,
17    /// Bytes the GPU can use without hurting performance.
18    pub recommended_working_set: u64,
19    /// Maximum threads per threadgroup along one dimension.
20    pub max_threads_per_threadgroup: usize,
21    /// The machine's physical memory in bytes (unified memory: shared by the CPU and the
22    /// GPU; `hw.memsize`, MLX's `device_info()["memory_size"]`).
23    pub physical_memory: u64,
24}
25
26/// Generation and class letter of an `applegpu_g<NN><class>` architecture string.
27#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
28pub(crate) fn parse_architecture(arch: &str) -> (Option<u32>, Option<char>) {
29    let Some(rest) = arch.strip_prefix("applegpu_g") else {
30        return (None, None);
31    };
32    let digits = rest.chars().take_while(char::is_ascii_digit).count();
33    let generation = rest[..digits].parse().ok();
34    let class = rest[digits..]
35        .chars()
36        .next()
37        .filter(char::is_ascii_lowercase);
38    (generation, class)
39}
40
41#[cfg(test)]
42mod tests {
43    use super::parse_architecture;
44
45    #[test]
46    fn parses_apple_gpu_architectures() {
47        assert_eq!(parse_architecture("applegpu_g16s"), (Some(16), Some('s')));
48        assert_eq!(parse_architecture("applegpu_g17p"), (Some(17), Some('p')));
49        assert_eq!(parse_architecture("applegpu_g13"), (Some(13), None));
50        assert_eq!(parse_architecture("Paravirtual"), (None, None));
51    }
52}
53
54/// Why the Neural Accelerator (NAX) kernels are not used ([`NaxStatus::Off`]).
55#[derive(Debug, Clone, Copy, PartialEq, Eq)]
56#[non_exhaustive]
57pub enum NaxReason {
58    /// `NUVEC_NAX=off`.
59    Override,
60    /// The OS is older than macOS 26.2 (Metal 4 TensorOps).
61    Os,
62    /// The GPU is not an M5-class one (MLX's rule: generation 17 or later, 18 for the `p`
63    /// class), or its generation is unknown.
64    Generation,
65    /// This build has no NAX kernel library (an SDK older than 26.2 or no Metal).
66    NoLibrary,
67    /// An eligible GPU, but nuvec's NAX kernels have no recorded M5 measurement, so `auto`
68    /// keeps them off. `NUVEC_NAX=on` forces them.
69    Unmeasured,
70    /// A NAX pipeline failed to build; the existing kernels are used from then on.
71    PipelineError,
72}
73
74/// Whether products run on the Neural Accelerator kernels ([`crate::Device::nax`]).
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76#[non_exhaustive]
77pub enum NaxStatus {
78    /// They run where MLX would route to its NAX kernels; `forced` when `NUVEC_NAX=on`
79    /// enabled them on a GPU MLX would not (TensorOps run on every Apple GPU of macOS 26.2,
80    /// not on matrix units).
81    Active {
82        /// Enabled by `NUVEC_NAX=on` rather than MLX's eligibility rule.
83        forced: bool,
84    },
85    /// They do not run.
86    Off(NaxReason),
87}
88
89impl NaxStatus {
90    /// Whether NAX kernels run.
91    pub fn is_active(&self) -> bool {
92        matches!(self, Self::Active { .. })
93    }
94}
95
96/// The `NUVEC_NAX` override: `auto` (default), `on` or `off`.
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
99pub(crate) enum NaxOverride {
100    Auto,
101    On,
102    Off,
103}
104
105#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
106impl NaxOverride {
107    /// From `NUVEC_NAX` (case-insensitive; anything else is `auto`, with a warning).
108    pub(crate) fn from_env() -> Self {
109        match std::env::var("NUVEC_NAX") {
110            Err(_) => Self::Auto,
111            Ok(v) => match v.to_ascii_lowercase().as_str() {
112                "" | "auto" => Self::Auto,
113                "on" | "1" => Self::On,
114                "off" | "0" => Self::Off,
115                other => {
116                    eprintln!("nuvec: ignoring NUVEC_NAX={other:?} (auto, on or off)");
117                    Self::Auto
118                }
119            },
120        }
121    }
122}
123
124#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
125/// MLX's NAX eligibility (`is_nax_available`): macOS 26.2 or later and GPU generation 17 or
126/// later (18 or later for the `p` class).
127pub(crate) fn nax_eligible(generation: Option<u32>, class: Option<char>, os: (u32, u32)) -> bool {
128    let min = if class == Some('p') { 18 } else { 17 };
129    os >= (26, 2) && generation.is_some_and(|g| g >= min)
130}
131
132#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
133/// The NAX decision: the override, the OS (TensorOps need macOS 26.2 even when forced),
134/// the library, then MLX's rule, with `auto` off on eligible GPUs (no recorded M5 measurement).
135pub(crate) fn nax_status(
136    generation: Option<u32>,
137    class: Option<char>,
138    os: (u32, u32),
139    overriding: NaxOverride,
140    library: bool,
141) -> NaxStatus {
142    let eligible = nax_eligible(generation, class, os);
143    match overriding {
144        NaxOverride::Off => NaxStatus::Off(NaxReason::Override),
145        _ if os < (26, 2) => NaxStatus::Off(NaxReason::Os),
146        _ if !library => NaxStatus::Off(NaxReason::NoLibrary),
147        NaxOverride::On => NaxStatus::Active { forced: !eligible },
148        NaxOverride::Auto if eligible => NaxStatus::Off(NaxReason::Unmeasured),
149        NaxOverride::Auto => NaxStatus::Off(NaxReason::Generation),
150    }
151}
152
153#[cfg(test)]
154mod nax_tests {
155    use super::*;
156
157    #[test]
158    fn the_nax_decision_follows_mlx_and_the_override() {
159        use NaxOverride::{Auto, Off, On};
160        let off = |r| NaxStatus::Off(r);
161        // MLX's rule over generations, classes and OS versions.
162        for generation in [None, Some(13), Some(16), Some(17), Some(18)] {
163            for class in ['g', 's', 'd', 'p', 'c'] {
164                for os in [(15, 6), (26, 0), (26, 1), (26, 2), (26, 6), (27, 0)] {
165                    let min = if class == 'p' { 18 } else { 17 };
166                    let eligible = os >= (26, 2) && generation.is_some_and(|g| g >= min);
167                    assert_eq!(nax_eligible(generation, Some(class), os), eligible);
168                    let auto = nax_status(generation, Some(class), os, Auto, true);
169                    let want = match (os >= (26, 2), eligible) {
170                        (false, _) => off(NaxReason::Os),
171                        (true, true) => off(NaxReason::Unmeasured),
172                        (true, false) => off(NaxReason::Generation),
173                    };
174                    assert_eq!(auto, want, "{generation:?} {class} {os:?}");
175                    let on = nax_status(generation, Some(class), os, On, true);
176                    let want = match os >= (26, 2) {
177                        true => NaxStatus::Active { forced: !eligible },
178                        false => off(NaxReason::Os),
179                    };
180                    assert_eq!(on, want);
181                    assert_eq!(
182                        nax_status(generation, Some(class), os, Off, true),
183                        off(NaxReason::Override)
184                    );
185                }
186            }
187        }
188        // No library: never active.
189        assert_eq!(
190            nax_status(Some(17), Some('s'), (26, 2), On, false),
191            off(NaxReason::NoLibrary)
192        );
193    }
194}
nuvec_gpu/info.rs

nuvec_gpu/
info.rs
