metrale_circuit/venn/
families.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: The kernel-family manifest, `kernels/<hw>/common/KERNEL_FAMILIES.toml`: which
4//! circuit ops each family implements, the parameters it varies (runtime, compile-time or
5//! policy), the points it instantiates, the evidence envelope (the points with microbench
6//! records) and the rules that rediscover the instantiated points from the kernel sources.
7//!
8//! Owner: metrale-circuit (venn).
9//! Invariants:
10//! - A parameter is an axis the family's code is structured to vary (a runtime argument, a
11//!   template or macro parameter, a policy or a per-point file copy), even when only one value
12//!   is instantiated. A format the code cannot vary without a rewrite is a constraint on the
13//!   op (`[[family.op]]`), not a parameter.
14//! - Every point states every compile-time and policy parameter, and nothing else: a runtime
15//!   parameter never defines a point. Evidence sits on an instantiated point.
16//! - Nothing defaults: an op, format, parameter, extractor or mode the loader does not know is a
17//!   typed error, and a parameter a node cannot supply is an error unless the manifest states
18//!   its `absent` value.
19
20use std::collections::{BTreeMap, BTreeSet};
21
22use super::compute::{ComputeUnit, FamilyCompute};
23use crate::format::Format;
24use crate::ir::{LayerKind, LinearRole, OpKind};
25use crate::rules::{KernelId, Mode};
26
27/// 2026-09-29: How a parameter reaches the kernel.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
29pub enum ParamKind {
30    /// 2026-09-29: A kernel argument: strides, counts, eps, scale pointers. Never defines a point.
31    Runtime,
32    /// 2026-09-29: Sizes registers, shared memory or unrolling: a template parameter or a
33    /// per-point copy (head_dim, group size, tiles).
34    Compile,
35    /// 2026-09-29: A policy the code plugs in: weight format, scale layout, activation
36    /// quantizer, activation epilogue, routing scoring.
37    Policy,
38}
39
40impl ParamKind {
41    /// 2026-09-29: The manifest and report spelling.
42    pub fn name(self) -> &'static str {
43        match self {
44            ParamKind::Runtime => "runtime",
45            ParamKind::Compile => "compile-time",
46            ParamKind::Policy => "policy",
47        }
48    }
49}
50
51/// 2026-09-29: Where a parameter's value is read from, for one node of one circuit.
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub enum Extract {
54    /// 2026-09-29: `dim:<name>`: the circuit's dim.
55    Dim(String),
56    /// 2026-09-29: `weight`: the node's weight format.
57    Weight,
58    /// 2026-09-29: `activation`: the format of the node's first input.
59    Activation,
60    /// 2026-09-29: `output`: the format of the node's first output.
61    Output,
62    /// 2026-09-29: `op`: the node's op name.
63    Op,
64    /// 2026-09-29: `consumer_op`: the op of the first node reading the node's first output.
65    ConsumerOp,
66    /// 2026-09-29: `param:<key>`: the node's template parameter.
67    Param(String),
68    /// 2026-09-29: `setting:<key>`: the instance's policy setting.
69    Setting(String),
70    /// 2026-09-29: `in_dim` / `out_dim`: the first input's / output's feature dimension.
71    InDim,
72    /// 2026-09-29: See [`Extract::InDim`].
73    OutDim,
74}
75
76impl Extract {
77    fn parse(s: &str) -> Option<Self> {
78        let (head, arg) = match s.split_once(':') {
79            Some((h, a)) if !a.is_empty() => (h, Some(a.to_string())),
80            Some(_) => return None,
81            None => (s, None),
82        };
83        match (head, arg) {
84            ("dim", Some(a)) => Some(Extract::Dim(a)),
85            ("param", Some(a)) => Some(Extract::Param(a)),
86            ("setting", Some(a)) => Some(Extract::Setting(a)),
87            ("weight", None) => Some(Extract::Weight),
88            ("activation", None) => Some(Extract::Activation),
89            ("output", None) => Some(Extract::Output),
90            ("op", None) => Some(Extract::Op),
91            ("consumer_op", None) => Some(Extract::ConsumerOp),
92            ("in_dim", None) => Some(Extract::InDim),
93            ("out_dim", None) => Some(Extract::OutDim),
94            _ => None,
95        }
96    }
97}
98
99/// 2026-09-29: One parameter of a family.
100#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct Param {
102    /// 2026-09-29: Name, e.g. `head_dim`.
103    pub name: String,
104    /// 2026-09-29: How it reaches the kernel.
105    pub kind: ParamKind,
106    /// 2026-09-29: Op base name to extractor. Ops of the family missing here do not carry it.
107    pub from: BTreeMap<String, Extract>,
108    /// 2026-09-29: The value when the extractor finds nothing (a dim or parameter the circuit
109    /// does not state), with the manifest's reason in a comment; `None` makes that an error.
110    pub absent: Option<String>,
111}
112
113/// 2026-09-29: An op a family implements, with the constraints a node must meet.
114#[derive(Debug, Clone, PartialEq, Eq)]
115pub struct OpSpec {
116    /// 2026-09-29: Op base name (`linear`, `paged_attention`, ...). 2026-10-02: or a quantizer
117    /// with its output format (`act_quant:nvfp4/g16`).
118    pub op: String,
119    /// 2026-09-29: For `linear`: the roles it takes; empty is every role.
120    pub roles: BTreeSet<LinearRole>,
121    /// 2026-09-29: Allowed weight formats; empty is any.
122    pub weight: BTreeSet<Format>,
123    /// 2026-09-29: Allowed first-input formats; empty is any.
124    pub activation: BTreeSet<Format>,
125    /// 2026-09-29: Node template parameters that must hold these values.
126    pub params: Values,
127    /// 2026-09-29: Ops one of which must read the node's first output, by base name
128    /// (`l2_norm`) or qualified name (`linear:q`); empty is no constraint. A kernel that fuses
129    /// its consumer (a conv with its L2 norm) implements the op only where that consumer
130    /// follows.
131    pub feeds: BTreeSet<String>,
132    /// 2026-09-29: Ops one of which must produce the node's first input, spelled as `feeds`;
133    /// empty is no constraint (a snapshot of the conv window, not of the SSM state).
134    pub after: BTreeSet<String>,
135    /// 2026-09-29: Ops one of which must also read the node's first input, spelled as
136    /// `feeds`; empty is no constraint. A conv kernel that writes its window snapshot and the
137    /// L2 norm of its output implements the snapshot only where the L2 norm reads beside it.
138    pub beside: BTreeSet<String>,
139}
140
141impl OpSpec {
142    /// 2026-10-02: The spec names `op`: by base name, or by its qualified name (an `act_quant`
143    /// with its format).
144    pub fn names(&self, op: &OpKind) -> bool {
145        self.op == op.base_name() || self.op == op.name()
146    }
147}
148
149/// 2026-09-29: How an instantiated point is realised in the sources.
150#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
151pub enum How {
152    /// 2026-09-29: A template or macro instantiation.
153    Instantiation,
154    /// 2026-09-29: A per-point copy of the source file (the parameterization target).
155    Copy,
156    /// 2026-09-29: A runtime branch inside one entry point.
157    Branch,
158}
159
160impl How {
161    /// 2026-09-29: The manifest spelling.
162    pub fn name(self) -> &'static str {
163        match self {
164            How::Instantiation => "instantiation",
165            How::Copy => "copy",
166            How::Branch => "branch",
167        }
168    }
169}
170
171/// 2026-09-29: Parameter name to value.
172pub type Values = BTreeMap<String, String>;
173
174/// 2026-09-29: One instantiated point.
175#[derive(Debug, Clone, PartialEq, Eq)]
176pub struct Point {
177    /// 2026-09-29: Every compile-time and policy parameter's value.
178    pub values: Values,
179    /// 2026-09-29: How it is realised.
180    pub how: How,
181    /// 2026-09-29: Repo-relative sources that realise it.
182    pub files: Vec<String>,
183    /// 2026-10-02: Its compute unit where it differs from the family's (`None`: the family's).
184    pub compute: Option<ComputeUnit>,
185    /// 2026-10-02: Its own pipeline declarations, where they differ from the family's
186    /// ([`crate::pipeline::declare`]).
187    pub pipeline: crate::pipeline::declare::ByOp,
188}
189
190/// 2026-09-29: Where an evidence record lives.
191#[derive(Debug, Clone, PartialEq, Eq)]
192pub enum EvidenceSource {
193    /// 2026-09-29: A `docs/kernel-perf/measurements.toml` row, keyed `<kernel> @ <regime>`.
194    Measurement(String),
195    /// 2026-09-29: A microbench not in measurements.toml: its method and result in words.
196    Microbench(String),
197}
198
199/// 2026-09-29: A microbench record at one point and the row counts it measured.
200#[derive(Debug, Clone, PartialEq, Eq)]
201pub struct Evidence {
202    /// 2026-09-29: The instantiated point measured.
203    pub point: Values,
204    /// 2026-09-29: The row counts measured; a record says nothing about any other count.
205    pub rows: BTreeSet<u64>,
206    /// 2026-09-29: The record.
207    pub source: EvidenceSource,
208}
209
210/// 2026-09-29: A rule that rediscovers instantiated points from the sources.
211#[derive(Debug, Clone, PartialEq, Eq)]
212pub enum Discover {
213    /// 2026-09-29: Each file matching `glob` realises `values` as a copy.
214    File {
215        /// 2026-09-29: Repo-relative glob (`*` matches any run).
216        glob: String,
217        /// 2026-09-29: The point's values.
218        values: Values,
219    },
220    /// 2026-09-29: Each `name(...)` invocation line in `file` realises one point: parameter
221    /// `p` is argument `args[p]` (0-based), mapped through `map[p]` when present.
222    Macro {
223        /// 2026-09-29: Repo-relative source.
224        file: String,
225        /// 2026-09-29: Macro or template name.
226        name: String,
227        /// 2026-09-29: Parameter to argument index.
228        args: BTreeMap<String, usize>,
229        /// 2026-09-29: Parameter to (argument text to value).
230        map: BTreeMap<String, BTreeMap<String, String>>,
231    },
232}
233
234/// 2026-09-29: One kernel family.
235#[derive(Debug, Clone, PartialEq, Eq)]
236pub struct Family {
237    /// 2026-09-29: Id.
238    pub id: String,
239    /// 2026-09-29: One line.
240    pub description: String,
241    /// 2026-09-29: Entry points, as FUSIONS.toml names them.
242    pub kernels: Vec<KernelId>,
243    /// 2026-09-29: FUSIONS.toml emitters of kernel-less groups (host copies) it stands for.
244    pub emitters: Vec<String>,
245    /// 2026-09-29: Inclusive range of rows one launch covers.
246    pub rows: (u64, u64),
247    /// 2026-09-29: Ops it implements.
248    pub ops: Vec<OpSpec>,
249    /// 2026-09-29: Its parameter space.
250    pub params: Vec<Param>,
251    /// 2026-09-29: Instantiated points.
252    pub points: Vec<Point>,
253    /// 2026-09-29: The evidence envelope.
254    pub evidence: Vec<Evidence>,
255    /// 2026-09-29: Discovery rules.
256    pub discover: Vec<Discover>,
257    /// 2026-10-02: The compute units it runs on ([`super::compute`]).
258    pub compute: FamilyCompute,
259    /// 2026-10-02: The numeric pipeline each op runs at, the family's and its kernels'
260    /// ([`crate::pipeline::declare`]).
261    pub pipeline: crate::pipeline::declare::FamilyPipelines,
262    /// 2026-10-02: Device scratch a launch needs beyond its edges (`crate::memory`).
263    pub workspace: Vec<Workspace>,
264}
265
266/// 2026-10-02: One workspace of a family (`[[family.workspace]]`).
267#[derive(Debug, Clone, PartialEq, Eq)]
268pub struct Workspace {
269    /// 2026-10-02: Name, unique in the family.
270    pub name: String,
271    /// 2026-10-02: Byte expressions over the circuit's dims, `n` (rows), `k` (the node's input
272    /// width) and `sm_count`; the largest is the workspace.
273    pub bytes: Vec<crate::dims::DimExpr>,
274    /// 2026-10-02: Held inside the legacy buffer arena.
275    pub arena: bool,
276    /// 2026-10-02: Where the engine allocates it, and why it has this size.
277    pub why: String,
278}
279
280impl Family {
281    /// 2026-09-29: One launch covers more than one row.
282    pub fn multi_row(&self) -> bool {
283        self.rows.1 > 1
284    }
285
286    /// 2026-09-29: The parameter named `name`.
287    pub fn param(&self, name: &str) -> Option<&Param> {
288        self.params.iter().find(|p| p.name == name)
289    }
290
291    /// 2026-10-02: Each point's values with its own pipeline declarations.
292    pub fn point_pipelines(&self) -> Vec<(&Values, &crate::pipeline::declare::ByOp)> {
293        self.points
294            .iter()
295            .map(|p| (&p.values, &p.pipeline))
296            .collect()
297    }
298}
299
300/// 2026-09-29: A fact about a legacy layer implementation the circuit cannot show: a mode in
301/// which the layer has no multi-row path and loops per sequence (or per row).
302#[derive(Debug, Clone, PartialEq, Eq)]
303pub struct LegacyPath {
304    /// 2026-09-29: Circuit arch.
305    pub arch: String,
306    /// 2026-09-29: Layer kind.
307    pub layer_kind: LayerKind,
308    /// 2026-09-29: Modes whose rows above one loop per sequence.
309    pub per_sequence: Vec<Mode>,
310    /// 2026-09-29: The sites (`block.node`) that loop; empty is every node of the layer kind.
311    pub sites: Vec<String>,
312    /// 2026-09-29: The loop runs only above this many rows (a batched path covers the rest).
313    pub rows_above: u64,
314    /// 2026-09-29: `path:line` of the loop or the trait default it falls back to.
315    pub cite: String,
316    /// 2026-09-29: Text the cited line holds (checked by the tests).
317    pub holds: String,
318    /// 2026-09-29: One line.
319    pub note: String,
320}
321
322/// 2026-09-29: The roofline constants of the hardware.
323#[derive(Debug, Clone, Copy, PartialEq)]
324pub struct Roofline {
325    /// 2026-09-29: DRAM bandwidth, GB/s.
326    pub dram_gbps: f64,
327    /// 2026-09-29: Dense BF16 tensor peak, TFLOPS.
328    pub bf16_tflops: f64,
329    /// 2026-09-29: FP8 tensor peak, TFLOPS.
330    pub fp8_tflops: f64,
331    /// 2026-09-29: NVFP4 tensor peak, TFLOPS.
332    pub nvfp4_tflops: f64,
333    /// 2026-09-29: KV context each sequence attends over in the estimate.
334    pub context_tokens: u64,
335}
336
337/// 2026-09-29: A parsed manifest.
338#[derive(Debug, Clone, PartialEq)]
339pub struct Families {
340    /// 2026-09-29: Hardware directory.
341    pub hardware: String,
342    /// 2026-09-29: Estimate constants.
343    pub roofline: Roofline,
344    /// 2026-09-29: Families, in file order.
345    pub families: Vec<Family>,
346    /// 2026-09-29: Legacy per-sequence facts.
347    pub legacy: Vec<LegacyPath>,
348}
349
350impl Families {
351    /// 2026-10-02: The compute unit `kernel` runs on: its family's answer; `None` when no
352    /// family lists it.
353    pub fn compute_of(&self, kernel: &KernelId) -> Option<&ComputeUnit> {
354        self.families
355            .iter()
356            .find(|f| f.kernels.contains(kernel))
357            .map(|f| f.compute.of(kernel))
358    }
359
360    /// 2026-09-29: The family that ran a node of `op` in a plan group launching `kernels`
361    /// (or, for a kernel-less group, emitted by `emitter`): the first, in manifest order, that
362    /// lists one of them and implements the op (and, for `linear`, the role). Formats are not
363    /// checked: a plan runs what its rules chose. `None` when no family lists them.
364    pub fn of_group(&self, kernels: &[KernelId], emitter: &str, op: &OpKind) -> Option<&Family> {
365        self.families.iter().find(|f| {
366            let listed = if kernels.is_empty() {
367                f.emitters.iter().any(|e| e == emitter)
368            } else {
369                kernels.iter().any(|k| f.kernels.contains(k))
370            };
371            listed
372                && f.ops.iter().any(|s| {
373                    s.names(op)
374                        && match op {
375                            OpKind::Linear(r) => s.roles.is_empty() || s.roles.contains(r),
376                            _ => true,
377                        }
378                })
379        })
380    }
381}
382
383/// 2026-09-29: Why the manifest did not load.
384#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
385pub enum FamilyError {
386    /// 2026-09-29: Not TOML, or not the schema.
387    #[error("KERNEL_FAMILIES.toml: {0}")]
388    Parse(String),
389    /// 2026-09-29: A field of one family.
390    #[error("family `{family}`: {detail}")]
391    Field {
392        /// 2026-09-29: Family id.
393        family: String,
394        /// 2026-09-29: What was wrong.
395        detail: String,
396    },
397    /// 2026-09-29: An op outside the circuit vocabulary.
398    #[error("family `{family}`: unknown op `{op}`")]
399    UnknownOp {
400        /// 2026-09-29: Family id.
401        family: String,
402        /// 2026-09-29: The op.
403        op: String,
404    },
405    /// 2026-09-29: A point, evidence record or discovery rule naming a parameter the family
406    /// does not declare, or a runtime one.
407    #[error("family `{family}`: `{param}` is not a compile-time or policy parameter of it")]
408    UnknownParam {
409        /// 2026-09-29: Family id.
410        family: String,
411        /// 2026-09-29: The parameter.
412        param: String,
413    },
414}
415
416#[path = "families_compute.rs"]
417mod compute_file;
418#[path = "families_file.rs"]
419mod file;
420#[path = "families_legacy.rs"]
421mod legacy_file;
422#[path = "families_workspace.rs"]
423mod workspace_file;
424
425/// 2026-09-29: Parse the manifest text.
426pub fn parse_families(text: &str) -> Result<Families, FamilyError> {
427    file::parse(text)
428}
429
430#[cfg(test)]
431#[path = "families_tests.rs"]
432mod families_tests;
metrale_circuit/venn/families.rs

metrale_circuit/venn/
families.rs
