metrale_circuit/pipeline/
require.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The REQUIRED pipeline of a node: what the circuit's formats, under the serving
4//! policy and the plan's stated routing, ask each step to compute at. This is the "what we
5//! need" view (`met circuit precision`) and the side every kernel declaration is checked
6//! against ([`super::check`]).
7//!
8//! Where each requirement comes from:
9//! - inputs and outputs: the formats the plan stores each edge in (the circuit's, or a rule's
10//!   `writes`); inside a group, the format a rule states the edge is handed on in (`holds`);
11//! - `act`: the input format, unless `activation_quantization` routes the projection's family
12//!   to a fixed format at this row count, or, under `adaptive` (today's routing, which the rules
13//!   encode), the rule states the format its kernels run the activation at;
14//! - `weight`: the node's weight format, its values widened to the activation's precision for a
15//!   16/32-bit activation, consumed natively (E4M3, E2M1) under an 8- or 4-bit one (an NVFP4
16//!   weight under FP8 activations by the exact E2M1 to E4M3 conversion);
17//! - `cache`: `kv_cache_dtype`; `state`: the state's declared format, a keyed one from its
18//!   setting (`ssm_h_storage` from `ssm_h_dtype`);
19//! - every accumulation, scale application, element-wise compute, softmax, routing score and
20//!   blend sum: FP32, the circuit's reference definition of those ops;
21//! - any step a plan's rule states (`steps` on a pattern element): the stated value. A rule
22//!   states where today's routing departs from the reference; the plan prints it, and the
23//!   kernel must declare the same.
24//!
25//! Owner: metrale-circuit (pipeline).
26//! Invariants:
27//! - Pure; a setting the requirement reads and the policy does not state is an error.
28//! - The base requirement (no plan) and a plan's requirement differ only where the plan states
29//!   something: a stored, held or stated format.
30
31use std::collections::BTreeMap;
32
33use super::vocab::steps_of;
34use super::{NodePipeline, Num, PipelineError, Step, StepKind, Value};
35use crate::format::Format;
36use crate::fuser::{EdgeState, FusionPlan};
37use crate::ir::{Circuit, EdgeIdx, Node, NodeIdx, OpKind};
38use crate::rules::Rule;
39use crate::state::{StateAccess, StateDtype, StateFormat};
40
41/// 2026-10-02: Everything a requirement reads besides the node.
42#[derive(Debug, Clone)]
43pub struct Need<'a> {
44    /// 2026-10-02: The circuit.
45    pub circuit: &'a Circuit,
46    /// 2026-10-02: The policy's settings.
47    pub settings: &'a BTreeMap<String, String>,
48    /// 2026-10-02: The row count the plan runs (the activation-quantization ladder reads it).
49    pub rows: u64,
50    /// 2026-10-02: Each edge's stored format.
51    pub stored: Vec<Format>,
52    /// 2026-10-02: Each node's group under the plan; all `None` without one.
53    pub group_of: Vec<Option<usize>>,
54    /// 2026-10-02: Each edge's state under the plan; all `None` without one.
55    pub edge_states: Vec<Option<EdgeState>>,
56    /// 2026-10-02: In-group hand-off formats the plan's rules state (`holds`), by edge.
57    pub held: BTreeMap<EdgeIdx, Format>,
58    /// 2026-10-02: Step values the plan's rules state (`steps`), by node.
59    pub stated: BTreeMap<NodeIdx, BTreeMap<StepKind, Value>>,
60}
61
62impl<'a> Need<'a> {
63    /// 2026-10-02: The requirement before any plan: the circuit's own formats.
64    pub fn base(circuit: &'a Circuit, settings: &'a BTreeMap<String, String>, rows: u64) -> Self {
65        Need {
66            circuit,
67            settings,
68            rows,
69            stored: circuit.edges.iter().map(|e| e.format).collect(),
70            group_of: vec![None; circuit.nodes.len()],
71            edge_states: vec![None; circuit.edges.len()],
72            held: BTreeMap::new(),
73            stated: BTreeMap::new(),
74        }
75    }
76
77    /// 2026-10-02: The requirement under `plan`, whose groups' rules are among `rules` (a
78    /// group whose rule is not there, a hardware planner's placeholder, states nothing).
79    pub fn of_plan(
80        circuit: &'a Circuit,
81        plan: &FusionPlan,
82        rules: &[Rule],
83        settings: &'a BTreeMap<String, String>,
84    ) -> Self {
85        let mut need = Need::base(circuit, settings, plan.rows);
86        need.stored.clone_from(&plan.edge_formats);
87        need.edge_states.clone_from(&plan.edge_states);
88        for (g, grp) in plan.groups.iter().enumerate() {
89            for &n in &grp.nodes {
90                need.group_of[n] = Some(g);
91            }
92            let Some(rule) = rules.iter().find(|r| r.id == grp.rule) else {
93                continue;
94            };
95            for (p, &n) in rule.pattern.iter().zip(&grp.nodes) {
96                if let Some(f) = p.holds {
97                    for &e in &circuit.nodes[n].outputs {
98                        need.held.insert(e, f);
99                    }
100                }
101                if !p.steps.is_empty() {
102                    need.stated.insert(n, p.steps.clone());
103                }
104            }
105        }
106        need
107    }
108
109    fn input(&self, node: NodeIdx, e: EdgeIdx) -> Format {
110        match self.held.get(&e) {
111            Some(f) if self.held_input(node, e) => *f,
112            _ => self.stored[e],
113        }
114    }
115
116    /// 2026-10-02: Node `node` reads edge `e` from a producer in its own group, in a format a
117    /// rule states that producer holds it in.
118    pub fn held_input(&self, node: NodeIdx, e: EdgeIdx) -> bool {
119        self.held.contains_key(&e)
120            && self.circuit.edges[e].producer.is_some_and(|p| {
121                self.group_of[p].is_some() && self.group_of[p] == self.group_of[node]
122            })
123    }
124
125    fn output(&self, e: EdgeIdx) -> Format {
126        match (self.edge_states[e], self.held.get(&e)) {
127            (Some(EdgeState::Fused(_)), Some(f)) => *f,
128            _ => self.stored[e],
129        }
130    }
131
132    /// 2026-10-02: Edge `e` is handed on inside its group in a format a rule states.
133    pub fn holds_in_group(&self, e: EdgeIdx) -> bool {
134        self.held.contains_key(&e) && matches!(self.edge_states[e], Some(EdgeState::Fused(_)))
135    }
136
137    /// 2026-10-02: The steps a rule states for node `n`.
138    pub fn stated_steps(&self, n: NodeIdx) -> impl Iterator<Item = StepKind> + '_ {
139        self.stated
140            .get(&n)
141            .into_iter()
142            .flat_map(|m| m.keys().copied())
143    }
144}
145
146/// 2026-10-02: The pipeline `need` requires of node `n`.
147pub fn required(need: &Need<'_>, n: NodeIdx) -> Result<NodePipeline, PipelineError> {
148    let node = &need.circuit.nodes[n];
149    let fail = |detail: String| PipelineError::Required {
150        node: node.id.clone(),
151        op: node.op.name(),
152        detail,
153    };
154    let inputs: Vec<Format> = node.inputs.iter().map(|&e| need.input(n, e)).collect();
155    let outputs: Vec<Format> = node.outputs.iter().map(|&e| need.output(e)).collect();
156    let stated = need.stated.get(&n);
157    let first = |what: &str| {
158        inputs
159            .first()
160            .copied()
161            .ok_or_else(|| fail(format!("{what} needs an input")))
162    };
163    let mut projection: Option<(Format, Format)> = None;
164    let mut steps = Vec::new();
165    for &kind in steps_of(&node.op) {
166        let said = stated.and_then(|m| m.get(&kind)).cloned();
167        let value = match (kind, said) {
168            (StepKind::Act, said) => {
169                let weight = node
170                    .weight
171                    .ok_or_else(|| fail("a projection without a weight format".into()))?;
172                let said = match said {
173                    Some(Value::Format(f)) => Some(f),
174                    Some(other) => return Err(fail(format!("a stated act `{}`", other.name()))),
175                    None => None,
176                };
177                let act = super::act_policy::required_act(
178                    need.settings,
179                    &node.op,
180                    need.rows,
181                    (first("a projection")?, said, weight),
182                )
183                .map_err(fail)?;
184                projection = Some((act, weight));
185                Value::Format(act)
186            }
187            (_, Some(v)) => v,
188            (StepKind::Gather, None) => match node.op {
189                OpKind::Embed => Value::Format(
190                    outputs
191                        .first()
192                        .copied()
193                        .ok_or_else(|| fail("an embedding writes its rows".into()))?,
194                ),
195                _ => Value::Format(first("a gather")?),
196            },
197            (StepKind::Weight | StepKind::Mma | StepKind::Scale, None) => {
198                let (act, weight) =
199                    projection.ok_or_else(|| fail("`act` precedes the weight steps".into()))?;
200                derived(kind, act, weight).map_err(fail)?
201            }
202            (
203                StepKind::Accumulate
204                | StepKind::Compute
205                | StepKind::Scores
206                | StepKind::Softmax
207                | StepKind::Score
208                | StepKind::Scatter
209                | StepKind::Reduce,
210                None,
211            ) => Value::Num(Num::F32),
212            (StepKind::Combine, None) => Value::OptNum((inputs.len() > 2).then_some(Num::F32)),
213            (StepKind::Move | StepKind::Compare, None) => Value::Format(first("a move")?),
214            (StepKind::Cache, None) => Value::Kv(
215                need.settings
216                    .get("kv_cache_dtype")
217                    .cloned()
218                    .ok_or_else(|| fail("the policy states no `kv_cache_dtype`".into()))?,
219            ),
220            (StepKind::State, None) => Value::State(state_of(need, node).map_err(fail)?),
221        };
222        steps.push(Step { kind, value });
223    }
224    Ok(NodePipeline {
225        inputs,
226        steps,
227        outputs,
228    })
229}
230
231/// 2026-10-02: The weight, multiply and scale requirement of an `act` x `weight` projection.
232fn derived(kind: StepKind, act: Format, weight: Format) -> Result<Value, String> {
233    let operand = match (act, weight) {
234        (Format::Bf16, _) => Num::Bf16,
235        (Format::F32, _) => Num::F32,
236        (Format::Fp8E4m3 { .. }, Format::Fp8E4m3 { .. } | Format::Nvfp4 { .. }) => Num::E4m3,
237        (Format::Nvfp4 { .. }, Format::Nvfp4 { .. }) => Num::E2m1,
238        (a, w) => {
239            return Err(format!(
240                "no multiply takes a {w} weight under a {a} activation"
241            ));
242        }
243    };
244    Ok(match kind {
245        StepKind::Weight => Value::Weight {
246            stored: weight,
247            operand,
248        },
249        StepKind::Mma => Value::Mma {
250            a: Num::of_format(act).ok_or_else(|| format!("an activation in {act}"))?,
251            b: operand,
252        },
253        _ => Value::OptNum((!(act.is_plain() && weight.is_plain())).then_some(Num::F32)),
254    })
255}
256
257/// 2026-10-02: The precision of the state `node` updates (or snapshots).
258fn state_of(need: &Need<'_>, node: &Node) -> Result<StateDtype, String> {
259    let want = match node.op {
260        OpKind::StateSnapshot => StateAccess::Snapshot,
261        _ => StateAccess::Update,
262    };
263    let (idx, _) = node
264        .state
265        .iter()
266        .find(|(_, a)| *a == want)
267        .ok_or_else(|| format!("the node touches no state as `{want:?}`"))?;
268    let decl = &need.circuit.states[*idx];
269    match &decl.format {
270        StateFormat::Fixed(d) => Ok(*d),
271        StateFormat::Keyed(key) => keyed_state(key, need.settings),
272    }
273}
274
275/// 2026-10-02: The precision a keyed state format holds under `settings`. The recurrent h-state
276/// (`ssm_h_storage`) holds the precision `--ssm-h-dtype` names: FP16 under `f16` and
277/// `f16-pool` (which differ only in the pool's size), FP32 under `f32`.
278pub fn keyed_state(key: &str, settings: &BTreeMap<String, String>) -> Result<StateDtype, String> {
279    let setting = match key {
280        "ssm_h_storage" => "ssm_h_dtype",
281        other => other,
282    };
283    let value = settings
284        .get(setting)
285        .ok_or_else(|| format!("the policy states no `{setting}`"))?;
286    match (key, value.as_str()) {
287        ("ssm_h_storage", "f32") => Ok(StateDtype::F32),
288        ("ssm_h_storage", "f16" | "f16-pool") => Ok(StateDtype::F16),
289        ("ssm_h_storage", other) => Err(format!(
290            "`ssm_h_dtype = {other}` names no h-state precision"
291        )),
292        (_, other) => StateDtype::parse(other)
293            .ok_or_else(|| format!("`{setting} = {other}` is no state dtype")),
294    }
295}
296
297#[cfg(test)]
298#[path = "require_tests.rs"]
299mod require_tests;
metrale_circuit/pipeline/require.rs

metrale_circuit/pipeline/
require.rs
