metrale_circuit/pipeline/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The numeric pipeline of a node: the formats it reads, the ordered steps inside it
4//! (activation and weight preparation, the multiply's operands, accumulation, scale application,
5//! element-wise compute, cache and state precision, routing and blending) each with its
6//! precision, and the formats it writes. One pipeline is REQUIRED per node of a plan
7//! ([`require`]: the circuit's formats under the serving policy) and one is DECLARED per kernel
8//! family, point or kernel in `KERNEL_FAMILIES.toml` ([`declare`]); [`check`] refuses a plan
9//! whose kernels declare anything but exactly what the plan requires.
10//!
11//! Owner: metrale-circuit (pipeline).
12//! Invariants:
13//! - The step vocabulary of every op is one table ([`vocab::steps_of`]); a pipeline names every
14//!   step of its op's vocabulary in that order, and nothing else.
15//! - Every value has one canonical spelling that parses back to itself (`Value::parse(k,
16//!   v.name()) == Ok(v)`), so the golden plans and the manifests read one text per value.
17//! - Nothing defaults: a step whose precision cannot be stated is an error, never a guess.
18//! - Precision is about values, not instructions: a BF16 activation held in an FP32 register is
19//!   BF16, an E4M3 weight widened exactly is the precision it widens to; which unit and atom
20//!   carry the multiply is the compute unit's (`venn::compute`).
21
22use std::fmt;
23
24use crate::format::Format;
25use crate::state::StateDtype;
26
27pub mod check;
28pub mod declare;
29pub mod query;
30pub mod require;
31pub mod vocab;
32
33#[path = "act_policy.rs"]
34mod act_policy;
35
36pub use check::{PlanPipelines, check_plan};
37pub use require::{Need, required};
38
39/// 2026-10-02: The element precision of a value inside a kernel (registers, operands, sums).
40#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
41pub enum Num {
42    /// 2026-10-02: bfloat16.
43    Bf16,
44    /// 2026-10-02: IEEE half.
45    F16,
46    /// 2026-10-02: float32.
47    F32,
48    /// 2026-10-02: FP8 E4M3 values (an FP8 activation or weight as multiplied).
49    E4m3,
50    /// 2026-10-02: FP4 E2M1 values (an NVFP4 activation or weight as multiplied).
51    E2m1,
52}
53
54const NUMS: [(Num, &str); 5] = [
55    (Num::Bf16, "bf16"),
56    (Num::F16, "f16"),
57    (Num::F32, "f32"),
58    (Num::E4m3, "e4m3"),
59    (Num::E2m1, "e2m1"),
60];
61
62impl Num {
63    /// 2026-10-02: The canonical spelling.
64    pub fn name(self) -> &'static str {
65        NUMS.iter()
66            .find(|(n, _)| *n == self)
67            .map_or("?", |(_, s)| s)
68    }
69
70    /// 2026-10-02: Parse the canonical spelling.
71    pub fn parse(s: &str) -> Option<Self> {
72        NUMS.iter().find(|(_, n)| *n == s).map(|(v, _)| *v)
73    }
74
75    /// 2026-10-02: The precision of a tensor format's values: an FP8 tensor is E4M3, an NVFP4 one
76    /// E2M1; `None` for `i32`, which no multiply or compute step reads as a number.
77    pub fn of_format(f: Format) -> Option<Self> {
78        match f {
79            Format::Bf16 => Some(Num::Bf16),
80            Format::F32 => Some(Num::F32),
81            Format::Fp8E4m3 { .. } => Some(Num::E4m3),
82            Format::Nvfp4 { .. } => Some(Num::E2m1),
83            Format::I32 => None,
84        }
85    }
86}
87
88/// 2026-10-02: A step of a node's pipeline. Each has one value shape ([`vocab::shape_of`]).
89#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
90pub enum StepKind {
91    /// 2026-10-02: Rows gathered by expert id (or embedding rows by token id): their format.
92    Gather,
93    /// 2026-10-02: The format the activation is rounded or quantized to before the multiply.
94    Act,
95    /// 2026-10-02: The stored weight format and the precision of its unscaled values as multiplied.
96    Weight,
97    /// 2026-10-02: The multiply's operand precisions, activation first.
98    Mma,
99    /// 2026-10-02: The precision dot products (or attention's PV sums) accumulate in.
100    Accumulate,
101    /// 2026-10-02: The precision weight and activation scales are applied in; none without scales.
102    Scale,
103    /// 2026-10-02: The precision of an element-wise op's or a reduction's arithmetic.
104    Compute,
105    /// 2026-10-02: Data moved without arithmetic: its format.
106    Move,
107    /// 2026-10-02: The KV-cache dtype read or written.
108    Cache,
109    /// 2026-10-02: Attention: the precision of the QK dot products.
110    Scores,
111    /// 2026-10-02: Attention: the precision of the softmax.
112    Softmax,
113    /// 2026-10-02: The precision of the recurrent state as stored and updated.
114    State,
115    /// 2026-10-02: Routing: the precision of the scoring function and of the routing weights.
116    Score,
117    /// 2026-10-02: Blend: the precision of the weighted sum of each token's top-k expert rows.
118    Scatter,
119    /// 2026-10-02: Blend: the precision the (gated) shared expert is added in; none without one.
120    Combine,
121    /// 2026-10-02: A cross-rank reduction's precision.
122    Reduce,
123    /// 2026-10-02: Greedy selection: the format of the values compared.
124    Compare,
125}
126
127/// 2026-10-02: A step's value.
128#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
129pub enum Value {
130    /// 2026-10-02: A tensor format (`gather`, `act`, `move`, `compare`).
131    Format(Format),
132    /// 2026-10-02: `<stored>-><operand>` (`weight`).
133    Weight {
134        /// 2026-10-02: The weight as loaded.
135        stored: Format,
136        /// 2026-10-02: Its unscaled values' precision as multiplied.
137        operand: Num,
138    },
139    /// 2026-10-02: `<a>*<b>` (`mma`).
140    Mma {
141        /// 2026-10-02: The activation operand.
142        a: Num,
143        /// 2026-10-02: The weight operand.
144        b: Num,
145    },
146    /// 2026-10-02: A precision.
147    Num(Num),
148    /// 2026-10-02: A precision, or `none` where the step does not apply (no scales, no shared
149    /// expert).
150    OptNum(Option<Num>),
151    /// 2026-10-02: A KV-cache dtype, spelled as `--kv-cache-dtype` spells it.
152    Kv(String),
153    /// 2026-10-02: A recurrent state's storage element.
154    State(StateDtype),
155}
156
157/// 2026-10-02: The spelling of a state dtype (the circuit TOML's, [`StateDtype::name`]).
158pub fn state_name(d: StateDtype) -> &'static str {
159    d.name()
160}
161
162impl Value {
163    /// 2026-10-02: The canonical spelling.
164    pub fn name(&self) -> String {
165        match self {
166            Value::Format(f) => f.name(),
167            Value::Weight { stored, operand } => format!("{}->{}", stored.name(), operand.name()),
168            Value::Mma { a, b } => format!("{}*{}", a.name(), b.name()),
169            Value::Num(n) => n.name().to_string(),
170            Value::OptNum(None) => "none".to_string(),
171            Value::OptNum(Some(n)) => n.name().to_string(),
172            Value::Kv(k) => k.clone(),
173            Value::State(d) => state_name(*d).to_string(),
174        }
175    }
176}
177
178/// 2026-10-02: One step and its value.
179#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
180pub struct Step {
181    /// 2026-10-02: Which step.
182    pub kind: StepKind,
183    /// 2026-10-02: Its value.
184    pub value: Value,
185}
186
187/// 2026-10-02: A node's pipeline: input formats in the node's input order, the op's steps in
188/// vocabulary order, output formats in the node's output order.
189#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
190pub struct NodePipeline {
191    /// 2026-10-02: The format each input is read in.
192    pub inputs: Vec<Format>,
193    /// 2026-10-02: The steps.
194    pub steps: Vec<Step>,
195    /// 2026-10-02: The format each output is written in, or handed on in inside its group.
196    pub outputs: Vec<Format>,
197}
198
199/// 2026-10-02: The parts of a pipeline a plan's rule states rather than the reference: steps,
200/// and the inputs and outputs whose in-group hand-off format it states (`holds`). Rendered
201/// ` (rule)`.
202#[derive(Debug, Clone, Default, PartialEq, Eq)]
203pub struct Stated {
204    /// 2026-10-02: Stated steps.
205    pub steps: std::collections::BTreeSet<StepKind>,
206    /// 2026-10-02: Indices of inputs read in a stated hand-off format.
207    pub inputs: std::collections::BTreeSet<usize>,
208    /// 2026-10-02: Indices of outputs handed on in a stated format.
209    pub outputs: std::collections::BTreeSet<usize>,
210}
211
212impl NodePipeline {
213    /// 2026-10-02: The pipeline as facts: `in <formats>`, `<step> <value>` per step, `out
214    /// <formats>`, each part a rule states followed by ` (rule)`; `-` for no inputs or outputs.
215    pub fn facts(&self, stated: &Stated) -> Vec<String> {
216        let mark = |on: bool| if on { " (rule)" } else { "" };
217        let list = |v: &[Format], marked: &dyn Fn(usize) -> bool| {
218            if v.is_empty() {
219                "-".to_string()
220            } else {
221                v.iter()
222                    .enumerate()
223                    .map(|(i, f)| format!("{}{}", f.name(), mark(marked(i))))
224                    .collect::<Vec<_>>()
225                    .join(",")
226            }
227        };
228        let mut out = vec![format!(
229            "in {}",
230            list(&self.inputs, &|i| stated.inputs.contains(&i))
231        )];
232        out.extend(self.steps.iter().map(|s| {
233            format!(
234                "{} {}{}",
235                vocab::step_name(s.kind),
236                s.value.name(),
237                mark(stated.steps.contains(&s.kind))
238            )
239        }));
240        out.push(format!(
241            "out {}",
242            list(&self.outputs, &|i| stated.outputs.contains(&i))
243        ));
244        out
245    }
246
247    /// 2026-10-02: `in -> [step value | ...] -> out`, ASCII ([`NodePipeline::facts`]).
248    pub fn render(&self, stated: &Stated) -> String {
249        let facts = self.facts(stated);
250        let (first, rest) = facts.split_first().map_or(("", &[][..]), |(f, r)| (f, r));
251        let (last, steps) = rest.split_last().map_or(("", &[][..]), |(l, s)| (l, s));
252        format!(
253            "{} -> [{}] -> {}",
254            first.trim_start_matches("in "),
255            steps.join(" | "),
256            last.trim_start_matches("out ")
257        )
258    }
259}
260
261impl fmt::Display for NodePipeline {
262    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
263        f.write_str(&self.render(&Default::default()))
264    }
265}
266
267/// 2026-10-02: What differs between a required and a declared pipeline, one entry per field
268/// (`in[1]`, `weight`, `out[0]`), as `field: required X, declared Y`.
269pub fn differences(required: &NodePipeline, declared: &NodePipeline) -> Vec<String> {
270    let mut out = Vec::new();
271    let fmts = |label: &str, r: &[Format], d: &[Format], out: &mut Vec<String>| {
272        if r.len() != d.len() {
273            out.push(format!(
274                "{label}: required {} formats, declared {}",
275                r.len(),
276                d.len()
277            ));
278            return;
279        }
280        for (i, (a, b)) in r.iter().zip(d).enumerate() {
281            if a != b {
282                out.push(format!("{label}[{i}]: required {a}, declared {b}"));
283            }
284        }
285    };
286    fmts("in", &required.inputs, &declared.inputs, &mut out);
287    for (r, d) in required.steps.iter().zip(&declared.steps) {
288        if r != d {
289            out.push(format!(
290                "{}: required {}, declared {}",
291                vocab::step_name(r.kind),
292                r.value.name(),
293                d.value.name()
294            ));
295        }
296    }
297    if required.steps.len() != declared.steps.len() {
298        out.push("steps: the two pipelines name different step lists".into());
299    }
300    fmts("out", &required.outputs, &declared.outputs, &mut out);
301    out
302}
303
304/// 2026-10-02: Why a pipeline could not be required, declared or matched.
305#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
306pub enum PipelineError {
307    /// 2026-10-02: The policy or the circuit does not state what a step needs.
308    #[error("node `{node}` ({op}): no required pipeline: {detail}")]
309    Required {
310        /// 2026-10-02: Node id.
311        node: String,
312        /// 2026-10-02: Its op.
313        op: String,
314        /// 2026-10-02: What is missing or contradictory.
315        detail: String,
316    },
317    /// 2026-10-02: A group's kernels are in no family, so their pipeline is undeclared.
318    #[error("node `{node}` ({op}) runs {kernels}, which no kernel family declares a pipeline for")]
319    Undeclared {
320        /// 2026-10-02: Node id.
321        node: String,
322        /// 2026-10-02: Its op.
323        op: String,
324        /// 2026-10-02: The group's kernels or emitter.
325        kernels: String,
326    },
327    /// 2026-10-02: A family's declaration cannot be resolved for a node.
328    #[error("node `{node}` ({op}) in family `{family}`: {detail}")]
329    Declared {
330        /// 2026-10-02: Node id.
331        node: String,
332        /// 2026-10-02: Its op.
333        op: String,
334        /// 2026-10-02: Family id.
335        family: String,
336        /// 2026-10-02: Why.
337        detail: String,
338    },
339    /// 2026-10-02: Kernels declare another pipeline than a plan requires, one entry per node.
340    #[error(
341        "{} node(s) run a kernel whose declared pipeline is not the required one:\n  {}",
342        .0.len(),
343        .0.iter().map(Mismatch::describe).collect::<Vec<_>>().join("\n  ")
344    )]
345    Mismatch(Vec<Mismatch>),
346}
347
348/// 2026-10-02: One node whose kernel declares another pipeline than its plan requires.
349#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
350pub struct Mismatch {
351    /// 2026-10-02: Node id.
352    pub node: String,
353    /// 2026-10-02: The rule of its group.
354    pub rule: String,
355    /// 2026-10-02: The family that declares the kernel's pipeline.
356    pub family: String,
357    /// 2026-10-02: What differs ([`differences`]).
358    pub diffs: Vec<String>,
359}
360
361impl Mismatch {
362    /// 2026-10-02: One line.
363    pub fn describe(&self) -> String {
364        format!(
365            "{} (rule `{}`, family `{}`): {}",
366            self.node,
367            self.rule,
368            self.family,
369            self.diffs.join("; ")
370        )
371    }
372}
373
374#[cfg(test)]
375#[path = "pipeline_tests.rs"]
376mod pipeline_tests;
metrale_circuit/pipeline/mod.rs

metrale_circuit/pipeline/
mod.rs
