fuel_graph/jit.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Kernel-seam JIT — the Fuel-side projection + structural matcher
3//! (kernel-seam-interop §5; fkc-fusion-patterns §3/§3a). The frozen grammar
4//! types ([`OpTag`], [`OpAttrs`], [`PatternNode`]) live in
5//! [`fuel_kernel_seam_types`] and are re-exported here; this module holds the
6//! two parts that need the graph: the `Op -> OpTag` projection [`op_to_tag`]
7//! and the structural matcher [`match_region`] (the declarative-fusion engine
8//! behind `PatternKind::Declarative`).
9//!
10//! One [`PatternNode`] serves the **JIT region** (Fuel -> synthesizer, "build a
11//! kernel for this subgraph"), a contract's `pattern:` **re-fuse rule**, and a
12//! synthesized op's **`decompose`** (the region re-emitted).
13
14use crate::{Graph, NodeId, Op};
15use std::collections::BTreeMap;
16
17pub use fuel_kernel_seam_types::{OpAttrs, OpTag, PatternNode};
18
19// ===========================================================================
20// op_to_tag — the Op -> OpTag projection (stays Fuel-side: it needs the graph Op)
21// ===========================================================================
22
23/// Project a graph [`Op`] to its functional [`OpTag`]. Returns `None` for
24/// in-place variants, structural/bookkeeping ops, and `Op::Fused` (a fused op
25/// isn't a region node — its *decomposition* is). Same for `Op::Scan` /
26/// `Op::ScanPlaceholder`: a scan isn't a region node either — its body
27/// (referenced via `inputs`, not a `PatternNode`) is what a region walk
28/// would see. A `None` op in a region is an honest "outside the vocabulary"
29/// miss, never a crash.
30pub fn op_to_tag(op: &Op) -> Option<OpTag> {
31 Some(match op {
32 Op::Add => OpTag::Add,
33 Op::Sub => OpTag::Sub,
34 Op::Mul => OpTag::Mul,
35 Op::Div => OpTag::Div,
36 Op::Maximum => OpTag::Maximum,
37 Op::Minimum => OpTag::Minimum,
38 Op::Pow => OpTag::Pow,
39 Op::Rem => OpTag::Rem,
40 Op::Neg => OpTag::Neg,
41 Op::Abs => OpTag::Abs,
42 Op::Sqr => OpTag::Sqr,
43 Op::Sqrt => OpTag::Sqrt,
44 Op::Rsqrt => OpTag::Rsqrt,
45 Op::Recip => OpTag::Recip,
46 Op::Exp => OpTag::Exp,
47 Op::Log => OpTag::Log,
48 Op::Sin => OpTag::Sin,
49 Op::Cos => OpTag::Cos,
50 Op::Tanh => OpTag::Tanh,
51 Op::Sigmoid => OpTag::Sigmoid,
52 Op::Silu => OpTag::Silu,
53 Op::GeluTanh => OpTag::GeluTanh,
54 Op::Gelu => OpTag::Gelu,
55 Op::Relu => OpTag::Relu,
56 Op::Erf => OpTag::Erf,
57 Op::Step => OpTag::Step,
58 Op::Floor => OpTag::Floor,
59 Op::Ceil => OpTag::Ceil,
60 Op::Round => OpTag::Round,
61 Op::Sign => OpTag::Sign,
62 Op::AddScalar(_) => OpTag::AddScalar,
63 Op::MulScalar(_) => OpTag::MulScalar,
64 Op::PowI(_) => OpTag::PowI,
65 Op::Clamp { .. } => OpTag::Clamp,
66 Op::Equal => OpTag::Equal,
67 Op::Ne => OpTag::Ne,
68 Op::Lt => OpTag::Lt,
69 Op::Le => OpTag::Le,
70 Op::Gt => OpTag::Gt,
71 Op::Ge => OpTag::Ge,
72 Op::Where => OpTag::Where,
73 Op::MaskedFill { .. } => OpTag::MaskedFill,
74 Op::SumAll => OpTag::SumAll,
75 Op::MaxAll => OpTag::MaxAll,
76 Op::MinAll => OpTag::MinAll,
77 Op::MeanAll => OpTag::MeanAll,
78 Op::SumDim(_) => OpTag::SumDim,
79 Op::MaxDim(_) => OpTag::MaxDim,
80 Op::MeanDim(_) => OpTag::MeanDim,
81 Op::ReduceSumTo(_) => OpTag::ReduceSumTo,
82 Op::ReduceMaxTo(_) => OpTag::ReduceMaxTo,
83 Op::CumSum { .. } => OpTag::CumSum,
84 Op::MatMul => OpTag::MatMul,
85 Op::Transpose => OpTag::Transpose,
86 Op::Permute(_) => OpTag::Permute,
87 Op::Reshape(_) => OpTag::Reshape,
88 Op::BroadcastTo(_) => OpTag::BroadcastTo,
89 Op::Unsqueeze { .. } => OpTag::Unsqueeze,
90 Op::Squeeze { .. } => OpTag::Squeeze,
91 Op::Cast(_) => OpTag::Cast,
92 Op::Slice { .. } => OpTag::Slice,
93 Op::Concat { .. } => OpTag::Concat,
94 Op::Flip { .. } => OpTag::Flip,
95 Op::Roll { .. } => OpTag::Roll,
96 Op::Pad { .. } => OpTag::Pad,
97 Op::Triu { .. } => OpTag::Triu,
98 Op::Tril { .. } => OpTag::Tril,
99 Op::IndexSelect { .. } => OpTag::IndexSelect,
100 Op::Gather { .. } => OpTag::Gather,
101 Op::IndexAdd { .. } => OpTag::IndexAdd,
102 Op::ScatterAdd { .. } => OpTag::ScatterAdd,
103 Op::LogSoftmaxLastDim => OpTag::LogSoftmaxLastDim,
104 Op::Iota { .. } => OpTag::Iota,
105 // In-place variants, structural / bookkeeping ops, and Op::Fused are
106 // not region nodes.
107 _ => return None,
108 })
109}
110
111// ===========================================================================
112// match_region — the structural declarative matcher (§3a)
113// ===========================================================================
114
115/// The §4.1 commutative ops, whose operands match order-independently (§3a.2a).
116/// `Maximum`/`Minimum` are DELIBERATELY ABSENT. They are commutative in
117/// value and NOT in bits: `chassis/binary.rs` returns the FIRST NaN operand
118/// (`if a.is_nan() { a } else if b.is_nan() { b }`), a documented lhs-first
119/// tie-break matching `torch.maximum`. Two distinct NaN payloads therefore
120/// swap with operand order -- measured, `max(0x7FC00001, 0x7FC00002)` gives
121/// `0x7FC00001` and the reverse gives `0x7FC00002` -- so folding `min(a,b)`
122/// with `min(b,a)` is an INVALID rewrite. See GAP-271.
123///
124/// The +-0 tie ALSO makes the fold invalid, but it is NOT A PROPERTY OF THIS
125/// SOURCE: it is decided per CALL SITE by the optimizer, because we delegate
126/// to Rust's `f32::max`, which disclaims the zero tie. Measured in ONE binary
127/// at opt-level 3 -- const-folded literals +0/+0, `black_box` -0/-0,
128/// `#[inline(never)]` -0/+0 (ORDER-DEPENDENT). Rely on the NaN-payload
129/// argument: it is STABLE across all four contexts, because that tie-break is
130/// explicit in our source. See `opt.rs::is_commutative` for the full note.
131fn is_commutative(op: OpTag) -> bool {
132 matches!(op, OpTag::Add | OpTag::Mul)
133}
134
135/// Project a graph [`Op`]'s load-bearing non-tensor parameters into an
136/// [`OpAttrs`] — the graph-side mirror of the region-side `attrs` a
137/// [`PatternNode::Op`] carries. The graph node stores these as *typed* `Op`
138/// payloads (`Op::Permute(Vec<usize>)`, `Op::AddScalar(f64)`, `Op::SumDim(usize)`,
139/// `Op::BroadcastTo(Shape)`, …); this reads them out into the flat `OpAttrs`
140/// surface so [`attrs_match`] can compare a pattern's `attrs` against them
141/// without the seam-types crate needing to know about the graph `Op` (it stays
142/// dependency-free). Ops with no attr payload project to `OpAttrs::default()`
143/// (all fields empty), which the wildcard rule treats as "no constraint".
144///
145/// **Scope (not a full `Op → OpAttrs → Op` round-trip).** This projection is
146/// **matcher-driven**: it fills only the fields `attrs_match` needs to
147/// discriminate the patterns that exist today. It is faithful (a projected
148/// field equals the `Op`'s value) for the ops it *does* project — including the
149/// Convergence-A additions Slice/Concat/Flip/Roll/Cast/Iota/Pad/MaskedFill — but
150/// it is **not** exhaustive: `tag_to_op` can *reconstruct* several axis-bearing
151/// ops (`CumSum`/`IndexSelect`/`Gather`/`IndexAdd`/`ScatterAdd`) from `axis` that
152/// this fn deliberately does **not** project yet (no matcher consumer), so a
153/// full round-trip through `op_to_attrs` would drop their axis. See the `_ => {}`
154/// arm. The re-emit path (`emit`) supplies those attrs directly from the region
155/// author, so this gap is projection-only, not an emit gap.
156fn op_to_attrs(op: &Op) -> OpAttrs {
157 let mut a = OpAttrs::default();
158 match op {
159 // Scalar-param ops → `scalars` (the region's slot snapshot; F1).
160 Op::AddScalar(v) | Op::MulScalar(v) => a.scalars = vec![*v],
161 Op::Clamp { min, max } => a.scalars = vec![*min, *max],
162 // PowI → i32 exponent on `scalars[0]` (Increment C carriers, A3): a
163 // BAKED constant (like MaskedFill's fill), completing the projection ↔
164 // `tag_to_op` round-trip and matching the §6.19 wire (which already
165 // serializes `scalars` for PowI). Exact for every i32 through f64.
166 Op::PowI(n) => a.scalars = vec![*n as f64],
167 // Dim-bearing ops → `axis`.
168 Op::SumDim(d) | Op::MaxDim(d) | Op::MeanDim(d) => a.axis = Some(*d as i64),
169 Op::Triu { diagonal } | Op::Tril { diagonal } => a.axis = Some(*diagonal),
170 // Permute/Transpose → absolute `perm` (F1/F2a). `Transpose` is the
171 // rank-2 last-two-axes special case; without the input rank on the op
172 // itself it projects to an empty perm (a wildcard) here — a `Permute`
173 // pattern is the discriminating form.
174 Op::Permute(axes) => a.perm = axes.iter().map(|&x| x as u8).collect(),
175 // Shape-target ops → `target_shape` (BroadcastTo + Reshape share it; F1).
176 Op::BroadcastTo(shape) | Op::Reshape(shape) => {
177 a.target_shape = shape.dims().iter().map(|&d| d as i64).collect()
178 }
179 // Squeeze/Unsqueeze → single-element `dims` (F1).
180 Op::Unsqueeze { dim } | Op::Squeeze { dim } => a.dims = vec![*dim as u8],
181 // --- Convergence Increment A: the full first-order set ---
182 // Slice → axis(dim) + start + len.
183 Op::Slice { dim, start, len } => {
184 a.axis = Some(*dim as i64);
185 a.slice_start = Some(*start as u64);
186 a.slice_len = Some(*len as u64);
187 }
188 // Concat/Flip → axis(dim).
189 Op::Concat { dim } | Op::Flip { dim } => a.axis = Some(*dim as i64),
190 // Roll → axis(dim) + signed shift.
191 Op::Roll { dim, shift } => {
192 a.axis = Some(*dim as i64);
193 a.roll_shift = Some(*shift);
194 }
195 // Cast → target dtype name (dep-free; fuel-graph maps back via FromStr).
196 Op::Cast(dt) => a.cast_dtype = Some(dt.as_str().to_string()),
197 // Iota len rides `target_shape` as the single-element output shape
198 // (mirrors how `target_shape` already serves BroadcastTo + Reshape).
199 Op::Iota { len } => a.target_shape = vec![*len as i64],
200 // Pad → per-axis amounts + mode code + constant value.
201 Op::Pad {
202 padding,
203 mode,
204 value,
205 } => {
206 a.pad_amounts = padding.iter().map(|&(b, e)| (b as u64, e as u64)).collect();
207 a.pad_mode = Some(match mode {
208 crate::PadMode::Constant => 0,
209 crate::PadMode::Reflect => 1,
210 crate::PadMode::Replicate => 2,
211 });
212 a.pad_value = Some(*value);
213 }
214 // MaskedFill → value on scalars[0] + its dtype on cast_dtype.
215 Op::MaskedFill { value } => {
216 a.scalars = vec![value.to_f64()];
217 a.cast_dtype = Some(value.dtype().as_str().to_string());
218 }
219 // Intentionally NOT projected yet (no matcher consumer): the axis-bearing
220 // ops `tag_to_op` can reconstruct from `attrs.axis` but no pattern needs
221 // to discriminate on today — `Op::CumSum`, `Op::IndexSelect`,
222 // `Op::Gather`, `Op::IndexAdd`, `Op::ScatterAdd`. Add an `axis` arm here
223 // only when a matcher requires it (deferred; the re-emit path already
224 // gets these attrs from the region author, not from this projection).
225 _ => {}
226 }
227 a
228}
229
230/// Compare a pattern node's `attrs` against the graph node's projected attrs
231/// with **wildcard-on-unset** semantics: an *empty/unset* field on the pattern
232/// is a wildcard (matches any graph value); a *set* field must equal the graph
233/// node's value. This is what keeps every existing attr-agnostic pattern (all
234/// authored with `OpAttrs::default()`) matching after attrs become comparable,
235/// while letting a layout/scalar pattern that *sets* a field discriminate.
236///
237/// `Vec` fields are unset ⇔ empty; `axis: Option` is unset ⇔ `None`. A set
238/// pattern field must equal the graph projection exactly (absolute perm, F2a).
239fn attrs_match(pattern: &OpAttrs, node: &OpAttrs) -> bool {
240 (pattern.scalars.is_empty() || pattern.scalars == node.scalars)
241 && (pattern.axis.is_none() || pattern.axis == node.axis)
242 && (pattern.perm.is_empty() || pattern.perm == node.perm)
243 && (pattern.target_shape.is_empty() || pattern.target_shape == node.target_shape)
244 && (pattern.dims.is_empty() || pattern.dims == node.dims)
245}
246
247/// Match a declarative region [`PatternNode`] against the graph rooted at
248/// `root` (the subgraph **sink**, §3a.1). Returns the region's external inputs
249/// in `bind`-index order on a match, or `None`. This is the structural core of
250/// the declarative matcher (`PatternKind::Declarative`); the §5 `guard:`/`extract:`
251/// layers and the `see_through`-set wrappers compose on top.
252///
253/// Implements: positional exact tensor-arity (scalar params are attributes, not
254/// operands — §3a.2); commutative-op order-independence (§3a.2a, by trying both
255/// orderings); the **interior sole-consumer guard** (§3a.4 — a matched Op that is
256/// neither the root nor a `bind` leaf must feed *only* the fusion, else fusing
257/// duplicates its computation and we decline); and the repeated-`bind`
258/// node-identity guard (§3.2). `consumers(n)` returns node `n`'s consumer count.
259pub fn match_region(
260 graph: &Graph,
261 root: NodeId,
262 pattern: &PatternNode,
263 consumers: &dyn Fn(NodeId) -> usize,
264) -> Option<Vec<NodeId>> {
265 match_region_extract(graph, root, pattern, consumers).map(|(binds, _)| binds)
266}
267
268/// [`match_region`] plus the **`extract:` layer** (FKC §5.3): alongside the
269/// bound inputs, return the matched region's live scalar values — one entry per
270/// scalar a **slot** pattern node left open (a scalar-carrying op whose pattern
271/// `attrs.scalars` is empty/wildcard), read from the matched graph node, in
272/// **pattern pre-order** (the canonical slot order the recipe's re-emit and a
273/// synthesized kernel's trailing `p{i}` launch args both use). A pattern node
274/// with *baked* scalars is a constant of the pattern (the attr guard enforced
275/// equality), not a slot, and extracts nothing. A slotless region extracts `[]`.
276pub fn match_region_extract(
277 graph: &Graph,
278 root: NodeId,
279 pattern: &PatternNode,
280 consumers: &dyn Fn(NodeId) -> usize,
281) -> Option<(Vec<NodeId>, Vec<f64>)> {
282 let mut binds: BTreeMap<u8, NodeId> = BTreeMap::new();
283 let mut scalars: Vec<f64> = Vec::new();
284 match_node(
285 graph,
286 root,
287 pattern,
288 true,
289 consumers,
290 &mut binds,
291 &mut scalars,
292 )?;
293 // Bind indices must form a contiguous [0, n) — exactly the region's inputs.
294 let n = binds.len() as u8;
295 if (0..n).all(|i| binds.contains_key(&i)) {
296 Some(((0..n).map(|i| binds[&i]).collect(), scalars))
297 } else {
298 None
299 }
300}
301
302fn match_node(
303 graph: &Graph,
304 node_id: NodeId,
305 pattern: &PatternNode,
306 is_root: bool,
307 consumers: &dyn Fn(NodeId) -> usize,
308 binds: &mut BTreeMap<u8, NodeId>,
309 scalars: &mut Vec<f64>,
310) -> Option<()> {
311 match pattern {
312 PatternNode::Any => Some(()),
313 PatternNode::Bind { index } => {
314 // Node-identity guard: a repeated index must bind the SAME node.
315 match binds.get(index) {
316 Some(&existing) if existing != node_id => None,
317 _ => {
318 binds.insert(*index, node_id);
319 Some(())
320 }
321 }
322 }
323 PatternNode::SeeThrough { then } => {
324 // The see_through-set skip is a follow-up; for now match `then`
325 // against this node directly.
326 match_node(graph, node_id, then, is_root, consumers, binds, scalars)
327 }
328 PatternNode::Op {
329 op,
330 operands,
331 attrs,
332 } => {
333 let node = graph.node(node_id);
334 if op_to_tag(&node.op) != Some(*op) {
335 return None;
336 }
337 // Attr guard (F1): a SET pattern attr must equal the graph node's
338 // projected value; an empty/unset pattern attr is a wildcard, so
339 // existing attr-agnostic patterns (all `OpAttrs::default()`) keep
340 // matching. Op-tag is checked first, so the projection is meaningful.
341 let node_attrs = op_to_attrs(&node.op);
342 if !attrs_match(attrs, &node_attrs) {
343 return None;
344 }
345 // Interior nodes (not the root, not a bind leaf) must be sole-consumer.
346 if !is_root && consumers(node_id) != 1 {
347 return None;
348 }
349 // Exact tensor-input arity (scalar/attribute params are not operands).
350 let inputs = &node.inputs;
351 if inputs.len() != operands.len() {
352 return None;
353 }
354 // The `extract:` layer (§5.3): empty pattern-scalars on a
355 // scalar-carrying op is a SLOT — record the live values from the
356 // matched node, pre-order (before descending into operands).
357 if attrs.scalars.is_empty() && !node_attrs.scalars.is_empty() {
358 scalars.extend_from_slice(&node_attrs.scalars);
359 }
360 if is_commutative(*op) && operands.len() == 2 {
361 // Try both orderings; commit the first that fully matches.
362 // Clone-commit covers `scalars` too: a failed first ordering
363 // must not leave partial extractions behind.
364 for (a, b) in [(0usize, 1usize), (1, 0)] {
365 let mut trial = binds.clone();
366 let mut trial_scalars = scalars.clone();
367 if match_node(
368 graph,
369 inputs[a],
370 &operands[0],
371 false,
372 consumers,
373 &mut trial,
374 &mut trial_scalars,
375 )
376 .is_some()
377 && match_node(
378 graph,
379 inputs[b],
380 &operands[1],
381 false,
382 consumers,
383 &mut trial,
384 &mut trial_scalars,
385 )
386 .is_some()
387 {
388 *binds = trial;
389 *scalars = trial_scalars;
390 return Some(());
391 }
392 }
393 return None;
394 }
395 for (child_pat, &child_id) in operands.iter().zip(inputs.iter()) {
396 match_node(graph, child_id, child_pat, false, consumers, binds, scalars)?;
397 }
398 Some(())
399 }
400 }
401}
402
403#[cfg(test)]
404mod tests {
405 use super::*;
406 use fuel_ir::DType;
407
408 #[test]
409 fn op_to_tag_projects_functional_ops_and_skips_structural() {
410 // Functional ops project; the GeluTanh/Gelu distinction is preserved.
411 assert_eq!(op_to_tag(&Op::Add), Some(OpTag::Add));
412 assert_eq!(op_to_tag(&Op::Gelu), Some(OpTag::Gelu));
413 assert_eq!(op_to_tag(&Op::GeluTanh), Some(OpTag::GeluTanh)); // tanh-approx, distinct
414 assert_ne!(op_to_tag(&Op::GeluTanh), op_to_tag(&Op::Gelu));
415 assert_eq!(op_to_tag(&Op::AddScalar(1.0)), Some(OpTag::AddScalar));
416 assert_eq!(op_to_tag(&Op::MatMul), Some(OpTag::MatMul));
417 // In-place + structural ops are not region nodes.
418 assert_eq!(op_to_tag(&Op::ReluInplace), None);
419 assert_eq!(op_to_tag(&Op::Const), None);
420 assert_eq!(op_to_tag(&Op::Release), None);
421 }
422
423 #[test]
424 fn op_to_attrs_projects_new_first_order_params() {
425 use fuel_ir::DType;
426 // Cast → cast_dtype name.
427 let a = op_to_attrs(&Op::Cast(DType::F16));
428 assert_eq!(a.cast_dtype.as_deref(), Some("f16"));
429 // Slice → axis(dim) + start + len.
430 let a = op_to_attrs(&Op::Slice {
431 dim: 2,
432 start: 3,
433 len: 5,
434 });
435 assert_eq!(
436 (a.axis, a.slice_start, a.slice_len),
437 (Some(2), Some(3), Some(5))
438 );
439 // Concat → axis(dim).
440 assert_eq!(op_to_attrs(&Op::Concat { dim: 1 }).axis, Some(1));
441 // Roll → axis(dim) + roll_shift.
442 let a = op_to_attrs(&Op::Roll { dim: 0, shift: -2 });
443 assert_eq!((a.axis, a.roll_shift), (Some(0), Some(-2)));
444 // Flip → axis(dim).
445 assert_eq!(op_to_attrs(&Op::Flip { dim: 1 }).axis, Some(1));
446 // Pad → amounts + mode + value.
447 let a = op_to_attrs(&Op::Pad {
448 padding: vec![(1, 1), (0, 2)],
449 mode: crate::PadMode::Constant,
450 value: 0.5,
451 });
452 assert_eq!(a.pad_amounts, vec![(1, 1), (0, 2)]);
453 assert_eq!((a.pad_mode, a.pad_value), (Some(0), Some(0.5)));
454 // Iota len rides target_shape.
455 assert_eq!(op_to_attrs(&Op::Iota { len: 7 }).target_shape, vec![7]);
456 // PowI → i32 exponent on scalars[0] (A3 carrier; the §6.19 wire arm
457 // already serializes `scalars` for PowI, so this completes the
458 // projection ↔ reconstruction round-trip). RED before A3 (fell in the
459 // `_ => {}` no-op arm).
460 assert_eq!(op_to_attrs(&Op::PowI(3)).scalars, vec![3.0]);
461 }
462
463 #[test]
464 fn max_dim_projects_tag_and_axis() {
465 // T4 (Increment C slice 1): additive OpTag::MaxDim — the D3 keepdim
466 // swap ({Max,Sum,Mean}Dim + Unsqueeze) needs MaxDim in the region
467 // vocabulary. Round-trip anchor: op_to_tag projects the tag,
468 // op_to_attrs projects the axis (same carrier as SumDim/MeanDim).
469 assert_eq!(op_to_tag(&Op::MaxDim(1)), Some(OpTag::MaxDim));
470 assert_eq!(op_to_attrs(&Op::MaxDim(2)).axis, Some(2));
471 }
472
473 // ---- the structural matcher (match_region) -------------------------------
474
475 fn consumer_counts(g: &Graph) -> std::collections::HashMap<NodeId, usize> {
476 let mut c = std::collections::HashMap::new();
477 for i in 0..g.len() {
478 for &inp in &g.node(NodeId(i)).inputs {
479 *c.entry(inp).or_insert(0) += 1;
480 }
481 }
482 c
483 }
484
485 fn leaf(g: &mut Graph, s: &fuel_ir::Shape) -> NodeId {
486 g.push(crate::Node {
487 op: Op::Const,
488 inputs: vec![],
489 shape: s.clone(),
490 dtype: DType::F32,
491 })
492 }
493 fn op1(g: &mut Graph, op: Op, x: NodeId, s: &fuel_ir::Shape) -> NodeId {
494 g.push(crate::Node {
495 op,
496 inputs: vec![x],
497 shape: s.clone(),
498 dtype: DType::F32,
499 })
500 }
501 fn op2(g: &mut Graph, op: Op, x: NodeId, y: NodeId, s: &fuel_ir::Shape) -> NodeId {
502 g.push(crate::Node {
503 op,
504 inputs: vec![x, y],
505 shape: s.clone(),
506 dtype: DType::F32,
507 })
508 }
509
510 fn relu_add_pattern() -> PatternNode {
511 PatternNode::Op {
512 op: OpTag::Relu,
513 attrs: OpAttrs::default(),
514 operands: vec![PatternNode::Op {
515 op: OpTag::Add,
516 attrs: OpAttrs::default(),
517 operands: vec![
518 PatternNode::Bind { index: 0 },
519 PatternNode::Bind { index: 1 },
520 ],
521 }],
522 }
523 }
524
525 #[test]
526 fn match_region_binds_relu_add() {
527 let mut g = Graph::new();
528 let s = fuel_ir::Shape::from_dims(&[2]);
529 let a = leaf(&mut g, &s);
530 let b = leaf(&mut g, &s);
531 let sum = op2(&mut g, Op::Add, a, b, &s);
532 let r = op1(&mut g, Op::Relu, sum, &s);
533 let counts = consumer_counts(&g);
534 let got = match_region(&g, r, &relu_add_pattern(), &|n| {
535 *counts.get(&n).unwrap_or(&0)
536 });
537 assert_eq!(got, Some(vec![a, b]));
538 }
539
540 #[test]
541 fn match_region_commutative_order_independent() {
542 // mul(c, relu(d)); the pattern puts relu on operand[0] — matches via
543 // commutativity (§3a.2a).
544 let mut g = Graph::new();
545 let s = fuel_ir::Shape::from_dims(&[2]);
546 let c = leaf(&mut g, &s);
547 let d = leaf(&mut g, &s);
548 let rd = op1(&mut g, Op::Relu, d, &s);
549 let prod = op2(&mut g, Op::Mul, c, rd, &s);
550 let pat = PatternNode::Op {
551 op: OpTag::Mul,
552 attrs: OpAttrs::default(),
553 operands: vec![
554 PatternNode::Op {
555 op: OpTag::Relu,
556 attrs: OpAttrs::default(),
557 operands: vec![PatternNode::Bind { index: 0 }],
558 },
559 PatternNode::Bind { index: 1 },
560 ],
561 };
562 let counts = consumer_counts(&g);
563 let got = match_region(&g, prod, &pat, &|n| *counts.get(&n).unwrap_or(&0));
564 assert_eq!(got, Some(vec![d, c]), "bind 0 = d (under relu), bind 1 = c");
565 }
566
567 #[test]
568 fn match_region_rejects_wrong_interior_op() {
569 // relu(mul(a, b)) does not match the relu(add(...)) pattern.
570 let mut g = Graph::new();
571 let s = fuel_ir::Shape::from_dims(&[2]);
572 let a = leaf(&mut g, &s);
573 let b = leaf(&mut g, &s);
574 let prod = op2(&mut g, Op::Mul, a, b, &s);
575 let r = op1(&mut g, Op::Relu, prod, &s);
576 let counts = consumer_counts(&g);
577 assert_eq!(
578 match_region(&g, r, &relu_add_pattern(), &|n| *counts
579 .get(&n)
580 .unwrap_or(&0)),
581 None
582 );
583 }
584
585 #[test]
586 fn match_region_declines_shared_interior() {
587 // sum feeds two consumers → fusing duplicates it → decline (§3a.4).
588 let mut g = Graph::new();
589 let s = fuel_ir::Shape::from_dims(&[2]);
590 let a = leaf(&mut g, &s);
591 let b = leaf(&mut g, &s);
592 let sum = op2(&mut g, Op::Add, a, b, &s);
593 let r1 = op1(&mut g, Op::Relu, sum, &s);
594 let _r2 = op1(&mut g, Op::Neg, sum, &s); // second consumer of `sum`
595 let counts = consumer_counts(&g);
596 assert_eq!(
597 match_region(&g, r1, &relu_add_pattern(), &|n| *counts
598 .get(&n)
599 .unwrap_or(&0)),
600 None
601 );
602 }
603
604 // ---- attr-comparison (F1: match_node compares OpAttrs) --------------------
605
606 /// A single-`Permute` region binding one input, with the given absolute
607 /// perm on `attrs.perm`. An empty `perm` (`&[]`) is the attr-agnostic
608 /// (wildcard) pattern — the shape every existing authored pattern has.
609 fn permute_pattern(perm: &[u8]) -> PatternNode {
610 PatternNode::Op {
611 op: OpTag::Permute,
612 attrs: OpAttrs {
613 perm: perm.to_vec(),
614 ..OpAttrs::default()
615 },
616 operands: vec![PatternNode::Bind { index: 0 }],
617 }
618 }
619
620 #[test]
621 fn match_node_discriminates_on_perm_attr() {
622 // A graph node that permutes with perm = [1, 0].
623 let mut g = Graph::new();
624 let s = fuel_ir::Shape::from_dims(&[2, 3]);
625 let x = leaf(&mut g, &s);
626 let p = op1(
627 &mut g,
628 Op::Permute(vec![1, 0]),
629 x,
630 &fuel_ir::Shape::from_dims(&[3, 2]),
631 );
632 let counts = consumer_counts(&g);
633 let cf = |n: NodeId| *counts.get(&n).unwrap_or(&0);
634
635 // The matching perm binds; the non-matching perm is rejected.
636 assert_eq!(
637 match_region(&g, p, &permute_pattern(&[1, 0]), &cf),
638 Some(vec![x]),
639 "perm=[1,0] pattern must match a [1,0] graph node",
640 );
641 assert_eq!(
642 match_region(&g, p, &permute_pattern(&[0, 2, 1]), &cf),
643 None,
644 "perm=[0,2,1] pattern must NOT match a [1,0] graph node (attr discrimination)",
645 );
646
647 // No-regression guard: an empty-attr (wildcard) pattern — the shape of
648 // every existing authored pattern — still matches regardless of the
649 // graph node's real perm. This is what keeps attr-agnostic patterns
650 // matching after attrs become comparable.
651 assert_eq!(
652 match_region(&g, p, &permute_pattern(&[]), &cf),
653 Some(vec![x]),
654 "empty-perm (wildcard) pattern must still match (no regression)",
655 );
656 }
657
658 // ---- scalar extraction (the `extract:` layer, §5.3) -----------------------
659
660 #[test]
661 fn match_region_extract_reads_live_scalars_in_pattern_pre_order() {
662 // Graph: add_scalar(mul_scalar(x, 2.5), 0.5) — two live scalars.
663 let mut g = Graph::new();
664 let s = fuel_ir::Shape::from_dims(&[2]);
665 let x = leaf(&mut g, &s);
666 let ms = op1(&mut g, Op::MulScalar(2.5), x, &s);
667 let asn = op1(&mut g, Op::AddScalar(0.5), ms, &s);
668 let counts = consumer_counts(&g);
669 let cf = |n: NodeId| *counts.get(&n).unwrap_or(&0);
670
671 // Slot template: both scalar attrs left empty (open slots).
672 let pat = PatternNode::Op {
673 op: OpTag::AddScalar,
674 attrs: OpAttrs::default(),
675 operands: vec![PatternNode::Op {
676 op: OpTag::MulScalar,
677 attrs: OpAttrs::default(),
678 operands: vec![PatternNode::Bind { index: 0 }],
679 }],
680 };
681 let (binds, scalars) = match_region_extract(&g, asn, &pat, &cf).expect("matches");
682 assert_eq!(binds, vec![x]);
683 // Pattern PRE-order: the AddScalar (root) slot before the MulScalar slot.
684 assert_eq!(scalars, vec![0.5, 2.5], "live values in pattern pre-order");
685 }
686
687 #[test]
688 fn baked_scalar_is_a_pattern_constant_not_a_slot() {
689 let mut g = Graph::new();
690 let s = fuel_ir::Shape::from_dims(&[2]);
691 let x = leaf(&mut g, &s);
692 let ms = op1(&mut g, Op::MulScalar(2.5), x, &s);
693 let counts = consumer_counts(&g);
694 let cf = |n: NodeId| *counts.get(&n).unwrap_or(&0);
695
696 let baked = |v: f64| PatternNode::Op {
697 op: OpTag::MulScalar,
698 attrs: OpAttrs {
699 scalars: vec![v],
700 ..OpAttrs::default()
701 },
702 operands: vec![PatternNode::Bind { index: 0 }],
703 };
704 // The equal baked value matches and extracts NOTHING…
705 let (binds, scalars) = match_region_extract(&g, ms, &baked(2.5), &cf).expect("matches");
706 assert_eq!(binds, vec![x]);
707 assert!(
708 scalars.is_empty(),
709 "baked value is a constant of the pattern, not a slot"
710 );
711 // …and a different baked value refuses to match at all (attr guard).
712 assert!(match_region_extract(&g, ms, &baked(3.0), &cf).is_none());
713 }
714
715 // ---- GAP-287 - the `to_canonical_bytes` unset-vs-zero collapse table ----
716 //
717 // `OpAttrs::to_canonical_bytes` reads 17 fields through `unwrap_or`
718 // defaults. Its doc used to justify them as harmless because "an op that
719 // reaches a given arm always has the field set (`op_to_attrs` /
720 // `tag_to_op` guarantee it)". THIS FILE IS THE CITED GUARANTOR AND IT DOES
721 // NOT GUARANTEE THAT: the `_ => {}` arm of `op_to_attrs` leaves `axis`
722 // unset for five axis-bearing ops, and `keepdim` is set by NOTHING in the
723 // repository. Two semantically distinct ops then serialize to identical
724 // bytes, on a format that exists to leave this repo. Note what checks it:
725 // ONE conformance fixture (the rank-2 matmul role-vector golden), covering
726 // the `MatMul` arm - the arm that has no `unwrap_or` at all. There is no
727 // second encoder; Baracuda confirmed the bytes and has no binary arm.
728 //
729 // SCOPE, stated because it is narrower than the 17: only the 12 sites
730 // reachable from the PROJECTION path are testable here. The other five
731 // (`Const`->const_bits, `RuntimeScalar`->slot_index, `ReducedCount`->axis,
732 // `ScanPlaceholder`->{scan_role, scan_index}) are unreachable from any
733 // `Op` -- `op_to_tag` emits none of those four tags -- so they are
734 // reachable only from hand-built regions and are pinned in
735 // fuel-kernel-seam-types instead.
736 //
737 // SOUNDNESS IS A PROPERTY OF (op, field), NOT OF A SITE: the single-axis
738 // arm serves eight tags and is sound for four of them.
739
740 /// Why a given (op, field) pair is or is not sound.
741 #[derive(PartialEq)]
742 enum Soundness {
743 /// `op_to_attrs` sets it, so the `unwrap_or` is unreachable.
744 Projected,
745 /// The default IS emitted and DIFFERS from the truth: distinct ops
746 /// collide. KISS-OPS-6.19 makes `axis` mandatory with no default, so
747 /// `unwrap_or(0)` fabricates a value the schema does not permit.
748 Lossy,
749 /// The field is never set AND the default is CORRECT, because no op
750 /// reaching this arm can have the other value. `keepdim` is the case:
751 /// `Op::SumDim`/`MaxDim`/`MeanDim` all document "the reduced dim is
752 /// REMOVED from the output" and `CumSum` is "same shape as input", so
753 /// every tag on this arm is keepdim=false. A keepdim=TRUE reduce is
754 /// `Op::ReduceSumTo(Shape)` -- a DIFFERENT tag on a DIFFERENT arm that
755 /// carries `target_shape` and has no `unwrap_or` at all.
756 ///
757 /// Kept as a row rather than dropped: "never set" has more than one
758 /// cause and only one of them is a defect. A table that shows only the
759 /// defect cannot be checked against the claim that the others are fine.
760 DeadField,
761 }
762
763 /// One (op, field) pair that [`OpAttrs::to_canonical_bytes`] reads through
764 /// an `unwrap_or` default.
765 struct Collapse {
766 op: Op,
767 field: &'static str,
768 soundness: Soundness,
769 is_set: fn(&OpAttrs) -> bool,
770 }
771
772 fn axis(a: &OpAttrs) -> bool {
773 a.axis.is_some()
774 }
775 fn keepdim(a: &OpAttrs) -> bool {
776 a.keepdim.is_some()
777 }
778 fn dtype(a: &OpAttrs) -> bool {
779 a.cast_dtype.is_some()
780 }
781 fn pad() -> Op {
782 Op::Pad {
783 padding: vec![(1, 1)],
784 mode: crate::PadMode::Constant,
785 value: 0.0,
786 }
787 }
788 fn slice() -> Op {
789 Op::Slice {
790 dim: 1,
791 start: 0,
792 len: 2,
793 }
794 }
795 fn row(
796 op: Op,
797 field: &'static str,
798 soundness: Soundness,
799 is_set: fn(&OpAttrs) -> bool,
800 ) -> Collapse {
801 Collapse {
802 op,
803 field,
804 soundness,
805 is_set,
806 }
807 }
808
809 fn collapse_table() -> Vec<Collapse> {
810 use Soundness::{DeadField, Lossy, Projected};
811 vec![
812 // -- sound: the producing arm sets the field --
813 row(slice(), "axis", Projected, axis),
814 row(slice(), "slice_start", Projected, |a| {
815 a.slice_start.is_some()
816 }),
817 row(slice(), "slice_len", Projected, |a| a.slice_len.is_some()),
818 row(Op::Roll { dim: 0, shift: 1 }, "axis", Projected, axis),
819 row(
820 Op::Roll { dim: 0, shift: 1 },
821 "roll_shift",
822 Projected,
823 |a| a.roll_shift.is_some(),
824 ),
825 row(Op::Cast(DType::F32), "cast_dtype", Projected, dtype),
826 row(pad(), "pad_mode", Projected, |a| a.pad_mode.is_some()),
827 row(pad(), "pad_value", Projected, |a| a.pad_value.is_some()),
828 row(
829 Op::MaskedFill {
830 value: crate::Scalar::F32(0.0),
831 },
832 "cast_dtype",
833 Projected,
834 dtype,
835 ),
836 row(Op::Concat { dim: 0 }, "axis", Projected, axis),
837 row(Op::Flip { dim: 0 }, "axis", Projected, axis),
838 row(Op::Triu { diagonal: 0 }, "axis", Projected, axis),
839 row(Op::Tril { diagonal: 0 }, "axis", Projected, axis),
840 row(Op::SumDim(0), "axis", Projected, axis),
841 row(Op::MaxDim(0), "axis", Projected, axis),
842 row(Op::MeanDim(0), "axis", Projected, axis),
843 // -- lossy: `_ => {}` leaves the field unset; the default is emitted --
844 row(Op::IndexSelect { dim: 1 }, "axis", Lossy, axis),
845 row(Op::Gather { dim: 1 }, "axis", Lossy, axis),
846 row(Op::IndexAdd { dim: 1 }, "axis", Lossy, axis),
847 row(Op::ScatterAdd { dim: 1 }, "axis", Lossy, axis),
848 row(Op::CumSum { dim: 1 }, "axis", Lossy, axis),
849 // `keepdim` is set by NOTHING repo-wide -- and unlike `axis` that is
850 // NOT a defect: every tag on this arm removes the reduced dim, so
851 // `false` is the true value. See `Soundness::DeadField`.
852 row(Op::SumDim(0), "keepdim", DeadField, keepdim),
853 row(Op::MaxDim(0), "keepdim", DeadField, keepdim),
854 row(Op::MeanDim(0), "keepdim", DeadField, keepdim),
855 row(Op::CumSum { dim: 1 }, "keepdim", DeadField, keepdim),
856 ]
857 }
858
859 /// Every row's op must actually project to a tag, or the row is a typo
860 /// asserting nothing. This is the table's own positive control.
861 #[test]
862 fn collapse_table_rows_all_reach_a_serializer_arm() {
863 for c in collapse_table() {
864 assert!(
865 op_to_tag(&c.op).is_some(),
866 "row ({:?}, {}) does not project to any OpTag, so it exercises \
867 no to_canonical_bytes arm and asserts nothing",
868 c.op,
869 c.field
870 );
871 }
872 }
873
874 /// LIVE GATE, the sound half. If a projection arm ever stops setting its
875 /// field, that field silently starts serializing a default and joins the
876 /// lossy set -- this test is what makes that visible.
877 #[test]
878 fn canonical_bytes_projected_fields_are_actually_set() {
879 for c in collapse_table()
880 .iter()
881 .filter(|c| c.soundness == Soundness::Projected)
882 {
883 let a = op_to_attrs(&c.op);
884 assert!(
885 (c.is_set)(&a),
886 "REGRESSION: op_to_attrs no longer sets `{}` for {:?}, so \
887 to_canonical_bytes now emits its unwrap_or default and this \
888 op collides with every other op on that arm (GAP-287)",
889 c.field,
890 c.op
891 );
892 }
893 }
894
895 /// PIN, not an endorsement: these five (op, field) pairs are DEFECTS,
896 /// held so that fixing one is visible. KISS-OPS-6.19 makes `axis` MANDATORY
897 /// with NO default, so `unwrap_or(0)` fabricates a value the schema does
898 /// not permit. Fixing the encoder is a wire-format change and needs the
899 /// external ask first; until then this records which collapses are live.
900 #[test]
901 fn canonical_bytes_lossy_collapses_are_pinned_gap287() {
902 for c in collapse_table()
903 .iter()
904 .filter(|c| c.soundness == Soundness::Lossy)
905 {
906 let a = op_to_attrs(&c.op);
907 assert!(
908 !(c.is_set)(&a),
909 "`{}` is now SET for {:?} -- the GAP-287 collapse was fixed. \
910 That is good news: move this row to `Soundness::Projected` and \
911 update the GAP-287 row and the to_canonical_bytes doc.",
912 c.field,
913 c.op
914 );
915 }
916 }
917
918 /// `keepdim` is unset everywhere AND that is CORRECT -- the separation this
919 /// table exists to make.
920 ///
921 /// I originally tiered `keepdim` as the WORST of these, reasoning "unset
922 /// universally, and it changes output rank". The first half is measured and
923 /// true; the second does not follow, and checking it refutes it. Every tag
924 /// on this arm removes the reduced dim (`CumSum` does not reduce at all),
925 /// so `false` is the true value, not a fabrication. keepdim=TRUE is
926 /// `Op::ReduceSumTo(Shape)` -- a different tag whose arm carries
927 /// `target_shape` and has no `unwrap_or`.
928 ///
929 /// A field being unset is only a defect if the default CAN differ from the
930 /// truth. This asserts the premise that makes the default correct, so that
931 /// adding a keepdim-varying reduce under one of these tags reddens here
932 /// rather than silently making the encoder wrong.
933 #[test]
934 fn keepdim_is_a_dead_field_not_a_lossy_collapse_gap287() {
935 for c in collapse_table()
936 .iter()
937 .filter(|c| c.soundness == Soundness::DeadField)
938 {
939 let a = op_to_attrs(&c.op);
940 assert!(
941 !(c.is_set)(&a),
942 "`{}` is now SET for {:?}. If a keepdim-varying reduce was added under this tag, unwrap_or(false) is no longer correct and this row moves to Soundness::Lossy.",
943 c.field,
944 c.op
945 );
946 // The premise: no op on this arm can have keepdim = true, because
947 // Fuel expresses that as a different TAG.
948 assert!(
949 !matches!(
950 op_to_tag(&c.op),
951 Some(OpTag::ReduceSumTo) | Some(OpTag::ReduceMaxTo)
952 ),
953 "{:?} now projects to a shape-target reduce tag, which is where keepdim=true lives -- the arm assumption has changed",
954 c.op
955 );
956 }
957 }
958}