fuel_graph/runtime_fused.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Runtime fused-op registration — the Tier-2 sidecar
3//! (`docs/specs/runtime-fused-op-registration.md`).
4//!
5//! A runtime-registered (JIT-synthesized or import-time) fused op **is** its
6//! region: its identity is a runtime [`FusedOpId`], its recipe is the §3
7//! [`PatternNode`] region kept here, and its `decompose` is that region
8//! re-emitted as primitives — so the recipe principle (total / never-panic /
9//! primitive→self) holds for free, since [`OpTag`] is the functional-primitive
10//! vocabulary only. No kernel field: the kernel binding lives in fuel-dispatch's
11//! `FusedKernelRegistry` (Tier-1 extensible); this sidecar holds only the
12//! graph-side recipe + the optimizer rules built from it.
13//!
14//! v1 scope: **same-shape elementwise** regions (the synthesizer's increment-1
15//! epilogues). Interior shape inference for broadcast/reduction regions is a
16//! follow-up — a re-emitted node takes its first operand's shape/dtype, exact
17//! for type-preserving same-shape ops and rejected-at-registration otherwise.
18
19use std::collections::HashMap;
20use std::sync::RwLock;
21
22use fuel_kernel_seam_types::{OpAttrs, OpTag, PatternNode, matmul_roles};
23
24use crate::registry::{FusedOpId, FusedOpParams};
25use crate::{Graph, Node, NodeId, Op};
26
27/// A runtime-registered fused op's metadata (the graph-side recipe).
28#[derive(Clone, Debug)]
29pub struct RuntimeFusedOpEntry {
30 /// The allocated runtime id (`>= FusedOpId::RUNTIME_FUSED_BASE`).
31 pub id: FusedOpId,
32 /// A human/telemetry name (e.g. `"jit::relu_add::sm89::<hash>"`).
33 pub name: String,
34 /// The §3 region (the subgraph sink) — the op's primitive recipe.
35 pub region: PatternNode,
36}
37
38/// A registration failure — never a panic (build-time validation).
39#[derive(Clone, Debug, PartialEq)]
40pub enum RuntimeFusedError {
41 /// The region's bind indices don't form a contiguous `[0, n)` (the op's
42 /// external inputs).
43 NonContiguousBinds(Vec<u8>),
44 /// The region carries an op with no primitive re-emission (outside the v1
45 /// re-emit vocabulary) — it could not decompose, so we refuse to register
46 /// it (the totality guard).
47 UnRepresentable(OpTag),
48 /// The region contains a matcher-only node (`Any`/`SeeThrough`) — a
49 /// concrete region must be `Op`/`Bind` only.
50 NonConcreteRegion,
51 /// The runtime id space (`u16` above `RUNTIME_FUSED_BASE`) is exhausted.
52 IdSpaceExhausted,
53 /// A shape-relative attr (D2) that can never resolve at ANY shape — a
54 /// STRUCTURAL authoring error caught at registration: a rel field and its
55 /// concrete sibling both set, a bind reference outside the region's bind
56 /// space, `axis_last` on an axis-less tag, or a `Param` reference (no
57 /// param threading until C-4). Value-dependent declines (a `Negative` or
58 /// symbolic-extent result at some particular shape) do NOT reject
59 /// registration — they surface at emit time as a decompose fixpoint (G2).
60 InvalidRelAttrs { tag: OpTag, error: RelAttrError },
61}
62
63static RUNTIME_FUSED_OPS: RwLock<Vec<RuntimeFusedOpEntry>> = RwLock::new(Vec::new());
64
65/// The recipe-identity index for runtime-registered ops: base-map content
66/// hash ([`crate::opt::base_map_hash`]) → the [`FusedOpId`] that first
67/// registered a region hashing to it. A **sibling** to `RUNTIME_FUSED_OPS`,
68/// not a reuse of [`crate::registry::FusedOpRegistry::by_pattern_hash`] —
69/// that field lives on the STATIC catalog (`FusedOpRegistry`, an
70/// `OnceLock`-frozen struct built at process startup for build-time-known
71/// ids `1..RUNTIME_FUSED_BASE`); runtime ops never populate a
72/// `FusedOpRegistry` instance at all, they live in this module's own
73/// `RUNTIME_FUSED_OPS` global with the disjoint `RUNTIME_FUSED_BASE..` id
74/// space, so `by_pattern_hash` is unreachable from here. This index is the
75/// natural home for runtime-region dedup: same lifetime/global-ness as
76/// `RUNTIME_FUSED_OPS`, cleared in the same breath by
77/// `clear_runtime_fused_for_tests`.
78///
79/// `HashMap::new()` isn't `const`, so this can't be a plain
80/// `static … : RwLock<HashMap<..>> = RwLock::new(HashMap::new())` the way
81/// `RUNTIME_FUSED_OPS` is a plain `RwLock::new(Vec::new())` — `Vec::new()`
82/// is `const`, `HashMap::new()` is not. `OnceLock` lazy-inits it instead
83/// (same pattern as `registry.rs`'s `static REGISTRY: OnceLock<..>` and the
84/// per-function `OnceLock` CPU-device singletons in `opt.rs`/`grad.rs`).
85fn hash_index() -> &'static RwLock<HashMap<u64, FusedOpId>> {
86 static IDX: std::sync::OnceLock<RwLock<HashMap<u64, FusedOpId>>> = std::sync::OnceLock::new();
87 IDX.get_or_init(|| RwLock::new(HashMap::new()))
88}
89
90/// Push `arity` uniform placeholder leaves (`Op::Const`, F32 `[1]`, no
91/// storage) onto `g` and return their ids. Uniform + storage-free is
92/// load-bearing: two independently-built graphs' leaves must hash
93/// IDENTICALLY under [`crate::opt::base_map_hash`] (which folds a const's
94/// shape/dtype and silently no-ops on an unpopulated storage slot) for the
95/// dedup comparison to be meaningful. Mirrors
96/// `fuel_dispatch::jit_ingest::push_placeholder_leaves` — that crate
97/// depends on this one (not the other way around), so the few-line helper
98/// is duplicated here rather than shared.
99fn push_placeholder_leaves(graph: &mut Graph, arity: usize) -> Vec<NodeId> {
100 (0..arity)
101 .map(|_| {
102 graph.push(Node {
103 op: Op::Const,
104 inputs: vec![],
105 shape: fuel_ir::Shape::from_dims(&[1]),
106 dtype: fuel_ir::DType::F32,
107 })
108 })
109 .collect()
110}
111
112/// `region`'s structural-identity hash: emit it onto placeholder leaves
113/// (via [`emit_region`]), lower to the primitive base map
114/// ([`crate::opt::lower_to_base_map`]), hash the result
115/// ([`crate::opt::base_map_hash`]). `None` on any structural failure (a
116/// poisoned lock, a rel-attr resolution decline at the placeholder shapes,
117/// or an empty lowering result) — the caller
118/// (`register_runtime_fused`) treats `None` as "hash unavailable" and skips
119/// dedup for this registration, never blocking it.
120///
121/// Every caller in this module runs this AFTER `validate_representable`
122/// already passed for the same region, so `emit_region`'s own panic risks
123/// (an unrepresentable `OpTag`, a `Bind` index out of range) are already
124/// ruled out here — `register_runtime_fused` still wraps the call in
125/// `catch_unwind` as the never-panic contract's last-resort guard for
126/// anything this doesn't anticipate.
127fn region_base_map_hash(region: &PatternNode) -> Option<u64> {
128 let n_binds = region.bind_indices().len();
129 let scalars = vec![0.0; count_scalar_slots(region)];
130 let graph: crate::SharedGraph = std::sync::Arc::new(RwLock::new(Graph::new()));
131 let sink = {
132 let mut g = graph.write().ok()?;
133 let inputs = push_placeholder_leaves(&mut g, n_binds);
134 // Fallible entry: a rel-attr region that declines at the rank-1 `[1]`
135 // placeholder shapes yields `None` — "hash unavailable", dedup skipped
136 // for this registration (allocate-fresh), never a panic.
137 try_emit_region(&mut g, region, &inputs, &scalars).ok()?
138 };
139 let roots = crate::opt::lower_to_base_map(&graph, &[sink]);
140 let root = *roots.first()?;
141 let g = graph.read().ok()?;
142 Some(crate::opt::base_map_hash(&g, root))
143}
144
145/// Register a runtime fused op for `region`, returning its runtime
146/// [`FusedOpId`]. Validates **before** allocating that the region's bind
147/// indices form the op's input list and that every op re-emits to
148/// primitives (totality) — a non-decomposable region is rejected, never
149/// registered.
150///
151/// **Dedup (recipe identity):** a region that is structurally identical
152/// (same [`crate::opt::base_map_hash`] over its primitive lowering) to an
153/// already-registered region resolves to the EXISTING [`FusedOpId`] instead
154/// of minting a duplicate — two calls with the same shape but different
155/// `name`s return the same id, and only the first call's `name`/region is
156/// kept in `RUNTIME_FUSED_OPS`. Never-panic: hashing runs inside
157/// `catch_unwind`; any failure (a poisoned lock, an unanticipated panic) is
158/// treated as "hash unavailable" and simply skips the dedup check —
159/// registration always proceeds to today's allocate-fresh path either way.
160pub fn register_runtime_fused(
161 name: impl Into<String>,
162 region: PatternNode,
163) -> Result<FusedOpId, RuntimeFusedError> {
164 let name = name.into();
165 validate_recipe(®ion)?;
166
167 let hash = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
168 region_base_map_hash(®ion)
169 }))
170 .unwrap_or(None);
171 if hash.is_none() {
172 eprintln!(
173 "register_runtime_fused: base-map hash unavailable for {name:?}; \
174 registering without dedup (allocate-fresh fallback)"
175 );
176 }
177
178 // Hold the hash index's write lock across the whole check-then-insert
179 // sequence below (not read-then-separately-write) so two concurrent
180 // registrations of the same NEW region can't both miss the lookup and
181 // each mint their own id: the second caller blocks on this lock and,
182 // once it acquires it, observes the first caller's insert.
183 let mut idx = hash_index().write().unwrap();
184 if let Some(h) = hash
185 && let Some(&existing) = idx.get(&h)
186 {
187 return Ok(existing);
188 }
189
190 // The Vec length under the write lock is the allocator: id = BASE + index,
191 // so the index is always `id - BASE` with no allocate/push race.
192 let mut w = RUNTIME_FUSED_OPS.write().unwrap();
193 let raw = FusedOpId::RUNTIME_FUSED_BASE as usize + w.len();
194 if raw > u16::MAX as usize {
195 return Err(RuntimeFusedError::IdSpaceExhausted);
196 }
197 let id = FusedOpId(raw as u16);
198 w.push(RuntimeFusedOpEntry { id, name, region });
199 drop(w);
200
201 if let Some(h) = hash {
202 idx.insert(h, id);
203 }
204
205 Ok(id)
206}
207
208/// The region (recipe) for a runtime fused op, or `None` if `id` is not a
209/// registered runtime op.
210pub fn runtime_region(id: FusedOpId) -> Option<PatternNode> {
211 if !id.is_runtime() {
212 return None;
213 }
214 let idx = (id.0 - FusedOpId::RUNTIME_FUSED_BASE) as usize;
215 RUNTIME_FUSED_OPS
216 .read()
217 .unwrap()
218 .get(idx)
219 .map(|e| e.region.clone())
220}
221
222/// A runtime op's name (telemetry / `op_short_name` routing).
223pub fn runtime_name(id: FusedOpId) -> Option<String> {
224 if !id.is_runtime() {
225 return None;
226 }
227 let idx = (id.0 - FusedOpId::RUNTIME_FUSED_BASE) as usize;
228 RUNTIME_FUSED_OPS
229 .read()
230 .unwrap()
231 .get(idx)
232 .map(|e| e.name.clone())
233}
234
235/// All registered runtime ops — the optimizer iterates this to build a fusion
236/// rule + a lowering rule per runtime op (`RuleRegistry::default_rules` /
237/// `lowering_only`).
238pub fn runtime_entries() -> Vec<RuntimeFusedOpEntry> {
239 RUNTIME_FUSED_OPS.read().unwrap().clone()
240}
241
242/// **TEST-ONLY.** Clear the metadata sidecar AND the recipe-identity
243/// `hash_index` in the same breath. Because the Vec length *is* the id
244/// allocator (`id = BASE + index`), clearing restarts allocation — any
245/// sidecar keyed by prior runtime ids MUST be cleared alongside it or a
246/// reused id resolves stale data. This was already true for
247/// `fuel_dispatch::runtime_fused_kernels::clear_runtime_fused_for_tests`'s
248/// kernel sidecar (call that one, not this, from dispatch-level tests) and
249/// is now ALSO true for `hash_index`: leaving a stale `hash → old_id`
250/// entry after a clear would let a later registration's dedup lookup
251/// return an id that no longer names the region it was hashed from (the
252/// slot at that index now holds whatever the NEXT registration after the
253/// clear pushed there). Adopting tests share one process, so callers must
254/// also serialize with any other adopting test (dd-shapes coordination,
255/// 2026-07-08: the hook alone races). `#[doc(hidden)] pub` rather than
256/// `#[cfg(test)]` because adopting tests live in downstream crates, which
257/// compile this crate without `cfg(test)`.
258#[doc(hidden)]
259pub fn clear_runtime_fused_for_tests() {
260 RUNTIME_FUSED_OPS.write().unwrap().clear();
261 hash_index().write().unwrap().clear();
262}
263
264// ---- the region → primitive re-emit (the runtime op's `decompose`) ---------
265
266/// Project a region [`OpTag`] (+ its [`OpAttrs`]) back to a primitive [`Op`].
267/// The inverse of `jit::op_to_tag`, over the **full first-order re-emit
268/// vocabulary** (Convergence Increment A): every non-basis-gap, non-`Fused`
269/// op — elementwise, comparison, `Where`, `Cast`, shape/layout
270/// (Transpose/Permute/Reshape/BroadcastTo/(Un)squeeze/Slice/Concat/Flip/Roll/
271/// Pad/Triu/Tril), reductions (SumDim/MaxDim/MeanDim/ReduceSumTo/ReduceMaxTo/
272/// CumSum/SumAll/MaxAll/MinAll/MeanAll), `MatMul`, `Iota`, and indexing (IndexSelect/
273/// Gather/IndexAdd/ScatterAdd) — plus the **`Op::Scan` multi-output structural
274/// terminal** (Increment C, B1): `Scan` (params ride the `scan_*` carriers, body
275/// sub-graph rides the operands), its `ScanPlaceholder` body holes, and the
276/// `View` slot projection (`view_slot`) that reads the scan's `output_views`
277/// bundle at emit time — and the scalar-carrying `MaskedFill` (Increment C
278/// carriers, A2 — fill value on `attrs.scalars[0]`, dtype on `cast_dtype` when
279/// present else a provisional F32 that `emit` re-resolves to operand[0]'s dtype)
280/// and `PowI` (A3 — the i32 exponent on `scalars[0]`). Structural params are
281/// decoded from the (extended) [`OpAttrs`]. Returns `None` (an honest miss,
282/// rejected at registration) for ops with no first-order re-emission: `Clamp`
283/// (no two-scalar carrier yet), fused/basis-gap tags, and any tag whose required
284/// attrs are unset (e.g. `Iota` with no `target_shape`, `Scan` with no
285/// `scan_bound`, `View` with no `view_slot`, `PowI` with no exponent, or
286/// `MaskedFill` with no fill value / a dtype `Scalar` cannot represent).
287fn tag_to_op(tag: OpTag, attrs: &OpAttrs) -> Option<Op> {
288 use OpTag as T;
289 use fuel_ir::DType;
290 use std::str::FromStr;
291 Some(match tag {
292 T::Add => Op::Add,
293 T::Sub => Op::Sub,
294 T::Mul => Op::Mul,
295 T::Div => Op::Div,
296 T::Maximum => Op::Maximum,
297 T::Minimum => Op::Minimum,
298 T::Pow => Op::Pow,
299 T::Rem => Op::Rem,
300 T::Neg => Op::Neg,
301 T::Abs => Op::Abs,
302 T::Sqr => Op::Sqr,
303 T::Sqrt => Op::Sqrt,
304 T::Rsqrt => Op::Rsqrt,
305 T::Recip => Op::Recip,
306 T::Exp => Op::Exp,
307 T::Log => Op::Log,
308 T::Sin => Op::Sin,
309 T::Cos => Op::Cos,
310 T::Tanh => Op::Tanh,
311 T::Sigmoid => Op::Sigmoid,
312 T::Silu => Op::Silu,
313 T::GeluTanh => Op::GeluTanh,
314 T::Gelu => Op::Gelu,
315 T::Relu => Op::Relu,
316 T::Erf => Op::Erf,
317 T::Step => Op::Step,
318 T::Floor => Op::Floor,
319 T::Ceil => Op::Ceil,
320 T::Round => Op::Round,
321 T::Sign => Op::Sign,
322 // Scalar-param ops: the value rides `attrs.scalars` (the slot snapshot;
323 // live-value substitution via the `extract:` path is a follow-up).
324 T::AddScalar => Op::AddScalar(*attrs.scalars.first()?),
325 T::MulScalar => Op::MulScalar(*attrs.scalars.first()?),
326 // PowI (Increment C carriers, A3): the i32 exponent rides
327 // `attrs.scalars[0]` as an f64 — an EXACT round-trip for every i32
328 // (`|n| < 2^53`), reconstructed via `as i32`. This is the same carrier
329 // the §6.19 wire already commits to (the `to_canonical_bytes` PowI arm
330 // serializes `scalars`) and mirrors `MaskedFill`'s baked-scalar posture:
331 // the exponent is a BAKED pattern constant, NOT an open cursor slot
332 // (`scalar_slot_arity(PowI) == 0`), so the recipe author supplies it and
333 // the emitter never draws it from the params projection. An unset value
334 // is an honest miss (`None`), never a defaulted `PowI(0)`.
335 T::PowI => Op::PowI(*attrs.scalars.first()? as i32),
336
337 // --- Convergence Increment A: the full first-order set ---
338 // Comparison (dtype→U8 handled by primitive_shape, not here).
339 T::Equal => Op::Equal,
340 T::Ne => Op::Ne,
341 T::Lt => Op::Lt,
342 T::Le => Op::Le,
343 T::Gt => Op::Gt,
344 T::Ge => Op::Ge,
345 // Ternary select.
346 T::Where => Op::Where,
347 // Dtype-changing: target dtype rides `cast_dtype` (the stable name).
348 T::Cast => Op::Cast(DType::from_str(attrs.cast_dtype.as_deref()?).ok()?),
349 // MatMul: the LOCKED role-vector contraction cell (§5/D5). Empty roles
350 // = the rank-polymorphic recipe form → implicit-accept (unchanged from
351 // today; recipes keep matmul implicit). Explicit roles must match the
352 // canonical cell EXACTLY — same-rank ≥ 2, leading Batch, lhs=[..,FreeM,
353 // ContractedK], rhs=[..,ContractedK,FreeN] — checked by role POSITION,
354 // not extent (so GQA-divisible batch stays all-Batch). Any other config
355 // (transposed / multi-ContractedK / FreeN-before-K / rank mismatch) is a
356 // SURFACED honest miss (`None`, rejected at registration), never a crash.
357 T::MatMul => {
358 if attrs.lhs_roles.is_empty() && attrs.rhs_roles.is_empty() {
359 Op::MatMul
360 } else {
361 let (canon_lhs, canon_rhs) =
362 matmul_roles(attrs.lhs_roles.len(), attrs.rhs_roles.len());
363 if attrs.lhs_roles.len() == attrs.rhs_roles.len()
364 && attrs.lhs_roles.len() >= 2
365 && attrs.lhs_roles == canon_lhs
366 && attrs.rhs_roles == canon_rhs
367 {
368 Op::MatMul
369 } else {
370 return None;
371 }
372 }
373 }
374 T::LogSoftmaxLastDim => Op::LogSoftmaxLastDim,
375 // Shape / layout.
376 T::Transpose => Op::Transpose,
377 T::Permute => Op::Permute(attrs.perm.iter().map(|&x| x as usize).collect()),
378 T::Reshape => Op::Reshape(shape_from_attr(attrs)?),
379 T::BroadcastTo => Op::BroadcastTo(shape_from_attr(attrs)?),
380 T::ReduceSumTo => Op::ReduceSumTo(shape_from_attr(attrs)?),
381 T::ReduceMaxTo => Op::ReduceMaxTo(shape_from_attr(attrs)?),
382 T::Unsqueeze => Op::Unsqueeze {
383 dim: *attrs.dims.first()? as usize,
384 },
385 T::Squeeze => Op::Squeeze {
386 dim: *attrs.dims.first()? as usize,
387 },
388 T::Slice => Op::Slice {
389 dim: attrs.axis? as usize,
390 start: attrs.slice_start? as usize,
391 len: attrs.slice_len? as usize,
392 },
393 T::Concat => Op::Concat {
394 dim: attrs.axis? as usize,
395 },
396 T::Flip => Op::Flip {
397 dim: attrs.axis? as usize,
398 },
399 T::Roll => Op::Roll {
400 dim: attrs.axis? as usize,
401 shift: attrs.roll_shift?,
402 },
403 T::Pad => Op::Pad {
404 padding: attrs
405 .pad_amounts
406 .iter()
407 .map(|&(b, e)| (b as usize, e as usize))
408 .collect(),
409 mode: match attrs.pad_mode? {
410 0 => crate::PadMode::Constant,
411 1 => crate::PadMode::Reflect,
412 2 => crate::PadMode::Replicate,
413 _ => return None,
414 },
415 value: attrs.pad_value.unwrap_or(0.0),
416 },
417 T::Triu => Op::Triu {
418 diagonal: attrs.axis?,
419 },
420 T::Tril => Op::Tril {
421 diagonal: attrs.axis?,
422 },
423 // Reductions (dim rides `axis`; keepdim reductions ride `target_shape`).
424 T::SumDim => Op::SumDim(attrs.axis? as usize),
425 T::MaxDim => Op::MaxDim(attrs.axis? as usize),
426 T::MeanDim => Op::MeanDim(attrs.axis? as usize),
427 T::SumAll => Op::SumAll,
428 T::MaxAll => Op::MaxAll,
429 T::MinAll => Op::MinAll,
430 T::MeanAll => Op::MeanAll,
431 T::CumSum => Op::CumSum {
432 dim: attrs.axis? as usize,
433 },
434 // Value source leaf (len rides `target_shape` as a 1-element shape).
435 T::Iota => Op::Iota {
436 len: *attrs.target_shape.first()? as usize,
437 },
438 // Indexing (dim rides `axis`).
439 T::IndexSelect => Op::IndexSelect {
440 dim: attrs.axis? as usize,
441 },
442 T::Gather => Op::Gather {
443 dim: attrs.axis? as usize,
444 },
445 T::IndexAdd => Op::IndexAdd {
446 dim: attrs.axis? as usize,
447 },
448 T::ScatterAdd => Op::ScatterAdd {
449 dim: attrs.axis? as usize,
450 },
451
452 // --- Op::Scan structural re-emit (Increment C, B1) ---
453 // A scan's body sub-graph rides the node's operands (the trailing
454 // inputs); this arm reconstructs only the `Op::Scan` params. The body
455 // holes re-emit through the `ScanPlaceholder` arm below. `Op::Scan`
456 // stays a base-map terminal (no native kernel, no `LoweringRule`) — the
457 // recipe just re-emits it so a decompose can round-trip through data.
458 T::Scan => Op::Scan {
459 // GAP-303: `?` (an honest miss) rather than `unwrap_or(1)`. A recipe
460 // that does not state its carry count is not a single-carry scan --
461 // it is a recipe this vocabulary cannot reconstruct, and defaulting
462 // would silently rebuild a DIFFERENT graph.
463 n_carries: attrs.scan_n_carries? as usize,
464 n_xs: attrs.scan_n_xs? as usize,
465 bound: attrs.scan_bound? as usize,
466 emit: match attrs.scan_emit? {
467 0 => crate::ScanEmit::All,
468 1 => crate::ScanEmit::Final,
469 _ => return None,
470 },
471 // Phase-1 `ScanPredicate` is a marker (the predicate DAG rides the
472 // trailing `pred_exit` operand); `Some(true)` re-emits the marker.
473 early_exit: if attrs.scan_early_exit.unwrap_or(false) {
474 Some(crate::ScanPredicate)
475 } else {
476 None
477 },
478 },
479 // A body hole of `Op::Scan` (a childless leaf): `role` + `index` ride
480 // the dedicated `scan_role`/`scan_index` carriers.
481 T::ScanPlaceholder => Op::ScanPlaceholder {
482 role: match attrs.scan_role? {
483 fuel_kernel_seam_types::SCAN_ROLE_CARRY => crate::ScanRole::Carry,
484 fuel_kernel_seam_types::SCAN_ROLE_ELEM => crate::ScanRole::Elem,
485 _ => return None,
486 },
487 index: attrs.scan_index? as usize,
488 },
489 // Multi-output slot projection. The slot's shape/dtype are NOT decoded
490 // here (they come from the producer's `output_views` bundle at emit
491 // time, mirroring `Graph::view`); this arm reconstructs only the op.
492 T::View => Op::View {
493 slot: attrs.view_slot?,
494 },
495
496 // Nested fused op carried AS-IS (Increment C, C-T2, mechanism 2a). The
497 // `fused_op` selector names the registry entry (mirroring `cast_dtype`'s
498 // name-string precedent); resolve it to a `FusedOpId` and reconstruct
499 // the param-less `Op::Fused(fid, params)` via the small `fid -> params`
500 // map. An unset selector, an unknown name, or a param-carrying id (not in
501 // the map — its per-instance params can't be recovered from a name alone)
502 // is a surfaced honest miss (`None`), never a crash. The nested node's
503 // shape/dtype are computed at emit time from the entry's shape/dtype
504 // rules (`primitive_shape` honest-misses `Fused`); this arm reconstructs
505 // only the op. Fuel-INTERNAL — a `Fused` node never reaches the §6.19 wire.
506 T::Fused => {
507 let fid =
508 crate::registry::default_registry().id_for_name(attrs.fused_op.as_deref()?)?;
509 Op::Fused(fid, fused_params_for(fid)?)
510 }
511
512 // MaskedFill (Increment C carriers, A2): the fill VALUE rides
513 // `attrs.scalars[0]` (the `op_to_attrs` projection); its dtype rides
514 // `cast_dtype` when present (a concrete round-trip), else a provisional
515 // F32 that `emit` re-resolves to operand[0]'s dtype (the
516 // dtype-polymorphic recipe path — the byte executor derives `fill_bytes`
517 // at the filled tensor's width, so the Scalar dtype must match). An
518 // unset value or a dtype `Scalar` cannot represent is an honest miss.
519 T::MaskedFill => {
520 let v = *attrs.scalars.first()?;
521 let dtype = match attrs.cast_dtype.as_deref() {
522 Some(name) => DType::from_str(name).ok()?,
523 None => DType::F32,
524 };
525 Op::MaskedFill {
526 value: masked_fill_scalar(v, dtype)?,
527 }
528 }
529
530 // Honest misses (rejected at registration): Clamp (no two-scalar
531 // carrier yet), and any tag whose required attrs are unset or that is
532 // added to OpTag later.
533 _ => return None,
534 })
535}
536
537/// The `fid -> params` map for the nested fused ops a recipe carries as-is
538/// (Increment C, C-T2, mechanism 2a): the two **param-less** attention-trio
539/// nested softmaxes — `SoftmaxLastDim` (forward, used by `flash_attn` /
540/// `paged_attn`) and `SoftmaxLastDimBackward` (used by `flash_attn_backward`).
541/// A param-carrying fused op is NOT recoverable from the name-only `fused_op`
542/// selector — its per-instance params (softmax_scale, block_size, …) have no
543/// carrier here — so any other id is a surfaced honest miss (`None`), never a
544/// crash. Extend this map only when a new *param-less* nested fused op needs to
545/// ride a recipe.
546fn fused_params_for(fid: FusedOpId) -> Option<FusedOpParams> {
547 use crate::registry::FusedOps;
548 Some(match fid {
549 FusedOps::SOFTMAX_LAST_DIM => FusedOpParams::SoftmaxLastDim,
550 FusedOps::SOFTMAX_LAST_DIM_BACKWARD => FusedOpParams::SoftmaxLastDimBackward,
551 _ => return None,
552 })
553}
554
555/// Reconstruct a `MaskedFill` fill [`fuel_ir::Scalar`] from its `f64` value +
556/// the FILLED tensor's dtype (the A2 carrier). `None` for a dtype a `Scalar`
557/// cannot represent (the sub-byte dummy quant formats) — an honest miss, never
558/// a panic (so [`fuel_ir::Scalar::from_f64`]'s dummy-dtype panic is unreachable
559/// through this guard). The value rides `attrs.scalars[0]`; the dtype is the
560/// filled tensor's (operand[0]) dtype, resolved at emit time.
561fn masked_fill_scalar(value: f64, dtype: fuel_ir::DType) -> Option<fuel_ir::Scalar> {
562 // Honest miss for any dtype without a scalar rep; the Result collapses
563 // the old hand-written dummy-dtype guard.
564 fuel_ir::Scalar::from_f64(value, dtype).ok()
565}
566
567/// Decode a target [`fuel_ir::Shape`] from `attrs.target_shape` (the shared
568/// LOGICAL-shape carrier for Reshape/BroadcastTo/ReduceSumTo/ReduceMaxTo).
569/// `None` for an unset (empty) target — an honest miss — UNLESS the
570/// [`OpAttrs::rank0_target`] marker is set, which denotes the intentional
571/// rank-0 (`[]`, scalar) shape (C-T1: a reduce-to-scalar loss tail such as
572/// FSCE's `ReduceSumTo([])`). The marker is what distinguishes an authored
573/// rank-0 target from the empty-`target_shape` wildcard/unset state.
574fn shape_from_attr(attrs: &OpAttrs) -> Option<fuel_ir::Shape> {
575 if attrs.target_shape.is_empty() {
576 if attrs.rank0_target {
577 return Some(fuel_ir::Shape::from_dims(&[]));
578 }
579 return None;
580 }
581 let dims: Vec<usize> = attrs.target_shape.iter().map(|&d| d as usize).collect();
582 Some(fuel_ir::Shape::from_dims(&dims))
583}
584
585/// How many scalar values `tag` consumes from `attrs.scalars` when re-emitted.
586/// The slot machinery (extraction, validation dummy-fill, decompose fill) is
587/// keyed on this; extend alongside `tag_to_op` when a new scalar-param op joins
588/// the v1 vocabulary.
589fn scalar_slot_arity(tag: OpTag) -> usize {
590 matches!(tag, OpTag::AddScalar | OpTag::MulScalar) as usize
591}
592
593// ---- shape-relative attr resolution (Increment C slice 1, T2/D2) -----------
594
595/// A shape-relative attr resolution failure — a typed decline, never a panic.
596/// The emit-integration caller (T3) surfaces any of these as a decompose
597/// fixpoint (`return id`, G2); the registration-validation caller rejects the
598/// region. The `field` names the CONCRETE sibling the rel field resolves into
599/// (`"target_shape"`, `"slice_start"`, `"slice_len"`, `"axis"`, `"dims"`).
600#[derive(Clone, Debug, PartialEq, Eq)]
601pub enum RelAttrError {
602 /// A rel field and its concrete sibling are BOTH set — ambiguous authoring
603 /// (rel XOR abs per field), refused rather than given a silent precedence.
604 RelAbsConflict { field: &'static str },
605 /// The underlying shape-expression evaluation declined (bind out of range,
606 /// axis out of range, divide-by-zero, `Param` with no param threading, …).
607 Expr(fuel_kernel_seam_types::shape_expr::ShapeExprError),
608 /// The expression evaluated over a SYMBOLIC bind extent → a surfaced gap
609 /// (§6.20-0004): a rel attr cannot resolve concrete at emit time.
610 SymbolicGap { field: &'static str },
611 /// The expression produced a negative value where a non-negative
612 /// extent/offset is required.
613 Negative { field: &'static str, value: i64 },
614 /// `axis_last` on a tag with no axis carrier (e.g. `Add`) — meaningless,
615 /// refused (build-time validation, never silently ignored).
616 AxisLastUnsupported { tag: OpTag },
617 /// `axis_last` with no child operand — no rank to resolve LAST against.
618 NoChildOperand,
619 /// The region's Bind-space broadcast **frame** is assembled by per-axis max
620 /// across ≥2 binds (`a[N,1] ⊗ b[1,M] → [N,M]`), so NO single operand carries
621 /// it — and `SameAs { operand }`, the §6.20 EXPRESSION kind's only
622 /// whole-shape constructor, therefore cannot express it. Accepting the
623 /// `SameAs` would SILENTLY resolve a `BroadcastTo` target to a PARTIAL frame
624 /// (`[N,1]` or `[1,M]`) and emit the wrong graph, so it is refused instead
625 /// (Baracuda's §6.20 finding). A **Dims-class gap**: `missing_ctor` is the
626 /// reserved wire tag of the constructor that WOULD express it
627 /// ([`fuel_kernel_seam_types::shape_expr::TAG_DIMS`] = `0x0B`, a §6.20-0002
628 /// extension-registry entrant — proposal filed KISS #80). `frame` is the
629 /// computed per-axis-max frame, for telemetry.
630 FrameNotExpressible {
631 field: &'static str,
632 frame: Vec<i64>,
633 missing_ctor: u8,
634 },
635}
636
637/// Whether `attrs` carries any shape-RELATIVE field (D2) — the emit fast-path
638/// guard: rel-free attrs skip resolution entirely (zero behavior change for
639/// existing concrete regions).
640fn has_rel_attrs(attrs: &OpAttrs) -> bool {
641 attrs.target_shape_rel.is_some()
642 || attrs.slice_start_rel.is_some()
643 || attrs.slice_len_rel.is_some()
644 || attrs.axis_last
645 || attrs.scalar_rel.is_some()
646}
647
648/// The ONE rel-XOR-abs mutual-exclusion oracle (shared by
649/// [`resolve_rel_attrs`] and the registration rel-probe — no second copy to
650/// drift). Returns the first conflicted field name in canonical field order
651/// (`target_shape`, `slice_start`, `slice_len`, then the `axis_last` carrier —
652/// `dims` for Squeeze/Unsqueeze, `axis` otherwise), or `None`. Note the
653/// `axis_last` arm reports a carrier conflict even for a tag the resolver
654/// would refuse as [`RelAttrError::AxisLastUnsupported`] — both are typed
655/// authoring declines, and both-set is checked first.
656fn rel_abs_conflict_field(tag: OpTag, attrs: &OpAttrs) -> Option<&'static str> {
657 if attrs.target_shape_rel.is_some() && !attrs.target_shape.is_empty() {
658 return Some("target_shape");
659 }
660 if attrs.slice_start_rel.is_some() && attrs.slice_start.is_some() {
661 return Some("slice_start");
662 }
663 if attrs.slice_len_rel.is_some() && attrs.slice_len.is_some() {
664 return Some("slice_len");
665 }
666 if attrs.scalar_rel.is_some() && !attrs.scalars.is_empty() {
667 return Some("scalars");
668 }
669 if attrs.axis_last {
670 match tag {
671 OpTag::Unsqueeze | OpTag::Squeeze => {
672 if !attrs.dims.is_empty() {
673 return Some("dims");
674 }
675 }
676 _ => {
677 if attrs.axis.is_some() {
678 return Some("axis");
679 }
680 }
681 }
682 }
683 None
684}
685
686/// The region's Bind-space broadcast **frame**: the NumPy right-aligned
687/// per-axis max over EVERY bind shape (`[N,1] ⊗ [1,M] → [N,M]`) — the shape an
688/// elementwise consumer of all the binds produces. `None` when the binds carry
689/// no joint elementwise frame at all: no binds, a SYMBOLIC extent (the frame is
690/// itself a gap — the `SameAs` arm already declines those as
691/// [`RelAttrError::SymbolicGap`]), or mutually broadcast-INcompatible binds (a
692/// matmul/gather region, where per-axis max is meaningless). Pure and total —
693/// unlike the graph-builder `compute_broadcast_shape`, incompatibility is
694/// `None`, never a panic.
695fn bind_broadcast_frame(bind_shapes: &[Vec<i64>]) -> Option<Vec<i64>> {
696 use fuel_kernel_seam_types::shape_expr::SYMBOLIC;
697 let rank = bind_shapes.iter().map(Vec::len).max()?;
698 let mut frame = vec![1i64; rank];
699 for s in bind_shapes {
700 if s.iter().any(|&e| e == SYMBOLIC || e < 0) {
701 return None;
702 }
703 let pad = rank - s.len(); // right-aligned: pad the shorter with leading 1s
704 for (i, &e) in s.iter().enumerate() {
705 let f = &mut frame[pad + i];
706 if *f == e || e == 1 {
707 continue;
708 }
709 if *f == 1 {
710 *f = e;
711 continue;
712 }
713 return None; // incompatible at this axis ⇒ no joint frame
714 }
715 }
716 Some(frame)
717}
718
719/// The SameAs **degradation guard** (I1, Baracuda's §6.20 finding). Even an
720/// ELEMENTWISE output shape is not always expressible as `SameAs(operand)`:
721/// when the region's broadcast frame is assembled by per-axis max across TWO
722/// binds (`a[N,1] ⊗ b[1,M] → [N,M]`) no single operand carries the full frame,
723/// so every `SameAs` spelling resolves to a PARTIAL frame. `BroadcastTo` is the
724/// recipe's frame carrier (its target IS the elementwise output shape), so a
725/// `SameAs` target is accepted only when SOME bind does carry the whole frame;
726/// otherwise the frame is surfaced as a typed Dims-class gap rather than
727/// silently resolved to one operand's partial shape.
728///
729/// Deliberately narrow — it fires ONLY when a joint frame exists and NO bind
730/// equals it. Sub-frame broadcasts (bind1 `[T,D]` inside a `[B,T,D]` region)
731/// and frame-less regions (matmul binds) are untouched, and all 5 slice-1
732/// migrated recipes are safe by construction (their `SameAs` target is bind 0,
733/// which carries the frame).
734fn same_as_frame_guard(
735 tag: OpTag,
736 se: &fuel_kernel_seam_types::shape_expr::ShapeExpr,
737 bind_shapes: &[Vec<i64>],
738) -> Result<(), RelAttrError> {
739 use fuel_kernel_seam_types::shape_expr::{ShapeExpr, TAG_DIMS};
740 // Exhaustive on purpose: the Dims/WithDim extension entrants (KISS #80)
741 // express the max-frame DIRECTLY (a whole-shape ctor), so they must NOT be
742 // routed through this partial-`SameAs`-frame guard — they return early.
743 match se {
744 ShapeExpr::SameAs { .. } => {}
745 ShapeExpr::WithDim { .. } | ShapeExpr::Dims(_) => return Ok(()),
746 }
747 if tag != OpTag::BroadcastTo {
748 return Ok(());
749 }
750 let Some(frame) = bind_broadcast_frame(bind_shapes) else {
751 return Ok(()); // no joint elementwise frame at play
752 };
753 if bind_shapes.contains(&frame) {
754 return Ok(()); // some operand carries the whole frame ⇒ expressible
755 }
756 Err(RelAttrError::FrameNotExpressible {
757 field: "target_shape",
758 frame,
759 missing_ctor: TAG_DIMS,
760 })
761}
762
763/// Resolve `attrs`' shape-RELATIVE fields (`target_shape_rel`,
764/// `slice_start_rel`/`slice_len_rel`, `axis_last` — D2) into their concrete
765/// siblings, returning a fully-concrete [`OpAttrs`] ready for the unchanged
766/// `tag_to_op` → `primitive_shape` path. Pure: no graph access.
767///
768/// * `bind_shapes` — the region's **Bind-space** shapes, `bind_shapes[i]` =
769/// `Bind { index: i }`'s shape. This is what `ShapeExpr::SameAs { operand }`
770/// and `Dim::Extent { operand, .. }` index (the recipe-interior reference
771/// convention, same as the merged KISS shape-oracle RFC's contract roles).
772/// * `child_shapes` — THIS op's direct operand shapes (the already-emitted
773/// children), which `axis_last` resolves its rank against — a region
774/// interior node's shape generally matches NO bind.
775///
776/// Evaluation reuses `shape_expr::eval_dim`/`eval_shape`/`resolve_axis` — the
777/// single §6.20 evaluator, no second one. `Dim::Param` declines with a typed
778/// [`fuel_kernel_seam_types::shape_expr::ShapeExprError::ParamOutOfRange`] until param threading lands (C-4);
779/// symbolic bind extents decline as [`RelAttrError::SymbolicGap`]. Rel fields
780/// are CLEARED in the output (rel+abs both set in the RESULT would trip the
781/// mutual-exclusion check on a second resolve).
782pub fn resolve_rel_attrs(
783 tag: OpTag,
784 attrs: &OpAttrs,
785 bind_shapes: &[Vec<i64>],
786 child_shapes: &[Vec<i64>],
787) -> Result<OpAttrs, RelAttrError> {
788 use fuel_kernel_seam_types::shape_expr::{self, Dim, DimValue, LAST, ShapeValue, resolve_axis};
789 // Mutual exclusion FIRST, for every field, before any evaluation — so a
790 // value-dependent decline in an earlier field can't mask a rel+abs
791 // authoring conflict in a later one (the registration probe relies on
792 // this completeness).
793 if let Some(field) = rel_abs_conflict_field(tag, attrs) {
794 return Err(RelAttrError::RelAbsConflict { field });
795 }
796 let mut out = attrs.clone();
797
798 // target_shape_rel → target_shape (SameAs over the Bind space).
799 if let Some(se) = &attrs.target_shape_rel {
800 match shape_expr::eval_shape(se, bind_shapes, &[]).map_err(RelAttrError::Expr)? {
801 ShapeValue::Concrete(s) => {
802 if let Some(&bad) = s.iter().find(|&&e| e < 0) {
803 return Err(RelAttrError::Negative {
804 field: "target_shape",
805 value: bad,
806 });
807 }
808 // I1: refuse a `SameAs` target whose region frame no operand
809 // carries — a silent PARTIAL frame otherwise (see
810 // [`same_as_frame_guard`]). Runs AFTER evaluation so the
811 // structural declines (`OperandOutOfRange`, `SymbolicGap`)
812 // keep their existing precedence.
813 same_as_frame_guard(tag, se, bind_shapes)?;
814 out.target_shape = s;
815 }
816 ShapeValue::Gap => {
817 return Err(RelAttrError::SymbolicGap {
818 field: "target_shape",
819 });
820 }
821 }
822 out.target_shape_rel = None;
823 }
824
825 // slice_{start,len}_rel → slice_{start,len} (DimExpr over the Bind space).
826 let eval_dim_field = |d: &Dim, field: &'static str| -> Result<u64, RelAttrError> {
827 match shape_expr::eval_dim(d, bind_shapes, &[]).map_err(RelAttrError::Expr)? {
828 DimValue::Concrete(v) if v < 0 => Err(RelAttrError::Negative { field, value: v }),
829 DimValue::Concrete(v) => Ok(v as u64),
830 DimValue::Gap => Err(RelAttrError::SymbolicGap { field }),
831 }
832 };
833 if let Some(d) = &attrs.slice_start_rel {
834 out.slice_start = Some(eval_dim_field(d, "slice_start")?);
835 out.slice_start_rel = None;
836 }
837 if let Some(d) = &attrs.slice_len_rel {
838 out.slice_len = Some(eval_dim_field(d, "slice_len")?);
839 out.slice_len_rel = None;
840 }
841
842 // scalar_rel → scalars (DimExpr over the Bind space; the reduced_count
843 // concept — `Dim::Extent { operand, axis: LAST }` = n = extent of that
844 // operand's last axis, the norm-backward `MulScalar(n)` divisor). Rides the
845 // SAME §6.20 `eval_dim`/`resolve_axis` evaluator as the slice_* rel fields,
846 // resolved against the Bind space (not this op's children — a scalar rides no
847 // operand's rank). Unlike an extent/offset, a scalar may be negative, so only
848 // a symbolic Gap declines (§6.20-0004); every concrete value flows through.
849 if let Some(d) = &attrs.scalar_rel {
850 match shape_expr::eval_dim(d, bind_shapes, &[]).map_err(RelAttrError::Expr)? {
851 DimValue::Concrete(v) => out.scalars = vec![v as f64],
852 DimValue::Gap => return Err(RelAttrError::SymbolicGap { field: "scalars" }),
853 }
854 out.scalar_rel = None;
855 }
856
857 // axis_last → the per-tag axis carrier, resolved against operand[0]'s rank.
858 if attrs.axis_last {
859 let rank = child_shapes
860 .first()
861 .ok_or(RelAttrError::NoChildOperand)?
862 .len();
863 use OpTag as T;
864 match tag {
865 // `axis`-carrier tags: this op's LAST = rank − 1 via the shared
866 // §6.20 resolver (typed AxisOutOfRange on a rank-0 operand).
867 T::SumDim
868 | T::MaxDim
869 | T::MeanDim
870 | T::CumSum
871 | T::Concat
872 | T::Flip
873 | T::Slice
874 | T::Roll
875 | T::IndexSelect
876 | T::Gather
877 | T::IndexAdd
878 | T::ScatterAdd => {
879 let a = resolve_axis(LAST, rank).map_err(RelAttrError::Expr)?;
880 out.axis = Some(a as i64);
881 }
882 // `dims`-carrier: Unsqueeze APPENDS — dim == rank (`primitive_shape`
883 // permits `dim == rank`; keepdim-restore spelling, D3).
884 T::Unsqueeze => {
885 out.dims = vec![rank as u8];
886 }
887 // `dims`-carrier: Squeeze drops the trailing axis = rank − 1.
888 T::Squeeze => {
889 let a = resolve_axis(LAST, rank).map_err(RelAttrError::Expr)?;
890 out.dims = vec![a as u8];
891 }
892 other => return Err(RelAttrError::AxisLastUnsupported { tag: other }),
893 }
894 out.axis_last = false;
895 }
896
897 Ok(out)
898}
899
900/// Count the region's open scalar **slots** in pattern pre-order — scalar-param
901/// ops whose `attrs.scalars` is empty (a baked value is a pattern constant, not
902/// a slot). This is the length of the `scalars` vec `match_region_extract`
903/// returns for a match, and of the `FusedOpParams::Runtime { scalars }` the
904/// fused node must carry for [`decompose_region`] to fill the re-emit.
905pub fn count_scalar_slots(node: &PatternNode) -> usize {
906 match node {
907 PatternNode::Op {
908 op,
909 operands,
910 attrs,
911 } => {
912 // A `scalar_rel` node is filled from an input SHAPE at emit time, NOT
913 // from the params cursor — so it is never an open slot (mirrors a
914 // baked-value node). Only an empty-`scalars`, no-`scalar_rel`
915 // scalar-param op is a cursor slot.
916 let own = if attrs.scalars.is_empty() && attrs.scalar_rel.is_none() {
917 scalar_slot_arity(*op)
918 } else {
919 0
920 };
921 own + operands.iter().map(count_scalar_slots).sum::<usize>()
922 }
923 _ => 0,
924 }
925}
926
927/// The ONE recipe-validation oracle, shared by [`register_runtime_fused`]
928/// (runtime Tier-2 registration) and the static-registry
929/// [`crate::registry::decompose_via_recipe`] bridge (T5): bind indices form a
930/// contiguous `[0, n)` AND every op re-emits to primitives (totality — incl.
931/// the rel-attr probe). A recipe carrying a semantics-absent op token (no
932/// primitive re-emission — the flip-withdrawal posture: unknown/non-registry
933/// tokens are surfaced honest-miss declines, never accepted, never a crash)
934/// is a typed [`RuntimeFusedError::UnRepresentable`] decline here.
935pub(crate) fn validate_recipe(region: &PatternNode) -> Result<(), RuntimeFusedError> {
936 let binds = region.bind_indices();
937 let n = binds.len() as u8;
938 if binds != (0..n).collect::<Vec<_>>() {
939 return Err(RuntimeFusedError::NonContiguousBinds(binds));
940 }
941 validate_representable(region)
942}
943
944fn validate_representable(region: &PatternNode) -> Result<(), RuntimeFusedError> {
945 let n_binds = region.bind_indices().len();
946 validate_node(region, n_binds)
947}
948
949fn validate_node(node: &PatternNode, n_binds: usize) -> Result<(), RuntimeFusedError> {
950 match node {
951 PatternNode::Op {
952 op,
953 operands,
954 attrs,
955 } => {
956 // A rel-attr op is a SHAPE-POLYMORPHIC template — probe-resolve it
957 // (T3, mirror of the scalar slot dummy-fill below) so the
958 // `tag_to_op` representability check can run on concrete attrs.
959 // Structural authoring errors reject the region with a typed
960 // decline; value-dependent declines at the probe shape register
961 // fine and surface at emit time as a decompose fixpoint.
962 let probed;
963 let attrs = if has_rel_attrs(attrs) {
964 probed = rel_probe(*op, attrs, n_binds)
965 .map_err(|error| RuntimeFusedError::InvalidRelAttrs { tag: *op, error })?;
966 &probed
967 } else {
968 attrs
969 };
970 // A scalar-param op with empty scalars is a SLOT template —
971 // validate re-emittability with a dummy fill (the live value is
972 // substituted from the fused node's `Runtime { scalars }` at
973 // decompose time).
974 let representable = if attrs.scalars.is_empty() && scalar_slot_arity(*op) > 0 {
975 let mut filled = attrs.clone();
976 filled.scalars = vec![0.0; scalar_slot_arity(*op)];
977 tag_to_op(*op, &filled).is_some()
978 } else {
979 tag_to_op(*op, attrs).is_some()
980 };
981 if !representable {
982 return Err(RuntimeFusedError::UnRepresentable(*op));
983 }
984 for o in operands {
985 validate_node(o, n_binds)?;
986 }
987 Ok(())
988 }
989 PatternNode::Bind { .. } => Ok(()),
990 PatternNode::Any | PatternNode::SeeThrough { .. } => {
991 Err(RuntimeFusedError::NonConcreteRegion)
992 }
993 }
994}
995
996/// The registration-time rel-attr probe: resolve `attrs` against a fixed
997/// `[2, 4]` probe shape (every bind + the child) through the ONE resolver.
998/// * `Ok(resolved)` — fully-concrete attrs for the `tag_to_op` probe.
999/// * `Err` — a STRUCTURAL authoring error that can never resolve at ANY
1000/// shape: rel+abs both set, a bind/`Param` reference out of range,
1001/// `axis_last` on an axis-less tag.
1002/// * A VALUE-dependent decline at the probe shape (`Negative`,
1003/// `AxisOutOfRange` against the probe rank, `DivideByZero` through a
1004/// derived extent, …) is NOT an authoring error — the attrs get a dummy
1005/// concrete fill instead (the emit-time resolver is the real gate; its
1006/// decline there is a G2 fixpoint).
1007fn rel_probe(tag: OpTag, attrs: &OpAttrs, n_binds: usize) -> Result<OpAttrs, RelAttrError> {
1008 use fuel_kernel_seam_types::shape_expr::ShapeExprError as E;
1009 let probe: Vec<i64> = vec![2, 4];
1010 let bind_shapes = vec![probe.clone(); n_binds];
1011 match resolve_rel_attrs(tag, attrs, &bind_shapes, std::slice::from_ref(&probe)) {
1012 Ok(resolved) => Ok(resolved),
1013 Err(
1014 e @ (RelAttrError::RelAbsConflict { .. }
1015 | RelAttrError::AxisLastUnsupported { .. }
1016 | RelAttrError::Expr(E::OperandOutOfRange { .. } | E::ParamOutOfRange { .. })),
1017 ) => Err(e),
1018 Err(_) => Ok(dummy_fill_rel(tag, attrs)),
1019 }
1020}
1021
1022/// Clear `attrs`' rel fields and dummy-fill their concrete siblings (only
1023/// where the sibling is unset — a rel+abs conflict never reaches here, the
1024/// probe rejects it first) so `tag_to_op` representability can be checked.
1025fn dummy_fill_rel(tag: OpTag, attrs: &OpAttrs) -> OpAttrs {
1026 let mut out = attrs.clone();
1027 if out.target_shape_rel.take().is_some() && out.target_shape.is_empty() {
1028 out.target_shape = vec![1];
1029 }
1030 if out.slice_start_rel.take().is_some() && out.slice_start.is_none() {
1031 out.slice_start = Some(0);
1032 }
1033 if out.slice_len_rel.take().is_some() && out.slice_len.is_none() {
1034 out.slice_len = Some(1);
1035 }
1036 if out.scalar_rel.take().is_some() && out.scalars.is_empty() {
1037 out.scalars = vec![1.0];
1038 }
1039 if out.axis_last {
1040 out.axis_last = false;
1041 match tag {
1042 OpTag::Unsqueeze | OpTag::Squeeze => {
1043 if out.dims.is_empty() {
1044 out.dims = vec![0];
1045 }
1046 }
1047 _ => {
1048 if out.axis.is_none() {
1049 out.axis = Some(0);
1050 }
1051 }
1052 }
1053 }
1054 out
1055}
1056
1057/// Decompose a runtime `Op::Fused(id, Runtime { .. })` node by re-emitting its
1058/// region as primitives, returning the new root (the re-emitted sink). If `id`
1059/// is not a registered runtime op the node is returned unchanged (a fixpoint —
1060/// no recipe, G2). The matched node's inputs are the region's bound external
1061/// inputs in bind-index order.
1062pub fn decompose_region(graph: &mut Graph, node_id: NodeId) -> NodeId {
1063 let (fid, node_scalars) = match &graph.node(node_id).op {
1064 Op::Fused(id, FusedOpParams::Runtime { scalars }) => (*id, scalars.clone()),
1065 Op::Fused(id, _) => (*id, Vec::new()),
1066 _ => return node_id,
1067 };
1068 let region = match runtime_region(fid) {
1069 Some(r) => r,
1070 None => return node_id,
1071 };
1072 // The node's live scalars must fill the region's slots exactly (pattern
1073 // pre-order, the same canon `match_region_extract` produced them in). A
1074 // mismatch is a malformed fused node — surfaced as a no-op fixpoint (the
1075 // lowering driver records no progress), never a crash (G2).
1076 if node_scalars.len() != count_scalar_slots(®ion) {
1077 return node_id;
1078 }
1079 let inputs = graph.node(node_id).inputs.clone();
1080 let bind_shapes = bind_operand_shapes(graph, &inputs);
1081 let mut cursor = node_scalars.as_slice();
1082 // A shape-relative attr that fails to resolve at THESE input shapes (a
1083 // symbolic extent, a negative result, …) is a typed decline surfaced as a
1084 // no-op fixpoint (G2) — same posture as the slot-count mismatch above,
1085 // never a panic. Any child nodes emitted before the decline stay in the
1086 // push-only graph as unreferenced dead nodes (inert).
1087 emit(
1088 graph,
1089 ®ion,
1090 &inputs,
1091 &bind_shapes,
1092 &mut cursor,
1093 &mut Vec::new(),
1094 )
1095 .unwrap_or(node_id)
1096}
1097
1098/// Re-emit a validated region on the given external input nodes (public entry
1099/// for callers holding a raw [`PatternNode`] + input [`NodeId`]s — e.g. the
1100/// reference realization during candidate-kernel verification, which has a raw
1101/// region and freshly-pushed `Op::Const` input nodes rather than a Fused node
1102/// already in the graph). `scalars` fill the region's open scalar slots in
1103/// pre-order (the canonical order `match_region_extract` recorded them in);
1104/// pass `&[]` for a parameterless region. Thin wrapper over the private
1105/// `emit`; the same re-emittability caveat applies (a non-re-emittable
1106/// `OpTag` panics inside `emit` — validated decomposes never carry one).
1107/// Second panic risk: `emit`'s scalar-cursor fill (`scalars.split_at(arity)`)
1108/// panics if `scalars` is shorter than the region's total open-slot count.
1109/// [`decompose_region`] (the node-driven caller) guards this with its own
1110/// length check before ever calling `emit`; `emit_region` deliberately does
1111/// NOT — it's a thin wrapper, so validating the length is the caller's job.
1112/// Callers must pass a `scalars` slice at least as long as the region's
1113/// open-slot count. Third (T3): a shape-RELATIVE attr (D2) that fails to
1114/// resolve at these input shapes panics through the wrapper's `expect` —
1115/// rel-attr callers use [`try_emit_region`], which surfaces it as a typed
1116/// [`RelAttrError`] instead.
1117pub fn emit_region(
1118 graph: &mut Graph,
1119 region: &PatternNode,
1120 inputs: &[NodeId],
1121 scalars: &[f64],
1122) -> NodeId {
1123 try_emit_region(graph, region, inputs, scalars).expect(
1124 "rel-attr resolution failed — emit_region callers pass concrete-attr or \
1125 shape-compatible pre-validated regions; fallible callers use try_emit_region",
1126 )
1127}
1128
1129/// The FALLIBLE re-emit entry (Increment C slice 1, T3): like [`emit_region`]
1130/// but surfacing a shape-relative attr resolution failure (D2) as a typed
1131/// [`RelAttrError`] instead of a panic. This is the resolving entry the
1132/// registry `decompose_via_recipe` bridge calls (any failure ⇒ `return id`,
1133/// the G2 fixpoint). Concrete-attr regions can never hit the `Err` arm — for
1134/// them this is exactly the legacy `emit_region`. The `emit_region` panic
1135/// caveats (non-re-emittable `OpTag`, short `scalars` slice) apply unchanged.
1136pub fn try_emit_region(
1137 graph: &mut Graph,
1138 region: &PatternNode,
1139 inputs: &[NodeId],
1140 scalars: &[f64],
1141) -> Result<NodeId, RelAttrError> {
1142 let bind_shapes = bind_operand_shapes(graph, inputs);
1143 let mut cursor = scalars;
1144 emit(
1145 graph,
1146 region,
1147 inputs,
1148 &bind_shapes,
1149 &mut cursor,
1150 &mut Vec::new(),
1151 )
1152}
1153
1154/// A graph [`fuel_ir::Shape`] as a §6.20 evaluator operand: per-axis extents
1155/// with a bounded-symbolic (`Extent::Range`) axis mapped to the
1156/// [`shape_expr::SYMBOLIC`] sentinel — so a rel attr over a symbolic extent
1157/// declines as [`RelAttrError::SymbolicGap`] (surfaced gap, §6.20-0004)
1158/// instead of silently resolving against the capacity bound.
1159fn shape_expr_operand(shape: &fuel_ir::Shape) -> Vec<i64> {
1160 use fuel_kernel_seam_types::shape_expr::SYMBOLIC;
1161 (0..shape.rank())
1162 .map(|a| {
1163 if shape.extent(a).is_dynamic() {
1164 SYMBOLIC
1165 } else {
1166 shape.dims()[a] as i64
1167 }
1168 })
1169 .collect()
1170}
1171
1172/// The region's **Bind-space** shapes (`bind_shapes[i]` = `inputs[i]`'s shape)
1173/// in §6.20 operand form — what `ShapeExpr::SameAs`/`Dim::Extent` index.
1174fn bind_operand_shapes(graph: &Graph, inputs: &[NodeId]) -> Vec<Vec<i64>> {
1175 inputs
1176 .iter()
1177 .map(|&id| shape_expr_operand(&graph.node(id).shape))
1178 .collect()
1179}
1180
1181/// The recursive re-emit core. `memo` is the per-emit-call identity-share
1182/// table (T5): a REPEATED slot-free subtree — the tree spelling of a DAG
1183/// recipe's shared interior (e.g. softmax's `e = Exp(..)`, consumed by both
1184/// the denominator reduce and the final Div) — emits ONCE, so the emitted
1185/// graph is the DAG, not a duplicated-compute tree. Lookup is by structural
1186/// equality (`PatternNode: PartialEq`; regions are tiny, a linear scan is
1187/// fine) and is sound because, within one call, `inputs`/`bind_shapes` are
1188/// fixed and emission is deterministic — equal subtrees emit equal nodes.
1189/// Subtrees with OPEN scalar slots are NEVER shared: each occurrence takes
1190/// its own value(s) from the pre-order cursor. (The flat-DAG node table with
1191/// real CSE is slice 3; this is only within-call identity-share.)
1192fn emit<'r>(
1193 graph: &mut Graph,
1194 node: &'r PatternNode,
1195 inputs: &[NodeId],
1196 bind_shapes: &[Vec<i64>],
1197 scalars: &mut &[f64],
1198 memo: &mut Vec<(&'r PatternNode, NodeId)>,
1199) -> Result<NodeId, RelAttrError> {
1200 match node {
1201 PatternNode::Bind { index } => Ok(inputs[*index as usize]),
1202 PatternNode::Op {
1203 op,
1204 operands,
1205 attrs,
1206 } => {
1207 // Identity-share: a slot-free subtree already emitted in THIS call
1208 // re-uses its node (see the fn doc). Checked before the cursor
1209 // fill — a slot-free subtree never moves the cursor, so a hit
1210 // cannot misalign later slots.
1211 let sharable = count_scalar_slots(node) == 0;
1212 if sharable && let Some(&(_, id)) = memo.iter().find(|(p, _)| *p == node) {
1213 return Ok(id);
1214 }
1215 // Fill an open scalar slot from the cursor in PRE-order (before
1216 // descending into operands) — the same canonical order
1217 // `match_region_extract` recorded the live values in. (T3 note:
1218 // children are now EMITTED before the attrs are USED, but the
1219 // cursor fill stays right here, before the descent — the cursor
1220 // order is authoring order, not emission order.)
1221 // A `scalar_rel` node is filled from an input shape by the rel-attr
1222 // resolver below (NOT the params cursor) — matching
1223 // `count_scalar_slots`, which does not count it as a slot, so it must
1224 // not consume a cursor value here either.
1225 let arity = scalar_slot_arity(*op);
1226 let filled;
1227 let attrs = if attrs.scalars.is_empty() && arity > 0 && attrs.scalar_rel.is_none() {
1228 let (take, rest) = scalars.split_at(arity);
1229 *scalars = rest;
1230 filled = OpAttrs {
1231 scalars: take.to_vec(),
1232 ..attrs.clone()
1233 };
1234 &filled
1235 } else {
1236 attrs
1237 };
1238 // Children FIRST (T3 reorder): their emitted shapes feed the
1239 // rel-attr resolver (`axis_last`'s rank, D4's pad decision).
1240 let mut child_ids = Vec::with_capacity(operands.len());
1241 for o in operands {
1242 child_ids.push(emit(graph, o, inputs, bind_shapes, scalars, memo)?);
1243 }
1244 let mut child_shapes: Vec<fuel_ir::Shape> = child_ids
1245 .iter()
1246 .map(|&c| graph.node(c).shape.clone())
1247 .collect();
1248 let child_dtypes: Vec<fuel_ir::DType> =
1249 child_ids.iter().map(|&c| graph.node(c).dtype).collect();
1250 // Shape-RELATIVE attrs (D2) resolve to fully-concrete siblings
1251 // against the region's Bind space + this op's operand shapes; the
1252 // unchanged tag_to_op → primitive_shape path then runs on the
1253 // result. A failure is a typed decline the caller surfaces
1254 // (`decompose_region` ⇒ fixpoint, `try_emit_region` ⇒ `Err`) —
1255 // never a panic. Nodes already pushed for the children stay in the
1256 // graph as unreferenced (dead) nodes: `Graph` is push-only and
1257 // base-map extraction walks from roots, so they are inert.
1258 let resolved;
1259 let attrs = if has_rel_attrs(attrs) {
1260 let child_ops: Vec<Vec<i64>> =
1261 child_shapes.iter().map(shape_expr_operand).collect();
1262 resolved = resolve_rel_attrs(*op, attrs, bind_shapes, &child_ops)?;
1263 &resolved
1264 } else {
1265 attrs
1266 };
1267 let mut prim =
1268 tag_to_op(*op, attrs).expect("region validated re-emittable at registration");
1269 // A `MaskedFill` fill Scalar must carry the FILLED tensor's
1270 // (operand[0]) dtype — the byte executor derives `fill_bytes` at
1271 // that width. A recipe authors the value dtype-polymorphically (no
1272 // `cast_dtype`), so re-resolve the Scalar at operand[0]'s emitted
1273 // dtype here, mirroring the imperative `Scalar::one(dtype)`. Concrete
1274 // round-trips already carry the matching dtype (identity); a dtype a
1275 // Scalar cannot represent leaves the provisional value untouched (an
1276 // inert dead node, never a panic — the executor rejects it later).
1277 if let Op::MaskedFill { value } = &prim
1278 && let Some(&dt) = child_dtypes.first()
1279 && value.dtype() != dt
1280 && let Some(fixed) = masked_fill_scalar(value.to_f64(), dt)
1281 {
1282 prim = Op::MaskedFill { value: fixed };
1283 }
1284 // D4: a `BroadcastTo` whose target rank EXCEEDS its operand's rank
1285 // first materializes the legacy `Reshape` pad (1-padded left,
1286 // right-aligned — byte-identical to `registry::rope`'s hand-built
1287 // broadcast prep, since `check_broadcast_compatible` is
1288 // right-aligned). Recipes stay free of rank-dependent nodes while
1289 // the emitted graph matches the legacy imperative builders.
1290 // Applied uniformly (rel-resolved AND absolute targets); an
1291 // equal-rank broadcast is unchanged (no pad).
1292 if let Op::BroadcastTo(target) = &prim
1293 && let Some(cs) = child_shapes.first()
1294 && target.rank() > cs.rank()
1295 {
1296 let mut padded: Vec<usize> = vec![1; target.rank() - cs.rank()];
1297 padded.extend_from_slice(cs.dims());
1298 let pad_shape = fuel_ir::Shape::from_dims(&padded);
1299 let pad = graph.push(Node {
1300 op: Op::Reshape(pad_shape.clone()),
1301 inputs: vec![child_ids[0]],
1302 shape: pad_shape.clone(),
1303 dtype: child_dtypes[0],
1304 });
1305 child_ids[0] = pad;
1306 child_shapes[0] = pad_shape;
1307 }
1308 // Shape/dtype for the emitted node. Most ops use the single source
1309 // of truth (`primitive_shape`) — a pure function of operand shapes,
1310 // correct for shape-changing/reducing/dtype-changing ops. Two
1311 // structural multi-output terminals are NOT such a function and are
1312 // resolved here with the graph in hand, mirroring `NodeHandle::scan` /
1313 // `Graph::view` (Increment C, B1):
1314 // * `Op::Scan` — the node's PRIMARY (slot-0) shape is the stacked
1315 // ys `[bound] ++ body_y`, and its 2-slot `output_views` bundle
1316 // (slot 0 = stacked ys, slot 1 = final carry) is attached AFTER
1317 // the push so downstream `Op::View`s can read it.
1318 // * `Op::View` — the slot's shape/dtype come from the producer's
1319 // `output_views[slot]` (set when the producing `Scan` emitted).
1320 // Both fall back to operand[0]'s shape/dtype on a malformed authored
1321 // region (never a panic) — the same posture as the primitive_shape
1322 // fallback (`.first()`, never `[0]`, guards a zero-operand leaf).
1323 fn fallback_sd(
1324 cs: &[fuel_ir::Shape],
1325 cd: &[fuel_ir::DType],
1326 ) -> (fuel_ir::Shape, fuel_ir::DType) {
1327 (
1328 cs.first()
1329 .cloned()
1330 .unwrap_or_else(|| fuel_ir::Shape::from_dims(&[])),
1331 cd.first().copied().unwrap_or(fuel_ir::DType::F32),
1332 )
1333 }
1334 let mut scan_bundle: Option<Vec<fuel_ir::storage::OutputViewSpec>> = None;
1335 let (s, d) = match &prim {
1336 Op::Scan {
1337 n_carries: _,
1338 n_xs,
1339 bound,
1340 early_exit,
1341 ..
1342 } => {
1343 // inputs = [init_carry, xs(n_xs), consts.., body_new_carry,
1344 // body_y, [pred_exit]] — the Phase-1 lax.scan encoding.
1345 let n_trailing = if early_exit.is_some() { 3 } else { 2 };
1346 let len = child_ids.len();
1347 if len >= 1 + *n_xs + n_trailing {
1348 let carry_shape = child_shapes[0].clone();
1349 let carry_dtype = child_dtypes[0];
1350 let body_y_i = len - n_trailing + 1;
1351 let y_shape = &child_shapes[body_y_i];
1352 let y_dtype = child_dtypes[body_y_i];
1353 let mut ys_dims = Vec::with_capacity(1 + y_shape.rank());
1354 ys_dims.push(*bound);
1355 ys_dims.extend_from_slice(y_shape.dims());
1356 let ys_shape = fuel_ir::Shape::from_dims(&ys_dims);
1357 scan_bundle = Some(vec![
1358 fuel_ir::storage::OutputViewSpec::contiguous(y_dtype, ys_shape.clone()),
1359 fuel_ir::storage::OutputViewSpec::contiguous(carry_dtype, carry_shape),
1360 ]);
1361 (ys_shape, y_dtype)
1362 } else {
1363 fallback_sd(&child_shapes, &child_dtypes)
1364 }
1365 }
1366 Op::View { slot } => child_ids
1367 .first()
1368 .and_then(|&producer| {
1369 graph
1370 .output_views(producer)
1371 .and_then(|v| v.get(*slot as usize))
1372 .map(|spec| (spec.shape.clone(), spec.dtype))
1373 })
1374 .unwrap_or_else(|| fallback_sd(&child_shapes, &child_dtypes)),
1375 // A nested fused node carried as-is (Increment C, C-T2, mechanism
1376 // 2a). `primitive_shape` honest-misses `Fused` — a fused op is not
1377 // a first-order shape-inferable primitive — so its output frame
1378 // comes from the named registry entry's OWN `shape_rule` /
1379 // `dtype_rule` (e.g. softmax passes both through). Falls back to
1380 // operand[0] on an unregistered id (never a panic), the same
1381 // posture as the primitive_shape fallback.
1382 Op::Fused(fid, params) => match crate::registry::default_registry().entry(*fid) {
1383 Some(e) => (
1384 (e.shape_rule)(&child_shapes, params),
1385 (e.dtype_rule)(&child_dtypes, params),
1386 ),
1387 None => fallback_sd(&child_shapes, &child_dtypes),
1388 },
1389 // A childless scan-body hole (`Op::ScanPlaceholder`). Its shape
1390 // is DECLARED on the recipe node (`target_shape` /
1391 // `target_shape_rel`, resolved above) rather than inferred from
1392 // (absent) operands, so the re-emitted body carries the same
1393 // per-step shapes the imperative `selective_scan`/`ssd_chunk_scan`
1394 // decompose stamps on its placeholders. This is load-bearing:
1395 // `unroll_scan` CLONES body interior nodes with their STORED
1396 // shapes, so a rank-0 fallback here would poison the unrolled
1397 // graph (e.g. `du = Mul(d_t, u_t)` would come out rank-0). Dtype =
1398 // the scan's working dtype = bind 0's (uniform-dtype scan inputs).
1399 // Both fall back to the operand/F32 default on a shapeless
1400 // authored placeholder (B1's round-trip recipes), never a panic.
1401 Op::ScanPlaceholder { .. } => (
1402 shape_from_attr(attrs)
1403 .unwrap_or_else(|| fallback_sd(&child_shapes, &child_dtypes).0),
1404 inputs
1405 .first()
1406 .map(|&b| graph.node(b).dtype)
1407 .unwrap_or_else(|| fallback_sd(&child_shapes, &child_dtypes).1),
1408 ),
1409 _ => crate::shape::primitive_shape(&prim, &child_shapes, &child_dtypes)
1410 .unwrap_or_else(|_| fallback_sd(&child_shapes, &child_dtypes)),
1411 };
1412 let out = graph.push(Node {
1413 op: prim,
1414 inputs: child_ids,
1415 shape: s,
1416 dtype: d,
1417 });
1418 // Attach the scan's 2-slot bundle (mirrors `NodeHandle::scan`). A
1419 // malformed spec (compose/validate failure) leaves the producer
1420 // bundle-less — a later `View` then falls back — never a panic.
1421 if let Some(specs) = scan_bundle
1422 && let Ok((_bytes, views)) = fuel_ir::storage::compose_bundle(&specs)
1423 {
1424 let _ = graph.set_output_views(out, std::sync::Arc::from(views.into_boxed_slice()));
1425 }
1426 if sharable {
1427 memo.push((node, out));
1428 }
1429 Ok(out)
1430 }
1431 PatternNode::Any | PatternNode::SeeThrough { .. } => {
1432 unreachable!("region validated concrete (Op/Bind) at registration")
1433 }
1434 }
1435}
1436
1437/// A [`crate::opt::LoweringRule`]-shaped `decompose` for runtime ops: re-emit
1438/// the region. The scalar `extract:` substitution rides on the NODE (its
1439/// `FusedOpParams::Runtime { scalars }` fills the region's open slots inside
1440/// [`decompose_region`]), so the rule-shaped `params` argument stays unused.
1441pub fn runtime_lowering_decompose(
1442 graph: &mut Graph,
1443 node_id: NodeId,
1444 _params: &FusedOpParams,
1445) -> NodeId {
1446 decompose_region(graph, node_id)
1447}
1448
1449#[cfg(test)]
1450mod tests {
1451 use super::*;
1452 use fuel_ir::{DType, Shape};
1453
1454 fn relu_add_region() -> PatternNode {
1455 PatternNode::Op {
1456 op: OpTag::Relu,
1457 attrs: OpAttrs::default(),
1458 operands: vec![PatternNode::Op {
1459 op: OpTag::Add,
1460 attrs: OpAttrs::default(),
1461 operands: vec![
1462 PatternNode::Bind { index: 0 },
1463 PatternNode::Bind { index: 1 },
1464 ],
1465 }],
1466 }
1467 }
1468
1469 /// Structurally DISTINCT from `relu_add_region()` (`Mul` inner op, not
1470 /// `Add`) — used only by
1471 /// `register_allocates_a_runtime_id_and_keeps_the_region`, whose
1472 /// assertion on `runtime_name` needs a region no OTHER test in this
1473 /// module also registers. Since Task 7's dedup (`register_runtime_fused`
1474 /// above) resolves any two structurally-identical regions registered
1475 /// anywhere in the process to the SAME id — and `RUNTIME_FUSED_OPS` /
1476 /// `hash_index` are process-global statics shared by every `#[test]` in
1477 /// this binary, which `cargo test` runs concurrently by default — a
1478 /// `runtime_name` assertion tied to one specific registration call would
1479 /// be racy against any other test using `relu_add_region()` (both
1480 /// `decompose_region_re_emits_relu_add` and
1481 /// `register_runtime_fused_dedups_structurally_identical_regions` do):
1482 /// whichever call reaches the shared hash slot FIRST wins the name, and
1483 /// thread scheduling decides which that is. Those other two tests never
1484 /// assert on `runtime_name`, so they're unaffected by dedup either way;
1485 /// this one needs its own hash to stay deterministic.
1486 fn relu_mul_region() -> PatternNode {
1487 PatternNode::Op {
1488 op: OpTag::Relu,
1489 attrs: OpAttrs::default(),
1490 operands: vec![PatternNode::Op {
1491 op: OpTag::Mul,
1492 attrs: OpAttrs::default(),
1493 operands: vec![
1494 PatternNode::Bind { index: 0 },
1495 PatternNode::Bind { index: 1 },
1496 ],
1497 }],
1498 }
1499 }
1500
1501 #[test]
1502 fn register_allocates_a_runtime_id_and_keeps_the_region() {
1503 let id = register_runtime_fused("test::relu_mul", relu_mul_region()).unwrap();
1504 assert!(id.is_runtime(), "allocated id is in the runtime range");
1505 assert_eq!(runtime_region(id), Some(relu_mul_region()));
1506 assert_eq!(runtime_name(id).as_deref(), Some("test::relu_mul"));
1507 }
1508
1509 #[test]
1510 fn register_runtime_fused_dedups_structurally_identical_regions() {
1511 let id1 = register_runtime_fused("dedup::a", relu_add_region()).unwrap();
1512 let id2 = register_runtime_fused("dedup::b", relu_add_region()).unwrap(); // same region, different name
1513 assert_eq!(
1514 id1, id2,
1515 "an identical region must resolve to the same FusedOpId, not a duplicate"
1516 );
1517 }
1518
1519 #[test]
1520 fn register_rejects_non_contiguous_binds() {
1521 // bind indices {0, 2} — missing 1.
1522 let region = PatternNode::Op {
1523 op: OpTag::Add,
1524 attrs: OpAttrs::default(),
1525 operands: vec![
1526 PatternNode::Bind { index: 0 },
1527 PatternNode::Bind { index: 2 },
1528 ],
1529 };
1530 assert_eq!(
1531 register_runtime_fused("bad", region),
1532 Err(RuntimeFusedError::NonContiguousBinds(vec![0, 2]))
1533 );
1534 }
1535
1536 #[test]
1537 fn register_rejects_unrepresentable_region() {
1538 // Convergence A made MatMul/shape/reduction ops representable, and
1539 // Increment C carriers made MaskedFill (A2) + PowI (A3) representable;
1540 // `Clamp` stays an honest miss (no two-scalar carrier yet), so it is the
1541 // current canonical still-unrepresentable tag.
1542 let region = PatternNode::Op {
1543 op: OpTag::Clamp,
1544 attrs: OpAttrs::default(),
1545 operands: vec![PatternNode::Bind { index: 0 }],
1546 };
1547 assert_eq!(
1548 register_runtime_fused("bad", region),
1549 Err(RuntimeFusedError::UnRepresentable(OpTag::Clamp))
1550 );
1551 }
1552
1553 #[test]
1554 fn tag_to_op_reconstructs_shape_changing_ops() {
1555 use fuel_ir::Shape;
1556 // Slice{dim:1,start:2,len:3}
1557 let attrs = OpAttrs {
1558 axis: Some(1),
1559 slice_start: Some(2),
1560 slice_len: Some(3),
1561 ..OpAttrs::default()
1562 };
1563 assert!(matches!(
1564 super::tag_to_op(OpTag::Slice, &attrs),
1565 Some(Op::Slice {
1566 dim: 1,
1567 start: 2,
1568 len: 3
1569 })
1570 ));
1571 // Concat{dim:0}
1572 let attrs = OpAttrs {
1573 axis: Some(0),
1574 ..OpAttrs::default()
1575 };
1576 assert!(matches!(
1577 super::tag_to_op(OpTag::Concat, &attrs),
1578 Some(Op::Concat { dim: 0 })
1579 ));
1580 // Reshape([6])
1581 let attrs = OpAttrs {
1582 target_shape: vec![6],
1583 ..OpAttrs::default()
1584 };
1585 assert_eq!(
1586 super::tag_to_op(OpTag::Reshape, &attrs),
1587 Some(Op::Reshape(Shape::from_dims(&[6])))
1588 );
1589 // BroadcastTo([2,3])
1590 let attrs = OpAttrs {
1591 target_shape: vec![2, 3],
1592 ..OpAttrs::default()
1593 };
1594 assert_eq!(
1595 super::tag_to_op(OpTag::BroadcastTo, &attrs),
1596 Some(Op::BroadcastTo(Shape::from_dims(&[2, 3])))
1597 );
1598 // ReduceMaxTo([2,1])
1599 let attrs = OpAttrs {
1600 target_shape: vec![2, 1],
1601 ..OpAttrs::default()
1602 };
1603 assert_eq!(
1604 super::tag_to_op(OpTag::ReduceMaxTo, &attrs),
1605 Some(Op::ReduceMaxTo(Shape::from_dims(&[2, 1])))
1606 );
1607 }
1608
1609 #[test]
1610 fn shape_target_ops_represent_rank0_via_the_marker() {
1611 // C-T1: an INTENTIONAL rank-0 (`[]`) reduce/reshape/broadcast target.
1612 // A rank-0 shape has empty `target_shape` — the same empty state as an
1613 // unset/wildcard target — so the `rank0_target` marker disambiguates.
1614 use fuel_ir::Shape;
1615 let empty = Shape::from_dims(&[]);
1616 // Marker SET → the concrete rank-0 shape (RED before the shape_from_attr
1617 // rank0 arm: an empty `target_shape` honest-missed to `None`).
1618 let rank0 = OpAttrs {
1619 rank0_target: true,
1620 ..OpAttrs::default()
1621 };
1622 assert_eq!(super::shape_from_attr(&rank0), Some(empty.clone()));
1623 assert_eq!(
1624 super::tag_to_op(OpTag::ReduceSumTo, &rank0),
1625 Some(Op::ReduceSumTo(empty.clone()))
1626 );
1627 assert_eq!(
1628 super::tag_to_op(OpTag::ReduceMaxTo, &rank0),
1629 Some(Op::ReduceMaxTo(empty.clone()))
1630 );
1631 assert_eq!(
1632 super::tag_to_op(OpTag::Reshape, &rank0),
1633 Some(Op::Reshape(empty.clone()))
1634 );
1635 assert_eq!(
1636 super::tag_to_op(OpTag::BroadcastTo, &rank0),
1637 Some(Op::BroadcastTo(empty))
1638 );
1639 // Marker UNSET + empty target_shape stays an honest miss (wildcard).
1640 assert_eq!(super::shape_from_attr(&OpAttrs::default()), None);
1641 assert_eq!(
1642 super::tag_to_op(OpTag::ReduceSumTo, &OpAttrs::default()),
1643 None
1644 );
1645 }
1646
1647 #[test]
1648 fn tag_to_op_reconstructs_reductions_dtype_and_matmul() {
1649 use fuel_ir::DType;
1650 assert!(matches!(
1651 super::tag_to_op(
1652 OpTag::MeanDim,
1653 &OpAttrs {
1654 axis: Some(1),
1655 ..OpAttrs::default()
1656 }
1657 ),
1658 Some(Op::MeanDim(1))
1659 ));
1660 assert!(matches!(
1661 super::tag_to_op(OpTag::MatMul, &OpAttrs::default()),
1662 Some(Op::MatMul)
1663 ));
1664 // Cast target dtype via name.
1665 let attrs = OpAttrs {
1666 cast_dtype: Some("f16".into()),
1667 ..OpAttrs::default()
1668 };
1669 assert_eq!(
1670 super::tag_to_op(OpTag::Cast, &attrs),
1671 Some(Op::Cast(DType::F16))
1672 );
1673 // Comparison.
1674 assert!(matches!(
1675 super::tag_to_op(OpTag::Lt, &OpAttrs::default()),
1676 Some(Op::Lt)
1677 ));
1678 }
1679
1680 #[test]
1681 fn tag_to_op_reconstructs_masked_fill() {
1682 // A2 (Increment C carriers): the MaskedFill re-emit carrier. Value on
1683 // `scalars[0]`, dtype on `cast_dtype` (the `op_to_attrs` projection) —
1684 // RED before A2 (MaskedFill was the `_ => return None` honest miss).
1685 use fuel_ir::Scalar;
1686 let attrs = OpAttrs {
1687 scalars: vec![-1.0],
1688 cast_dtype: Some("f16".into()),
1689 ..OpAttrs::default()
1690 };
1691 match super::tag_to_op(OpTag::MaskedFill, &attrs) {
1692 Some(Op::MaskedFill { value }) => {
1693 assert_eq!(value.dtype(), DType::F16, "dtype rides cast_dtype");
1694 assert_eq!(
1695 value,
1696 Scalar::from_f64(-1.0, DType::F16).unwrap(),
1697 "value rides scalars[0]"
1698 );
1699 }
1700 other => panic!("expected MaskedFill, got {other:?}"),
1701 }
1702 // No fill value ⇒ honest miss (None), never a defaulted 0.
1703 assert_eq!(
1704 super::tag_to_op(OpTag::MaskedFill, &OpAttrs::default()),
1705 None
1706 );
1707 // A dtype `Scalar` cannot represent (sub-byte dummy) ⇒ honest miss.
1708 let dummy = OpAttrs {
1709 scalars: vec![1.0],
1710 cast_dtype: Some("f4".into()),
1711 ..OpAttrs::default()
1712 };
1713 assert_eq!(super::tag_to_op(OpTag::MaskedFill, &dummy), None);
1714 // dtype-polymorphic authoring (no cast_dtype) reconstructs a provisional
1715 // F32 that `emit` later re-resolves to operand[0]'s dtype.
1716 let poly = OpAttrs {
1717 scalars: vec![1.0],
1718 ..OpAttrs::default()
1719 };
1720 assert!(matches!(
1721 super::tag_to_op(OpTag::MaskedFill, &poly),
1722 Some(Op::MaskedFill { .. })
1723 ));
1724 // The fill value is a BAKED pattern constant, not a cursor slot.
1725 assert_eq!(super::scalar_slot_arity(OpTag::MaskedFill), 0);
1726 }
1727
1728 #[test]
1729 fn tag_to_op_reconstructs_powi() {
1730 // A3 (Increment C carriers): the PowI i32-exponent re-emit carrier. The
1731 // exponent rides `scalars[0]` as an f64 (an EXACT round-trip for every
1732 // i32 — |n| < 2^53), reconstructed via `as i32`. RED before A3 (PowI was
1733 // in the `_ => return None` honest-miss set).
1734 assert!(matches!(
1735 super::tag_to_op(
1736 OpTag::PowI,
1737 &OpAttrs {
1738 scalars: vec![3.0],
1739 ..OpAttrs::default()
1740 }
1741 ),
1742 Some(Op::PowI(3))
1743 ));
1744 // Negative exponent (e.g. PowI(-1) = reciprocal — the exp==0 backward arm).
1745 assert!(matches!(
1746 super::tag_to_op(
1747 OpTag::PowI,
1748 &OpAttrs {
1749 scalars: vec![-1.0],
1750 ..OpAttrs::default()
1751 }
1752 ),
1753 Some(Op::PowI(-1))
1754 ));
1755 // No exponent ⇒ honest miss (None), never a defaulted PowI(0).
1756 assert_eq!(super::tag_to_op(OpTag::PowI, &OpAttrs::default()), None);
1757 // The exponent is a BAKED pattern constant (like MaskedFill's fill), not
1758 // a cursor slot — so it never draws from the params projection.
1759 assert_eq!(super::scalar_slot_arity(OpTag::PowI), 0);
1760 }
1761
1762 #[test]
1763 fn masked_fill_region_registers_and_emits_with_operand_dtype() {
1764 // Before A2 this region was `UnRepresentable(MaskedFill)`; now it
1765 // registers, and `emit` resolves the fill Scalar to operand[0]'s dtype
1766 // (dtype-polymorphic — the imperative `Scalar::one(dtype)` behaviour).
1767 // No `clear_runtime_fused_for_tests()` here: this test asserts on the
1768 // EMITTED op (dtype/value), never on the registry state or the minted
1769 // id, so the clear served nothing — and calling it unserialized in this
1770 // parallel #[test] binary raced every other registry-touching test by
1771 // wiping ids mid-run (GAP-219; the fn's own doc requires callers to
1772 // serialize). Registering a unique region without a clean slate is fine.
1773 let region = PatternNode::Op {
1774 op: OpTag::MaskedFill,
1775 attrs: OpAttrs {
1776 scalars: vec![1.0],
1777 ..OpAttrs::default()
1778 },
1779 operands: vec![
1780 PatternNode::Bind { index: 0 },
1781 PatternNode::Bind { index: 1 },
1782 ],
1783 };
1784 register_runtime_fused("mf::poly", region.clone())
1785 .expect("MaskedFill region registers (was UnRepresentable before A2)");
1786 // Emit over an F16 tensor + U8 mask ⇒ MaskedFill with an F16 fill Scalar.
1787 let mut g = Graph::new();
1788 let x = g.push(Node {
1789 op: Op::Const,
1790 inputs: vec![],
1791 shape: Shape::from_dims(&[2, 3]),
1792 dtype: DType::F16,
1793 });
1794 let mask = g.push(Node {
1795 op: Op::Const,
1796 inputs: vec![],
1797 shape: Shape::from_dims(&[2, 3]),
1798 dtype: DType::U8,
1799 });
1800 let root = emit_region(&mut g, ®ion, &[x, mask], &[]);
1801 match &g.node(root).op {
1802 Op::MaskedFill { value } => {
1803 assert_eq!(
1804 value.dtype(),
1805 DType::F16,
1806 "fill Scalar re-resolved to operand[0] dtype"
1807 );
1808 assert_eq!(
1809 value.to_f64(),
1810 1.0,
1811 "fill VALUE preserved through the re-resolution"
1812 );
1813 }
1814 other => panic!("expected MaskedFill, got {other:?}"),
1815 }
1816 assert_eq!(g.node(root).dtype, DType::F16);
1817 }
1818
1819 #[test]
1820 fn powi_region_registers_and_emits() {
1821 // Before A3 a bare PowI region was `UnRepresentable(PowI)`; now it
1822 // registers, and `emit` reconstructs the i32 exponent from `scalars[0]`.
1823 // No `clear_runtime_fused_for_tests()` here — see the masked_fill test
1824 // above: this test asserts on the emitted op, not registry state, and
1825 // an unserialized clear in a parallel binary raced other tests (GAP-219).
1826 let region = PatternNode::Op {
1827 op: OpTag::PowI,
1828 attrs: OpAttrs {
1829 scalars: vec![3.0],
1830 ..OpAttrs::default()
1831 },
1832 operands: vec![PatternNode::Bind { index: 0 }],
1833 };
1834 register_runtime_fused("powi::cube", region.clone())
1835 .expect("PowI region registers (was UnRepresentable before A3)");
1836 let mut g = Graph::new();
1837 let x = g.push(Node {
1838 op: Op::Const,
1839 inputs: vec![],
1840 shape: Shape::from_dims(&[2, 3]),
1841 dtype: DType::F32,
1842 });
1843 let root = emit_region(&mut g, ®ion, &[x], &[]);
1844 assert!(
1845 matches!(g.node(root).op, Op::PowI(3)),
1846 "emit reconstructs Op::PowI(3) from scalars[0]",
1847 );
1848 assert_eq!(
1849 g.node(root).shape,
1850 Shape::from_dims(&[2, 3]),
1851 "PowI is shape-preserving"
1852 );
1853 assert_eq!(g.node(root).dtype, DType::F32);
1854 }
1855
1856 #[test]
1857 fn tag_to_op_matmul_resolves_canonical_roles() {
1858 // T9 (D5): explicit CANONICAL role vectors resolve to Op::MatMul. The
1859 // resolver checks role POSITIONS against the locked cell, not extents.
1860 let attrs = OpAttrs {
1861 lhs_roles: vec![1, 3],
1862 rhs_roles: vec![3, 2],
1863 ..OpAttrs::default()
1864 };
1865 assert!(matches!(
1866 super::tag_to_op(OpTag::MatMul, &attrs),
1867 Some(Op::MatMul)
1868 ));
1869 // Rank-4 canonical (leading Batch dims) also resolves — GQA-divisible
1870 // batch extents stay all-Batch (positions, not extents).
1871 let attrs4 = OpAttrs {
1872 lhs_roles: vec![0, 0, 1, 3],
1873 rhs_roles: vec![0, 0, 3, 2],
1874 ..OpAttrs::default()
1875 };
1876 assert!(matches!(
1877 super::tag_to_op(OpTag::MatMul, &attrs4),
1878 Some(Op::MatMul)
1879 ));
1880 }
1881
1882 #[test]
1883 fn tag_to_op_matmul_empty_roles_implicit_accept() {
1884 // Empty roles = the rank-polymorphic recipe form → implicit-accept
1885 // (unchanged from today; recipes keep matmul implicit).
1886 assert!(matches!(
1887 super::tag_to_op(OpTag::MatMul, &OpAttrs::default()),
1888 Some(Op::MatMul)
1889 ));
1890 }
1891
1892 #[test]
1893 fn tag_to_op_matmul_rejects_noncanonical_roles() {
1894 // Non-canonical role configs are a SURFACED honest miss (typed decline at
1895 // registration), never a crash.
1896 // (1) transposed lhs = [ContractedK, FreeM] = [3,1] instead of [1,3].
1897 let transposed = OpAttrs {
1898 lhs_roles: vec![3, 1],
1899 rhs_roles: vec![3, 2],
1900 ..OpAttrs::default()
1901 };
1902 assert_eq!(super::tag_to_op(OpTag::MatMul, &transposed), None);
1903 // (2) multi-ContractedK on lhs = [3,3].
1904 let multi_k = OpAttrs {
1905 lhs_roles: vec![3, 3],
1906 rhs_roles: vec![3, 2],
1907 ..OpAttrs::default()
1908 };
1909 assert_eq!(super::tag_to_op(OpTag::MatMul, &multi_k), None);
1910 // (3) FreeN-before-K on rhs = [FreeN, ContractedK] = [2,3] instead of [3,2].
1911 let freen_before_k = OpAttrs {
1912 lhs_roles: vec![1, 3],
1913 rhs_roles: vec![2, 3],
1914 ..OpAttrs::default()
1915 };
1916 assert_eq!(super::tag_to_op(OpTag::MatMul, &freen_before_k), None);
1917 }
1918
1919 #[test]
1920 fn tag_to_op_reconstructs_max_dim() {
1921 // T4 (Increment C slice 1): OpTag::MaxDim → Op::MaxDim(axis), the
1922 // axis riding `attrs.axis` exactly like SumDim/MeanDim.
1923 assert!(matches!(
1924 super::tag_to_op(
1925 OpTag::MaxDim,
1926 &OpAttrs {
1927 axis: Some(1),
1928 ..OpAttrs::default()
1929 }
1930 ),
1931 Some(Op::MaxDim(1))
1932 ));
1933 // An unset axis is an honest miss (typed decline at registration),
1934 // never a defaulted axis.
1935 assert_eq!(super::tag_to_op(OpTag::MaxDim, &OpAttrs::default()), None);
1936 // Not a scalar-param op: zero scalar slots.
1937 assert_eq!(super::scalar_slot_arity(OpTag::MaxDim), 0);
1938 }
1939
1940 #[test]
1941 fn max_dim_axis_last_resolves_to_rank_minus_one() {
1942 // D3 consumer: migrated recipes spell keepdim as MaxDim(axis_last) +
1943 // Unsqueeze(append), so MaxDim must be an `axis`-carrier tag for the
1944 // rel-attr resolver (rank − 1 via the shared §6.20 LAST resolver).
1945 let attrs = OpAttrs {
1946 axis_last: true,
1947 ..OpAttrs::default()
1948 };
1949 let resolved =
1950 super::resolve_rel_attrs(OpTag::MaxDim, &attrs, &[vec![2, 3, 4]], &[vec![2, 3, 4]])
1951 .expect("axis_last must resolve on MaxDim, not AxisLastUnsupported");
1952 assert_eq!(resolved.axis, Some(2), "rank-3 operand → LAST = axis 2");
1953 assert!(
1954 !resolved.axis_last,
1955 "rel carrier must be cleared post-resolve"
1956 );
1957 }
1958
1959 #[test]
1960 fn tag_to_op_still_rejects_basis_gap_and_scan() {
1961 // qmatmul/conv flow through Op::Fused (no OpTag); Scan is higher-order.
1962 assert_eq!(
1963 super::tag_to_op(OpTag::Iota, &OpAttrs::default()),
1964 None,
1965 "Iota needs a len (target_shape) — empty attrs is a miss"
1966 );
1967 }
1968
1969 // ---- C-T2 (Increment C): the nested-fused re-emit carrier (mechanism 2a) --
1970
1971 /// `OpTag::Fused` reconstructs the param-less nested fused ops from the
1972 /// `fused_op` selector (name → fid → params), and honest-misses anything
1973 /// else. Born-red before the C-T2 `tag_to_op` arm: `Fused` falls to the
1974 /// `_ => return None` catch-all, so every `Some(..)` assertion below fails.
1975 #[test]
1976 fn tag_to_op_reconstructs_nested_fused() {
1977 use crate::registry::{FusedOpParams, FusedOps};
1978 let sm = OpAttrs {
1979 fused_op: Some("SoftmaxLastDim".into()),
1980 ..OpAttrs::default()
1981 };
1982 assert_eq!(
1983 super::tag_to_op(OpTag::Fused, &sm),
1984 Some(Op::Fused(
1985 FusedOps::SOFTMAX_LAST_DIM,
1986 FusedOpParams::SoftmaxLastDim
1987 )),
1988 "SoftmaxLastDim selector reconstructs the param-less nested fused op",
1989 );
1990 let smb = OpAttrs {
1991 fused_op: Some("SoftmaxLastDimBackward".into()),
1992 ..OpAttrs::default()
1993 };
1994 assert_eq!(
1995 super::tag_to_op(OpTag::Fused, &smb),
1996 Some(Op::Fused(
1997 FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
1998 FusedOpParams::SoftmaxLastDimBackward,
1999 )),
2000 "SoftmaxLastDimBackward selector reconstructs its param-less nested fused op",
2001 );
2002 // Honest misses (never a crash): unset selector, unknown name, and a
2003 // param-carrying fused op (not round-trippable through the name-only
2004 // selector — its params can't be reconstructed from the map).
2005 assert_eq!(
2006 super::tag_to_op(OpTag::Fused, &OpAttrs::default()),
2007 None,
2008 "an unset fused_op selector is an honest miss",
2009 );
2010 let bogus = OpAttrs {
2011 fused_op: Some("NotARealFusedOp".into()),
2012 ..OpAttrs::default()
2013 };
2014 assert_eq!(
2015 super::tag_to_op(OpTag::Fused, &bogus),
2016 None,
2017 "an unknown name is an honest miss"
2018 );
2019 let paged = OpAttrs {
2020 fused_op: Some("PagedAttn".into()),
2021 ..OpAttrs::default()
2022 };
2023 assert_eq!(
2024 super::tag_to_op(OpTag::Fused, &paged),
2025 None,
2026 "a registered but param-carrying fused op is an honest miss (not in the fid->params map)",
2027 );
2028 }
2029
2030 /// A hand-built recipe carrying a nested `Op::Fused(SOFTMAX_LAST_DIM)` node
2031 /// validates as re-emittable and round-trips through `emit`: the emitted node
2032 /// is the real nested fused op, with shape/dtype from the registry entry's
2033 /// `shape_rule`/`dtype_rule` (softmax passes both through) since
2034 /// `primitive_shape` honest-misses `Fused`. Born-red before the C-T2 arms:
2035 /// `validate_recipe` rejects the region `UnRepresentable`.
2036 #[test]
2037 fn emit_reconstructs_nested_fused_shape_and_dtype() {
2038 use crate::registry::{FusedOpParams, FusedOps};
2039 use fuel_ir::{DType, Shape};
2040 let region = PatternNode::Op {
2041 op: OpTag::Fused,
2042 attrs: OpAttrs {
2043 fused_op: Some("SoftmaxLastDim".into()),
2044 ..OpAttrs::default()
2045 },
2046 operands: vec![PatternNode::Bind { index: 0 }],
2047 };
2048 assert!(
2049 super::validate_recipe(®ion).is_ok(),
2050 "the nested-fused recipe validates as re-emittable (C-T2 tag_to_op arm)",
2051 );
2052 let mut g = Graph::new();
2053 let x = g.push(Node {
2054 op: Op::Const,
2055 inputs: vec![],
2056 shape: Shape::from_dims(&[2, 3, 4]),
2057 dtype: DType::F32,
2058 });
2059 let root = emit_region(&mut g, ®ion, &[x], &[]);
2060 assert_eq!(
2061 g.node(root).op,
2062 Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim),
2063 "the nested fused node is reconstructed as-is (2a)",
2064 );
2065 assert_eq!(
2066 g.node(root).shape,
2067 Shape::from_dims(&[2, 3, 4]),
2068 "shape from softmax shape_rule (passthrough), NOT primitive_shape (which misses Fused)",
2069 );
2070 assert_eq!(
2071 g.node(root).dtype,
2072 DType::F32,
2073 "dtype from softmax dtype_rule (passthrough)"
2074 );
2075 assert_eq!(
2076 g.node(root).inputs,
2077 vec![x],
2078 "the fused node's single operand is the bound input"
2079 );
2080 }
2081
2082 #[test]
2083 fn validate_representable_now_accepts_a_slice_region() {
2084 // Region: Concat{0}(Neg(Slice{...}(bind0)), bind0) — the rope rotate-half shape.
2085 let region = PatternNode::Op {
2086 op: OpTag::Concat,
2087 attrs: OpAttrs {
2088 axis: Some(0),
2089 ..OpAttrs::default()
2090 },
2091 operands: vec![
2092 PatternNode::Op {
2093 op: OpTag::Neg,
2094 attrs: OpAttrs::default(),
2095 operands: vec![PatternNode::Op {
2096 op: OpTag::Slice,
2097 attrs: OpAttrs {
2098 axis: Some(0),
2099 slice_start: Some(0),
2100 slice_len: Some(1),
2101 ..OpAttrs::default()
2102 },
2103 operands: vec![PatternNode::Bind { index: 0 }],
2104 }],
2105 },
2106 PatternNode::Bind { index: 0 },
2107 ],
2108 };
2109 assert!(
2110 super::validate_representable(®ion).is_ok(),
2111 "slice/concat region must now validate"
2112 );
2113 }
2114
2115 #[test]
2116 fn emit_gets_shape_right_for_a_reduction_region() {
2117 use fuel_ir::{DType, Shape};
2118 // Region: ReduceSumTo([2,1])(bind0). Input [2,5] → output [2,1].
2119 let region = PatternNode::Op {
2120 op: OpTag::ReduceSumTo,
2121 attrs: OpAttrs {
2122 target_shape: vec![2, 1],
2123 ..OpAttrs::default()
2124 },
2125 operands: vec![PatternNode::Bind { index: 0 }],
2126 };
2127 let mut g = Graph::new();
2128 let x = g.push(Node {
2129 op: Op::Const,
2130 inputs: vec![],
2131 shape: Shape::from_dims(&[2, 5]),
2132 dtype: DType::F32,
2133 });
2134 let root = emit_region(&mut g, ®ion, &[x], &[]);
2135 assert!(matches!(g.node(root).op, Op::ReduceSumTo(_)));
2136 assert_eq!(
2137 g.node(root).shape,
2138 Shape::from_dims(&[2, 1]),
2139 "emit must use the reduced shape, not operand[0]"
2140 );
2141 assert_eq!(g.node(root).dtype, DType::F32);
2142 }
2143
2144 #[test]
2145 fn emit_gets_dtype_right_for_a_cast_region() {
2146 use fuel_ir::{DType, Shape};
2147 // Region: Cast(F16)(bind0). Input F32 → output F16, same shape.
2148 let region = PatternNode::Op {
2149 op: OpTag::Cast,
2150 attrs: OpAttrs {
2151 cast_dtype: Some("f16".into()),
2152 ..OpAttrs::default()
2153 },
2154 operands: vec![PatternNode::Bind { index: 0 }],
2155 };
2156 let mut g = Graph::new();
2157 let x = g.push(Node {
2158 op: Op::Const,
2159 inputs: vec![],
2160 shape: Shape::from_dims(&[3, 3]),
2161 dtype: DType::F32,
2162 });
2163 let root = emit_region(&mut g, ®ion, &[x], &[]);
2164 assert!(matches!(g.node(root).op, Op::Cast(DType::F16)));
2165 assert_eq!(
2166 g.node(root).dtype,
2167 DType::F16,
2168 "emit must take Cast's target dtype, not operand[0]'s"
2169 );
2170 assert_eq!(g.node(root).shape, Shape::from_dims(&[3, 3]));
2171 }
2172
2173 #[test]
2174 fn emit_zero_operand_representable_region_is_panic_free() {
2175 // M-1 never-panic hardening: a MALFORMED region — a binary op given ZERO
2176 // operands. `validate_representable` accepts it (it checks
2177 // `tag_to_op(op).is_some()`, NOT arity), and `emit_region` is a public
2178 // raw-region entry (candidate-kernel verification) that does not
2179 // re-validate. `primitive_shape(Add, [], [])` errs, so the fallback runs
2180 // with an EMPTY child_shapes — it must NOT index-panic. emit stays total:
2181 // it returns a node (with a degenerate rank-0 shape), never a panic.
2182 let region = PatternNode::Op {
2183 op: OpTag::Add,
2184 attrs: OpAttrs::default(),
2185 operands: vec![],
2186 };
2187 let mut g = Graph::new();
2188 let root = emit_region(&mut g, ®ion, &[], &[]);
2189 assert!(
2190 matches!(g.node(root).op, Op::Add),
2191 "emit returns a node, not a panic"
2192 );
2193 }
2194
2195 #[test]
2196 fn decompose_region_re_emits_relu_add() {
2197 let id = register_runtime_fused("test::relu_add::decompose", relu_add_region()).unwrap();
2198 let mut g = Graph::new();
2199 let s = Shape::from_dims(&[4]);
2200 let a = g.push(Node {
2201 op: Op::Const,
2202 inputs: vec![],
2203 shape: s.clone(),
2204 dtype: DType::F32,
2205 });
2206 let b = g.push(Node {
2207 op: Op::Const,
2208 inputs: vec![],
2209 shape: s.clone(),
2210 dtype: DType::F32,
2211 });
2212 let fused = g.push(Node {
2213 op: Op::Fused(id, FusedOpParams::Runtime { scalars: vec![] }),
2214 inputs: vec![a, b],
2215 shape: s.clone(),
2216 dtype: DType::F32,
2217 });
2218
2219 let root = decompose_region(&mut g, fused);
2220
2221 // The re-emitted sink is Relu over Add(a, b) — the region, on primitives.
2222 assert!(matches!(g.node(root).op, Op::Relu));
2223 let add_id = g.node(root).inputs[0];
2224 assert!(matches!(g.node(add_id).op, Op::Add));
2225 assert_eq!(g.node(add_id).inputs, vec![a, b]);
2226 // Shapes propagated from the leaves (same-shape elementwise).
2227 assert_eq!(g.node(root).shape, s);
2228 assert_eq!(g.node(add_id).shape, s);
2229 }
2230
2231 // ---- scalar slots (the `extract:` substitution) ---------------------
2232
2233 /// tanh(mul_scalar(a)) with the scalar left OPEN (a slot template).
2234 fn tanh_mul_scalar_slot_region() -> PatternNode {
2235 PatternNode::Op {
2236 op: OpTag::Tanh,
2237 attrs: OpAttrs::default(),
2238 operands: vec![PatternNode::Op {
2239 op: OpTag::MulScalar,
2240 attrs: OpAttrs::default(), // empty scalars = an open slot
2241 operands: vec![PatternNode::Bind { index: 0 }],
2242 }],
2243 }
2244 }
2245
2246 #[test]
2247 fn slot_template_registers_and_counts() {
2248 // Born-red before slot support: validation rejected an AddScalar/
2249 // MulScalar pattern node with no baked value.
2250 let id =
2251 register_runtime_fused("test::tanh_mul_scalar::slot", tanh_mul_scalar_slot_region())
2252 .expect("a slot template is registrable");
2253 let region = runtime_region(id).unwrap();
2254 assert_eq!(count_scalar_slots(®ion), 1, "one open slot");
2255 }
2256
2257 #[test]
2258 fn decompose_fills_slots_from_the_node_scalars() {
2259 let id =
2260 register_runtime_fused("test::tanh_mul_scalar::fill", tanh_mul_scalar_slot_region())
2261 .unwrap();
2262 let mut g = Graph::new();
2263 let s = Shape::from_dims(&[4]);
2264 let a = g.push(Node {
2265 op: Op::Const,
2266 inputs: vec![],
2267 shape: s.clone(),
2268 dtype: DType::F32,
2269 });
2270 let fused = g.push(Node {
2271 op: Op::Fused(id, FusedOpParams::Runtime { scalars: vec![2.5] }),
2272 inputs: vec![a],
2273 shape: s.clone(),
2274 dtype: DType::F32,
2275 });
2276
2277 let root = decompose_region(&mut g, fused);
2278
2279 // tanh(mul_scalar(a, 2.5)) — the LIVE value filled the slot.
2280 assert!(matches!(g.node(root).op, Op::Tanh));
2281 let ms = g.node(root).inputs[0];
2282 assert!(
2283 matches!(g.node(ms).op, Op::MulScalar(v) if v == 2.5),
2284 "slot filled with the node's live scalar, got {:?}",
2285 g.node(ms).op,
2286 );
2287 assert_eq!(g.node(ms).inputs, vec![a]);
2288 }
2289
2290 #[test]
2291 fn decompose_slot_count_mismatch_is_a_fixpoint_not_a_crash() {
2292 let id = register_runtime_fused(
2293 "test::tanh_mul_scalar::mismatch",
2294 tanh_mul_scalar_slot_region(),
2295 )
2296 .unwrap();
2297 let mut g = Graph::new();
2298 let s = Shape::from_dims(&[4]);
2299 let a = g.push(Node {
2300 op: Op::Const,
2301 inputs: vec![],
2302 shape: s.clone(),
2303 dtype: DType::F32,
2304 });
2305 // One slot, but the node carries NO scalars — malformed; must be a
2306 // no-op fixpoint (G2), never a panic.
2307 let fused = g.push(Node {
2308 op: Op::Fused(id, FusedOpParams::Runtime { scalars: vec![] }),
2309 inputs: vec![a],
2310 shape: s.clone(),
2311 dtype: DType::F32,
2312 });
2313 assert_eq!(
2314 decompose_region(&mut g, fused),
2315 fused,
2316 "mismatch ⇒ fixpoint"
2317 );
2318 }
2319
2320 // ---- shape-derived scalar slot (A1: the reduced_count live-emission) -----
2321
2322 /// `mul_scalar(a)` whose scalar is a SHAPE-DERIVED value: `scalar_rel =
2323 /// Extent(operand 0, LAST)` — the reduced_count of bind 0's last axis, filled
2324 /// from the input shape at emit time (NOT from the params cursor). This is
2325 /// the narrow carrier a norm-backward's `MulScalar(n = dims[last])` needs.
2326 fn mul_scalar_reduced_count_region() -> PatternNode {
2327 use fuel_kernel_seam_types::shape_expr::{Dim, LAST};
2328 PatternNode::Op {
2329 op: OpTag::MulScalar,
2330 attrs: OpAttrs {
2331 scalar_rel: Some(Dim::Extent {
2332 operand: 0,
2333 axis: LAST,
2334 }),
2335 ..OpAttrs::default()
2336 },
2337 operands: vec![PatternNode::Bind { index: 0 }],
2338 }
2339 }
2340
2341 /// A1 born-red: a `scalar_rel` node is filled from an input SHAPE, so it is
2342 /// NOT a params-cursor slot (`count_scalar_slots == 0`) and its emitted value
2343 /// is the reduced_count of the resolving input's last axis. RED before the
2344 /// wiring: `count_scalar_slots` counts the empty-`scalars` MulScalar as one
2345 /// open slot, and `emit` ignores `scalar_rel` (fills from the cursor).
2346 #[test]
2347 fn scalar_rel_is_shape_derived_not_a_cursor_slot() {
2348 let region = mul_scalar_reduced_count_region();
2349 assert_eq!(
2350 count_scalar_slots(®ion),
2351 0,
2352 "a scalar_rel node is shape-derived, never a params-cursor slot",
2353 );
2354 // Emit over x[3,5]: n = last extent = 5, with NO cursor scalars.
2355 let mut g = Graph::new();
2356 let x = g.push(Node {
2357 op: Op::Const,
2358 inputs: vec![],
2359 shape: Shape::from_dims(&[3, 5]),
2360 dtype: DType::F32,
2361 });
2362 let root = try_emit_region(&mut g, ®ion, &[x], &[])
2363 .expect("scalar_rel resolves against the input shape");
2364 assert!(
2365 matches!(g.node(root).op, Op::MulScalar(v) if v == 5.0),
2366 "scalar_rel = Extent(0, LAST) over x[3,5] ⇒ MulScalar(5.0), got {:?}",
2367 g.node(root).op,
2368 );
2369 assert_eq!(g.node(root).inputs, vec![x]);
2370 // Rank-polymorphic: the SAME datum resolves to 7 over x[..,7].
2371 let mut g3 = Graph::new();
2372 let x3 = g3.push(Node {
2373 op: Op::Const,
2374 inputs: vec![],
2375 shape: Shape::from_dims(&[2, 4, 7]),
2376 dtype: DType::F32,
2377 });
2378 let root3 = try_emit_region(&mut g3, ®ion, &[x3], &[]).unwrap();
2379 assert!(
2380 matches!(g3.node(root3).op, Op::MulScalar(v) if v == 7.0),
2381 "same datum, x[2,4,7] ⇒ MulScalar(7.0), got {:?}",
2382 g3.node(root3).op,
2383 );
2384 }
2385
2386 /// A1: a `scalar_rel` set together with a non-empty `scalars` is a typed
2387 /// rel-XOR-abs authoring conflict at emit (never a silent precedence).
2388 #[test]
2389 fn scalar_rel_and_baked_scalar_together_is_a_typed_conflict() {
2390 use fuel_kernel_seam_types::shape_expr::{Dim, LAST};
2391 let region = PatternNode::Op {
2392 op: OpTag::MulScalar,
2393 attrs: OpAttrs {
2394 scalars: vec![2.0],
2395 scalar_rel: Some(Dim::Extent {
2396 operand: 0,
2397 axis: LAST,
2398 }),
2399 ..OpAttrs::default()
2400 },
2401 operands: vec![PatternNode::Bind { index: 0 }],
2402 };
2403 let mut g = Graph::new();
2404 let x = g.push(Node {
2405 op: Op::Const,
2406 inputs: vec![],
2407 shape: Shape::from_dims(&[3, 5]),
2408 dtype: DType::F32,
2409 });
2410 assert_eq!(
2411 try_emit_region(&mut g, ®ion, &[x], &[]),
2412 Err(RelAttrError::RelAbsConflict { field: "scalars" }),
2413 "scalar_rel XOR scalars — both set is a typed decline",
2414 );
2415 }
2416
2417 // ---- Task 5: byte-for-byte emit == registry::*::decompose parity --------
2418 //
2419 // The A.4 acceptance gate: express each hand-written decompose as a
2420 // PatternNode region, re-emit it via the grown `emit`, and assert the
2421 // result is structurally identical (op + shape + dtype at every node) to
2422 // the decompose-oracle output — the migration oracle. Since T5, `emit`
2423 // identity-shares repeated slot-free subtrees within one call, so a
2424 // shared oracle node compares against an equally-shared emitted node;
2425 // `assert_structural_eq` is recursive + order-sensitive (no commutative
2426 // canonicalization — stricter than `base_map_hash`), catching an
2427 // operand-swap the hash would mask.
2428
2429 fn op_node(op: OpTag, attrs: OpAttrs, operands: Vec<PatternNode>) -> PatternNode {
2430 PatternNode::Op {
2431 op,
2432 attrs,
2433 operands,
2434 }
2435 }
2436 fn bind(i: u8) -> PatternNode {
2437 PatternNode::Bind { index: i }
2438 }
2439
2440 /// Recursively assert two subgraphs are identical: same Op, shape, dtype,
2441 /// arity, and recursively-equal inputs. Shared leaves (same NodeId) match
2442 /// by identity. This is the "byte-for-byte" node-structure check.
2443 fn assert_structural_eq(g: &Graph, a: NodeId, b: NodeId) {
2444 if a == b {
2445 return; // shared leaf (bound external input)
2446 }
2447 let na = g.node(a);
2448 let nb = g.node(b);
2449 assert_eq!(na.op, nb.op, "op mismatch: {:?} vs {:?}", na.op, nb.op);
2450 assert_eq!(
2451 na.shape, nb.shape,
2452 "shape mismatch at {:?} vs {:?}",
2453 na.op, nb.op
2454 );
2455 assert_eq!(na.dtype, nb.dtype, "dtype mismatch at {:?}", na.op);
2456 assert_eq!(
2457 na.inputs.len(),
2458 nb.inputs.len(),
2459 "arity mismatch at {:?}",
2460 na.op
2461 );
2462 for (&ia, &ib) in na.inputs.iter().zip(nb.inputs.iter()) {
2463 assert_structural_eq(g, ia, ib);
2464 }
2465 }
2466
2467 /// FROZEN copy of the pre-migration imperative
2468 /// `registry::softmax_last_dim::decompose` (the legacy 7-node
2469 /// `ReduceMaxTo`/`ReduceSumTo` keepdim spelling), copied VERBATIM from
2470 /// that module @ `af4b7dd4` before T5 replaced the live body with the
2471 /// data recipe. Two consumers: the T5 numeric-parity oracle
2472 /// (`recipe_bridge` below) and `emit_matches_softmax_last_dim_decompose`
2473 /// (whose oracle was repointed here — the live decompose no longer emits
2474 /// this spelling).
2475 fn frozen_legacy_softmax_decompose(
2476 graph: &mut Graph,
2477 id: NodeId,
2478 _params: &FusedOpParams,
2479 ) -> NodeId {
2480 let (x_id, x_shape, dtype) = {
2481 let n = graph.node(id);
2482 (n.inputs[0], n.shape.clone(), n.dtype)
2483 };
2484 let dims = x_shape.dims().to_vec();
2485 let rank = dims.len();
2486 let last = rank - 1;
2487
2488 let mut keepdim_dims = dims.clone();
2489 keepdim_dims[last] = 1;
2490 let keepdim_shape = Shape::from_dims(&keepdim_dims);
2491
2492 let m_id = graph.push(Node {
2493 op: Op::ReduceMaxTo(keepdim_shape.clone()),
2494 inputs: vec![x_id],
2495 shape: keepdim_shape.clone(),
2496 dtype,
2497 });
2498 let mb_id = graph.push(Node {
2499 op: Op::BroadcastTo(x_shape.clone()),
2500 inputs: vec![m_id],
2501 shape: x_shape.clone(),
2502 dtype,
2503 });
2504 let s_id = graph.push(Node {
2505 op: Op::Sub,
2506 inputs: vec![x_id, mb_id],
2507 shape: x_shape.clone(),
2508 dtype,
2509 });
2510 let e_id = graph.push(Node {
2511 op: Op::Exp,
2512 inputs: vec![s_id],
2513 shape: x_shape.clone(),
2514 dtype,
2515 });
2516 let d_id = graph.push(Node {
2517 op: Op::ReduceSumTo(keepdim_shape.clone()),
2518 inputs: vec![e_id],
2519 shape: keepdim_shape,
2520 dtype,
2521 });
2522 let db_id = graph.push(Node {
2523 op: Op::BroadcastTo(x_shape.clone()),
2524 inputs: vec![d_id],
2525 shape: x_shape.clone(),
2526 dtype,
2527 });
2528
2529 graph.push(Node {
2530 op: Op::Div,
2531 inputs: vec![e_id, db_id],
2532 shape: x_shape,
2533 dtype,
2534 })
2535 }
2536
2537 /// FROZEN copy of the pre-migration imperative
2538 /// `registry::rope::decompose` (the legacy 11-node spelling with two
2539 /// leading-1-padded `Reshape` prep nodes), copied VERBATIM from that
2540 /// module @ `af4b7dd4` before T6 replaced the live body with the data
2541 /// recipe. Two consumers: `emit_matches_rope_decompose` (whose oracle was
2542 /// repointed here — the live decompose now emits the recipe, which at
2543 /// EQUAL rank elides the no-op prep `Reshape`, D4) and the T6
2544 /// numeric/structural parity oracle (`rope_recipe` below).
2545 fn frozen_legacy_rope_decompose(
2546 graph: &mut Graph,
2547 id: NodeId,
2548 _params: &FusedOpParams,
2549 ) -> NodeId {
2550 let (x_id, cos_id, sin_id, x_shape, dtype) = {
2551 let n = graph.node(id);
2552 (
2553 n.inputs[0],
2554 n.inputs[1],
2555 n.inputs[2],
2556 n.shape.clone(),
2557 n.dtype,
2558 )
2559 };
2560 let dims = x_shape.dims().to_vec();
2561 let rank = dims.len();
2562 let seq = dims[rank - 2];
2563 let d = dims[rank - 1];
2564 let half = d / 2;
2565 let last = rank - 1;
2566
2567 let mut broadcast_shape_dims: Vec<usize> = vec![1usize; rank];
2568 broadcast_shape_dims[rank - 2] = seq;
2569 broadcast_shape_dims[rank - 1] = d;
2570 let broadcast_shape = Shape::from_dims(&broadcast_shape_dims);
2571
2572 let cos_reshaped_id = graph.push(Node {
2573 op: Op::Reshape(broadcast_shape.clone()),
2574 inputs: vec![cos_id],
2575 shape: broadcast_shape.clone(),
2576 dtype,
2577 });
2578 let sin_reshaped_id = graph.push(Node {
2579 op: Op::Reshape(broadcast_shape.clone()),
2580 inputs: vec![sin_id],
2581 shape: broadcast_shape,
2582 dtype,
2583 });
2584 let cos_bcast_id = graph.push(Node {
2585 op: Op::BroadcastTo(x_shape.clone()),
2586 inputs: vec![cos_reshaped_id],
2587 shape: x_shape.clone(),
2588 dtype,
2589 });
2590 let sin_bcast_id = graph.push(Node {
2591 op: Op::BroadcastTo(x_shape.clone()),
2592 inputs: vec![sin_reshaped_id],
2593 shape: x_shape.clone(),
2594 dtype,
2595 });
2596
2597 let mut half_dims = dims.clone();
2598 half_dims[last] = half;
2599 let half_shape = Shape::from_dims(&half_dims);
2600
2601 let first_half_id = graph.push(Node {
2602 op: Op::Slice {
2603 dim: last,
2604 start: 0,
2605 len: half,
2606 },
2607 inputs: vec![x_id],
2608 shape: half_shape.clone(),
2609 dtype,
2610 });
2611 let second_half_id = graph.push(Node {
2612 op: Op::Slice {
2613 dim: last,
2614 start: half,
2615 len: half,
2616 },
2617 inputs: vec![x_id],
2618 shape: half_shape.clone(),
2619 dtype,
2620 });
2621 let neg_second_id = graph.push(Node {
2622 op: Op::Neg,
2623 inputs: vec![second_half_id],
2624 shape: half_shape,
2625 dtype,
2626 });
2627 let rotated_half_id = graph.push(Node {
2628 op: Op::Concat { dim: last },
2629 inputs: vec![neg_second_id, first_half_id],
2630 shape: x_shape.clone(),
2631 dtype,
2632 });
2633
2634 let left_id = graph.push(Node {
2635 op: Op::Mul,
2636 inputs: vec![x_id, cos_bcast_id],
2637 shape: x_shape.clone(),
2638 dtype,
2639 });
2640 let right_id = graph.push(Node {
2641 op: Op::Mul,
2642 inputs: vec![rotated_half_id, sin_bcast_id],
2643 shape: x_shape.clone(),
2644 dtype,
2645 });
2646
2647 graph.push(Node {
2648 op: Op::Add,
2649 inputs: vec![left_id, right_id],
2650 shape: x_shape,
2651 dtype,
2652 })
2653 }
2654
2655 /// FROZEN copy of the pre-migration imperative
2656 /// `registry::rms_norm_last_dim::decompose` (the legacy 7-node
2657 /// `MeanDim → Reshape(keepdim) → AddScalar(eps)` spelling), copied VERBATIM
2658 /// from that module before T7 replaced the live body with the data recipe.
2659 /// Consumer: the T7 numeric-parity oracle (`norm_recipe` below). The live
2660 /// decompose now emits the D3 shrink-via-swap spelling (`Unsqueeze` append
2661 /// in place of `Reshape(keepdim)`), so the parity test evaluates BOTH
2662 /// through the shared reference interpreter and asserts bit-exact
2663 /// equivalence (the swap is metadata-only).
2664 fn frozen_legacy_rms_norm_decompose(
2665 graph: &mut Graph,
2666 id: NodeId,
2667 params: &FusedOpParams,
2668 ) -> NodeId {
2669 let (x_id, x_shape, dtype) = {
2670 let n = graph.node(id);
2671 (n.inputs[0], n.shape.clone(), n.dtype)
2672 };
2673 let eps = match params {
2674 FusedOpParams::RmsNormLastDim { eps } => *eps,
2675 _ => return id,
2676 };
2677 let dims = x_shape.dims().to_vec();
2678 let rank = dims.len();
2679 let last = rank - 1;
2680
2681 let mut keepdim_dims = dims.clone();
2682 keepdim_dims[last] = 1;
2683 let keepdim_shape = Shape::from_dims(&keepdim_dims);
2684 let mut reduced_dims = dims.clone();
2685 reduced_dims.remove(last);
2686 let reduced_shape = Shape::from_dims(&reduced_dims);
2687
2688 let sq_id = graph.push(Node {
2689 op: Op::Sqr,
2690 inputs: vec![x_id],
2691 shape: x_shape.clone(),
2692 dtype,
2693 });
2694 let mean_id = graph.push(Node {
2695 op: Op::MeanDim(last),
2696 inputs: vec![sq_id],
2697 shape: reduced_shape,
2698 dtype,
2699 });
2700 let mean_kd_id = graph.push(Node {
2701 op: Op::Reshape(keepdim_shape.clone()),
2702 inputs: vec![mean_id],
2703 shape: keepdim_shape.clone(),
2704 dtype,
2705 });
2706 let denom_sq_id = graph.push(Node {
2707 op: Op::AddScalar(eps),
2708 inputs: vec![mean_kd_id],
2709 shape: keepdim_shape.clone(),
2710 dtype,
2711 });
2712 let denom_id = graph.push(Node {
2713 op: Op::Sqrt,
2714 inputs: vec![denom_sq_id],
2715 shape: keepdim_shape,
2716 dtype,
2717 });
2718 let denom_bcast_id = graph.push(Node {
2719 op: Op::BroadcastTo(x_shape.clone()),
2720 inputs: vec![denom_id],
2721 shape: x_shape.clone(),
2722 dtype,
2723 });
2724 graph.push(Node {
2725 op: Op::Div,
2726 inputs: vec![x_id, denom_bcast_id],
2727 shape: x_shape,
2728 dtype,
2729 })
2730 }
2731
2732 /// FROZEN copy of the pre-migration imperative
2733 /// `registry::layer_norm_last_dim::decompose` (the legacy 11-node spelling
2734 /// with two `Reshape(keepdim)` restores and the `centered` subterm shared
2735 /// between `Sqr` and the final `Div`), copied VERBATIM from that module
2736 /// before T7 replaced the live body with the data recipe. Two consumers:
2737 /// `emit_matches_layer_norm_last_dim_decompose` (whose oracle was repointed
2738 /// here — the live decompose now emits the `Unsqueeze` D3 spelling) and the
2739 /// T7 numeric-parity oracle (`norm_recipe` below).
2740 fn frozen_legacy_layer_norm_decompose(
2741 graph: &mut Graph,
2742 id: NodeId,
2743 params: &FusedOpParams,
2744 ) -> NodeId {
2745 let (x_id, x_shape, dtype) = {
2746 let n = graph.node(id);
2747 (n.inputs[0], n.shape.clone(), n.dtype)
2748 };
2749 let eps = match params {
2750 FusedOpParams::LayerNormLastDim { eps } => *eps,
2751 _ => return id,
2752 };
2753 let dims = x_shape.dims().to_vec();
2754 let rank = dims.len();
2755 let last = rank - 1;
2756
2757 let mut keepdim_dims = dims.clone();
2758 keepdim_dims[last] = 1;
2759 let keepdim_shape = Shape::from_dims(&keepdim_dims);
2760 let mut reduced_dims = dims.clone();
2761 reduced_dims.remove(last);
2762 let reduced_shape = Shape::from_dims(&reduced_dims);
2763
2764 let mean_id = graph.push(Node {
2765 op: Op::MeanDim(last),
2766 inputs: vec![x_id],
2767 shape: reduced_shape.clone(),
2768 dtype,
2769 });
2770 let mean_kd_id = graph.push(Node {
2771 op: Op::Reshape(keepdim_shape.clone()),
2772 inputs: vec![mean_id],
2773 shape: keepdim_shape.clone(),
2774 dtype,
2775 });
2776 let mean_bcast_id = graph.push(Node {
2777 op: Op::BroadcastTo(x_shape.clone()),
2778 inputs: vec![mean_kd_id],
2779 shape: x_shape.clone(),
2780 dtype,
2781 });
2782 let centered_id = graph.push(Node {
2783 op: Op::Sub,
2784 inputs: vec![x_id, mean_bcast_id],
2785 shape: x_shape.clone(),
2786 dtype,
2787 });
2788 let centered_sq_id = graph.push(Node {
2789 op: Op::Sqr,
2790 inputs: vec![centered_id],
2791 shape: x_shape.clone(),
2792 dtype,
2793 });
2794 let var_id = graph.push(Node {
2795 op: Op::MeanDim(last),
2796 inputs: vec![centered_sq_id],
2797 shape: reduced_shape,
2798 dtype,
2799 });
2800 let var_kd_id = graph.push(Node {
2801 op: Op::Reshape(keepdim_shape.clone()),
2802 inputs: vec![var_id],
2803 shape: keepdim_shape.clone(),
2804 dtype,
2805 });
2806 let var_eps_id = graph.push(Node {
2807 op: Op::AddScalar(eps),
2808 inputs: vec![var_kd_id],
2809 shape: keepdim_shape.clone(),
2810 dtype,
2811 });
2812 let denom_id = graph.push(Node {
2813 op: Op::Sqrt,
2814 inputs: vec![var_eps_id],
2815 shape: keepdim_shape,
2816 dtype,
2817 });
2818 let denom_bcast_id = graph.push(Node {
2819 op: Op::BroadcastTo(x_shape.clone()),
2820 inputs: vec![denom_id],
2821 shape: x_shape.clone(),
2822 dtype,
2823 });
2824 graph.push(Node {
2825 op: Op::Div,
2826 inputs: vec![centered_id, denom_bcast_id],
2827 shape: x_shape,
2828 dtype,
2829 })
2830 }
2831
2832 /// FROZEN copy of the pre-migration imperative
2833 /// `registry::layer_norm_last_dim_backward::decompose` (the legacy ~20-node
2834 /// `MeanDim`/`Reshape(keepdim)`/`BroadcastTo` recompute of
2835 /// `grad_x = istd · (g − mean(g) − xhat·mean(g·xhat))`), copied VERBATIM
2836 /// from that module @ `b967bdb1` before slice-2 replaced the live body with
2837 /// the data recipe. Sole consumer: the slice-2 numeric-parity oracle
2838 /// (`layer_norm_backward_recipe` below). Reads `inputs[0] = x`, `inputs[1] =
2839 /// g` (the upstream gradient) off the node — the order the autograd
2840 /// `BackwardKind::Fused(LAYER_NORM_LAST_DIM_BACKWARD)` edge emits.
2841 fn frozen_legacy_layer_norm_backward_decompose(
2842 graph: &mut Graph,
2843 id: NodeId,
2844 params: &FusedOpParams,
2845 ) -> NodeId {
2846 let (x_id, g_id, x_shape, dtype) = {
2847 let n = graph.node(id);
2848 (n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
2849 };
2850 let eps = match params {
2851 FusedOpParams::LayerNormLastDimBackward { eps } => *eps,
2852 // G2: total + never-panic — impossible params; return self.
2853 _ => return id,
2854 };
2855 let dims = x_shape.dims().to_vec();
2856 let last = dims.len() - 1;
2857 let mut kd = dims.clone();
2858 kd[last] = 1;
2859 let keepdim = Shape::from_dims(&kd);
2860 let mut rd = dims.clone();
2861 rd.remove(last);
2862 let reduced = Shape::from_dims(&rd);
2863
2864 // reduce-mean over the last dim, keepdim, broadcast back to x_shape.
2865 let mean_b = |graph: &mut Graph, src: NodeId| -> NodeId {
2866 let m = graph.push(Node {
2867 op: Op::MeanDim(last),
2868 inputs: vec![src],
2869 shape: reduced.clone(),
2870 dtype,
2871 });
2872 let m_kd = graph.push(Node {
2873 op: Op::Reshape(keepdim.clone()),
2874 inputs: vec![m],
2875 shape: keepdim.clone(),
2876 dtype,
2877 });
2878 graph.push(Node {
2879 op: Op::BroadcastTo(x_shape.clone()),
2880 inputs: vec![m_kd],
2881 shape: x_shape.clone(),
2882 dtype,
2883 })
2884 };
2885
2886 // xhat = (x − mean(x)) · istd ; istd = rsqrt(var + eps).
2887 let mean_x = mean_b(graph, x_id);
2888 let xc = graph.push(Node {
2889 op: Op::Sub,
2890 inputs: vec![x_id, mean_x],
2891 shape: x_shape.clone(),
2892 dtype,
2893 });
2894 let xc_sq = graph.push(Node {
2895 op: Op::Sqr,
2896 inputs: vec![xc],
2897 shape: x_shape.clone(),
2898 dtype,
2899 });
2900 let var = mean_b(graph, xc_sq);
2901 let var_eps = graph.push(Node {
2902 op: Op::AddScalar(eps),
2903 inputs: vec![var],
2904 shape: x_shape.clone(),
2905 dtype,
2906 });
2907 let istd = graph.push(Node {
2908 op: Op::Rsqrt,
2909 inputs: vec![var_eps],
2910 shape: x_shape.clone(),
2911 dtype,
2912 });
2913 let xhat = graph.push(Node {
2914 op: Op::Mul,
2915 inputs: vec![xc, istd],
2916 shape: x_shape.clone(),
2917 dtype,
2918 });
2919
2920 // grad_x = istd · (g − mean(g) − xhat·mean(g·xhat)).
2921 let mean_g = mean_b(graph, g_id);
2922 let g_xhat = graph.push(Node {
2923 op: Op::Mul,
2924 inputs: vec![g_id, xhat],
2925 shape: x_shape.clone(),
2926 dtype,
2927 });
2928 let mean_gxh = mean_b(graph, g_xhat);
2929 let t1 = graph.push(Node {
2930 op: Op::Sub,
2931 inputs: vec![g_id, mean_g],
2932 shape: x_shape.clone(),
2933 dtype,
2934 });
2935 let t2 = graph.push(Node {
2936 op: Op::Mul,
2937 inputs: vec![xhat, mean_gxh],
2938 shape: x_shape.clone(),
2939 dtype,
2940 });
2941 let inner = graph.push(Node {
2942 op: Op::Sub,
2943 inputs: vec![t1, t2],
2944 shape: x_shape.clone(),
2945 dtype,
2946 });
2947 graph.push(Node {
2948 op: Op::Mul,
2949 inputs: vec![istd, inner],
2950 shape: x_shape,
2951 dtype,
2952 })
2953 }
2954
2955 /// FROZEN copy of the pre-migration imperative
2956 /// `registry::rms_norm_last_dim_backward::decompose` (the legacy ~22-node
2957 /// `MeanDim`/`SumDim`/`Reshape(keepdim)`/`AddScalar`/`MulScalar(n)`/`Rsqrt`/
2958 /// `BroadcastTo`/`Mul`/`Sub`/`Div` spelling, with the `n = dims[last]`
2959 /// reduced-count baked as an absolute `MulScalar`), copied VERBATIM from that
2960 /// module @ `9d7a1380` before A1 replaced the live body with the shape-
2961 /// polymorphic data recipe. Sole consumer: the A1 numeric-parity oracle
2962 /// (`rms_norm_backward_recipe` below). Reads `inputs = [x, upstream]` + the
2963 /// output shape off the node — the convention the autograd
2964 /// `BackwardKind::Fused` edge emits.
2965 fn frozen_legacy_rms_norm_backward_decompose(
2966 graph: &mut Graph,
2967 id: NodeId,
2968 params: &FusedOpParams,
2969 ) -> NodeId {
2970 let (x_id, up_id, x_shape, dtype) = {
2971 let n = graph.node(id);
2972 (n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
2973 };
2974 let eps = match params {
2975 FusedOpParams::RmsNormLastDimBackward { eps } => *eps,
2976 // G2: total + never-panic — impossible params; return self.
2977 _ => return id,
2978 };
2979 let dims = x_shape.dims().to_vec();
2980 let last = dims.len() - 1;
2981 let n = dims[last] as f64;
2982 let mut kd = dims.clone();
2983 kd[last] = 1;
2984 let keepdim = Shape::from_dims(&kd);
2985 let mut rd = dims.clone();
2986 rd.remove(last);
2987 let reduced = Shape::from_dims(&rd);
2988
2989 // denom = mean(x²) + eps (keepdim).
2990 let sq = graph.push(Node {
2991 op: Op::Sqr,
2992 inputs: vec![x_id],
2993 shape: x_shape.clone(),
2994 dtype,
2995 });
2996 let mean = graph.push(Node {
2997 op: Op::MeanDim(last),
2998 inputs: vec![sq],
2999 shape: reduced.clone(),
3000 dtype,
3001 });
3002 let mean_kd = graph.push(Node {
3003 op: Op::Reshape(keepdim.clone()),
3004 inputs: vec![mean],
3005 shape: keepdim.clone(),
3006 dtype,
3007 });
3008 let denom_kd = graph.push(Node {
3009 op: Op::AddScalar(eps),
3010 inputs: vec![mean_kd],
3011 shape: keepdim.clone(),
3012 dtype,
3013 });
3014 // r_rms = rsqrt(denom), broadcast.
3015 let rrms_kd = graph.push(Node {
3016 op: Op::Rsqrt,
3017 inputs: vec![denom_kd],
3018 shape: keepdim.clone(),
3019 dtype,
3020 });
3021 let rrms_b = graph.push(Node {
3022 op: Op::BroadcastTo(x_shape.clone()),
3023 inputs: vec![rrms_kd],
3024 shape: x_shape.clone(),
3025 dtype,
3026 });
3027 // s = sum(g·x, last) (keepdim).
3028 let gx = graph.push(Node {
3029 op: Op::Mul,
3030 inputs: vec![up_id, x_id],
3031 shape: x_shape.clone(),
3032 dtype,
3033 });
3034 let s = graph.push(Node {
3035 op: Op::SumDim(last),
3036 inputs: vec![gx],
3037 shape: reduced,
3038 dtype,
3039 });
3040 let s_kd = graph.push(Node {
3041 op: Op::Reshape(keepdim.clone()),
3042 inputs: vec![s],
3043 shape: keepdim.clone(),
3044 dtype,
3045 });
3046 let s_b = graph.push(Node {
3047 op: Op::BroadcastTo(x_shape.clone()),
3048 inputs: vec![s_kd],
3049 shape: x_shape.clone(),
3050 dtype,
3051 });
3052 // term = x·s / (n·denom).
3053 let ndenom_kd = graph.push(Node {
3054 op: Op::MulScalar(n),
3055 inputs: vec![denom_kd],
3056 shape: keepdim.clone(),
3057 dtype,
3058 });
3059 let ndenom_b = graph.push(Node {
3060 op: Op::BroadcastTo(x_shape.clone()),
3061 inputs: vec![ndenom_kd],
3062 shape: x_shape.clone(),
3063 dtype,
3064 });
3065 let xs = graph.push(Node {
3066 op: Op::Mul,
3067 inputs: vec![x_id, s_b],
3068 shape: x_shape.clone(),
3069 dtype,
3070 });
3071 let term = graph.push(Node {
3072 op: Op::Div,
3073 inputs: vec![xs, ndenom_b],
3074 shape: x_shape.clone(),
3075 dtype,
3076 });
3077 // grad_x = r_rms · (g − term).
3078 let inner = graph.push(Node {
3079 op: Op::Sub,
3080 inputs: vec![up_id, term],
3081 shape: x_shape.clone(),
3082 dtype,
3083 });
3084 graph.push(Node {
3085 op: Op::Mul,
3086 inputs: vec![rrms_b, inner],
3087 shape: x_shape,
3088 dtype,
3089 })
3090 }
3091
3092 /// FROZEN copy of the pre-migration imperative
3093 /// `registry::reduce_max_to_backward::decompose` (the legacy 9-node
3094 /// `ReduceMaxTo`→`BroadcastTo`→`Equal(U8)`→`MaskedFill`→`ReduceSumTo`→`Div`→
3095 /// `BroadcastTo`→`Mul` spelling, sharing the single `mask_f`), copied VERBATIM
3096 /// from that module @ `94c69ec7` before A2 replaced the live body with the
3097 /// data recipe. Sole consumer: the A2 numeric-parity oracle
3098 /// (`reduce_max_to_backward_recipe` below). Reads `inputs[0] = x` and
3099 /// `inputs[1] = up` (the upstream gradient) off the node — the order autograd
3100 /// emits. The `MaskedFill` fill is `Scalar::one(dtype)` at the node's dtype.
3101 fn frozen_legacy_reduce_max_to_backward_decompose(
3102 graph: &mut Graph,
3103 id: NodeId,
3104 _params: &FusedOpParams,
3105 ) -> NodeId {
3106 use fuel_ir::Scalar;
3107 let (x_id, up_id, x_shape, dtype) = {
3108 let n = graph.node(id);
3109 (n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
3110 };
3111 let target = graph.node(up_id).shape.clone();
3112
3113 // y = per-window max, broadcast back to x's shape.
3114 let y = graph.push(Node {
3115 op: Op::ReduceMaxTo(target.clone()),
3116 inputs: vec![x_id],
3117 shape: target.clone(),
3118 dtype,
3119 });
3120 let y_b = graph.push(Node {
3121 op: Op::BroadcastTo(x_shape.clone()),
3122 inputs: vec![y],
3123 shape: x_shape.clone(),
3124 dtype,
3125 });
3126 // U8 mask = (x == max), then a float mask = MaskedFill(1.0 into zeros).
3127 let mask_u8 = graph.push(Node {
3128 op: Op::Equal,
3129 inputs: vec![x_id, y_b],
3130 shape: x_shape.clone(),
3131 dtype: DType::U8,
3132 });
3133 let zeros = graph.push(Node {
3134 op: Op::MulScalar(0.0),
3135 inputs: vec![x_id],
3136 shape: x_shape.clone(),
3137 dtype,
3138 });
3139 // A reduce-max backward over a non-real (packed/scale) dtype is not a
3140 // thing this frozen oracle handles; self-return the node unchanged
3141 // (decompose-fixpoint / surfaced-gap convention) rather than panic.
3142 let fill = match Scalar::one(dtype) {
3143 Ok(s) => s,
3144 Err(_) => return id,
3145 };
3146 let mask_f = graph.push(Node {
3147 op: Op::MaskedFill { value: fill },
3148 inputs: vec![zeros, mask_u8],
3149 shape: x_shape.clone(),
3150 dtype,
3151 });
3152 // ties = count per window; share = upstream / ties (fair share for ties).
3153 let ties = graph.push(Node {
3154 op: Op::ReduceSumTo(target.clone()),
3155 inputs: vec![mask_f],
3156 shape: target.clone(),
3157 dtype,
3158 });
3159 let share = graph.push(Node {
3160 op: Op::Div,
3161 inputs: vec![up_id, ties],
3162 shape: target,
3163 dtype,
3164 });
3165 let share_b = graph.push(Node {
3166 op: Op::BroadcastTo(x_shape.clone()),
3167 inputs: vec![share],
3168 shape: x_shape.clone(),
3169 dtype,
3170 });
3171 graph.push(Node {
3172 op: Op::Mul,
3173 inputs: vec![mask_f, share_b],
3174 shape: x_shape,
3175 dtype,
3176 })
3177 }
3178
3179 /// FROZEN copy of the pre-migration imperative
3180 /// `registry::powi_backward::decompose` (the legacy 3-node
3181 /// `PowI(exp-1)`→`MulScalar(exp)`→`Mul` spelling), copied VERBATIM from that
3182 /// module before A3 replaced the live body with the param-derived data
3183 /// recipe. Sole consumer: the A3 numeric-parity oracle (`powi_backward_recipe`
3184 /// below). Reads `inputs[0] = x` and `inputs[1] = up` (the upstream gradient)
3185 /// off the node — the order autograd emits — plus the output shape/dtype;
3186 /// `exp` rides the params.
3187 fn frozen_legacy_powi_backward_decompose(
3188 graph: &mut Graph,
3189 id: NodeId,
3190 params: &FusedOpParams,
3191 ) -> NodeId {
3192 let (x_id, up_id, shape, dtype) = {
3193 let n = graph.node(id);
3194 (n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
3195 };
3196 let exp = match params {
3197 FusedOpParams::PowIBackward { exp } => *exp,
3198 // G2: total + never-panic — impossible params; return self.
3199 _ => return id,
3200 };
3201 let pow = graph.push(Node {
3202 op: Op::PowI(exp - 1),
3203 inputs: vec![x_id],
3204 shape: shape.clone(),
3205 dtype,
3206 });
3207 let scaled = graph.push(Node {
3208 op: Op::MulScalar(exp as f64),
3209 inputs: vec![pow],
3210 shape: shape.clone(),
3211 dtype,
3212 });
3213 graph.push(Node {
3214 op: Op::Mul,
3215 inputs: vec![scaled, up_id],
3216 shape,
3217 dtype,
3218 })
3219 }
3220
3221 /// FROZEN copy of the pre-migration imperative
3222 /// `registry::fused_linear::decompose` (the legacy 3-node `MatMul` +
3223 /// `BroadcastTo(rank-1 bias)` + `Add` spelling, broadcasting the rank-1
3224 /// bias DIRECTLY with no leading-1 pad), copied VERBATIM from that module @
3225 /// `b967bdb1` before slice-2 replaced the live body with the WithDim data
3226 /// recipe. Sole consumer: the slice-2 numeric-parity oracle
3227 /// (`fused_linear_recipe` below). Reads `inputs = [a, b, bias]` + the
3228 /// matmul output shape off the node.
3229 fn frozen_legacy_fused_linear_decompose(
3230 graph: &mut Graph,
3231 id: NodeId,
3232 _params: &FusedOpParams,
3233 ) -> NodeId {
3234 let (a_id, b_id, bias_id, out_shape, dtype) = {
3235 let n = graph.node(id);
3236 (
3237 n.inputs[0],
3238 n.inputs[1],
3239 n.inputs[2],
3240 n.shape.clone(),
3241 n.dtype,
3242 )
3243 };
3244 let mm_id = graph.push(Node {
3245 op: Op::MatMul,
3246 inputs: vec![a_id, b_id],
3247 shape: out_shape.clone(),
3248 dtype,
3249 });
3250 let bias_bcst_id = graph.push(Node {
3251 op: Op::BroadcastTo(out_shape.clone()),
3252 inputs: vec![bias_id],
3253 shape: out_shape.clone(),
3254 dtype,
3255 });
3256 graph.push(Node {
3257 op: Op::Add,
3258 inputs: vec![mm_id, bias_bcst_id],
3259 shape: out_shape,
3260 dtype,
3261 })
3262 }
3263
3264 /// FROZEN copy of the pre-migration imperative
3265 /// `registry::softmax_last_dim_backward::decompose` (the legacy 5-node
3266 /// `Mul`/`ReduceSumTo(keepdim)`/`BroadcastTo`/`Sub`/`Mul` spelling), copied
3267 /// VERBATIM from that module @ `aa2eee3c` before T8 replaced the live body
3268 /// with the data recipe. Sole consumer: the T8 numeric-parity oracle
3269 /// (`softmax_backward_recipe` below). Reads `inputs[0] = s` (the forward
3270 /// softmax output) and `inputs[1] = g` (the upstream gradient) off the node
3271 /// — the same convention the autograd `BackwardKind::Fused` edge emits
3272 /// (`lib.rs` softmax-backward arm: `vec![id, up_id]`).
3273 fn frozen_legacy_softmax_backward_decompose(
3274 graph: &mut Graph,
3275 id: NodeId,
3276 _params: &FusedOpParams,
3277 ) -> NodeId {
3278 let (s_id, g_id, x_shape, dtype) = {
3279 let n = graph.node(id);
3280 (n.inputs[0], n.inputs[1], n.shape.clone(), n.dtype)
3281 };
3282 // keepdim shape: last dim → 1.
3283 let mut kd = x_shape.dims().to_vec();
3284 let last = kd.len() - 1;
3285 kd[last] = 1;
3286 let keepdim = Shape::from_dims(&kd);
3287
3288 let gs = graph.push(Node {
3289 op: Op::Mul,
3290 inputs: vec![g_id, s_id],
3291 shape: x_shape.clone(),
3292 dtype,
3293 });
3294 let summed = graph.push(Node {
3295 op: Op::ReduceSumTo(keepdim.clone()),
3296 inputs: vec![gs],
3297 shape: keepdim,
3298 dtype,
3299 });
3300 let summed_b = graph.push(Node {
3301 op: Op::BroadcastTo(x_shape.clone()),
3302 inputs: vec![summed],
3303 shape: x_shape.clone(),
3304 dtype,
3305 });
3306 let sub = graph.push(Node {
3307 op: Op::Sub,
3308 inputs: vec![g_id, summed_b],
3309 shape: x_shape.clone(),
3310 dtype,
3311 });
3312 graph.push(Node {
3313 op: Op::Mul,
3314 inputs: vec![s_id, sub],
3315 shape: x_shape,
3316 dtype,
3317 })
3318 }
3319
3320 #[test]
3321 fn emit_matches_softmax_last_dim_decompose() {
3322 use fuel_ir::{DType, Shape};
3323 let mut g = Graph::new();
3324 let sh = Shape::from_dims(&[2, 4]);
3325 let x = g.push(Node {
3326 op: Op::Const,
3327 inputs: vec![],
3328 shape: sh.clone(),
3329 dtype: DType::F32,
3330 });
3331 // Oracle: the FROZEN legacy builder (reads inputs[0] + shape + dtype
3332 // off the node). T5 repointed this from the live registry decompose —
3333 // which now emits the 9-node recipe spelling — so this test keeps
3334 // pinning the Increment-A guarantee it always pinned: the grown
3335 // `emit` reconstructs the LEGACY imperative structure from the legacy
3336 // region datum.
3337 let fused = g.push(Node {
3338 op: Op::Const,
3339 inputs: vec![x],
3340 shape: sh.clone(),
3341 dtype: DType::F32,
3342 });
3343 let oracle = frozen_legacy_softmax_decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
3344
3345 // keepdim shape [2,1]; full shape [2,4].
3346 let kd = OpAttrs {
3347 target_shape: vec![2, 1],
3348 ..OpAttrs::default()
3349 };
3350 let full = OpAttrs {
3351 target_shape: vec![2, 4],
3352 ..OpAttrs::default()
3353 };
3354 // e = Exp(Sub(x, BroadcastTo(ReduceMaxTo(x)))) — mirrors decompose order
3355 // `Sub{[x, mb]}` exactly; built fresh each call so numerator and the
3356 // denominator's ReduceSumTo input are identical subtrees.
3357 let softmax_e = |kd: &OpAttrs, full: &OpAttrs| {
3358 op_node(
3359 OpTag::Exp,
3360 OpAttrs::default(),
3361 vec![op_node(
3362 OpTag::Sub,
3363 OpAttrs::default(),
3364 vec![
3365 bind(0),
3366 op_node(
3367 OpTag::BroadcastTo,
3368 full.clone(),
3369 vec![op_node(OpTag::ReduceMaxTo, kd.clone(), vec![bind(0)])],
3370 ),
3371 ],
3372 )],
3373 )
3374 };
3375 // out = Div(e, BroadcastTo(ReduceSumTo(e))) — mirrors `Div{[e, db]}`.
3376 let region = op_node(
3377 OpTag::Div,
3378 OpAttrs::default(),
3379 vec![
3380 softmax_e(&kd, &full),
3381 op_node(
3382 OpTag::BroadcastTo,
3383 full.clone(),
3384 vec![op_node(
3385 OpTag::ReduceSumTo,
3386 kd.clone(),
3387 vec![softmax_e(&kd, &full)],
3388 )],
3389 ),
3390 ],
3391 );
3392 let emitted = emit_region(&mut g, ®ion, &[x], &[]);
3393 assert_structural_eq(&g, oracle, emitted);
3394 }
3395
3396 #[test]
3397 fn emit_matches_rope_decompose() {
3398 use fuel_ir::{DType, Shape};
3399 let mut g = Graph::new();
3400 let sh = Shape::from_dims(&[2, 4]); // seq=2, d=4, half=2
3401 let x = g.push(Node {
3402 op: Op::Const,
3403 inputs: vec![],
3404 shape: sh.clone(),
3405 dtype: DType::F32,
3406 });
3407 let cos = g.push(Node {
3408 op: Op::Const,
3409 inputs: vec![],
3410 shape: sh.clone(),
3411 dtype: DType::F32,
3412 });
3413 let sin = g.push(Node {
3414 op: Op::Const,
3415 inputs: vec![],
3416 shape: sh.clone(),
3417 dtype: DType::F32,
3418 });
3419 let fused = g.push(Node {
3420 op: Op::Const,
3421 inputs: vec![x, cos, sin],
3422 shape: sh.clone(),
3423 dtype: DType::F32,
3424 });
3425 // Oracle: the FROZEN legacy builder. T6 repointed this from the live
3426 // registry decompose — which now emits the data recipe (byte-identical
3427 // to legacy where a rank-raise occurs, but at EQUAL rank the recipe
3428 // elides the legacy's no-op prep `Reshape`, D4). This test keeps
3429 // pinning the Increment-A guarantee it always pinned: the grown `emit`
3430 // reconstructs the LEGACY imperative structure from a legacy-spelled
3431 // region datum.
3432 let oracle = frozen_legacy_rope_decompose(&mut g, fused, &FusedOpParams::Rope);
3433
3434 // decompose's broadcast_shape for rank-2 [2,4] is [seq,d] = [2,4]; half slices along last dim.
3435 let full = OpAttrs {
3436 target_shape: vec![2, 4],
3437 ..OpAttrs::default()
3438 };
3439 let sl_first = OpAttrs {
3440 axis: Some(1),
3441 slice_start: Some(0),
3442 slice_len: Some(2),
3443 ..OpAttrs::default()
3444 };
3445 let sl_second = OpAttrs {
3446 axis: Some(1),
3447 slice_start: Some(2),
3448 slice_len: Some(2),
3449 ..OpAttrs::default()
3450 };
3451 let cat = OpAttrs {
3452 axis: Some(1),
3453 ..OpAttrs::default()
3454 };
3455 let bcast_reshape = |full: &OpAttrs, i: u8| {
3456 op_node(
3457 OpTag::BroadcastTo,
3458 full.clone(),
3459 vec![op_node(OpTag::Reshape, full.clone(), vec![bind(i)])],
3460 )
3461 };
3462 // left = Mul(x, cos_bcast); right = Mul(rotated_half, sin_bcast); out = Add(left, right).
3463 // rotated_half = Concat{dim:1}(Neg(second_half), first_half).
3464 let rotated = op_node(
3465 OpTag::Concat,
3466 cat,
3467 vec![
3468 op_node(
3469 OpTag::Neg,
3470 OpAttrs::default(),
3471 vec![op_node(OpTag::Slice, sl_second, vec![bind(0)])],
3472 ),
3473 op_node(OpTag::Slice, sl_first, vec![bind(0)]),
3474 ],
3475 );
3476 let left = op_node(
3477 OpTag::Mul,
3478 OpAttrs::default(),
3479 vec![bind(0), bcast_reshape(&full, 1)],
3480 );
3481 let right = op_node(
3482 OpTag::Mul,
3483 OpAttrs::default(),
3484 vec![rotated, bcast_reshape(&full, 2)],
3485 );
3486 let region = op_node(OpTag::Add, OpAttrs::default(), vec![left, right]);
3487
3488 let emitted = emit_region(&mut g, ®ion, &[x, cos, sin], &[]);
3489 assert_structural_eq(&g, oracle, emitted);
3490 }
3491
3492 #[test]
3493 fn emit_matches_layer_norm_last_dim_decompose() {
3494 use fuel_ir::{DType, Shape};
3495 let mut g = Graph::new();
3496 let sh = Shape::from_dims(&[2, 4]); // last=1, reduced [2], keepdim [2,1]
3497 let x = g.push(Node {
3498 op: Op::Const,
3499 inputs: vec![],
3500 shape: sh.clone(),
3501 dtype: DType::F32,
3502 });
3503 let fused = g.push(Node {
3504 op: Op::Const,
3505 inputs: vec![x],
3506 shape: sh.clone(),
3507 dtype: DType::F32,
3508 });
3509 // Oracle: the FROZEN legacy builder (the `Reshape(keepdim)` spelling).
3510 // T7 repointed this from the live registry decompose — which now emits
3511 // the D3 `Unsqueeze` swap — so this test keeps pinning the Increment-A
3512 // guarantee it always pinned: the grown `emit` reconstructs the LEGACY
3513 // imperative structure from a legacy-spelled (Reshape) region datum.
3514 let oracle = frozen_legacy_layer_norm_decompose(
3515 &mut g,
3516 fused,
3517 &FusedOpParams::LayerNormLastDim { eps: 1e-5 },
3518 );
3519
3520 let kd = OpAttrs {
3521 target_shape: vec![2, 1],
3522 ..OpAttrs::default()
3523 };
3524 let full = OpAttrs {
3525 target_shape: vec![2, 4],
3526 ..OpAttrs::default()
3527 };
3528 let md = OpAttrs {
3529 axis: Some(1),
3530 ..OpAttrs::default()
3531 };
3532 let eps_attrs = OpAttrs {
3533 scalars: vec![1e-5],
3534 ..OpAttrs::default()
3535 }; // BAKED constant, not a slot
3536 // centered = Sub(x, BroadcastTo(Reshape(MeanDim(x)))) — shared subterm.
3537 let centered = op_node(
3538 OpTag::Sub,
3539 OpAttrs::default(),
3540 vec![
3541 bind(0),
3542 op_node(
3543 OpTag::BroadcastTo,
3544 full.clone(),
3545 vec![op_node(
3546 OpTag::Reshape,
3547 kd.clone(),
3548 vec![op_node(OpTag::MeanDim, md.clone(), vec![bind(0)])],
3549 )],
3550 ),
3551 ],
3552 );
3553 // denom_bcast = BroadcastTo(Sqrt(AddScalar(eps)(Reshape(MeanDim(Sqr(centered)))))).
3554 let denom_bcast = op_node(
3555 OpTag::BroadcastTo,
3556 full.clone(),
3557 vec![op_node(
3558 OpTag::Sqrt,
3559 OpAttrs::default(),
3560 vec![op_node(
3561 OpTag::AddScalar,
3562 eps_attrs,
3563 vec![op_node(
3564 OpTag::Reshape,
3565 kd.clone(),
3566 vec![op_node(
3567 OpTag::MeanDim,
3568 md.clone(),
3569 vec![op_node(
3570 OpTag::Sqr,
3571 OpAttrs::default(),
3572 vec![centered.clone()],
3573 )],
3574 )],
3575 )],
3576 )],
3577 )],
3578 );
3579 // out = Div(centered, denom_bcast).
3580 let region = op_node(OpTag::Div, OpAttrs::default(), vec![centered, denom_bcast]);
3581
3582 let emitted = emit_region(&mut g, ®ion, &[x], &[]);
3583 assert_structural_eq(&g, oracle, emitted);
3584 }
3585
3586 // ---- T2 (Increment C slice 1): shape-relative attr resolution ----------
3587 //
3588 // `resolve_rel_attrs` is the PURE resolver behind recipe polymorphism: it
3589 // turns the shape-relative interior fields (`target_shape_rel` /
3590 // `slice_{start,len}_rel` / `axis_last`, D2) into the concrete sibling
3591 // fields against the given bind/child shapes, reusing `shape_expr`'s
3592 // evaluator (no second evaluator). Every failure is a typed
3593 // [`RelAttrError`], never a panic.
3594
3595 mod resolve_rel {
3596 use super::super::{RelAttrError, resolve_rel_attrs};
3597 use fuel_kernel_seam_types::shape_expr::{
3598 Dim, LAST, SYMBOLIC, ShapeExpr, ShapeExprError, TAG_DIMS,
3599 };
3600 use fuel_kernel_seam_types::{OpAttrs, OpTag};
3601
3602 fn half_of_bind0_last() -> Dim {
3603 Dim::Div(
3604 Box::new(Dim::Extent {
3605 operand: 0,
3606 axis: LAST,
3607 }),
3608 Box::new(Dim::Const(2)),
3609 )
3610 }
3611
3612 #[test]
3613 fn same_as_bind0_tracks_the_bind_shape() {
3614 // The polymorphism seed: ONE recipe datum, two shapes, two targets.
3615 let attrs = OpAttrs {
3616 target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
3617 ..OpAttrs::default()
3618 };
3619 let r = resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &[vec![2, 3]], &[vec![2, 1]])
3620 .expect("resolves");
3621 assert_eq!(r.target_shape, vec![2, 3]);
3622 assert!(
3623 r.target_shape_rel.is_none(),
3624 "resolved attrs are fully concrete"
3625 );
3626 let r = resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &[vec![4, 5]], &[vec![4, 1]])
3627 .expect("resolves");
3628 assert_eq!(r.target_shape, vec![4, 5]);
3629 }
3630
3631 #[test]
3632 fn slice_bounds_track_the_bind_extent() {
3633 // start = len = Extent(bind0, LAST) / 2 — the rope-half worked example.
3634 let attrs = OpAttrs {
3635 axis: Some(1),
3636 slice_start_rel: Some(half_of_bind0_last()),
3637 slice_len_rel: Some(half_of_bind0_last()),
3638 ..OpAttrs::default()
3639 };
3640 let r = resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]])
3641 .expect("resolves at d=4");
3642 assert_eq!(r.slice_start, Some(2));
3643 assert_eq!(r.slice_len, Some(2));
3644 assert!(r.slice_start_rel.is_none() && r.slice_len_rel.is_none());
3645 let r = resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 8]], &[vec![2, 8]])
3646 .expect("resolves at d=8");
3647 assert_eq!(r.slice_start, Some(4));
3648 assert_eq!(r.slice_len, Some(4));
3649 }
3650
3651 #[test]
3652 fn axis_last_resolves_per_tag() {
3653 let attrs = OpAttrs {
3654 axis_last: true,
3655 ..OpAttrs::default()
3656 };
3657 // Reduce family (axis carrier): LAST = rank − 1.
3658 let r = resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[vec![2, 4]]).expect("rank 2");
3659 assert_eq!(r.axis, Some(1));
3660 assert!(!r.axis_last, "resolved attrs are fully concrete");
3661 let r =
3662 resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[vec![2, 3, 4]]).expect("rank 3");
3663 assert_eq!(r.axis, Some(2));
3664 // Concat rides the same axis carrier.
3665 let r =
3666 resolve_rel_attrs(OpTag::Concat, &attrs, &[], &[vec![2, 3, 4]]).expect("concat");
3667 assert_eq!(r.axis, Some(2));
3668 // Unsqueeze (dims carrier): APPEND — dim == rank (`primitive_shape`
3669 // permits `dim == rank`).
3670 let r =
3671 resolve_rel_attrs(OpTag::Unsqueeze, &attrs, &[], &[vec![2, 4]]).expect("unsqueeze");
3672 assert_eq!(r.dims, vec![2]);
3673 assert!(!r.axis_last);
3674 // Squeeze (dims carrier): LAST = rank − 1.
3675 let r =
3676 resolve_rel_attrs(OpTag::Squeeze, &attrs, &[], &[vec![2, 4, 1]]).expect("squeeze");
3677 assert_eq!(r.dims, vec![2]);
3678 }
3679
3680 #[test]
3681 fn bind_out_of_range_is_a_typed_decline() {
3682 let attrs = OpAttrs {
3683 target_shape_rel: Some(ShapeExpr::SameAs { operand: 3 }),
3684 ..OpAttrs::default()
3685 };
3686 assert_eq!(
3687 resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &[vec![2, 3]], &[vec![2, 3]]),
3688 Err(RelAttrError::Expr(ShapeExprError::OperandOutOfRange {
3689 operand: 3,
3690 operands: 1
3691 })),
3692 );
3693 let attrs = OpAttrs {
3694 axis: Some(1),
3695 slice_start_rel: Some(Dim::Extent {
3696 operand: 7,
3697 axis: LAST,
3698 }),
3699 ..OpAttrs::default()
3700 };
3701 assert_eq!(
3702 resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]]),
3703 Err(RelAttrError::Expr(ShapeExprError::OperandOutOfRange {
3704 operand: 7,
3705 operands: 1
3706 })),
3707 );
3708 }
3709
3710 #[test]
3711 fn rel_and_abs_both_set_is_a_typed_conflict() {
3712 // target_shape XOR target_shape_rel.
3713 let attrs = OpAttrs {
3714 target_shape: vec![2, 3],
3715 target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
3716 ..OpAttrs::default()
3717 };
3718 assert_eq!(
3719 resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &[vec![2, 3]], &[vec![2, 3]]),
3720 Err(RelAttrError::RelAbsConflict {
3721 field: "target_shape"
3722 }),
3723 );
3724 // slice_start XOR slice_start_rel.
3725 let attrs = OpAttrs {
3726 axis: Some(1),
3727 slice_start: Some(0),
3728 slice_start_rel: Some(half_of_bind0_last()),
3729 ..OpAttrs::default()
3730 };
3731 assert_eq!(
3732 resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]]),
3733 Err(RelAttrError::RelAbsConflict {
3734 field: "slice_start"
3735 }),
3736 );
3737 // slice_len XOR slice_len_rel.
3738 let attrs = OpAttrs {
3739 axis: Some(1),
3740 slice_len: Some(2),
3741 slice_len_rel: Some(half_of_bind0_last()),
3742 ..OpAttrs::default()
3743 };
3744 assert_eq!(
3745 resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]]),
3746 Err(RelAttrError::RelAbsConflict { field: "slice_len" }),
3747 );
3748 // axis XOR axis_last.
3749 let attrs = OpAttrs {
3750 axis: Some(0),
3751 axis_last: true,
3752 ..OpAttrs::default()
3753 };
3754 assert_eq!(
3755 resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[vec![2, 4]]),
3756 Err(RelAttrError::RelAbsConflict { field: "axis" }),
3757 );
3758 // dims XOR axis_last (Unsqueeze's carrier is `dims`).
3759 let attrs = OpAttrs {
3760 dims: vec![0],
3761 axis_last: true,
3762 ..OpAttrs::default()
3763 };
3764 assert_eq!(
3765 resolve_rel_attrs(OpTag::Unsqueeze, &attrs, &[], &[vec![2, 4]]),
3766 Err(RelAttrError::RelAbsConflict { field: "dims" }),
3767 );
3768 }
3769
3770 #[test]
3771 fn negative_result_is_a_typed_decline() {
3772 // 0 − 2 = −2: a negative slice offset is malformed, not a wrap.
3773 let neg = Dim::Sub(Box::new(Dim::Const(0)), Box::new(Dim::Const(2)));
3774 let attrs = OpAttrs {
3775 axis: Some(1),
3776 slice_start_rel: Some(neg),
3777 ..OpAttrs::default()
3778 };
3779 assert_eq!(
3780 resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]]),
3781 Err(RelAttrError::Negative {
3782 field: "slice_start",
3783 value: -2
3784 }),
3785 );
3786 }
3787
3788 #[test]
3789 fn symbolic_extent_is_a_surfaced_gap_decline() {
3790 // A symbolic bind extent → the expression evaluates to Gap → typed
3791 // decline (the emit caller surfaces it as a fixpoint, G2).
3792 let attrs = OpAttrs {
3793 axis: Some(1),
3794 slice_len_rel: Some(half_of_bind0_last()),
3795 ..OpAttrs::default()
3796 };
3797 assert_eq!(
3798 resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, SYMBOLIC]], &[vec![2, 4]]),
3799 Err(RelAttrError::SymbolicGap { field: "slice_len" }),
3800 );
3801 let attrs = OpAttrs {
3802 target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
3803 ..OpAttrs::default()
3804 };
3805 assert_eq!(
3806 resolve_rel_attrs(
3807 OpTag::BroadcastTo,
3808 &attrs,
3809 &[vec![2, SYMBOLIC]],
3810 &[vec![2, 4]]
3811 ),
3812 Err(RelAttrError::SymbolicGap {
3813 field: "target_shape"
3814 }),
3815 );
3816 }
3817
3818 #[test]
3819 fn two_operand_max_frame_declines_instead_of_a_partial_shape() {
3820 // I1 (Baracuda §6.20 finding): an ELEMENTWISE output frame is not
3821 // always expressible as `SameAs(operand)`. When the frame is
3822 // assembled by per-axis max across TWO binds — `a[2,1] ⊗ b[1,3] →
3823 // [2,3]` — NO single operand carries it, so BOTH spellings resolve
3824 // to a PARTIAL frame ([2,1] / [1,3]) and would silently emit the
3825 // wrong `BroadcastTo` target. Must be a typed decline.
3826 // The decline is typed and names the Dims-class constructor that
3827 // WOULD express it (reserved tag 0x0B, KISS #80) — a surfaced gap,
3828 // never a panic and never a wrong shape.
3829 let binds = vec![vec![2, 1], vec![1, 3]]; // frame = [2,3], carried by neither
3830 for operand in [0u8, 1] {
3831 let attrs = OpAttrs {
3832 target_shape_rel: Some(ShapeExpr::SameAs { operand }),
3833 ..OpAttrs::default()
3834 };
3835 let child = binds[operand as usize].clone();
3836 assert_eq!(
3837 resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &binds, &[child]),
3838 Err(RelAttrError::FrameNotExpressible {
3839 field: "target_shape",
3840 frame: vec![2, 3],
3841 missing_ctor: TAG_DIMS,
3842 }),
3843 "SameAs {{ operand: {operand} }} must not resolve to a PARTIAL frame",
3844 );
3845 }
3846 }
3847
3848 #[test]
3849 fn frame_guard_does_not_fire_when_an_operand_carries_the_frame() {
3850 // The guard is deliberately narrow. It must NOT degrade the cases
3851 // the 5 migrated recipes actually use.
3852 let attrs = OpAttrs {
3853 target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
3854 ..OpAttrs::default()
3855 };
3856 // (a) bind0 IS the frame (softmax/rope/rms-norm/layer-norm shape).
3857 let binds = vec![vec![2, 3, 4], vec![4]];
3858 let r = resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &binds, &[vec![4]])
3859 .expect("bind0 carries the frame");
3860 assert_eq!(r.target_shape, vec![2, 3, 4]);
3861 // (b) a SUB-frame target is legitimate: the frame is bind0's, and
3862 // naming bind1's smaller shape is an ordinary interior broadcast.
3863 let sub = OpAttrs {
3864 target_shape_rel: Some(ShapeExpr::SameAs { operand: 1 }),
3865 ..OpAttrs::default()
3866 };
3867 let r = resolve_rel_attrs(OpTag::BroadcastTo, &sub, &binds, &[vec![4]])
3868 .expect("sub-frame broadcast stays expressible");
3869 assert_eq!(r.target_shape, vec![4]);
3870 // (c) binds with NO joint elementwise frame (a matmul region) —
3871 // per-axis max is meaningless there, so the guard stays out.
3872 let mm = vec![vec![8, 4096], vec![4096, 1024]];
3873 let r = resolve_rel_attrs(OpTag::BroadcastTo, &attrs, &mm, &[vec![8, 1]])
3874 .expect("no joint frame ⇒ no guard");
3875 assert_eq!(r.target_shape, vec![8, 4096]);
3876 // (d) a NON-frame-carrier tag is untouched: only `BroadcastTo`'s
3877 // target IS the elementwise output frame.
3878 let two = vec![vec![2, 1], vec![1, 3]];
3879 let r = resolve_rel_attrs(OpTag::Reshape, &attrs, &two, &[vec![2, 1]])
3880 .expect("Reshape's target is not a frame claim");
3881 assert_eq!(r.target_shape, vec![2, 1]);
3882 }
3883
3884 #[test]
3885 fn axis_last_on_an_axisless_tag_or_without_a_child_declines() {
3886 let attrs = OpAttrs {
3887 axis_last: true,
3888 ..OpAttrs::default()
3889 };
3890 // Add has no axis carrier — axis_last is meaningless, a typed decline
3891 // (build-time validation posture: never silently ignore).
3892 assert_eq!(
3893 resolve_rel_attrs(OpTag::Add, &attrs, &[], &[vec![2, 4], vec![2, 4]]),
3894 Err(RelAttrError::AxisLastUnsupported { tag: OpTag::Add }),
3895 );
3896 // No child operand → no rank to resolve LAST against.
3897 assert_eq!(
3898 resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[]),
3899 Err(RelAttrError::NoChildOperand),
3900 );
3901 // Rank-0 child: LAST has no axis → the shared resolve_axis decline.
3902 assert_eq!(
3903 resolve_rel_attrs(OpTag::SumDim, &attrs, &[], &[vec![]]),
3904 Err(RelAttrError::Expr(ShapeExprError::AxisOutOfRange {
3905 axis: LAST,
3906 rank: 0
3907 })),
3908 );
3909 }
3910
3911 #[test]
3912 fn rel_free_attrs_pass_through_unchanged() {
3913 // The no-rel fast path: absolute attrs resolve to themselves.
3914 let attrs = OpAttrs {
3915 axis: Some(1),
3916 slice_start: Some(2),
3917 slice_len: Some(3),
3918 ..OpAttrs::default()
3919 };
3920 let r = resolve_rel_attrs(OpTag::Slice, &attrs, &[vec![2, 4]], &[vec![2, 4]])
3921 .expect("no-op");
3922 assert_eq!(r, attrs);
3923 }
3924 }
3925
3926 // ---- T3 (Increment C slice 1): resolving emit + D4 pad + rel validation ----
3927 //
3928 // The emit integration behind recipe polymorphism: children are emitted
3929 // FIRST (their shapes feed the rel-attr resolver), `resolve_rel_attrs`
3930 // produces fully-concrete attrs, then the unchanged tag_to_op →
3931 // primitive_shape path runs. A resolved `BroadcastTo` whose target rank
3932 // exceeds its operand's materializes the legacy `Reshape` pad (D4).
3933 // `validate_representable` accepts rel-attr regions via a probe-resolve
3934 // (mirror of the scalar slot dummy-fill) and rejects structural authoring
3935 // errors (rel+abs conflict, bind out of range) with a typed decline.
3936
3937 mod emit_rel {
3938 use super::super::*;
3939 use super::{assert_structural_eq, bind, op_node};
3940 use fuel_ir::{DType, Shape};
3941 use fuel_kernel_seam_types::shape_expr::{Dim, ShapeExpr};
3942
3943 fn cst(g: &mut Graph, dims: &[usize]) -> NodeId {
3944 g.push(Node {
3945 op: Op::Const,
3946 inputs: vec![],
3947 shape: Shape::from_dims(dims),
3948 dtype: DType::F32,
3949 })
3950 }
3951
3952 fn bcast_same_as_0() -> OpAttrs {
3953 OpAttrs {
3954 target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
3955 ..OpAttrs::default()
3956 }
3957 }
3958
3959 #[test]
3960 fn rel_region_emits_polymorphically_across_shapes() {
3961 // The headline polymorphism: ONE region datum —
3962 // Add(bind0, BroadcastTo{SameAs{0}}(bind1)) — emitted at two
3963 // different shapes produces the correct target BOTH times
3964 // (impossible with absolute attrs: a baked target matches exactly
3965 // one shape).
3966 let region = op_node(
3967 OpTag::Add,
3968 OpAttrs::default(),
3969 vec![
3970 bind(0),
3971 op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(1)]),
3972 ],
3973 );
3974 let mut g = Graph::new();
3975 let x1 = cst(&mut g, &[2, 3]);
3976 let t1 = cst(&mut g, &[1, 3]);
3977 let r1 = emit_region(&mut g, ®ion, &[x1, t1], &[]);
3978 assert!(matches!(g.node(r1).op, Op::Add));
3979 assert_eq!(g.node(r1).shape, Shape::from_dims(&[2, 3]));
3980 let b1 = g.node(r1).inputs[1];
3981 assert_eq!(g.node(b1).op, Op::BroadcastTo(Shape::from_dims(&[2, 3])));
3982 assert_eq!(g.node(b1).shape, Shape::from_dims(&[2, 3]));
3983
3984 // The SAME region datum at different shapes → a different target.
3985 let x2 = cst(&mut g, &[4, 5]);
3986 let t2 = cst(&mut g, &[1, 5]);
3987 let r2 = emit_region(&mut g, ®ion, &[x2, t2], &[]);
3988 assert_eq!(g.node(r2).shape, Shape::from_dims(&[4, 5]));
3989 let b2 = g.node(r2).inputs[1];
3990 assert_eq!(g.node(b2).op, Op::BroadcastTo(Shape::from_dims(&[4, 5])));
3991 }
3992
3993 #[test]
3994 fn two_operand_max_frame_region_declines_through_emit() {
3995 // I1 end-to-end: the ONE region spelling that WANTS the per-axis-max
3996 // frame — Mul(BroadcastTo(a[2,1]), BroadcastTo(b[1,3])) → [2,3],
3997 // which Fuel's primitive `Mul` requires explicitly (`primitive_shape`
3998 // takes in[0]'s shape, it does not broadcast). `SameAs` cannot name
3999 // [2,3], so the resolving emit surfaces the typed Dims-class gap
4000 // instead of emitting a partial-frame BroadcastTo.
4001 let region = op_node(
4002 OpTag::Mul,
4003 OpAttrs::default(),
4004 vec![
4005 op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(0)]),
4006 op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(1)]),
4007 ],
4008 );
4009 let mut g = Graph::new();
4010 let a = cst(&mut g, &[2, 1]);
4011 let b = cst(&mut g, &[1, 3]);
4012 assert_eq!(
4013 try_emit_region(&mut g, ®ion, &[a, b], &[]),
4014 Err(RelAttrError::FrameNotExpressible {
4015 field: "target_shape",
4016 frame: vec![2, 3],
4017 missing_ctor: fuel_kernel_seam_types::shape_expr::TAG_DIMS,
4018 }),
4019 );
4020 }
4021
4022 #[test]
4023 fn broadcast_rank_raise_materializes_the_legacy_reshape_pad() {
4024 // D4: a rank-1 bind1 broadcast to rank-3 bind0's shape — the
4025 // resolver must first push the legacy `Reshape` (1-padded left,
4026 // right-aligned; `registry::rope`'s hand-built broadcast prep).
4027 let region = op_node(
4028 OpTag::Mul,
4029 OpAttrs::default(),
4030 vec![
4031 bind(0),
4032 op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(1)]),
4033 ],
4034 );
4035 let mut g = Graph::new();
4036 let x = cst(&mut g, &[2, 3, 4]);
4037 let t = cst(&mut g, &[4]);
4038 // Hand-built legacy reference:
4039 // Reshape([1,1,4])(t) → BroadcastTo([2,3,4]) → Mul(x, ·).
4040 let pad_shape = Shape::from_dims(&[1, 1, 4]);
4041 let full = Shape::from_dims(&[2, 3, 4]);
4042 let pad = g.push(Node {
4043 op: Op::Reshape(pad_shape.clone()),
4044 inputs: vec![t],
4045 shape: pad_shape,
4046 dtype: DType::F32,
4047 });
4048 let bc = g.push(Node {
4049 op: Op::BroadcastTo(full.clone()),
4050 inputs: vec![pad],
4051 shape: full.clone(),
4052 dtype: DType::F32,
4053 });
4054 let reference = g.push(Node {
4055 op: Op::Mul,
4056 inputs: vec![x, bc],
4057 shape: full,
4058 dtype: DType::F32,
4059 });
4060
4061 let emitted = emit_region(&mut g, ®ion, &[x, t], &[]);
4062 assert_structural_eq(&g, reference, emitted);
4063 }
4064
4065 #[test]
4066 fn concrete_broadcast_rank_raise_also_pads() {
4067 // D4 applies uniformly: an ABSOLUTE rank-raising BroadcastTo also
4068 // materializes the pad (deterministic emission, matches the graph
4069 // builders' right-aligned rank-raising semantics). Equal-rank
4070 // broadcasts stay pad-free (pinned by the softmax parity oracle).
4071 let region = op_node(
4072 OpTag::BroadcastTo,
4073 OpAttrs {
4074 target_shape: vec![2, 3, 4],
4075 ..OpAttrs::default()
4076 },
4077 vec![bind(0)],
4078 );
4079 let mut g = Graph::new();
4080 let t = cst(&mut g, &[3, 4]);
4081 let emitted = emit_region(&mut g, ®ion, &[t], &[]);
4082 assert_eq!(
4083 g.node(emitted).op,
4084 Op::BroadcastTo(Shape::from_dims(&[2, 3, 4]))
4085 );
4086 let pad = g.node(emitted).inputs[0];
4087 assert_eq!(
4088 g.node(pad).op,
4089 Op::Reshape(Shape::from_dims(&[1, 3, 4])),
4090 "rank-raise inserts the legacy 1-padded-left Reshape",
4091 );
4092 assert_eq!(g.node(pad).inputs, vec![t]);
4093 }
4094
4095 #[test]
4096 fn scalar_cursor_fill_stays_pre_order_after_the_reorder() {
4097 // Risk-2 guard: children are now EMITTED first, but the scalar
4098 // cursor fill stays PRE-order (parent before descent) — the
4099 // canonical authoring order `match_region_extract` records.
4100 let region = op_node(
4101 OpTag::AddScalar,
4102 OpAttrs::default(),
4103 vec![op_node(OpTag::MulScalar, OpAttrs::default(), vec![bind(0)])],
4104 );
4105 let mut g = Graph::new();
4106 let x = cst(&mut g, &[4]);
4107 let root = emit_region(&mut g, ®ion, &[x], &[10.0, 20.0]);
4108 assert!(
4109 matches!(g.node(root).op, Op::AddScalar(v) if v == 10.0),
4110 "parent takes scalars[0] (pre-order), got {:?}",
4111 g.node(root).op,
4112 );
4113 let child = g.node(root).inputs[0];
4114 assert!(
4115 matches!(g.node(child).op, Op::MulScalar(v) if v == 20.0),
4116 "child takes scalars[1], got {:?}",
4117 g.node(child).op,
4118 );
4119 }
4120
4121 #[test]
4122 fn validate_accepts_a_rel_attr_region() {
4123 // Born-red: today tag_to_op(BroadcastTo, {empty target_shape}) →
4124 // None → UnRepresentable. The rel-probe must accept the template.
4125 let region = op_node(
4126 OpTag::Add,
4127 OpAttrs::default(),
4128 vec![
4129 bind(0),
4130 op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(1)]),
4131 ],
4132 );
4133 register_runtime_fused("t3::rel_bcast", region)
4134 .expect("a rel-attr region is registrable");
4135 }
4136
4137 #[test]
4138 fn validate_rejects_rel_abs_conflict_and_bind_out_of_range() {
4139 // rel+abs both set → a typed authoring reject, never a silent
4140 // precedence.
4141 let conflicted = op_node(
4142 OpTag::BroadcastTo,
4143 OpAttrs {
4144 target_shape: vec![2, 3],
4145 target_shape_rel: Some(ShapeExpr::SameAs { operand: 0 }),
4146 ..OpAttrs::default()
4147 },
4148 vec![bind(0)],
4149 );
4150 assert_eq!(
4151 register_runtime_fused("t3::conflict", conflicted),
4152 Err(RuntimeFusedError::InvalidRelAttrs {
4153 tag: OpTag::BroadcastTo,
4154 error: RelAttrError::RelAbsConflict {
4155 field: "target_shape"
4156 },
4157 }),
4158 );
4159 // A bind reference outside the region's bind space can never
4160 // resolve at ANY shape → a typed authoring reject.
4161 let oob = op_node(
4162 OpTag::BroadcastTo,
4163 OpAttrs {
4164 target_shape_rel: Some(ShapeExpr::SameAs { operand: 7 }),
4165 ..OpAttrs::default()
4166 },
4167 vec![bind(0)],
4168 );
4169 assert_eq!(
4170 register_runtime_fused("t3::oob", oob),
4171 Err(RuntimeFusedError::InvalidRelAttrs {
4172 tag: OpTag::BroadcastTo,
4173 error: RelAttrError::Expr(
4174 fuel_kernel_seam_types::shape_expr::ShapeExprError::OperandOutOfRange {
4175 operand: 7,
4176 operands: 1,
4177 },
4178 ),
4179 }),
4180 );
4181 }
4182
4183 #[test]
4184 fn decompose_rel_resolution_failure_is_a_fixpoint_not_a_crash() {
4185 // slice_start_rel = 0 − 2 → Negative at emit time. Registration
4186 // TOLERATES it (a value-dependent decline at the probe shape is
4187 // not an authoring error); the decompose-path resolution failure
4188 // surfaces as a no-op fixpoint (G2), never a panic.
4189 let neg = Dim::Sub(Box::new(Dim::Const(0)), Box::new(Dim::Const(2)));
4190 let region = op_node(
4191 OpTag::Slice,
4192 OpAttrs {
4193 axis: Some(1),
4194 slice_start_rel: Some(neg),
4195 slice_len: Some(1),
4196 ..OpAttrs::default()
4197 },
4198 vec![bind(0)],
4199 );
4200 let id = register_runtime_fused("t3::neg_slice", region)
4201 .expect("a value-dependent decline still registers");
4202 let mut g = Graph::new();
4203 let x = cst(&mut g, &[2, 4]);
4204 let fused = g.push(Node {
4205 op: Op::Fused(id, FusedOpParams::Runtime { scalars: vec![] }),
4206 inputs: vec![x],
4207 shape: Shape::from_dims(&[2, 4]),
4208 dtype: DType::F32,
4209 });
4210 assert_eq!(
4211 decompose_region(&mut g, fused),
4212 fused,
4213 "resolution decline ⇒ fixpoint"
4214 );
4215 }
4216
4217 #[test]
4218 fn try_emit_region_surfaces_typed_resolution_errors() {
4219 // Negative → typed error from the fallible entry, never a panic.
4220 let neg = Dim::Sub(Box::new(Dim::Const(0)), Box::new(Dim::Const(2)));
4221 let region = op_node(
4222 OpTag::Slice,
4223 OpAttrs {
4224 axis: Some(1),
4225 slice_start_rel: Some(neg),
4226 slice_len: Some(1),
4227 ..OpAttrs::default()
4228 },
4229 vec![bind(0)],
4230 );
4231 let mut g = Graph::new();
4232 let x = cst(&mut g, &[2, 4]);
4233 assert_eq!(
4234 try_emit_region(&mut g, ®ion, &[x], &[]),
4235 Err(RelAttrError::Negative {
4236 field: "slice_start",
4237 value: -2
4238 }),
4239 );
4240 // A graph-side SYMBOLIC bind extent maps to the §6.20 SYMBOLIC
4241 // sentinel → SymbolicGap (the surfaced-gap posture, §6.20-0004).
4242 let region = op_node(OpTag::BroadcastTo, bcast_same_as_0(), vec![bind(0)]);
4243 let dyn_shape = Shape::from_dims(&[2, 8]).with_dynamic_axis(1, 0, fuel_ir::SymId(0));
4244 let d = g.push(Node {
4245 op: Op::Const,
4246 inputs: vec![],
4247 shape: dyn_shape,
4248 dtype: DType::F32,
4249 });
4250 assert_eq!(
4251 try_emit_region(&mut g, ®ion, &[d], &[]),
4252 Err(RelAttrError::SymbolicGap {
4253 field: "target_shape"
4254 }),
4255 );
4256 }
4257 }
4258
4259 // ---- T5 (Increment C slice 1): identity-share of repeated subtrees ----
4260 //
4261 // A `PatternNode` recipe is a TREE; a DAG recipe (softmax's shared
4262 // `e = Exp(..)` interior, consumed by both the denominator reduce and the
4263 // final Div) is spelled by REPEATING the subtree. `emit` must emit a
4264 // repeated slot-free subtree ONCE per emit call (identity-share), so the
4265 // emitted graph is the DAG, not a duplicated-compute tree. Subtrees with
4266 // OPEN scalar slots are never shared — each occurrence takes its own
4267 // cursor value. (The flat-DAG node table with real CSE is slice 3.)
4268
4269 #[test]
4270 fn emit_shares_repeated_slot_free_subtrees() {
4271 let region = op_node(
4272 OpTag::Add,
4273 OpAttrs::default(),
4274 vec![
4275 op_node(OpTag::Exp, OpAttrs::default(), vec![bind(0)]),
4276 op_node(OpTag::Exp, OpAttrs::default(), vec![bind(0)]),
4277 ],
4278 );
4279 let mut g = Graph::new();
4280 let x = g.push(Node {
4281 op: Op::Const,
4282 inputs: vec![],
4283 shape: Shape::from_dims(&[4]),
4284 dtype: DType::F32,
4285 });
4286 let root = emit_region(&mut g, ®ion, &[x], &[]);
4287 let add = g.node(root);
4288 assert!(matches!(add.op, Op::Add));
4289 assert_eq!(
4290 add.inputs[0], add.inputs[1],
4291 "structurally-equal slot-free subtrees must share ONE emitted node",
4292 );
4293 }
4294
4295 #[test]
4296 fn emit_does_not_share_subtrees_with_open_scalar_slots() {
4297 // Two open MulScalar slots take DIFFERENT cursor values (pre-order
4298 // fill) — sharing them would silently drop the second live value.
4299 let region = op_node(
4300 OpTag::Add,
4301 OpAttrs::default(),
4302 vec![
4303 op_node(OpTag::MulScalar, OpAttrs::default(), vec![bind(0)]),
4304 op_node(OpTag::MulScalar, OpAttrs::default(), vec![bind(0)]),
4305 ],
4306 );
4307 let mut g = Graph::new();
4308 let x = g.push(Node {
4309 op: Op::Const,
4310 inputs: vec![],
4311 shape: Shape::from_dims(&[4]),
4312 dtype: DType::F32,
4313 });
4314 let root = emit_region(&mut g, ®ion, &[x], &[2.0, 3.0]);
4315 let a = g.node(root).inputs[0];
4316 let b = g.node(root).inputs[1];
4317 assert_ne!(a, b, "open-slot subtrees are never shared");
4318 assert!(matches!(g.node(a).op, Op::MulScalar(v) if v == 2.0));
4319 assert!(matches!(g.node(b).op, Op::MulScalar(v) if v == 3.0));
4320 }
4321
4322 #[test]
4323 fn emit_matches_cast_over_add_reference() {
4324 // Exercises the dtype path through assert_structural_eq: a hand-built
4325 // two-node reference `Cast(F16)(Add(a, b))` vs the emitted region.
4326 use fuel_ir::{DType, Shape};
4327 let mut g = Graph::new();
4328 let sh = Shape::from_dims(&[4]);
4329 let a = g.push(Node {
4330 op: Op::Const,
4331 inputs: vec![],
4332 shape: sh.clone(),
4333 dtype: DType::F32,
4334 });
4335 let b = g.push(Node {
4336 op: Op::Const,
4337 inputs: vec![],
4338 shape: sh.clone(),
4339 dtype: DType::F32,
4340 });
4341 // Reference graph.
4342 let add = g.push(Node {
4343 op: Op::Add,
4344 inputs: vec![a, b],
4345 shape: sh.clone(),
4346 dtype: DType::F32,
4347 });
4348 let reference = g.push(Node {
4349 op: Op::Cast(DType::F16),
4350 inputs: vec![add],
4351 shape: sh.clone(),
4352 dtype: DType::F16,
4353 });
4354
4355 let region = op_node(
4356 OpTag::Cast,
4357 OpAttrs {
4358 cast_dtype: Some("f16".into()),
4359 ..OpAttrs::default()
4360 },
4361 vec![op_node(
4362 OpTag::Add,
4363 OpAttrs::default(),
4364 vec![bind(0), bind(1)],
4365 )],
4366 );
4367 let emitted = emit_region(&mut g, ®ion, &[a, b], &[]);
4368 assert_structural_eq(&g, reference, emitted);
4369 }
4370
4371 // ---- T5 (Increment C slice 1): decompose_via_recipe bridge + the
4372 // softmax_last_dim pilot migration --------------------------------------
4373 //
4374 // The registry bridge (`crate::registry::decompose_via_recipe`, design
4375 // D6) makes a static entry's `decompose` a re-emit of portable
4376 // `PatternNode` DATA: node inputs are the binds, a per-entry projection
4377 // supplies the open-slot scalars, the resolving emit does the rest. ANY
4378 // failure — wrong params payload, a semantics-absent op token (the
4379 // flip-withdrawal posture: unknown/non-registry tokens are surfaced
4380 // honest-miss declines, NEVER accepted, NEVER a crash), a bind/arity or
4381 // slot-count mismatch, a rel-resolution decline at these shapes — returns
4382 // `id` (fixpoint, G2), never panics.
4383
4384 mod recipe_bridge {
4385 use super::super::*;
4386 use super::{bind, frozen_legacy_softmax_decompose, op_node};
4387 use crate::registry::{FusedOps, decompose_via_recipe};
4388 use fuel_ir::{DType, Shape};
4389 use std::collections::HashMap;
4390
4391 /// Tiny f64 reference interpreter over the primitive vocabulary the
4392 /// two softmax spellings use (Const leaves, last-axis reduces, keepdim
4393 /// restores, last-dim broadcast, elementwise). BOTH parity sides run
4394 /// through it, with in-order accumulation per row — so the bit-exact
4395 /// assert isolates recipe STRUCTURE; float noise can't differ between
4396 /// two evaluations of the same interpreter. (Not code evaluation: a
4397 /// closed match over our own `Op` enum — no dynamic execution.)
4398 fn eval(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
4399 let node = g.node(id);
4400 match &node.op {
4401 Op::Const => leaves.get(&id).expect("leaf data provided").clone(),
4402 Op::Exp => eval(g, node.inputs[0], leaves)
4403 .iter()
4404 .map(|v| v.exp())
4405 .collect(),
4406 Op::Sub => {
4407 let a = eval(g, node.inputs[0], leaves);
4408 let b = eval(g, node.inputs[1], leaves);
4409 a.iter().zip(&b).map(|(x, y)| x - y).collect()
4410 }
4411 Op::Div => {
4412 let a = eval(g, node.inputs[0], leaves);
4413 let b = eval(g, node.inputs[1], leaves);
4414 a.iter().zip(&b).map(|(x, y)| x / y).collect()
4415 }
4416 // Last-axis reduces — one arm per spelling pair, identical
4417 // in-order fold.
4418 Op::MaxDim(_) | Op::ReduceMaxTo(_) => {
4419 let input = eval(g, node.inputs[0], leaves);
4420 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
4421 input
4422 .chunks(last)
4423 .map(|row| row.iter().copied().fold(f64::NEG_INFINITY, f64::max))
4424 .collect()
4425 }
4426 Op::SumDim(_) | Op::ReduceSumTo(_) => {
4427 let input = eval(g, node.inputs[0], leaves);
4428 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
4429 input.chunks(last).map(|row| row.iter().sum()).collect()
4430 }
4431 // Metadata-only keepdim restores.
4432 Op::Unsqueeze { .. } | Op::Reshape(_) => eval(g, node.inputs[0], leaves),
4433 // Broadcast a keepdim/reduced tensor back along the last axis.
4434 Op::BroadcastTo(target) => {
4435 let input = eval(g, node.inputs[0], leaves);
4436 let out_n: usize = target.dims().iter().product();
4437 let last = *target.dims().last().unwrap();
4438 assert_eq!(
4439 input.len() * last,
4440 out_n,
4441 "broadcast is a last-dim repeat in these graphs",
4442 );
4443 input
4444 .iter()
4445 .flat_map(|&v| std::iter::repeat_n(v, last))
4446 .collect()
4447 }
4448 other => panic!("eval: unhandled op {other:?}"),
4449 }
4450 }
4451
4452 fn softmax_fused_node(g: &mut Graph, dims: &[usize]) -> (NodeId, NodeId) {
4453 let sh = Shape::from_dims(dims);
4454 let x = g.push(Node {
4455 op: Op::Const,
4456 inputs: vec![],
4457 shape: sh.clone(),
4458 dtype: DType::F32,
4459 });
4460 let fused = g.push(Node {
4461 op: Op::Fused(FusedOps::SOFTMAX_LAST_DIM, FusedOpParams::SoftmaxLastDim),
4462 inputs: vec![x],
4463 shape: sh,
4464 dtype: DType::F32,
4465 });
4466 (x, fused)
4467 }
4468
4469 /// T5 red (a): ONE recipe datum decomposes at BOTH rank 2 and rank 3
4470 /// (the polymorphism the baked-shape legacy body never had), and its
4471 /// numerics match the FROZEN legacy builder bit-exactly under the
4472 /// shared reference interpreter.
4473 #[test]
4474 fn softmax_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
4475 for dims in [vec![2usize, 4], vec![3, 5, 7]] {
4476 let mut g = Graph::new();
4477 let (x, fused) = softmax_fused_node(&mut g, &dims);
4478 let sh = Shape::from_dims(&dims);
4479 let new_root = crate::registry::softmax_last_dim::decompose(
4480 &mut g,
4481 fused,
4482 &FusedOpParams::SoftmaxLastDim,
4483 );
4484 assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
4485 assert_eq!(g.node(new_root).shape, sh, "softmax is shape-preserving");
4486 assert_eq!(g.node(new_root).dtype, DType::F32);
4487
4488 let legacy_root =
4489 frozen_legacy_softmax_decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
4490
4491 let n: usize = dims.iter().product();
4492 let data: Vec<f64> = (0..n)
4493 .map(|i| ((i as f64) * 0.37).sin() * 3.0 - 0.5)
4494 .collect();
4495 let mut leaves = HashMap::new();
4496 leaves.insert(x, data);
4497 let got = eval(&g, new_root, &leaves);
4498 let want = eval(&g, legacy_root, &leaves);
4499 assert_eq!(got.len(), want.len());
4500 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
4501 assert_eq!(
4502 a.to_bits(),
4503 b.to_bits(),
4504 "softmax[{i}] at {dims:?}: recipe={a} vs legacy={b}",
4505 );
4506 }
4507 }
4508 }
4509
4510 /// T5 red (b): the structural golden — the ratified D3 shrink-via-swap
4511 /// spelling, 9 op nodes with the `e = Exp(..)` interior SHARED (node
4512 /// identity) between the denominator reduce and the final Div.
4513 #[test]
4514 fn softmax_recipe_emits_the_nine_node_shared_spelling() {
4515 let mut g = Graph::new();
4516 let (x, fused) = softmax_fused_node(&mut g, &[2, 4]);
4517 let sh = Shape::from_dims(&[2, 4]);
4518 let root = crate::registry::softmax_last_dim::decompose(
4519 &mut g,
4520 fused,
4521 &FusedOpParams::SoftmaxLastDim,
4522 );
4523
4524 // out = Div(e, db)
4525 assert!(matches!(g.node(root).op, Op::Div));
4526 let e = g.node(root).inputs[0];
4527 let db = g.node(root).inputs[1];
4528 assert!(matches!(g.node(e).op, Op::Exp));
4529 assert_eq!(g.node(db).op, Op::BroadcastTo(sh.clone()));
4530 // db = BroadcastTo(Unsqueeze(SumDim(e))) — the SAME e node.
4531 let u2 = g.node(db).inputs[0];
4532 assert!(matches!(g.node(u2).op, Op::Unsqueeze { dim: 1 }));
4533 assert_eq!(g.node(u2).shape, Shape::from_dims(&[2, 1]));
4534 let d = g.node(u2).inputs[0];
4535 assert!(matches!(g.node(d).op, Op::SumDim(1)));
4536 assert_eq!(g.node(d).shape, Shape::from_dims(&[2]));
4537 assert_eq!(
4538 g.node(d).inputs[0],
4539 e,
4540 "the denominator reduces the SHARED Exp node — identity-share, not a duplicate",
4541 );
4542 // e = Exp(Sub(x, mb)); mb = BroadcastTo(Unsqueeze(MaxDim(x))).
4543 let s = g.node(e).inputs[0];
4544 assert!(matches!(g.node(s).op, Op::Sub));
4545 assert_eq!(g.node(s).inputs[0], x);
4546 let mb = g.node(s).inputs[1];
4547 assert_eq!(g.node(mb).op, Op::BroadcastTo(sh.clone()));
4548 let u1 = g.node(mb).inputs[0];
4549 assert!(matches!(g.node(u1).op, Op::Unsqueeze { dim: 1 }));
4550 let m = g.node(u1).inputs[0];
4551 assert!(matches!(g.node(m).op, Op::MaxDim(1)));
4552 assert_eq!(g.node(m).inputs[0], x);
4553 // 9 op nodes + the x leaf = 10 reachable (NO duplicated interior).
4554 assert_eq!(
4555 crate::topo_order_multi(&g, &[root]).len(),
4556 10,
4557 "MaxDim/Unsqueeze/Bcast/Sub/Exp/SumDim/Unsqueeze/Bcast/Div + leaf",
4558 );
4559 }
4560
4561 /// T5 red (c): totality — a wrong params payload is a typed decline
4562 /// surfaced as a fixpoint (G2), never a panic, and declines BEFORE any
4563 /// emission (no partial nodes).
4564 #[test]
4565 fn softmax_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
4566 let mut g = Graph::new();
4567 let (_x, fused) = softmax_fused_node(&mut g, &[2, 4]);
4568 let before = g.len();
4569 let out =
4570 crate::registry::softmax_last_dim::decompose(&mut g, fused, &FusedOpParams::Rope);
4571 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
4572 assert_eq!(g.len(), before, "declined before any emission");
4573 }
4574
4575 /// INJECTED item (2), the flip-withdrawal posture (Baracuda #68 /
4576 /// KISS-Ops closed registry): an op token with no registry semantics —
4577 /// the in-memory analog of the withdrawn reverse-scan "flip" spelling —
4578 /// must surface as a typed honest-miss decline: the node stays fused
4579 /// (fixpoint), NEVER accepted, NEVER a crash. The fabricated recipe
4580 /// stands in for a foreign token by carrying `OpTag::Clamp` — a tag
4581 /// with NO primitive re-emission today (`tag_to_op` → `None`), exactly
4582 /// the semantics-absent posture an unregistered op name resolves to.
4583 /// (This stand-in was `OpTag::PowI` until Increment C carriers A3 gave
4584 /// PowI its i32-exponent carrier; `Clamp` — still awaiting a two-scalar
4585 /// carrier — is the current canonical carrier-less token.) If/when the
4586 /// token registers, it becomes a NAMED-op resolution case — semantics
4587 /// arrive via registration, never via silent acceptance here.
4588 #[test]
4589 fn decompose_via_recipe_declines_an_unknown_token_recipe() {
4590 let fabricated = op_node(
4591 OpTag::Clamp,
4592 OpAttrs {
4593 scalars: vec![0.0, 1.0],
4594 ..OpAttrs::default()
4595 },
4596 vec![bind(0)],
4597 );
4598 let mut g = Graph::new();
4599 let (_x, fused) = softmax_fused_node(&mut g, &[2, 4]);
4600 let before = g.len();
4601 let out = decompose_via_recipe(&mut g, fused, &fabricated, Some(Vec::new()));
4602 assert_eq!(out, fused, "semantics-absent token ⇒ honest-miss fixpoint");
4603 assert_eq!(
4604 g.len(),
4605 before,
4606 "declined BEFORE any emission — no partial nodes"
4607 );
4608 }
4609
4610 /// The bridge's bind/input arity guard: a recipe over 2 binds cannot
4611 /// decompose a 1-input node — fixpoint, not a crash (and not a
4612 /// misbound emission).
4613 #[test]
4614 fn decompose_via_recipe_bind_arity_mismatch_is_a_fixpoint() {
4615 let recipe = op_node(OpTag::Add, OpAttrs::default(), vec![bind(0), bind(1)]);
4616 let mut g = Graph::new();
4617 let (_x, fused) = softmax_fused_node(&mut g, &[2, 4]);
4618 let before = g.len();
4619 let out = decompose_via_recipe(&mut g, fused, &recipe, Some(Vec::new()));
4620 assert_eq!(out, fused, "bind/input arity mismatch ⇒ fixpoint");
4621 assert_eq!(g.len(), before, "declined before any emission");
4622 }
4623 }
4624
4625 // rope migration (Increment C slice 1, T6) ------------------------------
4626 //
4627 // Rope's 11-node imperative body becomes a 9-node portable `PatternNode`
4628 // DATA recipe: cos/sin broadcasts carry `SameAs { operand: 0 }`, the two
4629 // half-slices carry `DimExpr` start/len over the Bind space (the
4630 // reference-doc worked example: `start=Const(0), len=Div(E,2)` /
4631 // `start=Div(E,2), len=Sub(E, Div(E,2))`), the last-axis Concat carries
4632 // `axis_last`, and the two leading-1-padded prep `Reshape`s are NOT in the
4633 // datum — the emit resolver MATERIALIZES them (D4) only where the
4634 // broadcast target out-ranks its operand. Consequence: at a rank-RAISING
4635 // broadcast (the real attention consumer: cos/sin `[seq,d]`, x `[..,seq,d]`
4636 // rank ≥ 3) emission is BYTE-IDENTICAL to legacy (11 nodes, both pads); at
4637 // EQUAL rank (x itself rank 2 = `[seq,d]`) the recipe emits the 9-node
4638 // form, eliding legacy's no-op `Reshape([seq,d]→[seq,d])` — numerically
4639 // identical, structurally leaner. D4 is shared with softmax/norms and MUST
4640 // NOT add reshapes at equal rank (that would break the softmax parity
4641 // oracle), so the equal-rank elision is intrinsic, not a defect.
4642 mod rope_recipe {
4643 use super::super::*;
4644 use super::frozen_legacy_rope_decompose;
4645 use crate::registry::{FusedOps, rope};
4646 use fuel_ir::{DType, Shape};
4647 use std::collections::HashMap;
4648
4649 /// Tiny f64 reference interpreter over the rope primitive vocabulary
4650 /// (Const leaves, metadata-only rank-pad Reshape, leading-dim
4651 /// BroadcastTo, last-axis Slice/Concat, Neg, elementwise Mul/Add).
4652 /// BOTH parity sides run through it in identical in-order arithmetic —
4653 /// so a bit-exact assert isolates recipe STRUCTURE. (Not code
4654 /// evaluation: a closed match over our own `Op` enum.)
4655 fn eval_rope(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
4656 let node = g.node(id);
4657 match &node.op {
4658 Op::Const => leaves.get(&id).expect("leaf data provided").clone(),
4659 Op::Neg => eval_rope(g, node.inputs[0], leaves)
4660 .iter()
4661 .map(|v| -v)
4662 .collect(),
4663 Op::Mul => {
4664 let a = eval_rope(g, node.inputs[0], leaves);
4665 let b = eval_rope(g, node.inputs[1], leaves);
4666 a.iter().zip(&b).map(|(x, y)| x * y).collect()
4667 }
4668 Op::Add => {
4669 let a = eval_rope(g, node.inputs[0], leaves);
4670 let b = eval_rope(g, node.inputs[1], leaves);
4671 a.iter().zip(&b).map(|(x, y)| x + y).collect()
4672 }
4673 // Metadata-only leading-1 rank-pad: same row-major order.
4674 Op::Reshape(_) => eval_rope(g, node.inputs[0], leaves),
4675 // Broadcast a `[1,..,1,seq,d]` inner block over leading dims:
4676 // the block (all of `input`, since the leading dims are 1)
4677 // tiles to fill the target — row-major `input[i % block]`.
4678 Op::BroadcastTo(target) => {
4679 let input = eval_rope(g, node.inputs[0], leaves);
4680 let out_n: usize = target.dims().iter().product();
4681 (0..out_n).map(|i| input[i % input.len()]).collect()
4682 }
4683 Op::Slice { dim, start, len } => {
4684 let input = eval_rope(g, node.inputs[0], leaves);
4685 let in_dims = g.node(node.inputs[0]).shape.dims().to_vec();
4686 let last = in_dims.len() - 1;
4687 assert_eq!(*dim, last, "rope slices along the last axis");
4688 let row = in_dims[last];
4689 input
4690 .chunks(row)
4691 .flat_map(|r| r[*start..*start + *len].to_vec())
4692 .collect()
4693 }
4694 Op::Concat { dim } => {
4695 let a = eval_rope(g, node.inputs[0], leaves);
4696 let b = eval_rope(g, node.inputs[1], leaves);
4697 let a_last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
4698 let b_last = *g.node(node.inputs[1]).shape.dims().last().unwrap();
4699 let last = g.node(node.inputs[0]).shape.dims().len() - 1;
4700 assert_eq!(*dim, last, "rope concats along the last axis");
4701 let mut out = Vec::with_capacity(a.len() + b.len());
4702 let mut ai = a.chunks(a_last);
4703 let mut bi = b.chunks(b_last);
4704 loop {
4705 match (ai.next(), bi.next()) {
4706 (Some(ra), Some(rb)) => {
4707 out.extend_from_slice(ra);
4708 out.extend_from_slice(rb);
4709 }
4710 (None, None) => break,
4711 _ => panic!("concat row-count mismatch"),
4712 }
4713 }
4714 out
4715 }
4716 other => panic!("eval_rope: unhandled op {other:?}"),
4717 }
4718 }
4719
4720 /// Build a fused Rope node over `x [..,seq,d]`, `cos [seq,d]`,
4721 /// `sin [seq,d]`. Returns `(x, cos, sin, fused)`.
4722 fn rope_fused_node(g: &mut Graph, x_dims: &[usize]) -> (NodeId, NodeId, NodeId, NodeId) {
4723 let rank = x_dims.len();
4724 let table_dims = [x_dims[rank - 2], x_dims[rank - 1]];
4725 let x_sh = Shape::from_dims(x_dims);
4726 let t_sh = Shape::from_dims(&table_dims);
4727 let x = g.push(Node {
4728 op: Op::Const,
4729 inputs: vec![],
4730 shape: x_sh.clone(),
4731 dtype: DType::F32,
4732 });
4733 let cos = g.push(Node {
4734 op: Op::Const,
4735 inputs: vec![],
4736 shape: t_sh.clone(),
4737 dtype: DType::F32,
4738 });
4739 let sin = g.push(Node {
4740 op: Op::Const,
4741 inputs: vec![],
4742 shape: t_sh,
4743 dtype: DType::F32,
4744 });
4745 let fused = g.push(Node {
4746 op: Op::Fused(FusedOps::ROPE, FusedOpParams::Rope),
4747 inputs: vec![x, cos, sin],
4748 shape: x_sh,
4749 dtype: DType::F32,
4750 });
4751 (x, cos, sin, fused)
4752 }
4753
4754 /// T6 red (a): ONE recipe datum decomposes at BOTH rank 2 and rank 4
4755 /// — the shape/rank polymorphism the baked-shape legacy body never had
4756 /// — and its numerics match the FROZEN legacy builder bit-exactly
4757 /// under the shared reference interpreter.
4758 #[test]
4759 fn rope_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
4760 for x_dims in [vec![2usize, 4], vec![1, 2, 3, 8]] {
4761 let mut g = Graph::new();
4762 let (x, cos, sin, fused) = rope_fused_node(&mut g, &x_dims);
4763 let x_sh = Shape::from_dims(&x_dims);
4764 let rank = x_dims.len();
4765 let seq = x_dims[rank - 2];
4766 let d = x_dims[rank - 1];
4767
4768 let new_root = rope::decompose(&mut g, fused, &FusedOpParams::Rope);
4769 assert_ne!(new_root, fused, "recipe decompose must fire at {x_dims:?}");
4770 assert_eq!(g.node(new_root).shape, x_sh, "rope is shape-preserving");
4771 assert_eq!(g.node(new_root).dtype, DType::F32);
4772
4773 let legacy_root = frozen_legacy_rope_decompose(&mut g, fused, &FusedOpParams::Rope);
4774
4775 // Distinct, deterministic leaf data for x / cos / sin.
4776 let x_n: usize = x_dims.iter().product();
4777 let t_n = seq * d;
4778 let x_data: Vec<f64> = (0..x_n)
4779 .map(|i| ((i as f64) * 0.31).sin() * 2.0 - 0.4)
4780 .collect();
4781 let cos_data: Vec<f64> =
4782 (0..t_n).map(|i| ((i as f64) * 0.17 + 0.5).cos()).collect();
4783 let sin_data: Vec<f64> =
4784 (0..t_n).map(|i| ((i as f64) * 0.23 - 0.2).sin()).collect();
4785 let mut leaves = HashMap::new();
4786 leaves.insert(x, x_data);
4787 leaves.insert(cos, cos_data);
4788 leaves.insert(sin, sin_data);
4789
4790 let got = eval_rope(&g, new_root, &leaves);
4791 let want = eval_rope(&g, legacy_root, &leaves);
4792 assert_eq!(got.len(), want.len(), "element count at {x_dims:?}");
4793 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
4794 assert_eq!(
4795 a.to_bits(),
4796 b.to_bits(),
4797 "rope[{i}] at {x_dims:?}: recipe={a} vs legacy={b}",
4798 );
4799 }
4800 }
4801 }
4802
4803 /// T6 red (b): at a rank-RAISING broadcast (rank 4 — the real attention
4804 /// consumer's shape), the recipe emission is BYTE-IDENTICAL to legacy:
4805 /// D4 materializes both leading-1-padded prep `Reshape`s, so the whole
4806 /// 11-node DAG matches node-for-node (op, shape, dtype, wiring). This
4807 /// is the byte-identity guarantee that retires all backend risk for the
4808 /// live rope-decompose path.
4809 #[test]
4810 fn rope_recipe_is_byte_identical_to_legacy_at_rank_raise() {
4811 let mut g = Graph::new();
4812 let (_x, _cos, _sin, fused) = rope_fused_node(&mut g, &[1, 2, 3, 8]);
4813 let recipe_root = rope::decompose(&mut g, fused, &FusedOpParams::Rope);
4814 let legacy_root = frozen_legacy_rope_decompose(&mut g, fused, &FusedOpParams::Rope);
4815 super::assert_structural_eq(&g, recipe_root, legacy_root);
4816 // Both leading-1 prep Reshapes are present (byte-identical, 11 ops).
4817 let reachable = crate::topo_order_multi(&g, &[recipe_root]);
4818 let reshapes = reachable
4819 .iter()
4820 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
4821 .count();
4822 assert_eq!(
4823 reshapes, 2,
4824 "rank-raise materializes both legacy prep Reshapes (D4)"
4825 );
4826 }
4827
4828 /// T6 red (c): at EQUAL rank (x itself `[seq,d]`) the recipe emits the
4829 /// 9-node form — the resolver adds NO pad `Reshape` (D4 pads only on a
4830 /// rank-raise), eliding legacy's no-op `Reshape([seq,d]→[seq,d])`.
4831 /// Numerically identical (covered by the parity test), structurally
4832 /// leaner: 9 op nodes + 3 leaves, zero `Reshape`.
4833 #[test]
4834 fn rope_recipe_elides_the_noop_prep_reshape_at_equal_rank() {
4835 let mut g = Graph::new();
4836 let (_x, _cos, _sin, fused) = rope_fused_node(&mut g, &[2, 4]);
4837 let root = rope::decompose(&mut g, fused, &FusedOpParams::Rope);
4838 assert_ne!(root, fused, "recipe decompose fires");
4839 let reachable = crate::topo_order_multi(&g, &[root]);
4840 let reshapes = reachable
4841 .iter()
4842 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
4843 .count();
4844 assert_eq!(
4845 reshapes, 0,
4846 "equal-rank broadcast needs no prep Reshape (D4)"
4847 );
4848 let op_nodes = reachable
4849 .iter()
4850 .filter(|&&n| !matches!(g.node(n).op, Op::Const))
4851 .count();
4852 assert_eq!(
4853 op_nodes, 9,
4854 "the 9-node rope recipe (2×Bcast/2×Slice/Neg/Concat/2×Mul/Add)"
4855 );
4856 }
4857
4858 /// T6 red (d): totality (G2) — a wrong params payload is a typed
4859 /// decline surfaced as a fixpoint, never a panic, and declines BEFORE
4860 /// any emission (no partial nodes). The legacy imperative body ignored
4861 /// `params` entirely and always decomposed; the recipe bridge gates on
4862 /// the projection, so a non-`Rope` payload now correctly no-ops.
4863 #[test]
4864 fn rope_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
4865 let mut g = Graph::new();
4866 let (_x, _cos, _sin, fused) = rope_fused_node(&mut g, &[2, 4]);
4867 let before = g.len();
4868 let out = rope::decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
4869 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
4870 assert_eq!(g.len(), before, "declined before any emission");
4871 }
4872 }
4873
4874 // rms_norm + layer_norm migration (Increment C slice 1, T7) --------------
4875 //
4876 // Both norms' imperative bodies become portable `PatternNode` DATA recipes.
4877 // Two forces at play beyond softmax/rope:
4878 //
4879 // * `eps` is an OPEN scalar slot. The recipe's `AddScalar` carries EMPTY
4880 // `scalars`, so it is a slot template; the per-entry projection
4881 // (`RmsNormLastDim { eps } → vec![eps]` / `LayerNormLastDim { eps } →
4882 // vec![eps]`) supplies the live value, and the resolving emit fills the
4883 // slot in pre-order. The eps-wiring tests below decompose the SAME op at
4884 // two eps values and assert the realized outputs DIFFER accordingly — the
4885 // proof that eps rides the projection→slot path, not a baked constant.
4886 //
4887 // * The keepdim restore is the RATIFIED D3 shrink-via-swap: `Reshape(keepdim)`
4888 // → `Unsqueeze(axis_last = append)` (a node-TYPE change, metadata-only, so
4889 // numerically bit-exact). `MeanDim(axis_last)` stays a rank-reducing mean.
4890 // The parity tests evaluate the new recipe emission and the FROZEN legacy
4891 // builder through one reference interpreter (which treats `Reshape` and
4892 // `Unsqueeze` identically) and assert bit-exact equivalence at two ranks.
4893 //
4894 // Neither norm ever trips D4 (the keepdim restore rebuilds rank BEFORE the
4895 // broadcast, so every `BroadcastTo` operand already matches its target's
4896 // rank — no leading-1 pad `Reshape` is materialized).
4897 mod norm_recipe {
4898 use super::super::*;
4899 use super::{frozen_legacy_layer_norm_decompose, frozen_legacy_rms_norm_decompose};
4900 use crate::registry::{FusedOps, layer_norm_last_dim, rms_norm_last_dim};
4901 use fuel_ir::{DType, Shape};
4902 use std::collections::HashMap;
4903
4904 /// Tiny f64 reference interpreter over the norm primitive vocabulary
4905 /// (Const leaves, `Sqr`, last-axis `MeanDim`, metadata-only keepdim
4906 /// restores `Unsqueeze`/`Reshape`, `AddScalar`, `Sqrt`, last-dim
4907 /// `BroadcastTo`, elementwise `Sub`/`Div`). BOTH parity sides run
4908 /// through it with identical in-order arithmetic, so a bit-exact assert
4909 /// isolates recipe STRUCTURE (the `Unsqueeze`-vs-`Reshape` swap can't
4910 /// perturb it). Not code evaluation: a closed match over our own `Op`.
4911 fn eval_norm(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
4912 let node = g.node(id);
4913 match &node.op {
4914 Op::Const => leaves.get(&id).expect("leaf data provided").clone(),
4915 Op::Sqr => eval_norm(g, node.inputs[0], leaves)
4916 .iter()
4917 .map(|v| v * v)
4918 .collect(),
4919 Op::Sqrt => eval_norm(g, node.inputs[0], leaves)
4920 .iter()
4921 .map(|v| v.sqrt())
4922 .collect(),
4923 Op::AddScalar(e) => eval_norm(g, node.inputs[0], leaves)
4924 .iter()
4925 .map(|v| v + e)
4926 .collect(),
4927 Op::Sub => {
4928 let a = eval_norm(g, node.inputs[0], leaves);
4929 let b = eval_norm(g, node.inputs[1], leaves);
4930 a.iter().zip(&b).map(|(x, y)| x - y).collect()
4931 }
4932 Op::Div => {
4933 let a = eval_norm(g, node.inputs[0], leaves);
4934 let b = eval_norm(g, node.inputs[1], leaves);
4935 a.iter().zip(&b).map(|(x, y)| x / y).collect()
4936 }
4937 // Last-axis mean — rank-reducing; identical fold both spellings.
4938 Op::MeanDim(_) => {
4939 let input = eval_norm(g, node.inputs[0], leaves);
4940 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
4941 input
4942 .chunks(last)
4943 .map(|row| row.iter().sum::<f64>() / last as f64)
4944 .collect()
4945 }
4946 // Metadata-only keepdim restores (the D3 swap and its legacy
4947 // twin evaluate identically here).
4948 Op::Unsqueeze { .. } | Op::Reshape(_) => eval_norm(g, node.inputs[0], leaves),
4949 // Broadcast a keepdim `[.., 1]` tensor back along the last axis.
4950 Op::BroadcastTo(target) => {
4951 let input = eval_norm(g, node.inputs[0], leaves);
4952 let out_n: usize = target.dims().iter().product();
4953 let last = *target.dims().last().unwrap();
4954 assert_eq!(
4955 input.len() * last,
4956 out_n,
4957 "broadcast is a last-dim repeat in these graphs",
4958 );
4959 input
4960 .iter()
4961 .flat_map(|&v| std::iter::repeat_n(v, last))
4962 .collect()
4963 }
4964 other => panic!("eval_norm: unhandled op {other:?}"),
4965 }
4966 }
4967
4968 /// Build a fused RmsNormLastDim node over `x [dims]`, carrying `eps`.
4969 /// Returns `(x, fused)`.
4970 fn rms_norm_fused_node(g: &mut Graph, dims: &[usize], eps: f64) -> (NodeId, NodeId) {
4971 let sh = Shape::from_dims(dims);
4972 let x = g.push(Node {
4973 op: Op::Const,
4974 inputs: vec![],
4975 shape: sh.clone(),
4976 dtype: DType::F32,
4977 });
4978 let fused = g.push(Node {
4979 op: Op::Fused(
4980 FusedOps::RMS_NORM_LAST_DIM,
4981 FusedOpParams::RmsNormLastDim { eps },
4982 ),
4983 inputs: vec![x],
4984 shape: sh,
4985 dtype: DType::F32,
4986 });
4987 (x, fused)
4988 }
4989
4990 /// Build a fused LayerNormLastDim node over `x [dims]`, carrying `eps`.
4991 /// Returns `(x, fused)`.
4992 fn layer_norm_fused_node(g: &mut Graph, dims: &[usize], eps: f64) -> (NodeId, NodeId) {
4993 let sh = Shape::from_dims(dims);
4994 let x = g.push(Node {
4995 op: Op::Const,
4996 inputs: vec![],
4997 shape: sh.clone(),
4998 dtype: DType::F32,
4999 });
5000 let fused = g.push(Node {
5001 op: Op::Fused(
5002 FusedOps::LAYER_NORM_LAST_DIM,
5003 FusedOpParams::LayerNormLastDim { eps },
5004 ),
5005 inputs: vec![x],
5006 shape: sh,
5007 dtype: DType::F32,
5008 });
5009 (x, fused)
5010 }
5011
5012 /// T7 red (a, rms): ONE recipe datum decomposes at BOTH rank 2 and rank
5013 /// 3 (the polymorphism the baked-shape legacy body never had), and its
5014 /// numerics match the FROZEN legacy builder bit-exactly.
5015 #[test]
5016 fn rms_norm_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
5017 for dims in [vec![2usize, 4], vec![3, 5, 7]] {
5018 let mut g = Graph::new();
5019 let (x, fused) = rms_norm_fused_node(&mut g, &dims, 1e-5);
5020 let sh = Shape::from_dims(&dims);
5021 let new_root = rms_norm_last_dim::decompose(
5022 &mut g,
5023 fused,
5024 &FusedOpParams::RmsNormLastDim { eps: 1e-5 },
5025 );
5026 assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
5027 assert_eq!(g.node(new_root).shape, sh, "rms_norm is shape-preserving");
5028 assert_eq!(g.node(new_root).dtype, DType::F32);
5029
5030 let legacy_root = frozen_legacy_rms_norm_decompose(
5031 &mut g,
5032 fused,
5033 &FusedOpParams::RmsNormLastDim { eps: 1e-5 },
5034 );
5035
5036 let n: usize = dims.iter().product();
5037 let data: Vec<f64> = (0..n)
5038 .map(|i| ((i as f64) * 0.37).sin() * 3.0 - 0.5)
5039 .collect();
5040 let mut leaves = HashMap::new();
5041 leaves.insert(x, data);
5042 let got = eval_norm(&g, new_root, &leaves);
5043 let want = eval_norm(&g, legacy_root, &leaves);
5044 assert_eq!(got.len(), want.len());
5045 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
5046 assert_eq!(
5047 a.to_bits(),
5048 b.to_bits(),
5049 "rms_norm[{i}] at {dims:?}: recipe={a} vs legacy={b}",
5050 );
5051 }
5052 }
5053 }
5054
5055 /// T7 red (a, layer): same polymorphism + bit-exact parity for the
5056 /// 11-node layer-norm recipe (with the `centered` subterm identity-
5057 /// shared between `Sqr` and the final `Div`).
5058 #[test]
5059 fn layer_norm_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
5060 for dims in [vec![2usize, 4], vec![3, 5, 7]] {
5061 let mut g = Graph::new();
5062 let (x, fused) = layer_norm_fused_node(&mut g, &dims, 1e-5);
5063 let sh = Shape::from_dims(&dims);
5064 let new_root = layer_norm_last_dim::decompose(
5065 &mut g,
5066 fused,
5067 &FusedOpParams::LayerNormLastDim { eps: 1e-5 },
5068 );
5069 assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
5070 assert_eq!(g.node(new_root).shape, sh, "layer_norm is shape-preserving");
5071 assert_eq!(g.node(new_root).dtype, DType::F32);
5072
5073 let legacy_root = frozen_legacy_layer_norm_decompose(
5074 &mut g,
5075 fused,
5076 &FusedOpParams::LayerNormLastDim { eps: 1e-5 },
5077 );
5078
5079 let n: usize = dims.iter().product();
5080 let data: Vec<f64> = (0..n)
5081 .map(|i| ((i as f64) * 0.29).cos() * 2.0 + 0.3)
5082 .collect();
5083 let mut leaves = HashMap::new();
5084 leaves.insert(x, data);
5085 let got = eval_norm(&g, new_root, &leaves);
5086 let want = eval_norm(&g, legacy_root, &leaves);
5087 assert_eq!(got.len(), want.len());
5088 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
5089 assert_eq!(
5090 a.to_bits(),
5091 b.to_bits(),
5092 "layer_norm[{i}] at {dims:?}: recipe={a} vs legacy={b}",
5093 );
5094 }
5095 }
5096 }
5097
5098 /// T7 red (structural, rms): the keepdim restore is the D3 shrink-via-
5099 /// swap — `Unsqueeze` append, NOT a baked `Reshape(keepdim)`. This is
5100 /// the crisp discriminator against the pre-migration imperative body.
5101 #[test]
5102 fn rms_norm_recipe_uses_the_unsqueeze_keepdim_swap() {
5103 let mut g = Graph::new();
5104 let (_x, fused) = rms_norm_fused_node(&mut g, &[2, 4], 1e-5);
5105 let root = rms_norm_last_dim::decompose(
5106 &mut g,
5107 fused,
5108 &FusedOpParams::RmsNormLastDim { eps: 1e-5 },
5109 );
5110 assert_ne!(root, fused, "recipe decompose fires");
5111 let reachable = crate::topo_order_multi(&g, &[root]);
5112 let unsqueezes = reachable
5113 .iter()
5114 .filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
5115 .count();
5116 let reshapes = reachable
5117 .iter()
5118 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
5119 .count();
5120 assert_eq!(
5121 unsqueezes, 1,
5122 "keepdim restored via Unsqueeze append (D3 swap)"
5123 );
5124 assert_eq!(reshapes, 0, "no baked keepdim Reshape after the D3 swap");
5125 }
5126
5127 /// T7 red (structural, layer): both keepdim restores are `Unsqueeze`
5128 /// appends; zero `Reshape` (the equal-rank broadcasts add no D4 pad).
5129 #[test]
5130 fn layer_norm_recipe_uses_the_unsqueeze_keepdim_swap() {
5131 let mut g = Graph::new();
5132 let (_x, fused) = layer_norm_fused_node(&mut g, &[2, 4], 1e-5);
5133 let root = layer_norm_last_dim::decompose(
5134 &mut g,
5135 fused,
5136 &FusedOpParams::LayerNormLastDim { eps: 1e-5 },
5137 );
5138 assert_ne!(root, fused, "recipe decompose fires");
5139 let reachable = crate::topo_order_multi(&g, &[root]);
5140 let unsqueezes = reachable
5141 .iter()
5142 .filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
5143 .count();
5144 let reshapes = reachable
5145 .iter()
5146 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
5147 .count();
5148 assert_eq!(
5149 unsqueezes, 2,
5150 "both keepdim restores via Unsqueeze append (D3 swap)"
5151 );
5152 assert_eq!(
5153 reshapes, 0,
5154 "no baked keepdim Reshape / no D4 pad after the swap"
5155 );
5156 // The `centered` Sub is SHARED (Sqr input == final Div numerator):
5157 // 11 op nodes + 1 leaf = 12 reachable, not the 12-op unshared tree.
5158 let op_nodes = reachable
5159 .iter()
5160 .filter(|&&n| !matches!(g.node(n).op, Op::Const))
5161 .count();
5162 assert_eq!(op_nodes, 11, "11 op nodes with `centered` identity-shared");
5163 }
5164
5165 /// T7 red (eps-wiring, rms): the eps rides the projection→open-slot
5166 /// path. Decomposing the SAME op at two eps values yields DIFFERENT
5167 /// realized outputs — impossible if eps were dropped or baked to a
5168 /// single constant. Small `x` (so `mean(x²) ≈ eps`) makes the eps
5169 /// choice materially move every element.
5170 #[test]
5171 fn rms_norm_recipe_eps_flows_through_the_open_slot() {
5172 let dims = [2usize, 4];
5173 let n: usize = dims.iter().product();
5174 let data: Vec<f64> = (0..n)
5175 .map(|i| 0.001 * (((i as f64) * 0.37).sin() + 1.2))
5176 .collect();
5177
5178 let realize = |eps: f64| -> Vec<f64> {
5179 let mut g = Graph::new();
5180 let (x, fused) = rms_norm_fused_node(&mut g, &dims, eps);
5181 let root = rms_norm_last_dim::decompose(
5182 &mut g,
5183 fused,
5184 &FusedOpParams::RmsNormLastDim { eps },
5185 );
5186 let mut leaves = HashMap::new();
5187 leaves.insert(x, data.clone());
5188 eval_norm(&g, root, &leaves)
5189 };
5190 let a = realize(1e-5);
5191 let b = realize(1e-6);
5192 assert_eq!(a.len(), b.len());
5193 assert_ne!(
5194 a, b,
5195 "different eps must change the output — proves projection→slot, not a baked constant",
5196 );
5197 }
5198
5199 /// T7 red (eps-wiring, layer): same proof for layer-norm's open slot.
5200 #[test]
5201 fn layer_norm_recipe_eps_flows_through_the_open_slot() {
5202 let dims = [2usize, 4];
5203 let n: usize = dims.iter().product();
5204 // Near-constant rows so the variance ≈ eps and the eps choice moves
5205 // the output materially.
5206 let data: Vec<f64> = (0..n)
5207 .map(|i| 1.0 + 0.001 * ((i as f64) * 0.37).sin())
5208 .collect();
5209
5210 let realize = |eps: f64| -> Vec<f64> {
5211 let mut g = Graph::new();
5212 let (x, fused) = layer_norm_fused_node(&mut g, &dims, eps);
5213 let root = layer_norm_last_dim::decompose(
5214 &mut g,
5215 fused,
5216 &FusedOpParams::LayerNormLastDim { eps },
5217 );
5218 let mut leaves = HashMap::new();
5219 leaves.insert(x, data.clone());
5220 eval_norm(&g, root, &leaves)
5221 };
5222 let a = realize(1e-5);
5223 let b = realize(1e-6);
5224 assert_eq!(a.len(), b.len());
5225 assert_ne!(
5226 a, b,
5227 "different eps must change the output — proves projection→slot, not a baked constant",
5228 );
5229 }
5230
5231 /// T7 red (totality, rms): a wrong params payload is a typed decline
5232 /// surfaced as a fixpoint (G2), never a panic, declining BEFORE any
5233 /// emission (no partial nodes).
5234 #[test]
5235 fn rms_norm_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
5236 let mut g = Graph::new();
5237 let (_x, fused) = rms_norm_fused_node(&mut g, &[2, 4], 1e-5);
5238 let before = g.len();
5239 let out = rms_norm_last_dim::decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
5240 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
5241 assert_eq!(g.len(), before, "declined before any emission");
5242 }
5243
5244 /// T7 red (totality, layer): same fixpoint posture for layer-norm.
5245 #[test]
5246 fn layer_norm_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
5247 let mut g = Graph::new();
5248 let (_x, fused) = layer_norm_fused_node(&mut g, &[2, 4], 1e-5);
5249 let before = g.len();
5250 let out = layer_norm_last_dim::decompose(&mut g, fused, &FusedOpParams::SoftmaxLastDim);
5251 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
5252 assert_eq!(g.len(), before, "declined before any emission");
5253 }
5254 }
5255
5256 // softmax_last_dim_backward migration (Increment C slice 1, T8) ----------
5257 //
5258 // The 5-node imperative backward body `s · (g − sum(g·s, last, keepdim))`
5259 // becomes a portable `PatternNode` DATA recipe. Bind space: `0 = s` (the
5260 // forward softmax output), `1 = g` (the upstream gradient) — the order the
5261 // autograd `BackwardKind::Fused(SOFTMAX_LAST_DIM_BACKWARD)` edge emits. The
5262 // keepdim restore is the ratified D3 shrink-via-swap
5263 // (`ReduceSumTo(keepdim)` → `SumDim(axis_last)` + `Unsqueeze(axis_last =
5264 // append)`, node-TYPE change, numerically bit-exact) and the broadcast
5265 // targets `SameAs { operand: 0 }` over the Bind space (D2). D4 never fires
5266 // (the `Unsqueeze` rebuilds rank BEFORE the broadcast, so the broadcast
5267 // operand already matches its target's rank — no leading-1 pad `Reshape`).
5268 // This activates the registry's backward-helper edge END-TO-END on a data
5269 // recipe for the first time.
5270 mod softmax_backward_recipe {
5271 use super::super::*;
5272 use super::frozen_legacy_softmax_backward_decompose;
5273 use crate::registry::{FusedOps, softmax_last_dim_backward};
5274 use fuel_ir::{DType, Shape};
5275 use std::collections::HashMap;
5276
5277 /// Tiny f64 reference interpreter over the softmax-backward primitive
5278 /// vocabulary (leaf-lookup FIRST, then `Mul`/`Sub`, last-axis
5279 /// `SumDim`/`ReduceSumTo`, metadata-only keepdim restore
5280 /// `Unsqueeze`/`Reshape`, last-dim `BroadcastTo`). Leaf-first lets ANY
5281 /// node stand in as a bound input — a `Const`, or the autograd path's
5282 /// forward-softmax (`Op::Fused`) and upstream nodes. BOTH parity sides
5283 /// run through it with identical in-order arithmetic, so a bit-exact
5284 /// assert isolates recipe STRUCTURE (the `SumDim`+`Unsqueeze`-vs-
5285 /// `ReduceSumTo` swap can't perturb it). Not code evaluation: a closed
5286 /// match over our own `Op`.
5287 fn eval_bwd(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
5288 if let Some(v) = leaves.get(&id) {
5289 return v.clone();
5290 }
5291 let node = g.node(id);
5292 match &node.op {
5293 Op::Mul => {
5294 let a = eval_bwd(g, node.inputs[0], leaves);
5295 let b = eval_bwd(g, node.inputs[1], leaves);
5296 a.iter().zip(&b).map(|(x, y)| x * y).collect()
5297 }
5298 Op::Sub => {
5299 let a = eval_bwd(g, node.inputs[0], leaves);
5300 let b = eval_bwd(g, node.inputs[1], leaves);
5301 a.iter().zip(&b).map(|(x, y)| x - y).collect()
5302 }
5303 // Last-axis sum — one arm per spelling pair, identical fold.
5304 Op::SumDim(_) | Op::ReduceSumTo(_) => {
5305 let input = eval_bwd(g, node.inputs[0], leaves);
5306 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
5307 input.chunks(last).map(|row| row.iter().sum()).collect()
5308 }
5309 // Metadata-only keepdim restores (the D3 swap and its legacy
5310 // twin evaluate identically here).
5311 Op::Unsqueeze { .. } | Op::Reshape(_) => eval_bwd(g, node.inputs[0], leaves),
5312 // Broadcast a keepdim/reduced tensor back along the last axis.
5313 Op::BroadcastTo(target) => {
5314 let input = eval_bwd(g, node.inputs[0], leaves);
5315 let out_n: usize = target.dims().iter().product();
5316 let last = *target.dims().last().unwrap();
5317 assert_eq!(
5318 input.len() * last,
5319 out_n,
5320 "broadcast is a last-dim repeat in these graphs",
5321 );
5322 input
5323 .iter()
5324 .flat_map(|&v| std::iter::repeat_n(v, last))
5325 .collect()
5326 }
5327 other => panic!("eval_bwd: unhandled op {other:?}"),
5328 }
5329 }
5330
5331 /// Build a fused SoftmaxLastDimBackward node over `s [dims]` (input 0,
5332 /// the forward output) and `g [dims]` (input 1, the upstream gradient).
5333 /// Returns `(s, g, fused)`.
5334 fn softmax_backward_fused_node(g: &mut Graph, dims: &[usize]) -> (NodeId, NodeId, NodeId) {
5335 let sh = Shape::from_dims(dims);
5336 let s = g.push(Node {
5337 op: Op::Const,
5338 inputs: vec![],
5339 shape: sh.clone(),
5340 dtype: DType::F32,
5341 });
5342 let up = g.push(Node {
5343 op: Op::Const,
5344 inputs: vec![],
5345 shape: sh.clone(),
5346 dtype: DType::F32,
5347 });
5348 let fused = g.push(Node {
5349 op: Op::Fused(
5350 FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
5351 FusedOpParams::SoftmaxLastDimBackward,
5352 ),
5353 inputs: vec![s, up],
5354 shape: sh,
5355 dtype: DType::F32,
5356 });
5357 (s, up, fused)
5358 }
5359
5360 /// T8 red (a): ONE recipe datum decomposes at BOTH rank 2 and rank 3
5361 /// (the polymorphism the baked-shape legacy body never had), and its
5362 /// numerics match the FROZEN legacy builder bit-exactly under the
5363 /// shared reference interpreter.
5364 #[test]
5365 fn softmax_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
5366 for dims in [vec![2usize, 4], vec![3, 5, 7]] {
5367 let mut g = Graph::new();
5368 let (s, up, fused) = softmax_backward_fused_node(&mut g, &dims);
5369 let sh = Shape::from_dims(&dims);
5370 let new_root = softmax_last_dim_backward::decompose(
5371 &mut g,
5372 fused,
5373 &FusedOpParams::SoftmaxLastDimBackward,
5374 );
5375 assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
5376 assert_eq!(
5377 g.node(new_root).shape,
5378 sh,
5379 "softmax backward is shape-preserving"
5380 );
5381 assert_eq!(g.node(new_root).dtype, DType::F32);
5382
5383 let legacy_root = frozen_legacy_softmax_backward_decompose(
5384 &mut g,
5385 fused,
5386 &FusedOpParams::SoftmaxLastDimBackward,
5387 );
5388
5389 let n: usize = dims.iter().product();
5390 let s_data: Vec<f64> = (0..n)
5391 .map(|i| ((i as f64) * 0.29).sin() * 0.5 + 0.5)
5392 .collect();
5393 let g_data: Vec<f64> = (0..n)
5394 .map(|i| ((i as f64) * 0.53).cos() * 2.0 - 0.3)
5395 .collect();
5396 let mut leaves = HashMap::new();
5397 leaves.insert(s, s_data);
5398 leaves.insert(up, g_data);
5399 let got = eval_bwd(&g, new_root, &leaves);
5400 let want = eval_bwd(&g, legacy_root, &leaves);
5401 assert_eq!(got.len(), want.len());
5402 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
5403 assert_eq!(
5404 a.to_bits(),
5405 b.to_bits(),
5406 "softmax_backward[{i}] at {dims:?}: recipe={a} vs legacy={b}",
5407 );
5408 }
5409 }
5410 }
5411
5412 /// T8 red (structural): the keepdim restore is the D3 shrink-via-swap —
5413 /// `SumDim(last)` + `Unsqueeze` append, NOT the baked
5414 /// `ReduceSumTo(keepdim)`. The crisp discriminator against the
5415 /// pre-migration imperative body; the backward root is the outer `Mul`.
5416 #[test]
5417 fn softmax_backward_recipe_uses_the_sumdim_unsqueeze_swap() {
5418 let mut g = Graph::new();
5419 let (_s, _up, fused) = softmax_backward_fused_node(&mut g, &[2, 4]);
5420 let root = softmax_last_dim_backward::decompose(
5421 &mut g,
5422 fused,
5423 &FusedOpParams::SoftmaxLastDimBackward,
5424 );
5425 assert_ne!(root, fused, "recipe decompose fires");
5426 assert!(
5427 matches!(g.node(root).op, Op::Mul),
5428 "backward root is the outer Mul"
5429 );
5430 let reachable = crate::topo_order_multi(&g, &[root]);
5431 let sumdims = reachable
5432 .iter()
5433 .filter(|&&n| matches!(g.node(n).op, Op::SumDim(_)))
5434 .count();
5435 let unsqueezes = reachable
5436 .iter()
5437 .filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
5438 .count();
5439 let reduce_sum_tos = reachable
5440 .iter()
5441 .filter(|&&n| matches!(g.node(n).op, Op::ReduceSumTo(_)))
5442 .count();
5443 assert_eq!(sumdims, 1, "the reduce is SumDim(last) — the D3 swap");
5444 assert_eq!(unsqueezes, 1, "keepdim restored via Unsqueeze append");
5445 assert_eq!(
5446 reduce_sum_tos, 0,
5447 "no baked keepdim ReduceSumTo after the swap"
5448 );
5449 }
5450
5451 /// T8 red (totality): a wrong params payload is a typed decline
5452 /// surfaced as a fixpoint (G2), never a panic, declining BEFORE any
5453 /// emission. (The pre-migration imperative body IGNORED params and
5454 /// always decomposed; the recipe bridge's `scalars(params)` projection
5455 /// is what makes a wrong payload a fixpoint.)
5456 #[test]
5457 fn softmax_backward_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
5458 let mut g = Graph::new();
5459 let (_s, _up, fused) = softmax_backward_fused_node(&mut g, &[2, 4]);
5460 let before = g.len();
5461 let out = softmax_last_dim_backward::decompose(&mut g, fused, &FusedOpParams::Rope);
5462 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
5463 assert_eq!(g.len(), before, "declined before any emission");
5464 }
5465
5466 /// T8 red (autograd path): the `BackwardKind::Fused(SOFTMAX_LAST_DIM_BACKWARD)`
5467 /// edge exercised END-TO-END. Build a softmax forward, backprop; the
5468 /// input gradient node is `Op::Fused(SOFTMAX_LAST_DIM_BACKWARD)` over
5469 /// `[y, upstream]`; decomposing it fires the MIGRATED recipe (the D3
5470 /// SumDim spelling) and matches the frozen legacy numerically — the
5471 /// "realize" leg, via the reference interpreter feeding synthetic leaf
5472 /// data on the two bound inputs (`y = s`, `upstream = g`).
5473 #[test]
5474 fn softmax_backward_reaches_the_recipe_through_autograd() {
5475 let dev: std::sync::Arc<dyn fuel_backend_contract::DynBackendDevice> =
5476 std::sync::Arc::new(fuel_cpu_backend::dyn_impl::CpuBackendDevice);
5477 let x = crate::NodeHandle::from_f32(
5478 vec![0.1f32, -0.2, 0.3, 0.4, -0.5, 0.6],
5479 Shape::from_dims(&[2, 3]),
5480 &dev,
5481 )
5482 .unwrap();
5483 let y = x.softmax_last_dim();
5484 let y_id = y.id();
5485 let grads = y.backward();
5486 let g_x = grads.get(&x).expect("softmax has an input gradient");
5487 let handle = g_x.graph();
5488 let bwd_id = g_x.id();
5489
5490 // The input-gradient node IS the registry backward fused op over
5491 // `[y, upstream]` (x feeds only the softmax, so no accumulation Add).
5492 let up_id = {
5493 let gr = handle.read().unwrap();
5494 let node = gr.node(bwd_id).clone();
5495 match node.op {
5496 Op::Fused(fid, params) => {
5497 assert_eq!(
5498 fid,
5499 FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
5500 "autograd emits the registry backward fused op",
5501 );
5502 assert!(matches!(params, FusedOpParams::SoftmaxLastDimBackward));
5503 }
5504 other => panic!("expected the backward fused op, got {other:?}"),
5505 }
5506 assert_eq!(
5507 node.inputs[0], y_id,
5508 "backward input 0 = the forward softmax output"
5509 );
5510 node.inputs[1]
5511 };
5512
5513 // Decompose the SAME autograd backward node both ways (push-only
5514 // graph — the fused node survives), then compare numerically.
5515 let (new_root, legacy_root, sh) = {
5516 let mut gr = handle.write().unwrap();
5517 let sh = gr.node(bwd_id).shape.clone();
5518 let new_root = softmax_last_dim_backward::decompose(
5519 &mut gr,
5520 bwd_id,
5521 &FusedOpParams::SoftmaxLastDimBackward,
5522 );
5523 let legacy_root = frozen_legacy_softmax_backward_decompose(
5524 &mut gr,
5525 bwd_id,
5526 &FusedOpParams::SoftmaxLastDimBackward,
5527 );
5528 (new_root, legacy_root, sh)
5529 };
5530
5531 let gr = handle.read().unwrap();
5532 assert_ne!(
5533 new_root, bwd_id,
5534 "the autograd backward node decomposes via the recipe"
5535 );
5536 assert_eq!(gr.node(new_root).shape, sh, "shape-preserving");
5537 let reachable = crate::topo_order_multi(&gr, &[new_root]);
5538 assert!(
5539 reachable
5540 .iter()
5541 .any(|&n| matches!(gr.node(n).op, Op::SumDim(_))),
5542 "the autograd path reaches the D3 SumDim spelling",
5543 );
5544
5545 // Numeric parity (leaf-first interpreter over `[y = s, up = g]`).
5546 let n: usize = sh.dims().iter().product();
5547 let s_data: Vec<f64> = (0..n)
5548 .map(|i| ((i as f64) * 0.31).sin() * 0.5 + 0.5)
5549 .collect();
5550 let g_data: Vec<f64> = (0..n).map(|i| ((i as f64) * 0.47).cos() - 0.1).collect();
5551 let mut leaves = HashMap::new();
5552 leaves.insert(y_id, s_data);
5553 leaves.insert(up_id, g_data);
5554 let got = eval_bwd(&gr, new_root, &leaves);
5555 let want = eval_bwd(&gr, legacy_root, &leaves);
5556 assert_eq!(got.len(), want.len());
5557 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
5558 assert_eq!(
5559 a.to_bits(),
5560 b.to_bits(),
5561 "autograd softmax_backward[{i}]: recipe={a} vs legacy={b}",
5562 );
5563 }
5564 }
5565 }
5566
5567 // layer_norm_last_dim_backward migration (Increment C slice 2, S2-1) ------
5568 //
5569 // The ~20-node imperative backward body
5570 // `grad_x = istd · (g − mean(g) − xhat·mean(g·xhat))` (mean over the last
5571 // dim, `xhat = (x − mean(x))·istd`, `istd = rsqrt(var + eps)`) becomes a
5572 // portable `PatternNode` DATA recipe. Bind space: `0 = x`, `1 = g` (the
5573 // upstream gradient) — the order the autograd
5574 // `BackwardKind::Fused(LAYER_NORM_LAST_DIM_BACKWARD)` edge emits. The four
5575 // keepdim restores are the ratified D3 shrink-via-swap (`Reshape(keepdim)`
5576 // → `MeanDim(axis_last)` shrink + `Unsqueeze(axis_last = append)`,
5577 // node-TYPE change, numerically bit-exact); each broadcast targets
5578 // `SameAs { operand: 0 }` (x's full shape) over the Bind space (D2). One
5579 // `eps` open scalar slot. D4 never fires at equal rank (the `Unsqueeze`
5580 // rebuilds rank before every broadcast). RISK-A: `xhat`/`istd` carry the
5581 // eps slot, so emit's slot-free identity-share does NOT dedup them — see
5582 // `..._open_slot_xhat_is_recomputed_...` below.
5583 mod layer_norm_backward_recipe {
5584 use super::super::*;
5585 use super::frozen_legacy_layer_norm_backward_decompose;
5586 use crate::registry::{FusedOps, layer_norm_last_dim_backward};
5587 use fuel_ir::{DType, Shape};
5588 use std::collections::HashMap;
5589
5590 /// Tiny f64 reference interpreter over the layer-norm-backward
5591 /// primitive vocabulary (leaf-lookup FIRST, then elementwise
5592 /// `Mul`/`Sub`, unary `Sqr`/`Rsqrt`, `AddScalar`, last-axis `MeanDim`,
5593 /// metadata-only keepdim restore `Unsqueeze`/`Reshape`, last-dim
5594 /// `BroadcastTo`). BOTH parity sides run through it with identical
5595 /// in-order arithmetic, so a bit-exact assert isolates recipe STRUCTURE
5596 /// — neither the `Reshape`→`Unsqueeze` D3 swap NOR the open-slot
5597 /// `xhat`/`istd` recompute can perturb it (recompute is deterministic).
5598 fn eval_lnb(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
5599 if let Some(v) = leaves.get(&id) {
5600 return v.clone();
5601 }
5602 let node = g.node(id);
5603 match &node.op {
5604 Op::Mul => {
5605 let a = eval_lnb(g, node.inputs[0], leaves);
5606 let b = eval_lnb(g, node.inputs[1], leaves);
5607 a.iter().zip(&b).map(|(x, y)| x * y).collect()
5608 }
5609 Op::Sub => {
5610 let a = eval_lnb(g, node.inputs[0], leaves);
5611 let b = eval_lnb(g, node.inputs[1], leaves);
5612 a.iter().zip(&b).map(|(x, y)| x - y).collect()
5613 }
5614 Op::Sqr => eval_lnb(g, node.inputs[0], leaves)
5615 .iter()
5616 .map(|v| v * v)
5617 .collect(),
5618 Op::Rsqrt => eval_lnb(g, node.inputs[0], leaves)
5619 .iter()
5620 .map(|v| 1.0 / v.sqrt())
5621 .collect(),
5622 Op::AddScalar(e) => eval_lnb(g, node.inputs[0], leaves)
5623 .iter()
5624 .map(|v| v + e)
5625 .collect(),
5626 // Last-axis mean — rank-reducing; identical fold both spellings.
5627 Op::MeanDim(_) => {
5628 let input = eval_lnb(g, node.inputs[0], leaves);
5629 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
5630 input
5631 .chunks(last)
5632 .map(|row| row.iter().sum::<f64>() / last as f64)
5633 .collect()
5634 }
5635 // Metadata-only keepdim restores (the D3 swap and its legacy
5636 // twin evaluate identically here).
5637 Op::Unsqueeze { .. } | Op::Reshape(_) => eval_lnb(g, node.inputs[0], leaves),
5638 // Broadcast a keepdim `[.., 1]` tensor back along the last axis.
5639 Op::BroadcastTo(target) => {
5640 let input = eval_lnb(g, node.inputs[0], leaves);
5641 let out_n: usize = target.dims().iter().product();
5642 let last = *target.dims().last().unwrap();
5643 assert_eq!(
5644 input.len() * last,
5645 out_n,
5646 "broadcast is a last-dim repeat in these graphs",
5647 );
5648 input
5649 .iter()
5650 .flat_map(|&v| std::iter::repeat_n(v, last))
5651 .collect()
5652 }
5653 other => panic!("eval_lnb: unhandled op {other:?}"),
5654 }
5655 }
5656
5657 /// Build a fused LayerNormLastDimBackward node over `x [dims]` (input 0)
5658 /// and `g [dims]` (input 1, the upstream gradient), carrying `eps`.
5659 /// Returns `(x, g, fused)`.
5660 fn ln_backward_fused_node(
5661 g: &mut Graph,
5662 dims: &[usize],
5663 eps: f64,
5664 ) -> (NodeId, NodeId, NodeId) {
5665 let sh = Shape::from_dims(dims);
5666 let x = g.push(Node {
5667 op: Op::Const,
5668 inputs: vec![],
5669 shape: sh.clone(),
5670 dtype: DType::F32,
5671 });
5672 let up = g.push(Node {
5673 op: Op::Const,
5674 inputs: vec![],
5675 shape: sh.clone(),
5676 dtype: DType::F32,
5677 });
5678 let fused = g.push(Node {
5679 op: Op::Fused(
5680 FusedOps::LAYER_NORM_LAST_DIM_BACKWARD,
5681 FusedOpParams::LayerNormLastDimBackward { eps },
5682 ),
5683 inputs: vec![x, up],
5684 shape: sh,
5685 dtype: DType::F32,
5686 });
5687 (x, up, fused)
5688 }
5689
5690 /// S2-1 (a): ONE recipe datum decomposes at BOTH rank 2 and rank 3 (the
5691 /// polymorphism the baked-shape legacy body never had), and its
5692 /// numerics match the FROZEN legacy builder bit-exactly under the
5693 /// shared reference interpreter.
5694 #[test]
5695 fn layer_norm_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
5696 for dims in [vec![2usize, 4], vec![3, 5, 7]] {
5697 let mut g = Graph::new();
5698 let (x, up, fused) = ln_backward_fused_node(&mut g, &dims, 1e-5);
5699 let sh = Shape::from_dims(&dims);
5700 let new_root = layer_norm_last_dim_backward::decompose(
5701 &mut g,
5702 fused,
5703 &FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
5704 );
5705 assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
5706 assert_eq!(
5707 g.node(new_root).shape,
5708 sh,
5709 "layer_norm backward is shape-preserving"
5710 );
5711 assert_eq!(g.node(new_root).dtype, DType::F32);
5712
5713 let legacy_root = frozen_legacy_layer_norm_backward_decompose(
5714 &mut g,
5715 fused,
5716 &FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
5717 );
5718
5719 let n: usize = dims.iter().product();
5720 let x_data: Vec<f64> = (0..n)
5721 .map(|i| ((i as f64) * 0.29).sin() * 2.0 + 0.3)
5722 .collect();
5723 let g_data: Vec<f64> = (0..n)
5724 .map(|i| ((i as f64) * 0.53).cos() * 1.5 - 0.2)
5725 .collect();
5726 let mut leaves = HashMap::new();
5727 leaves.insert(x, x_data);
5728 leaves.insert(up, g_data);
5729 let got = eval_lnb(&g, new_root, &leaves);
5730 let want = eval_lnb(&g, legacy_root, &leaves);
5731 assert_eq!(got.len(), want.len());
5732 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
5733 assert_eq!(
5734 a.to_bits(),
5735 b.to_bits(),
5736 "layer_norm_backward[{i}] at {dims:?}: recipe={a} vs legacy={b}",
5737 );
5738 }
5739 }
5740 }
5741
5742 /// S2-1 (structural, born-red): every keepdim restore is the D3
5743 /// shrink-via-swap `Unsqueeze` append, NOT the baked `Reshape(keepdim)`
5744 /// — the crisp discriminator against the pre-migration imperative body
5745 /// (which emits `Reshape`, so this test is RED until migration). Zero
5746 /// `Reshape` also confirms no D4 pad fires (equal rank throughout). The
5747 /// backward root is the outer `Mul(istd, inner)`.
5748 #[test]
5749 fn layer_norm_backward_recipe_uses_the_unsqueeze_keepdim_swap() {
5750 let mut g = Graph::new();
5751 let (_x, _up, fused) = ln_backward_fused_node(&mut g, &[2, 4], 1e-5);
5752 let root = layer_norm_last_dim_backward::decompose(
5753 &mut g,
5754 fused,
5755 &FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
5756 );
5757 assert_ne!(root, fused, "recipe decompose fires");
5758 assert!(
5759 matches!(g.node(root).op, Op::Mul),
5760 "backward root is the outer Mul(istd, inner)"
5761 );
5762 let reachable = crate::topo_order_multi(&g, &[root]);
5763 let unsqueezes = reachable
5764 .iter()
5765 .filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
5766 .count();
5767 let reshapes = reachable
5768 .iter()
5769 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
5770 .count();
5771 assert!(
5772 unsqueezes >= 1,
5773 "keepdim restored via Unsqueeze append (D3 swap)"
5774 );
5775 assert_eq!(
5776 reshapes, 0,
5777 "no baked keepdim Reshape / no D4 pad after the D3 swap"
5778 );
5779 }
5780
5781 /// S2-1 (RISK-A): `xhat = Mul(xc, istd)` and `istd = Rsqrt(AddScalar[
5782 /// eps](...))` carry the eps OPEN slot, so emit's slot-free
5783 /// identity-share does NOT dedup them — the recipe emits them once per
5784 /// occurrence, yielding a STRICTLY HEAVIER base map than the legacy DAG
5785 /// (which shares them). This is the ACCEPTED redundancy documented on
5786 /// `recipe()` (numerically identical — deterministic recompute — and no
5787 /// open-slot-sharing emit extension is built in this slice). The
5788 /// downstream full-optimizer CSE re-collapses the duplicates, which this
5789 /// pins as the mitigation.
5790 #[test]
5791 fn layer_norm_backward_recipe_open_slot_xhat_is_recomputed_then_cse_recollapses() {
5792 let count_ops = |g: &Graph, root: NodeId| -> usize {
5793 crate::topo_order_multi(g, &[root])
5794 .iter()
5795 .filter(|&&n| !matches!(g.node(n).op, Op::Const))
5796 .count()
5797 };
5798
5799 // Recipe emission (base map — no fused ops remain to lower).
5800 let mut gr = Graph::new();
5801 let (_x, _up, fused) = ln_backward_fused_node(&mut gr, &[2, 4], 1e-5);
5802 let recipe_root = layer_norm_last_dim_backward::decompose(
5803 &mut gr,
5804 fused,
5805 &FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
5806 );
5807 let recipe_n = count_ops(&gr, recipe_root);
5808
5809 // Legacy DAG (shares xhat/istd via let-bindings).
5810 let mut gl = Graph::new();
5811 let (_x2, _up2, fused2) = ln_backward_fused_node(&mut gl, &[2, 4], 1e-5);
5812 let legacy_root = frozen_legacy_layer_norm_backward_decompose(
5813 &mut gl,
5814 fused2,
5815 &FusedOpParams::LayerNormLastDimBackward { eps: 1e-5 },
5816 );
5817 let legacy_n = count_ops(&gl, legacy_root);
5818
5819 assert!(
5820 recipe_n > legacy_n,
5821 "RISK-A: emit's slot-free share does NOT dedup the eps-slot-carrying \
5822 xhat/istd — recipe base map {recipe_n} > legacy {legacy_n}",
5823 );
5824
5825 // Downstream full-optimizer CSE re-collapses the redundant recompute.
5826 let shared: crate::SharedGraph = std::sync::Arc::new(std::sync::RwLock::new(gr));
5827 let cse_roots = crate::opt::optimize(&shared, &[recipe_root]);
5828 let gr = shared.read().unwrap();
5829 let cse_n = count_ops(&gr, cse_roots[0]);
5830 assert!(
5831 cse_n < recipe_n,
5832 "downstream CSE re-collapses the redundant recompute: {cse_n} < {recipe_n}",
5833 );
5834 assert_eq!(
5835 cse_n, legacy_n,
5836 "CSE collapses the recipe base map back to the legacy DAG node count",
5837 );
5838 }
5839
5840 /// S2-1 (totality): a wrong params payload is a typed decline surfaced
5841 /// as a fixpoint (G2), never a panic, declining BEFORE any emission.
5842 #[test]
5843 fn layer_norm_backward_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
5844 let mut g = Graph::new();
5845 let (_x, _up, fused) = ln_backward_fused_node(&mut g, &[2, 4], 1e-5);
5846 let before = g.len();
5847 let out = layer_norm_last_dim_backward::decompose(&mut g, fused, &FusedOpParams::Rope);
5848 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
5849 assert_eq!(g.len(), before, "declined before any emission");
5850 }
5851 }
5852
5853 // rms_norm_last_dim_backward migration (Increment C carriers, A1) ---------
5854 //
5855 // The ~22-node imperative closed-form
5856 // `grad_x = r_rms · (g − x·s / (n·(mean_sq + eps)))` becomes a portable
5857 // `PatternNode` DATA recipe — the FIRST recipe to drive a SHAPE-DERIVED
5858 // scalar through live emit: `n = dims[last]` (the reduced_count) is a
5859 // `MulScalar(scalar_rel = Extent(0, LAST))`, resolved from x's shape at emit
5860 // (not baked, not a params slot). Sibling of the layer-norm-backward recipe:
5861 // `MeanDim(axis_last)` + `Unsqueeze(append)` keepdim (D3 swap) + broadcast to
5862 // `SameAs 0`; the `eps` `AddScalar` is an OPEN slot filled by the projection.
5863 // Binds: `0 = x`, `1 = upstream`. RISK-A: `denom_kd` (which carries the eps
5864 // slot) is consumed by BOTH the `Rsqrt` and the `MulScalar(n)`, so emit's
5865 // slot-free identity-share does NOT dedup it — the recipe emits it twice
5866 // (numerically identical deterministic recompute; the legacy DAG shared it),
5867 // and downstream CSE re-collapses. The parity oracle runs both sides through
5868 // a toy f64 interpreter (Reshape/Unsqueeze = metadata passthrough) so neither
5869 // the D3 swap nor the recompute can perturb the bit-exact assert.
5870 mod rms_norm_backward_recipe {
5871 use super::super::*;
5872 use super::frozen_legacy_rms_norm_backward_decompose;
5873 use crate::registry::{FusedOps, rms_norm_last_dim_backward};
5874 use fuel_ir::{DType, Shape};
5875 use std::collections::HashMap;
5876
5877 /// Tiny f64 reference interpreter over the rms-norm-backward primitive
5878 /// vocabulary (leaf-lookup FIRST, then elementwise `Mul`/`Sub`/`Div`,
5879 /// unary `Sqr`/`Rsqrt`, `AddScalar`/`MulScalar`, last-axis
5880 /// `MeanDim`/`SumDim`, metadata-only keepdim restore
5881 /// `Unsqueeze`/`Reshape`, last-dim `BroadcastTo`). BOTH parity sides run
5882 /// through it with identical in-order arithmetic, so a bit-exact assert
5883 /// isolates recipe STRUCTURE — neither the `Reshape`→`Unsqueeze` D3 swap
5884 /// NOR the open-slot `denom_kd` recompute can perturb it.
5885 fn eval_rnb(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
5886 if let Some(v) = leaves.get(&id) {
5887 return v.clone();
5888 }
5889 let node = g.node(id);
5890 match &node.op {
5891 Op::Mul => {
5892 let a = eval_rnb(g, node.inputs[0], leaves);
5893 let b = eval_rnb(g, node.inputs[1], leaves);
5894 a.iter().zip(&b).map(|(x, y)| x * y).collect()
5895 }
5896 Op::Sub => {
5897 let a = eval_rnb(g, node.inputs[0], leaves);
5898 let b = eval_rnb(g, node.inputs[1], leaves);
5899 a.iter().zip(&b).map(|(x, y)| x - y).collect()
5900 }
5901 Op::Div => {
5902 let a = eval_rnb(g, node.inputs[0], leaves);
5903 let b = eval_rnb(g, node.inputs[1], leaves);
5904 a.iter().zip(&b).map(|(x, y)| x / y).collect()
5905 }
5906 Op::Sqr => eval_rnb(g, node.inputs[0], leaves)
5907 .iter()
5908 .map(|v| v * v)
5909 .collect(),
5910 Op::Rsqrt => eval_rnb(g, node.inputs[0], leaves)
5911 .iter()
5912 .map(|v| 1.0 / v.sqrt())
5913 .collect(),
5914 Op::AddScalar(e) => eval_rnb(g, node.inputs[0], leaves)
5915 .iter()
5916 .map(|v| v + e)
5917 .collect(),
5918 Op::MulScalar(m) => eval_rnb(g, node.inputs[0], leaves)
5919 .iter()
5920 .map(|v| v * m)
5921 .collect(),
5922 // Last-axis reductions — rank-reducing; identical fold both
5923 // spellings.
5924 Op::MeanDim(_) => {
5925 let input = eval_rnb(g, node.inputs[0], leaves);
5926 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
5927 input
5928 .chunks(last)
5929 .map(|row| row.iter().sum::<f64>() / last as f64)
5930 .collect()
5931 }
5932 Op::SumDim(_) => {
5933 let input = eval_rnb(g, node.inputs[0], leaves);
5934 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
5935 input
5936 .chunks(last)
5937 .map(|row| row.iter().sum::<f64>())
5938 .collect()
5939 }
5940 // Metadata-only keepdim restores (the D3 swap and its legacy twin
5941 // evaluate identically here).
5942 Op::Unsqueeze { .. } | Op::Reshape(_) => eval_rnb(g, node.inputs[0], leaves),
5943 // Broadcast a keepdim `[.., 1]` tensor back along the last axis.
5944 Op::BroadcastTo(target) => {
5945 let input = eval_rnb(g, node.inputs[0], leaves);
5946 let out_n: usize = target.dims().iter().product();
5947 let last = *target.dims().last().unwrap();
5948 assert_eq!(
5949 input.len() * last,
5950 out_n,
5951 "broadcast is a last-dim repeat in these graphs",
5952 );
5953 input
5954 .iter()
5955 .flat_map(|&v| std::iter::repeat_n(v, last))
5956 .collect()
5957 }
5958 other => panic!("eval_rnb: unhandled op {other:?}"),
5959 }
5960 }
5961
5962 /// Build a fused RmsNormLastDimBackward node over `x [dims]` (input 0)
5963 /// and `g [dims]` (input 1, the upstream gradient), carrying `eps`.
5964 /// Returns `(x, g, fused)`.
5965 fn rms_backward_fused_node(
5966 g: &mut Graph,
5967 dims: &[usize],
5968 eps: f64,
5969 ) -> (NodeId, NodeId, NodeId) {
5970 let sh = Shape::from_dims(dims);
5971 let x = g.push(Node {
5972 op: Op::Const,
5973 inputs: vec![],
5974 shape: sh.clone(),
5975 dtype: DType::F32,
5976 });
5977 let up = g.push(Node {
5978 op: Op::Const,
5979 inputs: vec![],
5980 shape: sh.clone(),
5981 dtype: DType::F32,
5982 });
5983 let fused = g.push(Node {
5984 op: Op::Fused(
5985 FusedOps::RMS_NORM_LAST_DIM_BACKWARD,
5986 FusedOpParams::RmsNormLastDimBackward { eps },
5987 ),
5988 inputs: vec![x, up],
5989 shape: sh,
5990 dtype: DType::F32,
5991 });
5992 (x, up, fused)
5993 }
5994
5995 /// A1 (a): ONE recipe datum decomposes at BOTH rank 2 and rank 3 (the
5996 /// polymorphism the baked-shape legacy body never had — including the
5997 /// shape-derived `MulScalar(n)`), and its numerics match the FROZEN
5998 /// legacy builder bit-exactly under the shared reference interpreter.
5999 #[test]
6000 fn rms_norm_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
6001 for dims in [vec![2usize, 4], vec![3, 5, 7]] {
6002 let mut g = Graph::new();
6003 let (x, up, fused) = rms_backward_fused_node(&mut g, &dims, 1e-5);
6004 let sh = Shape::from_dims(&dims);
6005 let new_root = rms_norm_last_dim_backward::decompose(
6006 &mut g,
6007 fused,
6008 &FusedOpParams::RmsNormLastDimBackward { eps: 1e-5 },
6009 );
6010 assert_ne!(new_root, fused, "recipe decompose must fire at {dims:?}");
6011 assert_eq!(
6012 g.node(new_root).shape,
6013 sh,
6014 "rms_norm backward is shape-preserving"
6015 );
6016 assert_eq!(g.node(new_root).dtype, DType::F32);
6017
6018 let legacy_root = frozen_legacy_rms_norm_backward_decompose(
6019 &mut g,
6020 fused,
6021 &FusedOpParams::RmsNormLastDimBackward { eps: 1e-5 },
6022 );
6023
6024 let n: usize = dims.iter().product();
6025 let x_data: Vec<f64> = (0..n)
6026 .map(|i| ((i as f64) * 0.41).sin() * 2.0 + 0.7)
6027 .collect();
6028 let g_data: Vec<f64> = (0..n)
6029 .map(|i| ((i as f64) * 0.61).cos() * 1.3 - 0.15)
6030 .collect();
6031 let mut leaves = HashMap::new();
6032 leaves.insert(x, x_data);
6033 leaves.insert(up, g_data);
6034 let got = eval_rnb(&g, new_root, &leaves);
6035 let want = eval_rnb(&g, legacy_root, &leaves);
6036 assert_eq!(got.len(), want.len());
6037 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
6038 assert_eq!(
6039 a.to_bits(),
6040 b.to_bits(),
6041 "rms_norm_backward[{i}] at {dims:?}: recipe={a} vs legacy={b}",
6042 );
6043 }
6044 }
6045 }
6046
6047 /// A1 (structural, born-red): every keepdim restore is the D3
6048 /// shrink-via-swap `Unsqueeze` append, NOT the baked `Reshape(keepdim)`
6049 /// (the crisp discriminator against the pre-migration imperative body,
6050 /// which emits `Reshape` — so this is RED until migration), and the
6051 /// reduced-count divisor is a live `MulScalar(n = dims[last])` resolved
6052 /// from x's shape by the `scalar_rel` carrier (evidence the shape-derived
6053 /// scalar reached emit). The backward root is the outer `Mul(rrms, inner)`.
6054 #[test]
6055 fn rms_norm_backward_recipe_uses_the_unsqueeze_keepdim_swap_and_shape_scalar() {
6056 for dims in [vec![2usize, 4], vec![3, 5, 7]] {
6057 let mut g = Graph::new();
6058 let (_x, _up, fused) = rms_backward_fused_node(&mut g, &dims, 1e-5);
6059 let root = rms_norm_last_dim_backward::decompose(
6060 &mut g,
6061 fused,
6062 &FusedOpParams::RmsNormLastDimBackward { eps: 1e-5 },
6063 );
6064 assert_ne!(root, fused, "recipe decompose fires at {dims:?}");
6065 assert!(
6066 matches!(g.node(root).op, Op::Mul),
6067 "backward root is the outer Mul(rrms, inner)"
6068 );
6069 let reachable = crate::topo_order_multi(&g, &[root]);
6070 let unsqueezes = reachable
6071 .iter()
6072 .filter(|&&n| matches!(g.node(n).op, Op::Unsqueeze { .. }))
6073 .count();
6074 let reshapes = reachable
6075 .iter()
6076 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
6077 .count();
6078 assert!(
6079 unsqueezes >= 1,
6080 "keepdim restored via Unsqueeze append (D3 swap) at {dims:?}"
6081 );
6082 assert_eq!(
6083 reshapes, 0,
6084 "no baked keepdim Reshape / no D4 pad after the D3 swap at {dims:?}"
6085 );
6086 // The reduced-count divisor resolved to n = dims[last] from x's
6087 // shape (the scalar_rel carrier), at whatever rank.
6088 let n = *dims.last().unwrap() as f64;
6089 let mul_scalars: Vec<f64> = reachable
6090 .iter()
6091 .filter_map(|&nid| match g.node(nid).op {
6092 Op::MulScalar(v) => Some(v),
6093 _ => None,
6094 })
6095 .collect();
6096 assert!(
6097 mul_scalars.contains(&n),
6098 "a MulScalar(n={n}) resolved from x's last-axis extent at {dims:?}, got {mul_scalars:?}",
6099 );
6100 }
6101 }
6102 }
6103
6104 // fused_linear migration (Increment C slice 2, S2-2) ----------------------
6105 //
6106 // The 3-node imperative `Add(MatMul(a, b), BroadcastTo(rank-1 bias))`
6107 // becomes a portable `PatternNode` DATA recipe, and the FIRST recipe to
6108 // drive `WithDim` through live emit: the bias broadcast targets
6109 // `WithDim { operand: 0, axis: LAST, dim: Extent { operand: 1, axis: LAST } }`
6110 // — a's shape with its LAST axis (K) replaced by b's LAST extent (N) = the
6111 // matmul output `[..batch, M, N]`. This is Fuel-INTERNAL shape resolution
6112 // (not §6.19 wire emission), so it is NOT gated on KISS #86. Binds:
6113 // `0 = a`, `1 = b`, `2 = bias`. The rank-1 bias rank-raises to the matmul
6114 // output, so emit materializes a leading-1 pad `Reshape` (D4) the imperative
6115 // body never had — the numeric-parity oracle below runs both sides through a
6116 // toy interpreter (Reshape = metadata passthrough) to isolate it. The
6117 // `canonical_pattern` see-through of that pad (RISK-C) + the lower→fuse
6118 // round-trip live in `registry/fused_linear.rs`'s own tests.
6119 mod fused_linear_recipe {
6120 use super::super::*;
6121 use super::frozen_legacy_fused_linear_decompose;
6122 use crate::registry::{FusedOps, fused_linear};
6123 use fuel_ir::{DType, Shape};
6124 use std::collections::HashMap;
6125
6126 /// Right-aligned NumPy broadcast of `input` (shape `in_shape`) to
6127 /// `target` — a size-1 or padded leading dim contributes stride 0.
6128 /// General on purpose: fused_linear's bias broadcast tiles along the
6129 /// LEADING `M`/batch dims (`[1,N] → [M,N]`), unlike the norm recipes'
6130 /// last-dim keepdim repeat.
6131 fn broadcast(input: &[f64], in_shape: &[usize], target: &[usize]) -> Vec<f64> {
6132 let rank = target.len();
6133 let pad = rank - in_shape.len();
6134 let mut real_strides = vec![0isize; in_shape.len()];
6135 let mut s = 1isize;
6136 for i in (0..in_shape.len()).rev() {
6137 real_strides[i] = s;
6138 s *= in_shape[i] as isize;
6139 }
6140 let mut in_strides = vec![0isize; rank];
6141 for (i, stride) in in_strides.iter_mut().enumerate() {
6142 if i >= pad {
6143 let id = i - pad;
6144 *stride = if in_shape[id] == 1 {
6145 0
6146 } else {
6147 real_strides[id]
6148 };
6149 }
6150 }
6151 let out_n: usize = target.iter().product();
6152 let mut out = Vec::with_capacity(out_n);
6153 let mut idx = vec![0usize; rank];
6154 for _ in 0..out_n {
6155 let fi: isize = (0..rank).map(|i| idx[i] as isize * in_strides[i]).sum();
6156 out.push(input[fi as usize]);
6157 for i in (0..rank).rev() {
6158 idx[i] += 1;
6159 if idx[i] < target[i] {
6160 break;
6161 }
6162 idx[i] = 0;
6163 }
6164 }
6165 out
6166 }
6167
6168 /// Deterministic batched f64 matmul `a[..batch, M, K] · b[..batch, K, N]`
6169 /// (same-rank ≥ 2, aligned batch). BOTH parity sides share the SAME
6170 /// `MatMul(a, b)` binds, so identical accumulation order ⇒ identical
6171 /// bits.
6172 fn matmul(a: &[f64], ash: &[usize], b: &[f64], bsh: &[usize]) -> Vec<f64> {
6173 let r = ash.len();
6174 let (m, k, n) = (ash[r - 2], ash[r - 1], bsh[r - 1]);
6175 let batch: usize = ash[..r - 2].iter().product();
6176 let (a_bs, b_bs, o_bs) = (m * k, k * n, m * n);
6177 let mut out = vec![0.0f64; batch * o_bs];
6178 for bi in 0..batch {
6179 for i in 0..m {
6180 for j in 0..n {
6181 let mut acc = 0.0f64;
6182 for kk in 0..k {
6183 acc += a[bi * a_bs + i * k + kk] * b[bi * b_bs + kk * n + j];
6184 }
6185 out[bi * o_bs + i * n + j] = acc;
6186 }
6187 }
6188 }
6189 out
6190 }
6191
6192 /// Tiny f64 reference interpreter over the fused-linear primitive
6193 /// vocabulary (leaf-lookup FIRST, then `MatMul`, metadata-only
6194 /// `Reshape` pad, right-aligned `BroadcastTo`, elementwise `Add`). BOTH
6195 /// parity sides run through it with identical arithmetic, so a bit-exact
6196 /// assert isolates recipe STRUCTURE — the D4 pad `Reshape` (recipe only)
6197 /// is a passthrough here.
6198 fn eval_fl(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
6199 if let Some(v) = leaves.get(&id) {
6200 return v.clone();
6201 }
6202 let node = g.node(id);
6203 match &node.op {
6204 Op::MatMul => {
6205 let a = eval_fl(g, node.inputs[0], leaves);
6206 let b = eval_fl(g, node.inputs[1], leaves);
6207 let ash: Vec<usize> = g.node(node.inputs[0]).shape.dims().to_vec();
6208 let bsh: Vec<usize> = g.node(node.inputs[1]).shape.dims().to_vec();
6209 matmul(&a, &ash, &b, &bsh)
6210 }
6211 // Metadata-only leading-1 pad (recipe's D4) — values unchanged.
6212 Op::Reshape(_) => eval_fl(g, node.inputs[0], leaves),
6213 Op::BroadcastTo(target) => {
6214 let input = eval_fl(g, node.inputs[0], leaves);
6215 let in_shape: Vec<usize> = g.node(node.inputs[0]).shape.dims().to_vec();
6216 broadcast(&input, &in_shape, target.dims())
6217 }
6218 Op::Add => {
6219 let a = eval_fl(g, node.inputs[0], leaves);
6220 let b = eval_fl(g, node.inputs[1], leaves);
6221 a.iter().zip(&b).map(|(x, y)| x + y).collect()
6222 }
6223 other => panic!("eval_fl: unhandled op {other:?}"),
6224 }
6225 }
6226
6227 /// Build a fused FusedLinear node over `a [a_dims]`, `b [b_dims]`, and a
6228 /// rank-1 `bias [N]` (N = b's last dim). Returns `(a, b, bias, fused)`.
6229 fn fused_linear_fused_node(
6230 g: &mut Graph,
6231 a_dims: &[usize],
6232 b_dims: &[usize],
6233 ) -> (NodeId, NodeId, NodeId, NodeId) {
6234 let ar = a_dims.len();
6235 let n = b_dims[b_dims.len() - 1];
6236 let mut out_dims = a_dims[..ar - 2].to_vec();
6237 out_dims.push(a_dims[ar - 2]);
6238 out_dims.push(n);
6239 let a = g.push(Node {
6240 op: Op::Const,
6241 inputs: vec![],
6242 shape: Shape::from_dims(a_dims),
6243 dtype: DType::F32,
6244 });
6245 let b = g.push(Node {
6246 op: Op::Const,
6247 inputs: vec![],
6248 shape: Shape::from_dims(b_dims),
6249 dtype: DType::F32,
6250 });
6251 let bias = g.push(Node {
6252 op: Op::Const,
6253 inputs: vec![],
6254 shape: Shape::from_dims(&[n]),
6255 dtype: DType::F32,
6256 });
6257 let fused = g.push(Node {
6258 op: Op::Fused(FusedOps::FUSED_LINEAR, FusedOpParams::FusedLinear),
6259 inputs: vec![a, b, bias],
6260 shape: Shape::from_dims(&out_dims),
6261 dtype: DType::F32,
6262 });
6263 (a, b, bias, fused)
6264 }
6265
6266 /// S2-2 (a): ONE recipe datum decomposes at BOTH rank-2 and rank-3
6267 /// (batched) activations — the WithDim-derived broadcast target tracks
6268 /// the matmul output shape — and its numerics match the FROZEN legacy
6269 /// builder bit-exactly under the shared reference interpreter.
6270 #[test]
6271 fn fused_linear_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
6272 for (a_dims, b_dims) in [
6273 (vec![2usize, 3], vec![3usize, 4]),
6274 (vec![2, 2, 3], vec![2, 3, 4]),
6275 ] {
6276 let mut g = Graph::new();
6277 let (a, b, bias, fused) = fused_linear_fused_node(&mut g, &a_dims, &b_dims);
6278 let out_sh = g.node(fused).shape.clone();
6279 let new_root = fused_linear::decompose(&mut g, fused, &FusedOpParams::FusedLinear);
6280 assert_ne!(new_root, fused, "recipe decompose fires at a={a_dims:?}");
6281 assert_eq!(
6282 g.node(new_root).shape,
6283 out_sh,
6284 "output = matmul output shape"
6285 );
6286 assert_eq!(g.node(new_root).dtype, DType::F32);
6287
6288 let legacy_root = frozen_legacy_fused_linear_decompose(
6289 &mut g,
6290 fused,
6291 &FusedOpParams::FusedLinear,
6292 );
6293
6294 let an: usize = a_dims.iter().product();
6295 let bn: usize = b_dims.iter().product();
6296 let bias_n = b_dims[b_dims.len() - 1];
6297 let a_data: Vec<f64> = (0..an)
6298 .map(|i| ((i as f64) * 0.31).sin() * 1.5 - 0.2)
6299 .collect();
6300 let b_data: Vec<f64> = (0..bn)
6301 .map(|i| ((i as f64) * 0.47).cos() * 0.8 + 0.1)
6302 .collect();
6303 let bias_data: Vec<f64> = (0..bias_n)
6304 .map(|i| ((i as f64) * 0.7).sin() * 2.0)
6305 .collect();
6306 let mut leaves = HashMap::new();
6307 leaves.insert(a, a_data);
6308 leaves.insert(b, b_data);
6309 leaves.insert(bias, bias_data);
6310 let got = eval_fl(&g, new_root, &leaves);
6311 let want = eval_fl(&g, legacy_root, &leaves);
6312 assert_eq!(got.len(), want.len());
6313 for (i, (x, y)) in got.iter().zip(want.iter()).enumerate() {
6314 assert_eq!(
6315 x.to_bits(),
6316 y.to_bits(),
6317 "fused_linear[{i}] a={a_dims:?}: recipe={x} vs legacy={y}",
6318 );
6319 }
6320 }
6321 }
6322
6323 /// S2-2 (structural, born-red): the WithDim-derived broadcast target
6324 /// rank-raises the rank-1 bias to the matmul output `[M,N]`, so emit
6325 /// materializes exactly one leading-1 pad `Reshape` (the D4 path, driven
6326 /// LIVE by WithDim). The imperative body broadcast the rank-1 bias
6327 /// directly (zero Reshape) — so this is RED until migration. The bias
6328 /// broadcast target equals the matmul output shape.
6329 #[test]
6330 fn fused_linear_recipe_drives_withdim_broadcast_with_a_d4_pad() {
6331 let mut g = Graph::new();
6332 let (_a, _b, _bias, fused) = fused_linear_fused_node(&mut g, &[2, 3], &[3, 4]);
6333 let root = fused_linear::decompose(&mut g, fused, &FusedOpParams::FusedLinear);
6334 assert_ne!(root, fused, "recipe decompose fires");
6335 assert!(
6336 matches!(g.node(root).op, Op::Add),
6337 "root is Add(mm, bias_bcst)"
6338 );
6339 let reachable = crate::topo_order_multi(&g, &[root]);
6340 let reshapes = reachable
6341 .iter()
6342 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
6343 .count();
6344 assert_eq!(
6345 reshapes, 1,
6346 "WithDim rank-raise materializes exactly one D4 pad Reshape on the bias",
6347 );
6348 let bcst = reachable
6349 .iter()
6350 .find(|&&n| matches!(g.node(n).op, Op::BroadcastTo(_)))
6351 .expect("bias broadcast present");
6352 assert_eq!(
6353 g.node(*bcst).shape.dims(),
6354 &[2, 4],
6355 "bias broadcast target = the matmul output shape (WithDim resolved)",
6356 );
6357 }
6358
6359 /// S2-2 (totality): a wrong params payload is a typed decline surfaced
6360 /// as a fixpoint (G2), never a panic, declining BEFORE any emission.
6361 #[test]
6362 fn fused_linear_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
6363 let mut g = Graph::new();
6364 let (_a, _b, _bias, fused) = fused_linear_fused_node(&mut g, &[2, 3], &[3, 4]);
6365 let before = g.len();
6366 let out = fused_linear::decompose(&mut g, fused, &FusedOpParams::Rope);
6367 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
6368 assert_eq!(g.len(), before, "declined before any emission");
6369 }
6370 }
6371
6372 /// Increment C, B1 — the `Op::Scan` structural re-emit machinery. These
6373 /// tests prove a hand-built `Op::Scan` recipe (an `OpTag::Scan` node whose
6374 /// body sub-graph — `ScanPlaceholder` holes + body ops — rides its trailing
6375 /// operands, per the Phase-1 lax.scan encoding) round-trips through
6376 /// `tag_to_op`/`emit` to a graph whose base map is bit-identical to the
6377 /// same scan built imperatively (the way `selective_scan::decompose` emits
6378 /// it). `Op::Scan` stays a base-map terminal — the recipe just RE-EMITS it.
6379 ///
6380 /// Structure-preserving: the recipe mirrors the imperative
6381 /// `unroll_scan`-shaped scan node exactly (same op types, same params, same
6382 /// body structure, same const-sharing via a shared `Bind`), so the
6383 /// `base_map_hash` equality is a structural-identity proof. Neither
6384 /// `Op::Scan` nor `Op::ScanPlaceholder` folds its own node shape into
6385 /// `base_map_hash` (op_key tags 210/211 fold only their params + the body
6386 /// via child recursion), so emit's fallback interior shapes are inert for
6387 /// the identity — the base map matches by construction.
6388 mod scan_recipe_roundtrip {
6389 use super::super::*;
6390 use super::{bind, op_node};
6391 use crate::opt::base_map_hash;
6392 use crate::{ScanEmit, ScanPredicate, ScanRole};
6393 use fuel_ir::storage::{OutputViewSpec, compose_bundle};
6394 use fuel_ir::{DType, Shape};
6395 use fuel_kernel_seam_types::{SCAN_ROLE_CARRY, SCAN_ROLE_ELEM};
6396 use std::sync::Arc;
6397
6398 fn placeholder(role: u8, index: u32) -> PatternNode {
6399 PatternNode::Op {
6400 op: OpTag::ScanPlaceholder,
6401 attrs: OpAttrs {
6402 scan_role: Some(role),
6403 scan_index: Some(index),
6404 ..OpAttrs::default()
6405 },
6406 operands: vec![],
6407 }
6408 }
6409
6410 /// The recipe: `scan(n_xs=1, bound=3, emit=All)` over
6411 /// `[init(B0), xs0(B1), const(B2), body_new_carry, body_y]` where
6412 /// `body_new_carry = Add(carry_hole, elem_hole)` and
6413 /// `body_y = Mul(body_new_carry, const)`. The `Add(carry,elem)` sub-tree
6414 /// appears in BOTH body exits — emit's slot-free identity-share collapses
6415 /// them to one node, matching the imperative `sum` referenced twice.
6416 fn scan_recipe() -> PatternNode {
6417 let sum = || {
6418 op_node(
6419 OpTag::Add,
6420 OpAttrs::default(),
6421 vec![
6422 placeholder(SCAN_ROLE_CARRY, 0),
6423 placeholder(SCAN_ROLE_ELEM, 0),
6424 ],
6425 )
6426 };
6427 PatternNode::Op {
6428 op: OpTag::Scan,
6429 attrs: OpAttrs {
6430 scan_n_carries: Some(1),
6431 scan_n_xs: Some(1),
6432 scan_bound: Some(3),
6433 scan_emit: Some(0), // All
6434 scan_early_exit: None,
6435 ..OpAttrs::default()
6436 },
6437 operands: vec![
6438 bind(0), // init_carry
6439 bind(1), // xs0
6440 bind(2), // const
6441 sum(), // body_new_carry
6442 op_node(OpTag::Mul, OpAttrs::default(), vec![sum(), bind(2)]), // body_y
6443 ],
6444 }
6445 }
6446
6447 /// Build the reference scan imperatively (the shape/layout
6448 /// `selective_scan::decompose` + `unroll_scan` assume) and return the
6449 /// scan NodeId. Leaves are unpopulated `Op::Const`s — same shape/dtype
6450 /// as the emitted recipe's, so they hash identically.
6451 fn reference_scan(g: &mut Graph) -> NodeId {
6452 let cs = Shape::from_dims(&[1]);
6453 let xs = Shape::from_dims(&[3, 1]);
6454 let init = g.push(Node {
6455 op: Op::Const,
6456 inputs: vec![],
6457 shape: cs.clone(),
6458 dtype: DType::F32,
6459 });
6460 let xs0 = g.push(Node {
6461 op: Op::Const,
6462 inputs: vec![],
6463 shape: xs,
6464 dtype: DType::F32,
6465 });
6466 let konst = g.push(Node {
6467 op: Op::Const,
6468 inputs: vec![],
6469 shape: cs.clone(),
6470 dtype: DType::F32,
6471 });
6472 let carry = g.push(Node {
6473 op: Op::ScanPlaceholder {
6474 role: ScanRole::Carry,
6475 index: 0,
6476 },
6477 inputs: vec![],
6478 shape: cs.clone(),
6479 dtype: DType::F32,
6480 });
6481 let elem = g.push(Node {
6482 op: Op::ScanPlaceholder {
6483 role: ScanRole::Elem,
6484 index: 0,
6485 },
6486 inputs: vec![],
6487 shape: cs.clone(),
6488 dtype: DType::F32,
6489 });
6490 let sum = g.push(Node {
6491 op: Op::Add,
6492 inputs: vec![carry, elem],
6493 shape: cs.clone(),
6494 dtype: DType::F32,
6495 });
6496 let y = g.push(Node {
6497 op: Op::Mul,
6498 inputs: vec![sum, konst],
6499 shape: cs.clone(),
6500 dtype: DType::F32,
6501 });
6502 g.push(Node {
6503 op: Op::Scan {
6504 n_carries: 1,
6505 n_xs: 1,
6506 bound: 3,
6507 emit: ScanEmit::All,
6508 early_exit: None,
6509 },
6510 inputs: vec![init, xs0, konst, sum, y],
6511 shape: Shape::from_dims(&[3, 1]),
6512 dtype: DType::F32,
6513 })
6514 }
6515
6516 /// `tag_to_op` reconstructs the `Op::Scan` params + `ScanPlaceholder`
6517 /// role/index verbatim from the carriers.
6518 #[test]
6519 fn tag_to_op_reconstructs_scan_and_placeholder() {
6520 let scan_attrs = OpAttrs {
6521 scan_n_carries: Some(1),
6522 scan_n_xs: Some(1),
6523 scan_bound: Some(3),
6524 scan_emit: Some(0),
6525 scan_early_exit: None,
6526 ..OpAttrs::default()
6527 };
6528 assert!(matches!(
6529 tag_to_op(OpTag::Scan, &scan_attrs),
6530 Some(Op::Scan {
6531 n_carries: 1,
6532 n_xs: 1,
6533 bound: 3,
6534 emit: ScanEmit::All,
6535 early_exit: None
6536 }),
6537 ));
6538 let fin = OpAttrs {
6539 scan_emit: Some(1),
6540 scan_early_exit: Some(true),
6541 ..scan_attrs.clone()
6542 };
6543 assert!(matches!(
6544 tag_to_op(OpTag::Scan, &fin),
6545 Some(Op::Scan {
6546 emit: ScanEmit::Final,
6547 early_exit: Some(ScanPredicate),
6548 ..
6549 }),
6550 ));
6551 // A scan node with no bound is an honest miss (unset required attr).
6552 assert!(tag_to_op(OpTag::Scan, &OpAttrs::default()).is_none());
6553 let carry = OpAttrs {
6554 scan_role: Some(SCAN_ROLE_CARRY),
6555 scan_index: Some(0),
6556 ..OpAttrs::default()
6557 };
6558 assert!(matches!(
6559 tag_to_op(OpTag::ScanPlaceholder, &carry),
6560 Some(Op::ScanPlaceholder {
6561 role: ScanRole::Carry,
6562 index: 0
6563 }),
6564 ));
6565 let elem = OpAttrs {
6566 scan_role: Some(SCAN_ROLE_ELEM),
6567 scan_index: Some(2),
6568 ..OpAttrs::default()
6569 };
6570 assert!(matches!(
6571 tag_to_op(OpTag::ScanPlaceholder, &elem),
6572 Some(Op::ScanPlaceholder {
6573 role: ScanRole::Elem,
6574 index: 2
6575 }),
6576 ));
6577 assert!(tag_to_op(OpTag::ScanPlaceholder, &OpAttrs::default()).is_none());
6578 }
6579
6580 /// The recipe is a valid, representable region: contiguous binds `[0,1,2]`
6581 /// and every op (incl. `Scan`/`ScanPlaceholder`) re-emits to a primitive.
6582 #[test]
6583 fn scan_recipe_validates_as_representable() {
6584 assert_eq!(scan_recipe().bind_indices(), vec![0, 1, 2]);
6585 assert!(
6586 validate_recipe(&scan_recipe()).is_ok(),
6587 "scan recipe is a total, representable region"
6588 );
6589 }
6590
6591 /// THE round-trip: emit the recipe onto fresh input leaves and assert its
6592 /// base map is bit-identical to the imperatively-built scan.
6593 #[test]
6594 fn scan_recipe_reemits_to_the_same_base_map_as_the_imperative_scan() {
6595 // Reference.
6596 let mut gr = Graph::new();
6597 let ref_scan = reference_scan(&mut gr);
6598 let want = base_map_hash(&gr, ref_scan);
6599
6600 // Emitted from the recipe onto identical leaves.
6601 let mut ge = Graph::new();
6602 let cs = Shape::from_dims(&[1]);
6603 let init = ge.push(Node {
6604 op: Op::Const,
6605 inputs: vec![],
6606 shape: cs.clone(),
6607 dtype: DType::F32,
6608 });
6609 let xs0 = ge.push(Node {
6610 op: Op::Const,
6611 inputs: vec![],
6612 shape: Shape::from_dims(&[3, 1]),
6613 dtype: DType::F32,
6614 });
6615 let konst = ge.push(Node {
6616 op: Op::Const,
6617 inputs: vec![],
6618 shape: cs,
6619 dtype: DType::F32,
6620 });
6621 let root = emit_region(&mut ge, &scan_recipe(), &[init, xs0, konst], &[]);
6622 assert!(
6623 matches!(
6624 ge.node(root).op,
6625 Op::Scan {
6626 n_carries: 1,
6627 n_xs: 1,
6628 bound: 3,
6629 emit: ScanEmit::All,
6630 early_exit: None
6631 }
6632 ),
6633 "recipe re-emits an Op::Scan terminal"
6634 );
6635 let got = base_map_hash(&ge, root);
6636
6637 assert_eq!(
6638 got, want,
6639 "scan recipe re-emit base map == imperative scan base map"
6640 );
6641 }
6642
6643 // ---- The View / output_views bundle half (B1) ----------------------
6644 //
6645 // A scan recipe whose root is an `Op::View{slot}` must re-emit the scan,
6646 // attach the 2-slot `output_views` bundle (slot 0 = stacked ys, slot 1 =
6647 // final carry — the SelectiveScan bundle contract), and project the
6648 // slot — round-tripping to the same base map AND the correct slot shape
6649 // as the imperative scan+view. Real shapes throughout (init/const via
6650 // Binds carry real shapes; `body_y = Mul(const, sum)` takes the const's
6651 // shape) so the emitted bundle is byte-faithful.
6652
6653 /// `view(slot)(scan(n_xs=1, bound=3, emit=All))` over binds `[init[2],
6654 /// xs0[3,2], const[2]]`. `body_y = Mul(const, sum)` is const-first so its
6655 /// re-emitted shape is `[2]` (the const's), making slot 0 = `[3,2]` and
6656 /// slot 1 = `[2]`.
6657 fn view_scan_recipe(slot: u32) -> PatternNode {
6658 let sum = || {
6659 op_node(
6660 OpTag::Add,
6661 OpAttrs::default(),
6662 vec![
6663 placeholder(SCAN_ROLE_CARRY, 0),
6664 placeholder(SCAN_ROLE_ELEM, 0),
6665 ],
6666 )
6667 };
6668 let scan = PatternNode::Op {
6669 op: OpTag::Scan,
6670 attrs: OpAttrs {
6671 scan_n_carries: Some(1),
6672 scan_n_xs: Some(1),
6673 scan_bound: Some(3),
6674 scan_emit: Some(0),
6675 scan_early_exit: None,
6676 ..OpAttrs::default()
6677 },
6678 operands: vec![
6679 bind(0), // init_carry [2]
6680 bind(1), // xs0 [3,2]
6681 bind(2), // const [2]
6682 sum(), // body_new_carry
6683 op_node(OpTag::Mul, OpAttrs::default(), vec![bind(2), sum()]), // body_y = Mul(const, sum)
6684 ],
6685 };
6686 op_node(
6687 OpTag::View,
6688 OpAttrs {
6689 view_slot: Some(slot),
6690 ..OpAttrs::default()
6691 },
6692 vec![scan],
6693 )
6694 }
6695
6696 /// The imperative reference: the scan + its composed 2-slot bundle + the
6697 /// slot projection (mirrors `NodeHandle::scan` + `Graph::view`).
6698 fn reference_view(g: &mut Graph, slot: u32) -> NodeId {
6699 let two = Shape::from_dims(&[2]);
6700 let init = g.push(Node {
6701 op: Op::Const,
6702 inputs: vec![],
6703 shape: two.clone(),
6704 dtype: DType::F32,
6705 });
6706 let xs0 = g.push(Node {
6707 op: Op::Const,
6708 inputs: vec![],
6709 shape: Shape::from_dims(&[3, 2]),
6710 dtype: DType::F32,
6711 });
6712 let konst = g.push(Node {
6713 op: Op::Const,
6714 inputs: vec![],
6715 shape: two.clone(),
6716 dtype: DType::F32,
6717 });
6718 let carry = g.push(Node {
6719 op: Op::ScanPlaceholder {
6720 role: ScanRole::Carry,
6721 index: 0,
6722 },
6723 inputs: vec![],
6724 shape: two.clone(),
6725 dtype: DType::F32,
6726 });
6727 let elem = g.push(Node {
6728 op: Op::ScanPlaceholder {
6729 role: ScanRole::Elem,
6730 index: 0,
6731 },
6732 inputs: vec![],
6733 shape: two.clone(),
6734 dtype: DType::F32,
6735 });
6736 let sum = g.push(Node {
6737 op: Op::Add,
6738 inputs: vec![carry, elem],
6739 shape: two.clone(),
6740 dtype: DType::F32,
6741 });
6742 let body_y = g.push(Node {
6743 op: Op::Mul,
6744 inputs: vec![konst, sum],
6745 shape: two.clone(),
6746 dtype: DType::F32,
6747 });
6748 let ys = Shape::from_dims(&[3, 2]);
6749 let scan = g.push(Node {
6750 op: Op::Scan {
6751 n_carries: 1,
6752 n_xs: 1,
6753 bound: 3,
6754 emit: ScanEmit::All,
6755 early_exit: None,
6756 },
6757 inputs: vec![init, xs0, konst, sum, body_y],
6758 shape: ys.clone(),
6759 dtype: DType::F32,
6760 });
6761 let specs = vec![
6762 OutputViewSpec::contiguous(DType::F32, ys),
6763 OutputViewSpec::contiguous(DType::F32, two),
6764 ];
6765 let (_bytes, views) = compose_bundle(&specs).expect("compose_bundle");
6766 g.set_output_views(scan, Arc::from(views.into_boxed_slice()))
6767 .expect("set_output_views");
6768 let (sh, dt) = {
6769 let v = g.output_views(scan).expect("bundle set");
6770 let s = &v[slot as usize];
6771 (s.shape.clone(), s.dtype)
6772 };
6773 g.push(Node {
6774 op: Op::View { slot },
6775 inputs: vec![scan],
6776 shape: sh,
6777 dtype: dt,
6778 })
6779 }
6780
6781 fn roundtrip_view(slot: u32, want_shape: &[usize]) {
6782 // Reference.
6783 let mut gr = Graph::new();
6784 let ref_view = reference_view(&mut gr, slot);
6785 let want = base_map_hash(&gr, ref_view);
6786
6787 // Emitted from the recipe onto identical leaves.
6788 let mut ge = Graph::new();
6789 let two = Shape::from_dims(&[2]);
6790 let init = ge.push(Node {
6791 op: Op::Const,
6792 inputs: vec![],
6793 shape: two.clone(),
6794 dtype: DType::F32,
6795 });
6796 let xs0 = ge.push(Node {
6797 op: Op::Const,
6798 inputs: vec![],
6799 shape: Shape::from_dims(&[3, 2]),
6800 dtype: DType::F32,
6801 });
6802 let konst = ge.push(Node {
6803 op: Op::Const,
6804 inputs: vec![],
6805 shape: two,
6806 dtype: DType::F32,
6807 });
6808 let root = emit_region(&mut ge, &view_scan_recipe(slot), &[init, xs0, konst], &[]);
6809
6810 // The View projects the requested slot with the bundle-derived shape.
6811 assert!(
6812 matches!(ge.node(root).op, Op::View { slot: s } if s == slot),
6813 "recipe root re-emits Op::View{{slot={slot}}}"
6814 );
6815 assert_eq!(
6816 ge.node(root).shape.dims(),
6817 want_shape,
6818 "view slot {slot} shape comes from the re-attached output_views bundle"
6819 );
6820 // The producing scan carries the faithful 2-slot bundle.
6821 let producer = ge.node(root).inputs[0];
6822 assert!(
6823 matches!(ge.node(producer).op, Op::Scan { .. }),
6824 "producer is the scan"
6825 );
6826 assert_eq!(
6827 ge.output_views(producer).map(|v| v.len()),
6828 Some(2),
6829 "emit re-attaches the 2-slot bundle so slot>=1 views resolve"
6830 );
6831
6832 let got = base_map_hash(&ge, root);
6833 assert_eq!(
6834 got, want,
6835 "view+scan recipe re-emit base map == imperative scan+view base map"
6836 );
6837 }
6838
6839 /// Slot 0 (stacked ys): its shape equals the scan node's primary shape,
6840 /// so even the child-fallback would find it — but it must come through
6841 /// the bundle for the layout/offset to be right.
6842 #[test]
6843 fn scan_view_slot0_reemits_to_the_same_base_map() {
6844 roundtrip_view(0, &[3, 2]);
6845 }
6846
6847 /// Slot 1 (final carry / last_state) is the load-bearing case: its shape
6848 /// (`[2]`) lives ONLY in the bundle and differs from the scan's primary
6849 /// (`[3,2]`). `op_key(View) = None` folds the node shape into the hash,
6850 /// so a wrong slot-1 shape would break the round-trip — this proves the
6851 /// bundle is re-attached and read.
6852 #[test]
6853 fn scan_view_slot1_reemits_to_the_same_base_map() {
6854 roundtrip_view(1, &[2]);
6855 }
6856
6857 /// A view recipe validates as a representable region (contiguous binds,
6858 /// every op — incl. `View`/`Scan`/`ScanPlaceholder` — re-emittable).
6859 #[test]
6860 fn view_scan_recipe_validates_as_representable() {
6861 assert_eq!(view_scan_recipe(1).bind_indices(), vec![0, 1, 2]);
6862 assert!(validate_recipe(&view_scan_recipe(1)).is_ok());
6863 // tag_to_op reconstructs the View op from `view_slot`.
6864 assert!(matches!(
6865 tag_to_op(
6866 OpTag::View,
6867 &OpAttrs {
6868 view_slot: Some(1),
6869 ..OpAttrs::default()
6870 }
6871 ),
6872 Some(Op::View { slot: 1 }),
6873 ));
6874 assert!(tag_to_op(OpTag::View, &OpAttrs::default()).is_none());
6875 }
6876
6877 /// B2 emit carrier (born-red without it): a `ScanPlaceholder` that
6878 /// declares its per-step shape via `target_shape_rel` emits with THAT
6879 /// shape, not the rank-0 `primitive_shape` fallback. This is
6880 /// load-bearing for the SSM scan recipes — `unroll_scan` clones body
6881 /// interior nodes with their STORED shapes, so a rank-0 placeholder
6882 /// would poison the unrolled body (e.g. `du = Mul(d_t, u_t)` rank-0).
6883 /// Dtype comes from bind 0 (the uniform-dtype scan inputs).
6884 #[test]
6885 fn scan_placeholder_recipe_emits_its_declared_shape() {
6886 use fuel_kernel_seam_types::shape_expr::{Dim, ShapeExpr};
6887 let dims2 = || {
6888 ShapeExpr::Dims(vec![
6889 Dim::Extent {
6890 operand: 0,
6891 axis: 0,
6892 },
6893 Dim::Extent {
6894 operand: 0,
6895 axis: 1,
6896 },
6897 ])
6898 };
6899 let ph = |index: u32| PatternNode::Op {
6900 op: OpTag::ScanPlaceholder,
6901 attrs: OpAttrs {
6902 scan_role: Some(SCAN_ROLE_ELEM),
6903 scan_index: Some(index),
6904 target_shape_rel: Some(dims2()),
6905 ..OpAttrs::default()
6906 },
6907 operands: vec![],
6908 };
6909 // Mul takes operand[0]'s shape, so its shape mirrors the placeholder.
6910 // A `Bind { 0 }` operand supplies the frame the placeholder's Dims
6911 // reference (operand 0) — so the recipe carries a valid bind.
6912 let recipe = op_node(OpTag::Mul, OpAttrs::default(), vec![ph(0), bind(0)]);
6913 assert!(validate_recipe(&recipe).is_ok());
6914 let mut g = Graph::new();
6915 let bind0 = g.push(Node {
6916 op: Op::Const,
6917 inputs: vec![],
6918 shape: Shape::from_dims(&[2, 3]),
6919 dtype: DType::F32,
6920 });
6921 let root = emit_region(&mut g, &recipe, &[bind0], &[]);
6922 let ph_id = g.node(root).inputs[0];
6923 assert!(matches!(
6924 g.node(ph_id).op,
6925 Op::ScanPlaceholder {
6926 role: ScanRole::Elem,
6927 index: 0
6928 }
6929 ));
6930 assert_eq!(
6931 g.node(ph_id).shape.dims(),
6932 &[2, 3],
6933 "placeholder emits its declared target_shape_rel shape, not rank-0",
6934 );
6935 assert_eq!(
6936 g.node(ph_id).dtype,
6937 DType::F32,
6938 "placeholder dtype = bind 0's"
6939 );
6940 assert_eq!(
6941 g.node(root).shape.dims(),
6942 &[2, 3],
6943 "Mul(placeholder,..) inherits the declared shape"
6944 );
6945 }
6946 }
6947
6948 // reduce_max_to_backward migration (Increment C carriers, A2) -------------
6949 //
6950 // The 9-node imperative fair-share max subgradient (ReduceMaxTo→BroadcastTo→
6951 // Equal(U8)→MaskedFill→ReduceSumTo→Div→BroadcastTo→Mul, sharing the single
6952 // `mask_f`) becomes a portable `PatternNode` DATA recipe, and the FIRST recipe
6953 // to carry a `MaskedFill` — driven LIVE by the A2 fill-Scalar re-emit carrier
6954 // (the fill dtype re-resolves to x's dtype, matching `Scalar::one(dtype)`).
6955 // A DIRECT structural mirror (no D3 keepdim swap, no D4 pad), so the parity
6956 // oracle runs both sides through a toy f64 interpreter and asserts BIT-EXACT
6957 // structure — the discriminating power is confirmed by the sabotage-calibrated
6958 // negative in `reduce_max_to_backward_recipe_shares_mask_and_has_no_reshape`
6959 // (identity-share ⇒ ONE MaskedFill node across its two uses). Binds:
6960 // `0 = x`, `1 = up`. The reduce/count targets are `SameAs 1` (upstream's
6961 // shape); the broadcasts are `SameAs 0` (x's shape).
6962 mod reduce_max_to_backward_recipe {
6963 use super::super::*;
6964 use super::frozen_legacy_reduce_max_to_backward_decompose;
6965 use crate::registry::{FusedOps, reduce_max_to_backward};
6966 use fuel_ir::{DType, Shape};
6967 use std::collections::HashMap;
6968
6969 /// Right-aligned NumPy broadcast of `input` (shape `in_shape`) to
6970 /// `target` — a size-1 or padded leading dim contributes stride 0. General
6971 /// on purpose (both `bcast_x` sites tile a `[..,1]` keepdim back to x).
6972 fn broadcast(input: &[f64], in_shape: &[usize], target: &[usize]) -> Vec<f64> {
6973 let rank = target.len();
6974 let pad = rank - in_shape.len();
6975 let mut real_strides = vec![0isize; in_shape.len()];
6976 let mut s = 1isize;
6977 for i in (0..in_shape.len()).rev() {
6978 real_strides[i] = s;
6979 s *= in_shape[i] as isize;
6980 }
6981 let mut in_strides = vec![0isize; rank];
6982 for (i, stride) in in_strides.iter_mut().enumerate() {
6983 if i >= pad {
6984 let id = i - pad;
6985 *stride = if in_shape[id] == 1 {
6986 0
6987 } else {
6988 real_strides[id]
6989 };
6990 }
6991 }
6992 let out_n: usize = target.iter().product();
6993 let mut out = Vec::with_capacity(out_n);
6994 let mut idx = vec![0usize; rank];
6995 for _ in 0..out_n {
6996 let fi: isize = (0..rank).map(|i| idx[i] as isize * in_strides[i]).sum();
6997 out.push(input[fi as usize]);
6998 for i in (0..rank).rev() {
6999 idx[i] += 1;
7000 if idx[i] < target[i] {
7001 break;
7002 }
7003 idx[i] = 0;
7004 }
7005 }
7006 out
7007 }
7008
7009 /// Tiny f64 reference interpreter over the reduce-max-backward primitive
7010 /// vocabulary (leaf-lookup FIRST, then elementwise `Mul`/`Div`,
7011 /// `MulScalar`, `Equal` (1.0/0.0), the `MaskedFill` fill from the op's
7012 /// `Scalar`, last-axis `ReduceMaxTo`/`ReduceSumTo`, right-aligned
7013 /// `BroadcastTo`). BOTH parity sides run through it with identical in-order
7014 /// arithmetic, so a bit-exact assert isolates recipe STRUCTURE. Test dims
7015 /// reduce the LAST axis (`up = [.., 1]`), so the reduces are a last-axis
7016 /// chunk fold.
7017 fn eval_rmb(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
7018 if let Some(v) = leaves.get(&id) {
7019 return v.clone();
7020 }
7021 let node = g.node(id);
7022 match &node.op {
7023 Op::Mul => {
7024 let a = eval_rmb(g, node.inputs[0], leaves);
7025 let b = eval_rmb(g, node.inputs[1], leaves);
7026 a.iter().zip(&b).map(|(x, y)| x * y).collect()
7027 }
7028 Op::Div => {
7029 let a = eval_rmb(g, node.inputs[0], leaves);
7030 let b = eval_rmb(g, node.inputs[1], leaves);
7031 a.iter().zip(&b).map(|(x, y)| x / y).collect()
7032 }
7033 Op::MulScalar(m) => eval_rmb(g, node.inputs[0], leaves)
7034 .iter()
7035 .map(|v| v * m)
7036 .collect(),
7037 Op::Equal => {
7038 let a = eval_rmb(g, node.inputs[0], leaves);
7039 let b = eval_rmb(g, node.inputs[1], leaves);
7040 a.iter()
7041 .zip(&b)
7042 .map(|(x, y)| if x == y { 1.0 } else { 0.0 })
7043 .collect()
7044 }
7045 Op::MaskedFill { value } => {
7046 let input = eval_rmb(g, node.inputs[0], leaves);
7047 let mask = eval_rmb(g, node.inputs[1], leaves);
7048 let fill = value.to_f64();
7049 input
7050 .iter()
7051 .zip(&mask)
7052 .map(|(x, m)| if *m != 0.0 { fill } else { *x })
7053 .collect()
7054 }
7055 // Last-axis reduce-to-shape (up = [.., 1]); identical fold both
7056 // parity sides.
7057 Op::ReduceMaxTo(_) => {
7058 let input = eval_rmb(g, node.inputs[0], leaves);
7059 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
7060 input
7061 .chunks(last)
7062 .map(|row| row.iter().cloned().fold(f64::NEG_INFINITY, f64::max))
7063 .collect()
7064 }
7065 Op::ReduceSumTo(_) => {
7066 let input = eval_rmb(g, node.inputs[0], leaves);
7067 let last = *g.node(node.inputs[0]).shape.dims().last().unwrap();
7068 input
7069 .chunks(last)
7070 .map(|row| row.iter().sum::<f64>())
7071 .collect()
7072 }
7073 Op::BroadcastTo(target) => {
7074 let input = eval_rmb(g, node.inputs[0], leaves);
7075 let in_shape: Vec<usize> = g.node(node.inputs[0]).shape.dims().to_vec();
7076 broadcast(&input, &in_shape, target.dims())
7077 }
7078 other => panic!("eval_rmb: unhandled op {other:?}"),
7079 }
7080 }
7081
7082 /// Build a fused ReduceMaxToBackward node over `x [x_dims]` (input 0) and
7083 /// `up [up_dims]` (input 1, the upstream gradient = the forward reduce
7084 /// target). Returns `(x, up, fused)`.
7085 fn rmb_fused_node(
7086 g: &mut Graph,
7087 x_dims: &[usize],
7088 up_dims: &[usize],
7089 ) -> (NodeId, NodeId, NodeId) {
7090 let xsh = Shape::from_dims(x_dims);
7091 let x = g.push(Node {
7092 op: Op::Const,
7093 inputs: vec![],
7094 shape: xsh.clone(),
7095 dtype: DType::F32,
7096 });
7097 let up = g.push(Node {
7098 op: Op::Const,
7099 inputs: vec![],
7100 shape: Shape::from_dims(up_dims),
7101 dtype: DType::F32,
7102 });
7103 let fused = g.push(Node {
7104 op: Op::Fused(
7105 FusedOps::REDUCE_MAX_TO_BACKWARD,
7106 FusedOpParams::ReduceMaxToBackward,
7107 ),
7108 inputs: vec![x, up],
7109 shape: xsh,
7110 dtype: DType::F32,
7111 });
7112 (x, up, fused)
7113 }
7114
7115 /// A2 (a): ONE recipe datum decomposes at BOTH rank 2 and rank 3 (the
7116 /// shape polymorphism the baked-shape legacy body never had), and its
7117 /// numerics match the FROZEN legacy builder bit-exactly under the shared
7118 /// reference interpreter — the MaskedFill carrier driving live emission.
7119 #[test]
7120 fn reduce_max_to_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
7121 for (x_dims, up_dims) in [
7122 (vec![2usize, 4], vec![2usize, 1]),
7123 (vec![3, 5, 7], vec![3, 5, 1]),
7124 ] {
7125 let mut g = Graph::new();
7126 let (x, up, fused) = rmb_fused_node(&mut g, &x_dims, &up_dims);
7127 let xsh = Shape::from_dims(&x_dims);
7128 let new_root = reduce_max_to_backward::decompose(
7129 &mut g,
7130 fused,
7131 &FusedOpParams::ReduceMaxToBackward,
7132 );
7133 assert_ne!(new_root, fused, "recipe decompose must fire at {x_dims:?}");
7134 assert_eq!(
7135 g.node(new_root).shape,
7136 xsh,
7137 "reduce_max backward is x-shaped"
7138 );
7139 assert_eq!(g.node(new_root).dtype, DType::F32);
7140
7141 let legacy_root = frozen_legacy_reduce_max_to_backward_decompose(
7142 &mut g,
7143 fused,
7144 &FusedOpParams::ReduceMaxToBackward,
7145 );
7146
7147 let xn: usize = x_dims.iter().product();
7148 let upn: usize = up_dims.iter().product();
7149 let x_data: Vec<f64> = (0..xn)
7150 .map(|i| ((i as f64) * 0.37).sin() * 2.0 + 0.3)
7151 .collect();
7152 let up_data: Vec<f64> = (0..upn)
7153 .map(|i| ((i as f64) * 0.53).cos() * 1.1 - 0.2)
7154 .collect();
7155 let mut leaves = HashMap::new();
7156 leaves.insert(x, x_data);
7157 leaves.insert(up, up_data);
7158 let got = eval_rmb(&g, new_root, &leaves);
7159 let want = eval_rmb(&g, legacy_root, &leaves);
7160 assert_eq!(got.len(), want.len());
7161 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
7162 assert_eq!(
7163 a.to_bits(),
7164 b.to_bits(),
7165 "reduce_max_to_backward[{i}] at {x_dims:?}: recipe={a} vs legacy={b}",
7166 );
7167 }
7168 }
7169 }
7170
7171 /// A2 (structural): the single `mask_f` is IDENTITY-SHARED across its two
7172 /// use sites (`ReduceSumTo` and the final `Mul`) — exactly ONE MaskedFill
7173 /// node reachable, the same single-compute DAG the imperative body had (a
7174 /// duplicated emit would break this). A direct structural mirror: no D3/D4
7175 /// `Reshape`, and the MaskedFill fill Scalar re-resolved to x's (F32)
7176 /// dtype (evidence the A2 carrier reached emit).
7177 #[test]
7178 fn reduce_max_to_backward_recipe_shares_mask_and_has_no_reshape() {
7179 for (x_dims, up_dims) in [
7180 (vec![2usize, 4], vec![2usize, 1]),
7181 (vec![3, 5, 7], vec![3, 5, 1]),
7182 ] {
7183 let mut g = Graph::new();
7184 let (_x, _up, fused) = rmb_fused_node(&mut g, &x_dims, &up_dims);
7185 let root = reduce_max_to_backward::decompose(
7186 &mut g,
7187 fused,
7188 &FusedOpParams::ReduceMaxToBackward,
7189 );
7190 assert_ne!(root, fused, "recipe decompose fires at {x_dims:?}");
7191 assert!(
7192 matches!(g.node(root).op, Op::Mul),
7193 "root is Mul(mask_f, share_b)"
7194 );
7195 let reachable = crate::topo_order_multi(&g, &[root]);
7196 let masked = reachable
7197 .iter()
7198 .filter(|&&n| matches!(g.node(n).op, Op::MaskedFill { .. }))
7199 .count();
7200 assert_eq!(
7201 masked, 1,
7202 "mask_f identity-shared across its two uses at {x_dims:?} (ONE MaskedFill)",
7203 );
7204 let reshapes = reachable
7205 .iter()
7206 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
7207 .count();
7208 assert_eq!(
7209 reshapes, 0,
7210 "direct structural mirror — no D3/D4 Reshape at {x_dims:?}"
7211 );
7212 let fill_dtype = reachable
7213 .iter()
7214 .find_map(|&n| match &g.node(n).op {
7215 Op::MaskedFill { value } => Some(value.dtype()),
7216 _ => None,
7217 })
7218 .expect("a MaskedFill is present");
7219 assert_eq!(
7220 fill_dtype,
7221 DType::F32,
7222 "the A2 carrier re-resolved the fill Scalar to x's dtype at {x_dims:?}",
7223 );
7224 }
7225 }
7226
7227 /// A2 (totality): a wrong params payload is a typed decline surfaced as a
7228 /// fixpoint (G2), never a panic, declining BEFORE any emission.
7229 #[test]
7230 fn reduce_max_to_backward_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
7231 let mut g = Graph::new();
7232 let (_x, _up, fused) = rmb_fused_node(&mut g, &[2, 4], &[2, 1]);
7233 let before = g.len();
7234 let out = reduce_max_to_backward::decompose(&mut g, fused, &FusedOpParams::Rope);
7235 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
7236 assert_eq!(g.len(), before, "declined before any emission");
7237 }
7238 }
7239
7240 // powi_backward migration (Increment C carriers, A3) ----------------------
7241 //
7242 // The 3-node imperative closed-form gradient `grad_x = exp · x^(exp-1) ·
7243 // upstream` (PowI(exp-1)→MulScalar(exp)→Mul) becomes a portable `PatternNode`
7244 // DATA recipe, and the FIRST recipe to carry a `PowI` — driven LIVE by the A3
7245 // i32-exponent re-emit carrier (the exponent rides scalars[0]) — and the first
7246 // whose STRUCTURE is param-derived (both `exp` and `exp-1` are constant-folded
7247 // into the datum at build time, the minimal "param-derived const in the
7248 // recipe" C-4 posture, NOT a restructure to dodge the missing thread). A
7249 // DIRECT structural mirror (no D3 keepdim swap, no D4 pad), so the parity
7250 // oracle runs both sides through a toy f64 interpreter and asserts BIT-EXACT
7251 // structure across ranks AND exponents. Binds: `0 = x`, `1 = up`.
7252 mod powi_backward_recipe {
7253 use super::super::*;
7254 use super::frozen_legacy_powi_backward_decompose;
7255 use crate::registry::{FusedOps, powi_backward};
7256 use fuel_ir::{DType, Shape};
7257 use std::collections::HashMap;
7258
7259 /// Tiny f64 reference interpreter over the powi-backward primitive
7260 /// vocabulary (leaf-lookup FIRST, then elementwise `Mul`, `MulScalar`,
7261 /// and integer-power `PowI` via `f64::powi`). BOTH parity sides run
7262 /// through it with identical in-order arithmetic, so a bit-exact assert
7263 /// isolates recipe STRUCTURE — the exponent transported by the A3 carrier
7264 /// and the baked `MulScalar`. `x` and `up` are the same shape (`PowI` is
7265 /// elementwise), so every op is a flat elementwise map.
7266 fn eval_pb(g: &Graph, id: NodeId, leaves: &HashMap<NodeId, Vec<f64>>) -> Vec<f64> {
7267 if let Some(v) = leaves.get(&id) {
7268 return v.clone();
7269 }
7270 let node = g.node(id);
7271 match &node.op {
7272 Op::Mul => {
7273 let a = eval_pb(g, node.inputs[0], leaves);
7274 let b = eval_pb(g, node.inputs[1], leaves);
7275 a.iter().zip(&b).map(|(x, y)| x * y).collect()
7276 }
7277 Op::MulScalar(m) => eval_pb(g, node.inputs[0], leaves)
7278 .iter()
7279 .map(|v| v * m)
7280 .collect(),
7281 Op::PowI(n) => eval_pb(g, node.inputs[0], leaves)
7282 .iter()
7283 .map(|v| v.powi(*n))
7284 .collect(),
7285 other => panic!("eval_pb: unhandled op {other:?}"),
7286 }
7287 }
7288
7289 /// Build a fused PowIBackward node over `x [dims]` (input 0) and `up
7290 /// [dims]` (input 1, the upstream gradient — same shape as x), carrying
7291 /// `exp`. Returns `(x, up, fused)`.
7292 fn pb_fused_node(g: &mut Graph, dims: &[usize], exp: i32) -> (NodeId, NodeId, NodeId) {
7293 let sh = Shape::from_dims(dims);
7294 let x = g.push(Node {
7295 op: Op::Const,
7296 inputs: vec![],
7297 shape: sh.clone(),
7298 dtype: DType::F32,
7299 });
7300 let up = g.push(Node {
7301 op: Op::Const,
7302 inputs: vec![],
7303 shape: sh.clone(),
7304 dtype: DType::F32,
7305 });
7306 let fused = g.push(Node {
7307 op: Op::Fused(FusedOps::POWI_BACKWARD, FusedOpParams::PowIBackward { exp }),
7308 inputs: vec![x, up],
7309 shape: sh,
7310 dtype: DType::F32,
7311 });
7312 (x, up, fused)
7313 }
7314
7315 /// A3 (a): ONE recipe datum decomposes at BOTH rank 2 and rank 3 and
7316 /// across several exponents, and its numerics match the FROZEN legacy
7317 /// builder bit-exactly under the shared reference interpreter — the A3
7318 /// PowI carrier driving live emission of `Op::PowI(exp-1)`.
7319 #[test]
7320 fn powi_backward_recipe_decompose_is_polymorphic_and_matches_frozen_legacy() {
7321 // x strictly positive so exp==0's PowI(-1) reciprocal is finite/clean;
7322 // the exponents span +, the exp==1 identity, exp==0 (MulScalar(0)=0),
7323 // and a negative exponent (PowI(-3)).
7324 for dims in [vec![2usize, 4], vec![3, 5, 7]] {
7325 for exp in [3i32, 2, 5, 1, 0, -2] {
7326 let mut g = Graph::new();
7327 let (x, up, fused) = pb_fused_node(&mut g, &dims, exp);
7328 let sh = Shape::from_dims(&dims);
7329 let new_root = powi_backward::decompose(
7330 &mut g,
7331 fused,
7332 &FusedOpParams::PowIBackward { exp },
7333 );
7334 assert_ne!(
7335 new_root, fused,
7336 "recipe decompose must fire at {dims:?} exp={exp}"
7337 );
7338 assert_eq!(
7339 g.node(new_root).shape,
7340 sh,
7341 "powi backward is shape-preserving"
7342 );
7343 assert_eq!(g.node(new_root).dtype, DType::F32);
7344
7345 let legacy_root = frozen_legacy_powi_backward_decompose(
7346 &mut g,
7347 fused,
7348 &FusedOpParams::PowIBackward { exp },
7349 );
7350
7351 let n: usize = dims.iter().product();
7352 let x_data: Vec<f64> = (0..n)
7353 .map(|i| ((i as f64) * 0.41).sin() * 0.5 + 1.5)
7354 .collect();
7355 let up_data: Vec<f64> = (0..n)
7356 .map(|i| ((i as f64) * 0.61).cos() * 1.3 - 0.15)
7357 .collect();
7358 let mut leaves = HashMap::new();
7359 leaves.insert(x, x_data);
7360 leaves.insert(up, up_data);
7361 let got = eval_pb(&g, new_root, &leaves);
7362 let want = eval_pb(&g, legacy_root, &leaves);
7363 assert_eq!(got.len(), want.len());
7364 for (i, (a, b)) in got.iter().zip(want.iter()).enumerate() {
7365 assert_eq!(
7366 a.to_bits(),
7367 b.to_bits(),
7368 "powi_backward[{i}] at {dims:?} exp={exp}: recipe={a} vs legacy={b}",
7369 );
7370 }
7371 }
7372 }
7373 }
7374
7375 /// A3 (structural): a DIRECT structural mirror — the root is `Mul(scaled,
7376 /// up)`, the A3 carrier resolves the exponent to `Op::PowI(exp-1)`, the
7377 /// scale is a baked `Op::MulScalar(exp as f64)`, and NO D3/D4 `Reshape` is
7378 /// materialized (evidence the i32 carrier + baked scale reached emit at the
7379 /// right values).
7380 #[test]
7381 fn powi_backward_recipe_is_a_direct_mirror_with_carrier_exponent() {
7382 for exp in [3i32, 0, -2] {
7383 let mut g = Graph::new();
7384 let (_x, _up, fused) = pb_fused_node(&mut g, &[2, 4], exp);
7385 let root =
7386 powi_backward::decompose(&mut g, fused, &FusedOpParams::PowIBackward { exp });
7387 assert_ne!(root, fused, "recipe decompose fires at exp={exp}");
7388 assert!(
7389 matches!(g.node(root).op, Op::Mul),
7390 "root is Mul(scaled, up)"
7391 );
7392 let reachable = crate::topo_order_multi(&g, &[root]);
7393 // The A3 carrier reconstructed exactly one Op::PowI(exp-1).
7394 let powis: Vec<i32> = reachable
7395 .iter()
7396 .filter_map(|&n| match g.node(n).op {
7397 Op::PowI(k) => Some(k),
7398 _ => None,
7399 })
7400 .collect();
7401 assert_eq!(
7402 powis,
7403 vec![exp - 1],
7404 "the A3 carrier resolved Op::PowI(exp-1) at exp={exp}"
7405 );
7406 // The scale is a single baked MulScalar(exp as f64).
7407 let mul_scalars: Vec<f64> = reachable
7408 .iter()
7409 .filter_map(|&n| match g.node(n).op {
7410 Op::MulScalar(v) => Some(v),
7411 _ => None,
7412 })
7413 .collect();
7414 assert_eq!(
7415 mul_scalars,
7416 vec![exp as f64],
7417 "baked MulScalar(exp) at exp={exp}"
7418 );
7419 // Direct mirror — no keepdim/pad Reshape.
7420 let reshapes = reachable
7421 .iter()
7422 .filter(|&&n| matches!(g.node(n).op, Op::Reshape(_)))
7423 .count();
7424 assert_eq!(
7425 reshapes, 0,
7426 "direct structural mirror — no D3/D4 Reshape at exp={exp}"
7427 );
7428 }
7429 }
7430
7431 /// A3 (totality): a wrong params payload is a typed decline surfaced as a
7432 /// fixpoint (G2), never a panic, declining BEFORE any emission.
7433 #[test]
7434 fn powi_backward_recipe_wrong_params_is_a_fixpoint_not_a_crash() {
7435 let mut g = Graph::new();
7436 let (_x, _up, fused) = pb_fused_node(&mut g, &[2, 4], 3);
7437 let before = g.len();
7438 let out = powi_backward::decompose(&mut g, fused, &FusedOpParams::Rope);
7439 assert_eq!(out, fused, "wrong params ⇒ typed decline ⇒ fixpoint");
7440 assert_eq!(g.len(), before, "declined before any emission");
7441 }
7442 }
7443}