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    /// There is no Metal device.
73    NoDevice,
74}
75
76/// Whether products run on the Neural Accelerator kernels ([`crate::Device::nax`]).
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78#[non_exhaustive]
79pub enum NaxStatus {
80    /// They run where MLX would route to its NAX kernels; `forced` when `NUVEC_NAX=on`
81    /// enabled them on a GPU MLX would not (TensorOps run on every Apple GPU of macOS 26.2,
82    /// not on matrix units).
83    Active {
84        /// Enabled by `NUVEC_NAX=on` rather than MLX's eligibility rule.
85        forced: bool,
86    },
87    /// They do not run.
88    Off(NaxReason),
89}
90
91impl NaxStatus {
92    /// Whether NAX kernels run.
93    pub fn is_active(&self) -> bool {
94        matches!(self, Self::Active { .. })
95    }
96}
97
98/// The `NUVEC_NAX` override: `auto` (default), `on` or `off`.
99#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
101pub(crate) enum NaxOverride {
102    Auto,
103    On,
104    Off,
105}
106
107#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
108impl NaxOverride {
109    /// From `NUVEC_NAX` (case-insensitive; anything else is `auto`, with a warning).
110    pub(crate) fn from_env() -> Self {
111        match std::env::var("NUVEC_NAX") {
112            Err(_) => Self::Auto,
113            Ok(v) => match v.to_ascii_lowercase().as_str() {
114                "" | "auto" => Self::Auto,
115                "on" | "1" => Self::On,
116                "off" | "0" => Self::Off,
117                other => {
118                    eprintln!("nuvec: ignoring NUVEC_NAX={other:?} (auto, on or off)");
119                    Self::Auto
120                }
121            },
122        }
123    }
124}
125
126#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
127/// MLX's NAX eligibility (`is_nax_available`): macOS 26.2 or later and GPU generation 17 or
128/// later (18 or later for the `p` class).
129pub(crate) fn nax_eligible(generation: Option<u32>, class: Option<char>, os: (u32, u32)) -> bool {
130    let min = if class == Some('p') { 18 } else { 17 };
131    os >= (26, 2) && generation.is_some_and(|g| g >= min)
132}
133
134#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
135/// The NAX decision: the override, the OS (TensorOps need macOS 26.2 even when forced),
136/// the library, then MLX's rule, with `auto` off on eligible GPUs (no recorded M5 measurement).
137pub(crate) fn nax_status(
138    generation: Option<u32>,
139    class: Option<char>,
140    os: (u32, u32),
141    overriding: NaxOverride,
142    library: bool,
143) -> NaxStatus {
144    let eligible = nax_eligible(generation, class, os);
145    match overriding {
146        NaxOverride::Off => NaxStatus::Off(NaxReason::Override),
147        _ if os < (26, 2) => NaxStatus::Off(NaxReason::Os),
148        _ if !library => NaxStatus::Off(NaxReason::NoLibrary),
149        NaxOverride::On => NaxStatus::Active { forced: !eligible },
150        NaxOverride::Auto if eligible => NaxStatus::Off(NaxReason::Unmeasured),
151        NaxOverride::Auto => NaxStatus::Off(NaxReason::Generation),
152    }
153}
154
155#[cfg(test)]
156mod nax_tests {
157    use super::*;
158
159    #[test]
160    fn the_nax_decision_follows_mlx_and_the_override() {
161        use NaxOverride::{Auto, Off, On};
162        let off = |r| NaxStatus::Off(r);
163        // MLX's rule over generations, classes and OS versions.
164        for generation in [None, Some(13), Some(16), Some(17), Some(18)] {
165            for class in ['g', 's', 'd', 'p', 'c'] {
166                for os in [(15, 6), (26, 0), (26, 1), (26, 2), (26, 6), (27, 0)] {
167                    let min = if class == 'p' { 18 } else { 17 };
168                    let eligible = os >= (26, 2) && generation.is_some_and(|g| g >= min);
169                    assert_eq!(nax_eligible(generation, Some(class), os), eligible);
170                    let auto = nax_status(generation, Some(class), os, Auto, true);
171                    let want = match (os >= (26, 2), eligible) {
172                        (false, _) => off(NaxReason::Os),
173                        (true, true) => off(NaxReason::Unmeasured),
174                        (true, false) => off(NaxReason::Generation),
175                    };
176                    assert_eq!(auto, want, "{generation:?} {class} {os:?}");
177                    let on = nax_status(generation, Some(class), os, On, true);
178                    let want = match os >= (26, 2) {
179                        true => NaxStatus::Active { forced: !eligible },
180                        false => off(NaxReason::Os),
181                    };
182                    assert_eq!(on, want);
183                    assert_eq!(
184                        nax_status(generation, Some(class), os, Off, true),
185                        off(NaxReason::Override)
186                    );
187                }
188            }
189        }
190        // No library: never active.
191        assert_eq!(
192            nax_status(Some(17), Some('s'), (26, 2), On, false),
193            off(NaxReason::NoLibrary)
194        );
195    }
196}
197
198/// The machine's physical memory in bytes, with or without a GPU: `NSProcessInfo` on
199/// macOS, `MemTotal` of `/proc/meminfo` on Linux, 0 where neither is available.
200pub fn physical_memory() -> u64 {
201    #[cfg(target_os = "macos")]
202    return objc2_foundation::NSProcessInfo::processInfo().physicalMemory();
203    #[cfg(not(target_os = "macos"))]
204    std::fs::read_to_string("/proc/meminfo")
205        .ok()
206        .and_then(|s| {
207            let line = s.lines().find(|l| l.starts_with("MemTotal:"))?;
208            let kib: u64 = line.split_whitespace().nth(1)?.parse().ok()?;
209            Some(kib * 1024)
210        })
211        .unwrap_or(0)
212}
nuvec_gpu/info.rs

nuvec_gpu/
info.rs
