metrale_circuit/hardware/
device.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: `kernels/DEVICES.toml`: the GPU SKUs a circuit is planned for, the kernel class
4//! each builds, its roofline and the tensor-core instruction kinds it runs natively, plus the
5//! source guards that compile a kernel only where an instruction kind exists.
6//!
7//! Owner: metrale-circuit (hardware).
8//! Invariants:
9//! - Nothing defaults: every key the planner reads is required; an unknown top-level key, guard
10//!   key, instruction family, `native_mma` key, source id or a duplicate device id is a typed
11//!   error. A profile's other spec keys are carried for the roadmap and not read.
12//! - A peak the device lacks (`fp4_block_scale` on Hopper, `0.0` in the research file) is
13//!   `None` here, and `native_mma` must agree with the dense peaks.
14
15use std::collections::{BTreeMap, BTreeSet};
16
17use serde::Deserialize;
18
19use super::HwError;
20
21/// 2026-09-30: The tensor-core instruction family.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
23pub enum MmaFamily {
24    /// 2026-09-30: Warp-level `mma.sync` (SM 8.x-12.x; the GB10 path).
25    MmaSync,
26    /// 2026-09-30: Hopper warpgroup `wgmma` (operand B from shared memory).
27    Wgmma,
28    /// 2026-09-30: Blackwell datacentre `tcgen05` (accumulators in tensor memory).
29    Tcgen05,
30}
31
32impl MmaFamily {
33    /// 2026-09-30: The registry spelling.
34    pub fn name(self) -> &'static str {
35        match self {
36            MmaFamily::MmaSync => "mma_sync",
37            MmaFamily::Wgmma => "wgmma",
38            MmaFamily::Tcgen05 => "tcgen05",
39        }
40    }
41
42    fn parse(s: &str) -> Option<Self> {
43        [MmaFamily::MmaSync, MmaFamily::Wgmma, MmaFamily::Tcgen05]
44            .into_iter()
45            .find(|f| f.name() == s)
46    }
47}
48
49/// 2026-09-30: The operand format class an MMA multiplies.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
51pub enum MmaKind {
52    /// 2026-09-30: BF16 x BF16.
53    Bf16,
54    /// 2026-09-30: E4M3/E5M2 x E4M3/E5M2.
55    Fp8,
56    /// 2026-09-30: INT8 x INT8.
57    Int8,
58    /// 2026-09-30: NVFP4: E2M1 with an E4M3 scale per 16 elements, both operands.
59    Fp4BlockScale,
60    /// 2026-09-30: MX formats (UE8M0 scales per 32): FP8, FP6 or FP4 operands.
61    Mxf8f6f4,
62    /// 2026-09-30: FP16 x FP16.
63    Fp16,
64    /// 2026-09-30: TF32 x TF32.
65    Tf32,
66    /// 2026-09-30: FP64 tensor.
67    Fp64,
68    /// 2026-09-30: FP4 x FP8 without block scales.
69    Fp4Fp8Unscaled,
70    /// 2026-09-30: FP4 x FP8 with NVFP4 (E4M3 per 16) scales: W4A8 in one instruction.
71    Fp4Fp8Nvfp4,
72    /// 2026-09-30: FP4 x BF16.
73    Fp4Bf16,
74}
75
76impl MmaKind {
77    /// 2026-09-30: The registry spelling.
78    pub fn name(self) -> &'static str {
79        match self {
80            MmaKind::Bf16 => "bf16",
81            MmaKind::Fp8 => "fp8",
82            MmaKind::Int8 => "int8",
83            MmaKind::Fp4BlockScale => "fp4_block_scale",
84            MmaKind::Mxf8f6f4 => "mxf8f6f4",
85            MmaKind::Fp16 => "fp16",
86            MmaKind::Tf32 => "tf32",
87            MmaKind::Fp64 => "fp64",
88            MmaKind::Fp4Fp8Unscaled => "fp4_fp8_unscaled",
89            MmaKind::Fp4Fp8Nvfp4 => "fp4_fp8_nvfp4_scaled",
90            MmaKind::Fp4Bf16 => "fp4_bf16",
91        }
92    }
93
94    fn parse(s: &str) -> Option<Self> {
95        [
96            MmaKind::Bf16,
97            MmaKind::Fp8,
98            MmaKind::Int8,
99            MmaKind::Fp4BlockScale,
100            MmaKind::Mxf8f6f4,
101            MmaKind::Fp16,
102            MmaKind::Tf32,
103            MmaKind::Fp64,
104            MmaKind::Fp4Fp8Unscaled,
105            MmaKind::Fp4Fp8Nvfp4,
106            MmaKind::Fp4Bf16,
107        ]
108        .into_iter()
109        .find(|k| k.name() == s)
110    }
111
112    /// 2026-09-30: The kind a research `native_mma` key names; `None` for an unknown key.
113    fn of_native_key(key: &str) -> Option<Self> {
114        Some(match key {
115            "bf16_bf16" => MmaKind::Bf16,
116            "fp16_fp16" => MmaKind::Fp16,
117            "tf32_tf32" => MmaKind::Tf32,
118            "fp8_fp8" => MmaKind::Fp8,
119            "int8_int8" => MmaKind::Int8,
120            "fp64" => MmaKind::Fp64,
121            "fp4_fp4_nvfp4_block16" => MmaKind::Fp4BlockScale,
122            "fp8_fp8_mx_block32" | "fp6_fp6" | "fp4_fp4_mxfp4_block32" => MmaKind::Mxf8f6f4,
123            "fp4_fp8_unscaled" => MmaKind::Fp4Fp8Unscaled,
124            "fp4_fp8_nvfp4_scaled" => MmaKind::Fp4Fp8Nvfp4,
125            "fp4_bf16" => MmaKind::Fp4Bf16,
126            _ => return None,
127        })
128    }
129}
130
131/// 2026-09-30: One instruction kind, `<family>.<kind>`.
132#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
133pub struct Instr {
134    /// 2026-09-30: Family.
135    pub family: MmaFamily,
136    /// 2026-09-30: Operand kind.
137    pub kind: MmaKind,
138}
139
140impl Instr {
141    /// 2026-09-30: Parse `mma_sync.fp8`.
142    pub fn parse(s: &str) -> Result<Self, HwError> {
143        let bad = || HwError::Registry(format!("`{s}` is not <mma_sync|wgmma|tcgen05>.<kind>"));
144        let (f, k) = s.split_once('.').ok_or_else(bad)?;
145        Ok(Instr {
146            family: MmaFamily::parse(f).ok_or_else(bad)?,
147            kind: MmaKind::parse(k).ok_or_else(bad)?,
148        })
149    }
150
151    /// 2026-09-30: The registry spelling.
152    pub fn name(self) -> String {
153        format!("{}.{}", self.family.name(), self.kind.name())
154    }
155}
156
157/// 2026-09-30: Dense tensor peaks, TFLOPS; a kind the device has no MMA for is `None`.
158#[derive(Debug, Clone, Copy, PartialEq)]
159pub struct Peaks {
160    /// 2026-09-30: BF16.
161    pub bf16: f64,
162    /// 2026-09-30: FP8.
163    pub fp8: Option<f64>,
164    /// 2026-09-30: NVFP4 block-scaled.
165    pub fp4_block_scale: Option<f64>,
166}
167
168/// 2026-09-30: One SKU.
169#[derive(Debug, Clone, PartialEq)]
170pub struct Device {
171    /// 2026-09-30: `h100-sxm`.
172    pub id: String,
173    /// 2026-09-30: Marketing name.
174    pub name: String,
175    /// 2026-09-30: Kernel class, `kernels/<class>/`.
176    pub class: String,
177    /// 2026-09-30: Build arch (`sm_90a`), equal to the class's HARDWARE.toml.
178    pub arch: String,
179    /// 2026-09-30: Compute capability.
180    pub compute_capability: String,
181    /// 2026-09-30: SM count.
182    pub sms: u32,
183    /// 2026-09-30: Visible device memory, bytes (`memory_gib_visible`, else the datasheet GB,
184    /// which tracks the visible GiB).
185    pub memory_bytes: f64,
186    /// 2026-09-30: Memory technology.
187    pub memory_type: String,
188    /// 2026-09-30: Datasheet DRAM bandwidth, GB/s.
189    pub bandwidth_gbps: f64,
190    /// 2026-09-30: L2, MB, when published.
191    pub l2_mb: Option<f64>,
192    /// 2026-09-30: Shared memory per SM, KB.
193    pub smem_per_sm_kb: u32,
194    /// 2026-09-30: Largest portable thread-block cluster.
195    pub cluster_max: u32,
196    /// 2026-09-30: Tensor memory per SM, KB (tcgen05 accumulators); 0 without.
197    pub tmem_per_sm_kb: u32,
198    /// 2026-09-30: The family fast kernels are written for.
199    pub mma_family: MmaFamily,
200    /// 2026-09-30: Instruction kinds executed natively (`mma_family` x the `native_mma` pairs).
201    pub native: BTreeSet<Instr>,
202    /// 2026-09-30: Dense datasheet peaks.
203    pub peaks: Peaks,
204    /// 2026-09-30: Where measured achievable figures live, when measured.
205    pub measured: Option<String>,
206    /// 2026-09-30: Share of memory the planner may use.
207    pub usable_fraction: f64,
208    /// 2026-09-30: Why that share.
209    pub usable_why: String,
210    /// 2026-09-30: Keys computed by the research, not published.
211    pub derived: Vec<String>,
212    /// 2026-09-30: Source ids (`S1`, ...), resolved in [`Registry::sources`].
213    pub sources: Vec<String>,
214}
215
216impl Device {
217    /// 2026-09-30: Any family runs `kind` natively.
218    pub fn runs(&self, kind: MmaKind) -> bool {
219        self.native.iter().any(|i| i.kind == kind)
220    }
221}
222
223/// 2026-09-30: Whether a guarded region compiles when its macro is defined or when it is not.
224#[derive(Debug, Clone, Copy, PartialEq, Eq)]
225pub enum Polarity {
226    /// 2026-09-30: `#ifdef M` / `#if defined(M)`: compiled only where `M` is defined.
227    Ifdef,
228    /// 2026-09-30: `#ifndef M`: compiled only where `M` is not defined.
229    Ifndef,
230}
231
232/// 2026-09-30: A source guard and the instruction kind it stands for.
233#[derive(Debug, Clone, PartialEq, Eq)]
234pub struct Guard {
235    /// 2026-09-30: The macro.
236    pub macro_name: String,
237    /// 2026-09-30: Its polarity.
238    pub polarity: Polarity,
239    /// 2026-09-30: What a kernel inside the region needs natively.
240    pub requires: Instr,
241}
242
243/// 2026-09-30: The parsed registry.
244#[derive(Debug, Clone, PartialEq)]
245pub struct Registry {
246    /// 2026-09-30: Devices, in file order.
247    pub devices: Vec<Device>,
248    /// 2026-09-30: Guards, in file order.
249    pub guards: Vec<Guard>,
250    /// 2026-09-30: Source id to citation.
251    pub sources: BTreeMap<String, String>,
252}
253
254impl Registry {
255    /// 2026-09-30: The device `id`; an unknown id lists the known ones.
256    pub fn device(&self, id: &str) -> Result<&Device, HwError> {
257        self.devices
258            .iter()
259            .find(|d| d.id == id)
260            .ok_or_else(|| HwError::UnknownDevice {
261                id: id.to_string(),
262                known: self.devices.iter().map(|d| d.id.clone()).collect(),
263            })
264    }
265}
266
267#[derive(Deserialize)]
268#[serde(deny_unknown_fields)]
269struct File {
270    schema_version: u32,
271    compiled: String,
272    sources: BTreeMap<String, String>,
273    guard: Vec<GuardFile>,
274    device: Vec<DeviceFile>,
275}
276
277#[derive(Deserialize)]
278#[serde(deny_unknown_fields)]
279struct GuardFile {
280    #[serde(rename = "macro")]
281    macro_name: String,
282    polarity: String,
283    requires: String,
284}
285
286#[derive(Deserialize)]
287struct PeaksFile {
288    bf16: f64,
289    fp8: f64,
290    fp4_nvfp4: f64,
291}
292
293// 2026-09-30: The planner's keys of one research profile. The other spec keys (clocks, NVLink,
294// PCIe, TDP, sparse peaks, ridges, ...) are carried for the roadmap and not read here.
295#[derive(Deserialize)]
296struct DeviceFile {
297    id: String,
298    name: String,
299    class: String,
300    arch: String,
301    compute_capability: String,
302    sm_count: u32,
303    memory_type: String,
304    memory_gb_datasheet: f64,
305    memory_gib_visible: Option<f64>,
306    memory_bandwidth_gbps: f64,
307    l2_mb: Option<f64>,
308    smem_per_sm_kb: u32,
309    tmem_per_sm_kb: u32,
310    mma_family: String,
311    max_cluster: u32,
312    derived: Vec<String>,
313    sources: Vec<String>,
314    peak_tflops: PeaksFile,
315    native_mma: BTreeMap<String, bool>,
316    usable_fraction: f64,
317    usable_why: String,
318    measured: Option<String>,
319}
320
321/// 2026-09-30: Parse `kernels/DEVICES.toml`.
322pub fn parse_devices(text: &str) -> Result<Registry, HwError> {
323    let f: File = toml::from_str(text).map_err(|e| HwError::Registry(e.to_string()))?;
324    if f.schema_version != 1 {
325        return Err(HwError::Registry(format!(
326            "schema_version {} (this build reads 1; compiled {})",
327            f.schema_version, f.compiled
328        )));
329    }
330    let guards = f
331        .guard
332        .into_iter()
333        .map(|g| {
334            let polarity = match g.polarity.as_str() {
335                "ifdef" => Polarity::Ifdef,
336                "ifndef" => Polarity::Ifndef,
337                other => {
338                    return Err(HwError::Registry(format!(
339                        "guard {}: polarity `{other}` (ifdef | ifndef)",
340                        g.macro_name
341                    )));
342                }
343            };
344            Ok(Guard {
345                requires: Instr::parse(&g.requires)?,
346                macro_name: g.macro_name,
347                polarity,
348            })
349        })
350        .collect::<Result<Vec<_>, _>>()?;
351    let mut devices: Vec<Device> = Vec::with_capacity(f.device.len());
352    for d in f.device {
353        devices.push(device_of(d, &devices, &f.sources)?);
354    }
355    Ok(Registry {
356        devices,
357        guards,
358        sources: f.sources,
359    })
360}
361
362fn device_of(
363    d: DeviceFile,
364    seen: &[Device],
365    sources: &BTreeMap<String, String>,
366) -> Result<Device, HwError> {
367    let field = |detail: String| HwError::Registry(format!("device {}: {detail}", d.id));
368    if seen.iter().any(|x| x.id == d.id) {
369        return Err(field("duplicate id".into()));
370    }
371    if let Some(s) = d.sources.iter().find(|s| !sources.contains_key(*s)) {
372        return Err(field(format!("source `{s}` is not in [sources]")));
373    }
374    let mma_family = MmaFamily::parse(&d.mma_family)
375        .ok_or_else(|| field(format!("mma_family `{}`", d.mma_family)))?;
376    let mut native = BTreeSet::new();
377    for (key, on) in &d.native_mma {
378        let kind = MmaKind::of_native_key(key)
379            .ok_or_else(|| field(format!("native_mma key `{key}` is no instruction kind")))?;
380        if *on {
381            native.insert(Instr {
382                family: mma_family,
383                kind,
384            });
385        }
386    }
387    if !(d.usable_fraction > 0.0 && d.usable_fraction <= 1.0) {
388        return Err(field(format!("usable_fraction {}", d.usable_fraction)));
389    }
390    let p = &d.peak_tflops;
391    let listed = |k: MmaKind| native.iter().any(|i| i.kind == k);
392    for (kind, peak) in [(MmaKind::Fp8, p.fp8), (MmaKind::Fp4BlockScale, p.fp4_nvfp4)] {
393        if listed(kind) != (peak > 0.0) {
394            return Err(field(format!(
395                "`{}` is {} in native_mma but its dense peak is {peak}",
396                kind.name(),
397                if listed(kind) { "true" } else { "false" }
398            )));
399        }
400    }
401    let nonzero = |v: f64| (v > 0.0).then_some(v);
402    Ok(Device {
403        peaks: Peaks {
404            bf16: p.bf16,
405            fp8: nonzero(p.fp8),
406            fp4_block_scale: nonzero(p.fp4_nvfp4),
407        },
408        memory_bytes: d.memory_gib_visible.unwrap_or(d.memory_gb_datasheet) * (1u64 << 30) as f64,
409        id: d.id,
410        name: d.name,
411        class: d.class,
412        arch: d.arch,
413        compute_capability: d.compute_capability,
414        sms: d.sm_count,
415        memory_type: d.memory_type,
416        bandwidth_gbps: d.memory_bandwidth_gbps,
417        l2_mb: d.l2_mb,
418        smem_per_sm_kb: d.smem_per_sm_kb,
419        cluster_max: d.max_cluster,
420        tmem_per_sm_kb: d.tmem_per_sm_kb,
421        mma_family,
422        native,
423        measured: d.measured,
424        usable_fraction: d.usable_fraction,
425        usable_why: d.usable_why,
426        derived: d.derived,
427        sources: d.sources,
428    })
429}
metrale_circuit/hardware/device.rs

metrale_circuit/hardware/
device.rs
