metrale_circuit/venn/
compute.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The compute unit a kernel runs on: tensor cores (with the MMA atom and the
4//! operand format it issues), CUDA cores, or memory (data movement, element-wise work and
5//! reductions). Declared per family in `KERNEL_FAMILIES.toml` (`compute`, `mma`), per point where
6//! a family's points differ, and per kernel where one family's kernels differ; read by the
7//! tensor-core policy (`hardware::tc_policy`) and shown on every planned group.
8//!
9//! Owner: metrale-circuit (venn).
10//! Invariants:
11//! - Nothing defaults: every family states its unit; `tensor_core` names its atom, the other
12//!   units name none.
13//! - A kernel's unit is its `kernel_compute` entry, else its family's. A family whose points
14//!   state different units names every kernel's unit, so no kernel's unit is inferred.
15
16use std::collections::BTreeMap;
17
18use crate::rules::KernelId;
19
20/// 2026-10-02: Where a kernel's arithmetic runs.
21#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
22pub enum ComputeUnit {
23    /// 2026-10-02: Tensor cores, issuing `atom` (e.g. `mma.sync.m16n8k32.e4m3`).
24    TensorCore {
25        /// 2026-10-02: The MMA instruction shape and operand format.
26        atom: String,
27    },
28    /// 2026-10-02: CUDA-core FMA arithmetic (GEMV dot products, recurrences, convolutions).
29    CudaCore,
30    /// 2026-10-02: Data movement, element-wise work and reductions.
31    Memory,
32}
33
34impl ComputeUnit {
35    /// 2026-10-02: `tensor_core`, `cuda_core` or `memory`.
36    pub fn kind(&self) -> &'static str {
37        match self {
38            ComputeUnit::TensorCore { .. } => "tensor_core",
39            ComputeUnit::CudaCore => "cuda_core",
40            ComputeUnit::Memory => "memory",
41        }
42    }
43
44    /// 2026-10-02: The unit as a report shows it: `tensor_core <atom>`, `cuda_core`, `memory`.
45    pub fn name(&self) -> String {
46        match self {
47            ComputeUnit::TensorCore { atom } => format!("tensor_core {atom}"),
48            other => other.kind().to_string(),
49        }
50    }
51
52    /// 2026-10-02: The unit as one word, for a plan line: `tensor_core:<atom>`, `cuda_core`,
53    /// `memory`.
54    pub fn tag(&self) -> String {
55        match self {
56            ComputeUnit::TensorCore { atom } => format!("tensor_core:{atom}"),
57            other => other.kind().to_string(),
58        }
59    }
60
61    /// 2026-10-02: Runs on tensor cores.
62    pub fn is_tensor_core(&self) -> bool {
63        matches!(self, ComputeUnit::TensorCore { .. })
64    }
65
66    /// 2026-10-02: Parse a `compute` value and its `mma` atom.
67    pub fn parse(compute: &str, mma: Option<&str>) -> Result<Self, String> {
68        match (compute, mma.map(str::trim)) {
69            ("tensor_core", Some(atom)) if !atom.is_empty() => Ok(ComputeUnit::TensorCore {
70                atom: atom.to_string(),
71            }),
72            ("tensor_core", _) => Err("compute `tensor_core` needs its `mma` atom".into()),
73            ("cuda_core", None) => Ok(ComputeUnit::CudaCore),
74            ("memory", None) => Ok(ComputeUnit::Memory),
75            ("cuda_core" | "memory", Some(_)) => Err(format!(
76                "compute `{compute}` issues no MMA, so it names no `mma`"
77            )),
78            (other, _) => Err(format!(
79                "compute `{other}` (tensor_core | cuda_core | memory)"
80            )),
81        }
82    }
83}
84
85/// 2026-10-02: A family's units: its own, and the kernels that differ from it.
86#[derive(Debug, Clone, PartialEq, Eq)]
87pub struct FamilyCompute {
88    /// 2026-10-02: The family's unit.
89    pub unit: ComputeUnit,
90    /// 2026-10-02: Kernels whose unit differs from the family's.
91    pub kernels: BTreeMap<KernelId, ComputeUnit>,
92}
93
94impl FamilyCompute {
95    /// 2026-10-02: The unit `kernel` runs on.
96    pub fn of(&self, kernel: &KernelId) -> &ComputeUnit {
97        self.kernels.get(kernel).unwrap_or(&self.unit)
98    }
99}
100
101/// 2026-10-02: The unit a plan group runs its node on, from its kernels' units: tensor cores
102/// when any kernel issues MMAs (the others quantize, stage or reduce around it), else CUDA
103/// cores when any kernel computes, else memory. `None` for a group without kernels.
104pub fn group_unit<'a>(units: impl IntoIterator<Item = &'a ComputeUnit>) -> Option<ComputeUnit> {
105    let mut best: Option<&ComputeUnit> = None;
106    for u in units {
107        let rank = |x: &ComputeUnit| match x {
108            ComputeUnit::TensorCore { .. } => 2,
109            ComputeUnit::CudaCore => 1,
110            ComputeUnit::Memory => 0,
111        };
112        if best.is_none_or(|b| rank(u) > rank(b)) {
113            best = Some(u);
114        }
115    }
116    best.cloned()
117}
118
119#[cfg(test)]
120#[path = "compute_tests.rs"]
121mod compute_tests;
metrale_circuit/venn/compute.rs

metrale_circuit/venn/
compute.rs
