metrale_circuit/
rules.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: Fusion rules: `kernels/<hw>/common/FUSIONS.toml`. Each rule names an op chain,
4//! the kernel that runs it, where it applies (rows, modes, caps, policy settings) and how its
5//! numerics relate to the reference.
6//!
7//! Owner: metrale-circuit.
8//! Invariants:
9//! - A `bit_identical` rule names its microtest, and a `differs` rule its opt-in lever; a
10//!   rule missing either is a load error, as is a lever or microtest on a `reference` rule.
11//! - Rule ids are unique; `rows` is a non-empty inclusive range; `modes` is non-empty.
12//! - Every rule cites the dispatch site that makes today's choice (`cite`).
13
14use std::collections::{BTreeMap, BTreeSet};
15
16use crate::format::Format;
17use crate::ir::{LayerKind, LinearRole, OpKind};
18
19/// 2026-09-28: Which forward a plan is for.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
21pub enum Mode {
22    /// 2026-09-28: One sequence, one row.
23    Decode,
24    /// 2026-09-28: `n` sequences, one row each, padded to the batch ladder.
25    MultiSeq,
26    /// 2026-09-28: `K` draft rows of one sequence (MTP verify).
27    Verify,
28    /// 2026-09-28: The MTP draft head proposing one token for each of `n` sequences.
29    Draft,
30    /// 2026-09-30: `n` sequences' MTP verifies in one forward (`decode_verify_batched`):
31    /// `R = Σ k` rows, seq-major, planned for a row table ([`crate::runs::RowTable`]).
32    VerifyBatch,
33}
34
35impl Mode {
36    /// 2026-09-28: Every mode, in plan-file order.
37    /// 2026-09-30: [`Mode::VerifyBatch`] is not among them: its plans are keyed by a row table,
38    /// not a row count.
39    pub const ALL: [Mode; 4] = [Mode::Decode, Mode::MultiSeq, Mode::Verify, Mode::Draft];
40
41    /// 2026-09-28: The spelling in rules, file names and the CLI.
42    pub fn name(self) -> &'static str {
43        match self {
44            Mode::Decode => "decode",
45            Mode::MultiSeq => "multi_seq",
46            Mode::Verify => "verify",
47            Mode::Draft => "draft",
48            Mode::VerifyBatch => "verify_batch",
49        }
50    }
51
52    /// 2026-09-28: Parse [`Mode::name`].
53    pub fn parse(s: &str) -> Option<Self> {
54        Mode::ALL
55            .into_iter()
56            .chain([Mode::VerifyBatch])
57            .find(|m| m.name() == s)
58    }
59}
60
61/// 2026-09-28: A kernel entry point as the engine looks it up: the module name (the source
62/// stem after KERNEL.toml `[modules]` renames, e.g. `norm` for `rms_norm.cu`) and the
63/// function.
64#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
65pub struct KernelId {
66    /// 2026-09-28: Module, e.g. `norm`.
67    pub module: String,
68    /// 2026-09-28: Function, e.g. `residual_add_rms_norm`.
69    pub func: String,
70}
71
72impl std::fmt::Display for KernelId {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        write!(f, "{}::{}", self.module, self.func)
75    }
76}
77
78/// 2026-09-28: How a rule's output relates to the reference numerics.
79#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
80pub enum Numerics {
81    /// 2026-09-28: Byte-identical to the unfused chain, proven by the named microtest.
82    BitIdentical {
83        /// 2026-09-28: The microtest.
84        microtest: String,
85    },
86    /// 2026-09-28: Today's default: the reference definition, grandfathered.
87    Reference,
88    /// 2026-09-28: Changes numerics; selected only when `lever` is opted in.
89    Differs {
90        /// 2026-09-28: The opt-in lever.
91        lever: String,
92    },
93}
94
95impl Numerics {
96    /// 2026-09-28: `bit_identical`, `reference` or `differs`.
97    pub fn class(&self) -> &'static str {
98        match self {
99            Numerics::BitIdentical { .. } => "bit_identical",
100            Numerics::Reference => "reference",
101            Numerics::Differs { .. } => "differs",
102        }
103    }
104}
105
106/// 2026-09-28: One element of a rule's op chain.
107#[derive(Debug, Clone, PartialEq, Eq)]
108pub struct PatternOp {
109    /// 2026-09-28: The op the node must be. For `linear` with a `roles` list, any of
110    /// [`PatternOp::roles`] matches and this holds the first.
111    pub op: OpKind,
112    /// 2026-09-28: Linear roles any of which matches; empty means `op` exactly.
113    pub roles: BTreeSet<LinearRole>,
114    /// 2026-09-28: The node's layer must be of this kind, when set.
115    pub layer_kind: Option<LayerKind>,
116    /// 2026-09-28: The node's template-local id must be this, when set.
117    pub local: Option<String>,
118    /// 2026-09-28: The node's weight format must be this, when set.
119    pub weight: Option<Format>,
120    /// 2026-09-28: The format the kernel reads its first input as. Checked after fusion
121    /// against the format the producing group writes, so a rule set whose kernels disagree
122    /// on a materialised edge is refused.
123    pub input: Option<Format>,
124    /// 2026-09-28: The format the kernel stores this element's outputs in, when it is not
125    /// the circuit's declared format (the verify conv stores BF16 where decode keeps F32).
126    pub writes: Option<Format>,
127    /// 2026-09-28: The kernel writes this element's outputs even though a later element
128    /// reads them, so they may have readers outside the group.
129    pub keep: bool,
130    /// 2026-09-30: The kernel stores this element's outputs in memory, where a later launch of
131    /// the group reads them back (a per-run schedule's conv rows), so they are materialized even
132    /// when every reader is in the group.
133    pub stored: bool,
134    /// 2026-09-28: The element joins the group by sitting in the same block as the first
135    /// element rather than by reading an output of an earlier one (a shared expert beside
136    /// the routed experts).
137    pub sibling: bool,
138    /// 2026-10-02: The format the kernel hands this element's outputs to later members of the
139    /// group in, where it is not the format the plan stores (a sum kept in FP32 registers);
140    /// `None` hands them on as stored ([`crate::pipeline::require`]).
141    pub holds: Option<Format>,
142    /// 2026-10-02: Steps of this element's node that the routing runs at another value than
143    /// the reference requirement (an activation quantized inside the group, a sum rounded to
144    /// BF16), stated so the plan shows them and the kernel must declare them.
145    pub steps: BTreeMap<crate::pipeline::StepKind, crate::pipeline::Value>,
146}
147
148/// 2026-09-28: How many times a group's kernels launch per step.
149#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
150pub enum Repeat {
151    /// 2026-09-28: Once.
152    Once,
153    /// 2026-09-28: Once per row (a per-sequence loop).
154    PerRow,
155    /// 2026-09-28: Once per `n`-row chunk.
156    Chunk(u64),
157    /// 2026-09-29: Once per row but the last.
158    PerRowButLast,
159    /// 2026-09-30: Per run of a batched verify, as the rule's selectors say
160    /// ([`Rule::runs`], [`crate::fuser::Group::runs`]); not a function of the rows alone.
161    PerRun,
162}
163
164impl Repeat {
165    /// 2026-09-28: Launches of each kernel for `rows` rows. 2026-09-30: `None` for
166    /// [`Repeat::PerRun`], whose count is its runs' ([`crate::fuser::Group::launch_count`]).
167    pub fn count(self, rows: u64) -> Option<u64> {
168        match self {
169            Repeat::Once => Some(1),
170            Repeat::PerRow => Some(rows),
171            Repeat::Chunk(c) => Some(rows.div_ceil(c)),
172            Repeat::PerRowButLast => Some(rows.saturating_sub(1)),
173            Repeat::PerRun => None,
174        }
175    }
176
177    /// 2026-09-28: `once`, `per_row` or `chunk<n>`.
178    pub fn name(self) -> String {
179        match self {
180            Repeat::Once => "once".into(),
181            Repeat::PerRow => "per_row".into(),
182            Repeat::Chunk(c) => format!("chunk{c}"),
183            Repeat::PerRowButLast => "per_row_but_last".into(),
184            Repeat::PerRun => "per_run".into(),
185        }
186    }
187
188    fn parse(s: &str) -> Option<Self> {
189        match s {
190            "once" => Some(Repeat::Once),
191            "per_row" => Some(Repeat::PerRow),
192            "per_row_but_last" => Some(Repeat::PerRowButLast),
193            "per_run" => Some(Repeat::PerRun),
194            _ => s
195                .strip_prefix("chunk")?
196                .parse::<u64>()
197                .ok()
198                .filter(|&c| c > 0)
199                .map(Repeat::Chunk),
200        }
201    }
202}
203
204/// 2026-09-28: One fusion rule.
205#[derive(Debug, Clone, PartialEq, Eq)]
206pub struct Rule {
207    /// 2026-09-28: Unique id.
208    pub id: String,
209    /// 2026-09-28: The op chain, each element reading an output of the one before.
210    pub pattern: Vec<PatternOp>,
211    /// 2026-09-28: The kernels the emitter launches for the chain, in order. Empty for a
212    /// chain that runs as a copy-engine transfer (the embedding row copy).
213    pub kernels: Vec<KernelId>,
214    /// 2026-09-28: How often the kernels launch per step.
215    pub repeat: Repeat,
216    /// 2026-09-29: Copy-engine transfers the emitter issues besides its kernels, and how
217    /// often; `None` for none (the verify conv's state snapshots are `per_row_but_last`).
218    pub copies: Option<Repeat>,
219    /// 2026-09-28: The executor's emitter for this kernel.
220    pub emitter: String,
221    /// 2026-09-28: Inclusive row range.
222    pub rows: (u64, u64),
223    /// 2026-09-28: Modes it applies in.
224    pub modes: BTreeSet<Mode>,
225    /// 2026-09-28: Kernel capability bits it needs.
226    pub requires: BTreeSet<String>,
227    /// 2026-09-28: Policy settings that must hold (`kv_cache_dtype = "bf16"`).
228    pub when: BTreeMap<String, String>,
229    /// 2026-09-28: Numerics class.
230    pub numerics: Numerics,
231    /// 2026-09-28: Higher wins.
232    pub priority: i64,
233    /// 2026-09-28: The dispatch site(s) that make this choice today.
234    pub cite: String,
235    /// 2026-09-30: A `per_run` rule's selectors (`[[rule.run]]`), in file order; empty for any
236    /// other repeat.
237    pub runs: Vec<crate::runs::RunSelect>,
238}
239
240/// 2026-09-28: Why FUSIONS.toml did not load.
241#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
242pub enum RuleError {
243    /// 2026-09-28: Not valid TOML or not the rules shape.
244    #[error("FUSIONS.toml: {0}")]
245    Parse(String),
246    /// 2026-09-28: Two rules share an id.
247    #[error("duplicate rule id `{0}`")]
248    DuplicateId(String),
249    /// 2026-09-30: A malformed `[[runtime]]` route ([`crate::runtime`]).
250    #[error("runtime `{route}`: {detail}")]
251    Runtime {
252        /// 2026-09-30: Route id.
253        route: String,
254        /// 2026-09-30: What was wrong.
255        detail: String,
256    },
257    /// 2026-09-28: A pattern op the vocabulary does not have.
258    #[error("rule `{rule}`: {detail}")]
259    Op {
260        /// 2026-09-28: Rule id.
261        rule: String,
262        /// 2026-09-28: What was wrong.
263        detail: String,
264    },
265    /// 2026-09-28: `bit_identical` without `microtest`.
266    #[error("rule `{0}` is bit_identical but names no microtest")]
267    MissingMicrotest(String),
268    /// 2026-09-28: `differs` without `lever`.
269    #[error("rule `{0}` is differs but names no opt-in lever")]
270    MissingLever(String),
271    /// 2026-09-28: A `microtest` or `lever` on a class that takes none, or an unknown class.
272    #[error("rule `{rule}`: {detail}")]
273    Numerics {
274        /// 2026-09-28: Rule id.
275        rule: String,
276        /// 2026-09-28: What was wrong.
277        detail: String,
278    },
279    /// 2026-09-28: Empty pattern, empty modes, unknown mode, bad row range or empty cite.
280    #[error("rule `{rule}`: {detail}")]
281    Shape {
282        /// 2026-09-28: Rule id.
283        rule: String,
284        /// 2026-09-28: What was wrong.
285        detail: String,
286    },
287}
288
289#[path = "rules_parse.rs"]
290mod parse;
291pub(crate) use parse::parse_file;
292pub use parse::parse_rules;
293
294#[cfg(test)]
295#[path = "rules_tests.rs"]
296mod rules_tests;
metrale_circuit/rules.rs

metrale_circuit/
rules.rs
