fuel_graph/

registry.rs

// SPDX-License-Identifier: MIT OR Apache-2.0
//! FusedOpRegistry — metadata side. Phase 7.6 step 1 (skeleton).
//!
//! Architecture v1.0 splits op identity into two layers:
//! - the closed [`crate::Op`] enum carries primitive variants exhaustively;
//! - one arm — `Op::Fused(FusedOpId, FusedOpParams)` — delegates to an
//!   open registry of fused-op entries populated at process startup.
//!
//! This module holds the *graph-side metadata* half of the registry:
//! identity ([`FusedOpId`], [`FusedOps`]), pattern + decomposition + backward
//! identity, and shape/dtype rules. The *kernel side* (BackendImpl with
//! its function-pointer KernelRef, cost estimate, PrecisionGuarantee, caps,
//! revision hash) lives in `fuel-storage::fused` because it carries
//! `KernelRef`, which fuel-graph cannot import without inverting the
//! existing fuel-storage → fuel-graph dependency direction. The two halves
//! are joined at runtime by [`FusedOpId`].
//!
//! See [docs/fused-op-registry.md](../../docs/fused-op-registry.md) for the
//! full design and [docs/architecture/03-ir.md](../../docs/architecture/03-ir.md)
//! for the architectural commitment.
//!
//! ## Status (step 1)
//!
//! Types only. No callers; no behavior change. Subsequent steps:
//! - Step 2: extend `Op` with `Op::Fused(FusedOpId, FusedOpParams)` arm.
//! - Step 3: migrate SoftmaxLastDim end-to-end through the registry.
//! - Step 4: migrate the remaining 12 fused ops.
//! - Step 5: drop the per-fused-op `Op` variants once nothing emits them.

use crate::{Graph, NodeId};
use fuel_ir::storage::OutputViewSpec;
use fuel_ir::{DType, DynScalar, Shape};
use fuel_kernel_seam_types::PatternNode;
use std::collections::HashMap;

pub mod awq_matmul;
pub mod causal_conv1d;
pub mod conv2d;
pub mod conv_transpose_2d;
pub mod flash_attn;
pub mod flash_attn_backward;
pub mod fused_linear;
pub mod fused_softmax_cross_entropy;
pub mod inplace_affine;
pub mod layer_norm_last_dim;
pub mod layer_norm_last_dim_backward;
pub mod nf4_matmul;
pub mod paged_attn;
pub mod powi_backward;
pub mod qmatmul;
pub mod reduce_max_to_backward;
pub mod rms_norm_last_dim;
pub mod rms_norm_last_dim_backward;
pub mod rope;
pub mod selective_scan;
pub mod softmax_last_dim;
pub mod softmax_last_dim_backward;
pub mod ssd_chunk_scan;

/// Stable identifier for a registered fused op. Indexes into
/// `FusedOpRegistry::entries`. Newtype over `u16` (~65K capacity is
/// plenty; today's catalog is 13-14 entries).
///
/// Constants for the well-known ids are exposed via [`FusedOps`]
/// associated constants for ergonomic pattern matching in rule code.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct FusedOpId(pub u16);

/// Reserved sentinel: an unallocated id. Never appears in a populated
/// registry; useful as a placeholder during step-2-and-step-3 migration
/// where a `BackwardKind::Fused(id)` may be written before the backward's
/// own id is assigned.
impl FusedOpId {
    pub const UNASSIGNED: FusedOpId = FusedOpId(0);

    /// The first id reserved for **runtime-registered** (e.g. JIT-synthesized
    /// or import-time) fused ops. Static build-time ids are dense
    /// `1..RUNTIME_FUSED_BASE`; runtime ids are allocated from
    /// `RUNTIME_FUSED_BASE..` so the two spaces never collide and
    /// [`Self::is_runtime`] is the routing bit. See
    /// `docs/specs/runtime-fused-op-registration.md`.
    pub const RUNTIME_FUSED_BASE: u16 = 0x8000;

    /// Whether this id names a runtime-registered fused op — its metadata lives
    /// in the `runtime_fused` sidecar, not the static [`default_registry`].
    pub fn is_runtime(self) -> bool {
        self.0 >= Self::RUNTIME_FUSED_BASE
    }
}

/// Per-fused-op metadata. Identity, pattern, decomposition, backward,
/// shape/dtype rules. The kernel-side metadata (per-backend KernelRef +
/// cost + PrecisionGuarantee + caps) lives in `fuel-storage::fused` and is
/// joined to this entry by [`FusedOpEntry::id`].
pub struct FusedOpEntry {
    /// Stable id for this fused op.
    pub id: FusedOpId,
    /// Stable human-readable name. Shows up in op_short_name, error
    /// messages, telemetry.
    pub name: &'static str,
    /// Categorical bucket. Used by telemetry and for cost-model defaults.
    pub family: FusedOpFamily,
    /// Canonical primitive subgraph this fused op represents. Used by
    /// fusion rules to recognize the pattern in a base map.
    pub pattern: SubgraphPattern,
    /// Decompose a fused-op node into its primitive subgraph. Used by
    /// lowering rules (and by autograd when [`Self::backward`] is
    /// `BackwardKind::Decompose`).
    ///
    /// Contract: the function appends primitive nodes to `graph` and
    /// returns the [`NodeId`] of the new root that replaces the fused
    /// node identified by the second argument. The fused node itself
    /// remains in the arena; the driver rewrites consumer edges to point
    /// at the returned id.
    pub decompose: fn(&mut Graph, NodeId, &FusedOpParams) -> NodeId,
    /// Backward identity for this fused op.
    pub backward: BackwardKind,
    /// Output shape rule, computed from input shapes and params.
    /// For a multi-output op, this must return slot 0's shape (which
    /// is also the primary `Node::shape` of the producer).
    pub shape_rule: fn(&[Shape], &FusedOpParams) -> Shape,
    /// Output dtype rule, computed from input dtypes and params.
    /// For a multi-output op, this must return slot 0's dtype (which
    /// is also the primary `Node::dtype` of the producer).
    pub dtype_rule: fn(&[DType], &FusedOpParams) -> DType,
    /// Multi-output authoring contract (Option C, Session 2).
    ///
    /// `None` for single-output ops (the default) — the op produces
    /// one logical output whose shape/dtype come from `shape_rule`
    /// and `dtype_rule`, and the realized storage is a plain
    /// non-bundled [`fuel_backend_contract::Storage`].
    ///
    /// `Some(fn)` for multi-output ops: returns the per-slot specs
    /// from input shapes + dtypes + params. The Graph and the realized
    /// storage allocator both consume the same Vec — `set_output_views`
    /// on `Graph::push` and `allocate_bundled_storage` at realization
    /// time — so the slot count and per-slot shapes/dtypes are a
    /// single source of truth.
    ///
    /// Contract for multi-output ops:
    /// - `output_views(...)[0].dtype` MUST equal `dtype_rule(...)`.
    /// - `output_views(...)[0].shape` MUST equal `shape_rule(...)`.
    /// - The returned Vec must be non-empty; a "multi-output op" with
    ///   one slot is just a single-output op (drop `output_views`).
    ///
    /// `Graph::push_fused` (the multi-output-aware emit path that
    /// Session 2 introduces) checks these at graph-build time. The
    /// raw `Graph::push` + `Graph::set_output_views` route preserves
    /// the same invariants via `set_output_views`'s validation pass.
    #[expect(
        clippy::type_complexity,
        reason = "optional output-view rule fn-pointer; a type alias would add indirection for a single field"
    )]
    pub output_views: Option<fn(&[Shape], &[DType], &FusedOpParams) -> Vec<OutputViewSpec>>,

    /// Which input index this fused op DESTROYS (mutates in place) on execution,
    /// or `None` if it is non-destructive (produces a fresh output, every input
    /// stays readable). **Mandatory** — declared here, beside `decompose`/
    /// `pattern`, so a fused op cannot be registered without stating it (the
    /// recipe principle extended: a third required property of the same object).
    /// `Op::destructive_input` reads this via `default_registry().entry(id)`
    /// instead of a separate hand-maintained match, so the fact lives with the
    /// thing and cannot drift out of sync. A destructive op needs the scheduler
    /// to pin it after every non-destructive reader of the destroyed input
    /// (`opt::derive_ordering`); a wrong `None` here is a silent data race. There
    /// is deliberately no default: `Option<usize>` is the *value* (which input,
    /// or none), but the field itself must be set at every construction site.
    /// GAP-076.
    pub destructive_input: Option<usize>,
}

/// Categorical bucket for a fused op. Drives telemetry grouping and
/// some cost-model defaults; not load-bearing for dispatch.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum FusedOpFamily {
    Forward,
    Backward,
    Quantized,
    Attention,
    Norm,
}

/// Per-instance parameters carried by `Op::Fused(id, params)` nodes.
/// Step 1 ships only the SoftmaxLastDim variant — the proof-of-concept
/// migration target. Step 4 extends this enum with the remaining 12
/// fused-op variants, one per migrated op.
///
/// The variant tag is implicitly aligned with the registry entry's
/// [`FusedOpId`] via the convention that registering an id places its
/// matching variant here. CSE and `op_key` derive a hashable key from
/// the tuple `(FusedOpId, FusedOpParams)` so identical fused-op nodes
/// with identical params dedupe.
#[derive(Debug, Clone, PartialEq)]
pub enum FusedOpParams {
    /// SoftmaxLastDim has no per-instance parameters; the last-dim axis
    /// is implicit in the input shape.
    SoftmaxLastDim,
    /// FusedLinear ((a @ b) + bias). No per-instance parameters: the
    /// matmul shape and bias rank are implicit in the three inputs.
    FusedLinear,
    /// RmsNormLastDim — carries the epsilon term used to stabilize the
    /// division by RMS magnitude. The last-dim axis itself is implicit
    /// in the input shape.
    RmsNormLastDim { eps: f64 },
    /// LayerNormLastDim — carries the epsilon term used to stabilize
    /// the division by sqrt(variance + eps). Like RmsNormLastDim, the
    /// last-dim axis is implicit in the input shape.
    LayerNormLastDim { eps: f64 },
    /// Rope — rotary position embedding with caller-supplied cos/sin
    /// tables. No per-instance parameters: `seq` and `head_dim` are
    /// recovered from the inputs at execution time (see
    /// `fuel-storage::pipelined::op_to_op_params`).
    Rope,
    /// Conv2D — 2-D cross-correlation with stride / padding / groups.
    /// Spatial dilation is always 1 today; the field is omitted from
    /// the param payload until backward's dilation-as-stride trick
    /// forces the primitive variant to grow one (per the comments
    /// in `NodeHandle::backward`'s `Op::Conv2D` arm).
    Conv2D {
        stride: (usize, usize),
        padding: (usize, usize),
        groups: usize,
    },
    /// SoftmaxLastDimBackward — `s * (g - sum(g * s, last, keepdim))`.
    /// Two inputs (forward_y, upstream); parameterless. Higher-order
    /// gradients panic per `NodeHandle::backward`'s MVP behavior.
    SoftmaxLastDimBackward,
    /// LayerNormLastDimBackward — recomputes mean/variance from the
    /// original x. Two inputs (x, upstream) + eps.
    LayerNormLastDimBackward { eps: f64 },
    /// RmsNormLastDimBackward — `r_rms · (g − x·s / (n·(mean_sq + eps)))`.
    /// Two inputs (x, upstream) + eps.
    RmsNormLastDimBackward { eps: f64 },
    /// ReduceMaxToBackward — routes upstream to argmax positions
    /// (fair-share on ties). Two inputs (x, upstream); parameterless.
    ReduceMaxToBackward,
    /// PowIBackward — `grad_x = exp · x^(exp-1) · upstream`. Two
    /// inputs `(x, upstream)`. Carries the same `exp` as the forward
    /// `Op::PowI(exp)` so the backward kernel can evaluate the
    /// derivative coefficient directly.
    PowIBackward { exp: i32 },
    /// ConvTranspose2D — fractionally-strided 2-D convolution.
    /// Two inputs `[x, weight]` (no bias). Carries the full
    /// stride/padding/output_padding/dilation/groups bundle.
    ConvTranspose2D {
        stride: (usize, usize),
        padding: (usize, usize),
        output_padding: (usize, usize),
        dilation: (usize, usize),
        groups: usize,
    },
    /// FlashAttn — multi-head scaled-dot-product attention with
    /// FlashAttention-shaped kernel hooks. 4-5 inputs (q, k, v,
    /// optional alibi).
    ///
    /// `k_len` (Phase D symbolic extents): `None` ⇒ the attended K/V
    /// length is the full K extent `k.dims()[2]` (today's behavior).
    /// `Some(dyn)` ⇒ the K/V buffers are a fixed **capacity**
    /// `[B, Hkv, max_seq, D]` and only the first `dyn.resolve(env)` rows
    /// are attended (the live prefix), with the causal mask
    /// bottom-right-aligned at offset `k_len - Sq` (the standard FA2
    /// convention, = `cached_len` in decode). This is the runtime
    /// `k_len` decoupled from K's allocated shape that backs flash
    /// decode over a persistent capacity KV-cache. Build via
    /// [`NodeHandle::flash_attn_dyn`](crate::NodeHandle::flash_attn_dyn).
    FlashAttn {
        softmax_scale: f32,
        causal: bool,
        window_size_left: Option<usize>,
        window_size_right: Option<usize>,
        softcap: Option<f32>,
        k_len: Option<DynScalar>,
    },
    /// PagedAttn — paged-cache attention. 5-6 inputs (q, k_cache,
    /// v_cache, block_table, context_lens, optional alibi).
    PagedAttn {
        softmax_scale: f32,
        block_size: usize,
        softcap: Option<f32>,
    },
    /// QMatMul — quantized matrix multiply `C = A @ dequant(W_Q)`.
    /// Two inputs `[a, w_q_bytes]`. The U32-typed w_q_bytes carries
    /// a Q-type block stream; `quant_type` selects the dequant
    /// implementation; `k` / `n` are the logical weight dims.
    QMatMul {
        quant_type: crate::QuantType,
        k: usize,
        n: usize,
    },
    /// InplaceAffine — `x = mul·x + add`, mutating input 0. Single
    /// input. Destructive on index 0 (marked via
    /// `Op::destructive_input`). Phase 1 of the in-place ops
    /// infrastructure.
    InplaceAffine { mul: f64, add: f64 },
    /// SsdChunkScan — Mamba-2's State-Space Duality chunked scan
    /// (forward only). Five inputs:
    /// - `x`:   `[batch, seqlen, heads, head_dim]`
    /// - `dt`:  `[batch, seqlen, heads]`
    /// - `a`:   `[heads]` — per-head scalar log A.
    /// - `b`/`c`: `[batch, seqlen, heads, state_dim]`
    ///
    /// Output: `y: [batch, seqlen, heads, head_dim]`.
    ///
    /// `chunk_size` is the SSD block size (typically 256 in Mamba-2).
    /// It controls GPU parallelism granularity but doesn't affect the
    /// mathematical result — the CPU kernel runs a sequential scan
    /// regardless. Any `chunk_size ∈ [1, seqlen]` that divides
    /// `seqlen` is valid.
    SsdChunkScan { chunk_size: usize },
    /// Nf4Matmul — bitsandbytes-style 4-bit NormalFloat quantized
    /// matrix multiply. Three inputs:
    /// - `activations`: `[..., M, K]`
    /// - `w_packed`:    `[N, K/2]` U8 (two 4-bit codes per byte,
    ///   K must be even).
    /// - `absmax`:      `[N, K/block_size]` F32 (per-output-row,
    ///   per-block scales).
    ///
    /// Output: `[..., M, N]` matching the activations' dtype.
    ///
    /// `block_size` is the NF4 quantization block size (typically
    /// 64 in bitsandbytes). The NF4 NormalFloat lookup table is
    /// baked into the kernel — not a runtime input.
    Nf4Matmul { block_size: usize },
    /// AwqMatmul — AWQ (mit-han-lab) asymmetric int4 weight-only
    /// matmul: `dequant(q) = (q - zero_point) * scale`, per-group.
    /// Four inputs, the quantized three in the baracuda W4A16 kernel
    /// layout (OC-major):
    /// - `activations`: `[..., M, K]`
    /// - `weight`:      `[N, K/8]` U32 (eight int4 codes per word,
    ///   sequential nibble order, K must be a multiple of 8).
    /// - `zeros`:       `[N, zeros_w]` U32 (group `g` is nibble
    ///   `g % 8` of word `g / 8`); `zeros_w` is
    ///   [`crate::registry::awq_matmul::awq_zeros_w`].
    /// - `scales`:      `[N, 8 * zeros_w]` F16 (entry `g` is group
    ///   `g`'s scale).
    ///
    /// Output: `[..., M, N]` matching the activations' dtype.
    ///
    /// `group_size` must be 64 or 128 and evenly divide `K`.
    /// `split_k_iters` is a CUDA-kernel launch-tuning parameter (unused
    /// by the decompose path; carried so the fused CUDA dispatch does
    /// not need a second params shape).
    ///
    /// See the module doc on [`crate::registry::awq_matmul`] for the
    /// layout rationale and the decompose recipe.
    AwqMatmul {
        group_size: usize,
        split_k_iters: i32,
    },
    /// FlashAttnBackward — shared params for the Q/K/V backward
    /// variants. Carries the same shape parameters as the forward
    /// `FlashAttn` so the recompute pass produces identical scores.
    /// All three FusedOpId-marked variants (FLASH_ATTN_BACKWARD_Q/K/V)
    /// use this single variant; the FusedOpId distinguishes which
    /// gradient is being computed.
    FlashAttnBackward {
        softmax_scale: f32,
        causal: bool,
        window_size_left: Option<usize>,
        window_size_right: Option<usize>,
        softcap: Option<f32>,
    },
    /// SelectiveScan — Mamba-1's selective state-space scan
    /// (forward only). Five inputs:
    /// - `u`:     `[batch, seqlen, dim]` — input sequence.
    /// - `delta`: `[batch, seqlen, dim]` — per-step state-update rate.
    /// - `a`:     `[dim, dstate]` — recurrence matrix log.
    /// - `b`:     `[batch, seqlen, dstate]` — selective input matrix.
    /// - `c`:     `[batch, seqlen, dstate]` — selective output matrix.
    ///
    /// Output: `y: [batch, seqlen, dim]` — the scanned sequence.
    ///
    /// `delta_softplus` toggles the standard softplus(delta)
    /// activation before use (matches baracuda's `selective_scan_*_run`
    /// flag). The optional inputs (d_skip, z, delta_bias) from the
    /// full Mamba algorithm are NOT exposed in v1 — they're
    /// mechanical extensions when a consumer needs them.
    ///
    /// Multi-output Option C (item 3, 2026-06-01): the op now always
    /// produces a bundled output with two slots — slot 0 `y` (the
    /// scanned sequence) and slot 1 `last_state` (the final
    /// h-state). Consumers project via `Op::View` / `Op::ViewOwned`.
    /// The earlier `return_state: bool` flag retired here.
    SelectiveScan { delta_softplus: bool },
    /// CausalConv1d — depthwise 1-D convolution with causal masking
    /// (left-pad-only) + optional fused SiLU. Three inputs:
    /// - `x`: `[batch, channels, seq + (kernel - 1)]` — caller pre-pads
    ///   with `kernel - 1` zeros on the left for the causal mask.
    /// - `weight`: `[channels, 1, kernel]` — depthwise, one filter per
    ///   channel (groups == channels in conv terminology).
    /// - `bias`: `[channels]`.
    ///
    /// Output: `[batch, channels, seq]`, same dtype as inputs.
    ///
    /// Designed to mirror baracuda's `causal_conv1d_*_run` signature
    /// (which already carries a `use_silu` flag) and the Mamba-1 /
    /// Mamba-2 prefill convolution. See
    /// [[match-external-convention-for-well-known-ops]] for the
    /// "match baracuda's contract exactly" reasoning.
    CausalConv1d { use_silu: bool },
    /// FusedSoftmaxCrossEntropy — fused softmax + negative log-likelihood
    /// over class-indexed targets, the standard PyTorch / Liger-Kernel
    /// training-time loss. Two inputs:
    /// - `logits: [..., V]` (F32) — pre-computed unnormalized class scores
    /// - `targets: [...]` (I64) — class indices, matching the rest of
    ///   PyTorch / baracuda CE convention (see
    ///   [[match-external-convention-for-well-known-ops]])
    ///
    /// Output is F32 regardless of input dtype (loss values are
    /// accumulated in F32 for stability):
    /// - `Reduction::Mean` / `Reduction::Sum` → scalar `[]`
    /// - `Reduction::None` → same shape as `targets` (`[...]`)
    ///
    /// `ignore_index` rows are dropped from the loss accumulator and
    /// from the Mean denominator. The conventional sentinel is `-100`.
    FusedSoftmaxCrossEntropy {
        reduction: Reduction,
        ignore_index: i64,
    },
    /// A **runtime-registered** fused op (JIT-synthesized or import-time). The
    /// op's identity is the runtime [`FusedOpId`] in `Op::Fused`; its recipe is
    /// the region in the `runtime_fused` sidecar. This variant carries only the
    /// extracted scalar args (the `extract:` slots, in slot order) — a
    /// parameterless runtime op carries `scalars: []`. See
    /// `docs/specs/runtime-fused-op-registration.md`.
    Runtime { scalars: Vec<f64> },
}

/// Reduction mode for losses with per-sample outputs. Matches PyTorch's
/// `reduction` parameter shape (`'mean' | 'sum' | 'none'`).
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum Reduction {
    /// Sum the per-sample losses and divide by the count of contributing
    /// samples (rows not masked by `ignore_index`). Standard training-time
    /// scalar.
    Mean,
    /// Sum the per-sample losses without dividing.
    Sum,
    /// No reduction — return the per-sample losses with `targets.shape`.
    None,
}

impl Reduction {
    /// Stable per-variant key for hashable encodings (`FusedOpParams::key`
    /// and the kernel-side `OpParams` payload). The variant tag space is
    /// local to whichever enum carries this discriminant.
    pub(crate) fn key(self) -> i64 {
        match self {
            Reduction::Mean => 0,
            Reduction::Sum => 1,
            Reduction::None => 2,
        }
    }
}

/// Hashable key for [`FusedOpParams`]. Used by `op_key`/CSE so that two
/// `Op::Fused(id, params)` nodes with identical params dedupe.
///
/// The encoding is variant-tag + payload as a `Vec<u64>` (bit patterns
/// for floats, repr for ints/usize). Mirrors the existing
/// [`crate::opt`]-style `OpKey` encoding so future extensions slot in
/// without rebuilding the CSE map.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FusedOpParamsKey {
    pub tag: u16,
    pub bits: Vec<u64>,
    pub ints: Vec<i64>,
}

impl FusedOpParams {
    /// Hashable encoding. Step 4 adds variants per-fused-op as each
    /// migrates; the tag uniquely identifies the variant and the
    /// bits/ints slots carry any payload (none for parameterless ops
    /// like SoftmaxLastDim and FusedLinear).
    pub fn key(&self) -> FusedOpParamsKey {
        match self {
            FusedOpParams::SoftmaxLastDim => FusedOpParamsKey {
                tag: 1,
                bits: Vec::new(),
                ints: Vec::new(),
            },
            FusedOpParams::FusedLinear => FusedOpParamsKey {
                tag: 2,
                bits: Vec::new(),
                ints: Vec::new(),
            },
            FusedOpParams::RmsNormLastDim { eps } => FusedOpParamsKey {
                tag: 3,
                // Encode eps as its raw bit pattern so CSE on two
                // RmsNormLastDim nodes with the same eps dedupes
                // (and two with different eps don't).
                bits: vec![eps.to_bits()],
                ints: Vec::new(),
            },
            FusedOpParams::LayerNormLastDim { eps } => FusedOpParamsKey {
                tag: 4,
                bits: vec![eps.to_bits()],
                ints: Vec::new(),
            },
            FusedOpParams::Rope => FusedOpParamsKey {
                tag: 5,
                bits: Vec::new(),
                ints: Vec::new(),
            },
            FusedOpParams::Conv2D {
                stride,
                padding,
                groups,
            } => FusedOpParamsKey {
                tag: 6,
                bits: Vec::new(),
                // Stride.{0,1}, padding.{0,1}, groups. Five i64 slots
                // — enough that two Conv2D nodes with different
                // strides or paddings hash to distinct CSE keys.
                ints: vec![
                    stride.0 as i64,
                    stride.1 as i64,
                    padding.0 as i64,
                    padding.1 as i64,
                    *groups as i64,
                ],
            },
            FusedOpParams::SoftmaxLastDimBackward => FusedOpParamsKey {
                tag: 7,
                bits: Vec::new(),
                ints: Vec::new(),
            },
            FusedOpParams::LayerNormLastDimBackward { eps } => FusedOpParamsKey {
                tag: 8,
                bits: vec![eps.to_bits()],
                ints: Vec::new(),
            },
            FusedOpParams::RmsNormLastDimBackward { eps } => FusedOpParamsKey {
                tag: 9,
                bits: vec![eps.to_bits()],
                ints: Vec::new(),
            },
            FusedOpParams::ReduceMaxToBackward => FusedOpParamsKey {
                tag: 10,
                bits: Vec::new(),
                ints: Vec::new(),
            },
            FusedOpParams::ConvTranspose2D {
                stride,
                padding,
                output_padding,
                dilation,
                groups,
            } => FusedOpParamsKey {
                tag: 11,
                bits: Vec::new(),
                // 9 i64 slots: stride.{0,1}, padding.{0,1},
                // output_padding.{0,1}, dilation.{0,1}, groups.
                ints: vec![
                    stride.0 as i64,
                    stride.1 as i64,
                    padding.0 as i64,
                    padding.1 as i64,
                    output_padding.0 as i64,
                    output_padding.1 as i64,
                    dilation.0 as i64,
                    dilation.1 as i64,
                    *groups as i64,
                ],
            },
            FusedOpParams::FlashAttn {
                softmax_scale,
                causal,
                window_size_left,
                window_size_right,
                softcap,
                k_len,
            } => {
                // Encode the optional dynamic k_len as (tag, value) so
                // two flash nodes whose attended-length symbol differs
                // are NOT CSE-merged. None → (0,0); Concrete(v) → (1,v);
                // Sym(s) → (2, s.0).
                let (k_tag, k_val): (i64, i64) = match k_len {
                    None => (0, 0),
                    Some(DynScalar::Concrete(v)) => (1, *v as i64),
                    Some(DynScalar::Sym(s)) => (2, s.0 as i64),
                };
                FusedOpParamsKey {
                    tag: 12,
                    // softmax_scale f32 → u64 bit pattern; softcap
                    // f32 → u64 bit pattern (NaN sentinel for None so
                    // two FlashAttn nodes with no softcap still hash
                    // equal).
                    bits: vec![
                        softmax_scale.to_bits() as u64,
                        softcap.map(|s| s.to_bits() as u64).unwrap_or(u64::MAX),
                    ],
                    // causal as i64; window sizes carry MIN as the
                    // "None" sentinel so a real Some(0) still distinguishes.
                    ints: vec![
                        if *causal { 1 } else { 0 },
                        window_size_left.map(|w| w as i64).unwrap_or(i64::MIN),
                        window_size_right.map(|w| w as i64).unwrap_or(i64::MIN),
                        k_tag,
                        k_val,
                    ],
                }
            }
            FusedOpParams::PagedAttn {
                softmax_scale,
                block_size,
                softcap,
            } => FusedOpParamsKey {
                tag: 13,
                bits: vec![
                    softmax_scale.to_bits() as u64,
                    softcap.map(|s| s.to_bits() as u64).unwrap_or(u64::MAX),
                ],
                ints: vec![*block_size as i64],
            },
            FusedOpParams::QMatMul { quant_type, k, n } => FusedOpParamsKey {
                tag: 14,
                bits: Vec::new(),
                // QuantType encoded as stable per-variant i64; k, n
                // as usize.
                ints: vec![quant_type_key(*quant_type), *k as i64, *n as i64],
            },
            FusedOpParams::PowIBackward { exp } => FusedOpParamsKey {
                tag: 15,
                bits: Vec::new(),
                ints: vec![*exp as i64],
            },
            FusedOpParams::InplaceAffine { mul, add } => FusedOpParamsKey {
                tag: 16,
                bits: vec![mul.to_bits(), add.to_bits()],
                ints: Vec::new(),
            },
            FusedOpParams::FusedSoftmaxCrossEntropy {
                reduction,
                ignore_index,
            } => FusedOpParamsKey {
                tag: 17,
                bits: Vec::new(),
                ints: vec![reduction.key(), *ignore_index],
            },
            FusedOpParams::CausalConv1d { use_silu } => FusedOpParamsKey {
                tag: 18,
                bits: Vec::new(),
                ints: vec![*use_silu as i64],
            },
            FusedOpParams::SelectiveScan { delta_softplus } => FusedOpParamsKey {
                tag: 19,
                bits: Vec::new(),
                ints: vec![*delta_softplus as i64],
            },
            FusedOpParams::SsdChunkScan { chunk_size } => FusedOpParamsKey {
                tag: 20,
                bits: Vec::new(),
                ints: vec![*chunk_size as i64],
            },
            FusedOpParams::Nf4Matmul { block_size } => FusedOpParamsKey {
                tag: 21,
                bits: Vec::new(),
                ints: vec![*block_size as i64],
            },
            FusedOpParams::AwqMatmul {
                group_size,
                split_k_iters,
            } => FusedOpParamsKey {
                tag: 23,
                bits: Vec::new(),
                ints: vec![*group_size as i64, *split_k_iters as i64],
            },
            FusedOpParams::FlashAttnBackward {
                softmax_scale,
                causal,
                window_size_left,
                window_size_right,
                softcap,
            } => FusedOpParamsKey {
                // Distinct tag from forward FlashAttn (12) so the
                // forward and backward nodes can't collide in CSE.
                tag: 22,
                bits: vec![
                    softmax_scale.to_bits() as u64,
                    softcap.map(|s| s.to_bits() as u64).unwrap_or(u64::MAX),
                ],
                ints: vec![
                    if *causal { 1 } else { 0 },
                    window_size_left.map(|w| w as i64).unwrap_or(i64::MIN),
                    window_size_right.map(|w| w as i64).unwrap_or(i64::MIN),
                ],
            },
            // Runtime-registered fused ops share one tag (0xF000); the runtime
            // FusedOpId in `Op::Fused` disambiguates distinct ops, and the
            // scalar bits disambiguate instances of the same op.
            FusedOpParams::Runtime { scalars } => FusedOpParamsKey {
                tag: 0xF000,
                bits: scalars.iter().map(|s| s.to_bits()).collect(),
                ints: Vec::new(),
            },
        }
    }
}

/// Stable per-variant key for [`crate::QuantType`]. Used by
/// [`FusedOpParams::key`] so CSE on two `Op::Fused(QMATMUL, _)`
/// nodes with the same `(quant_type, k, n)` collapses correctly.
fn quant_type_key(q: crate::QuantType) -> i64 {
    use crate::QuantType::*;
    match q {
        Q4_0 => 1,
        Q4_1 => 2,
        Q5_0 => 3,
        Q5_1 => 4,
        Q8_0 => 5,
        Q8_1 => 6,
        Q2K => 7,
        Q3K => 8,
        Q4_K_M => 9,
        Q5K => 10,
        Q6K => 11,
    }
}

/// What kind of backward this fused op has.
///
/// - `Fused(id)` — emit a fused-backward op with the given id. Used by
///   ops whose backward is awkward to express as a primitive
///   decomposition (e.g. SoftmaxLastDimBackward's fused subtract-and-
///   project formula).
/// - `Decompose` — autograd derives the backward from the primitive
///   decomposition (lower the fused node, then run autograd over the
///   primitives).
/// - `NotDifferentiable` — backward returns Err / panics cleanly.
///   Matches today's QMatMul / ArgMaxDim treatment.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum BackwardKind {
    Fused(FusedOpId),
    Decompose,
    NotDifferentiable,
}

/// Canonical primitive-subgraph pattern that a fused op represents.
/// Used by fusion rules.
///
/// Two reasonable shapes; the registry supports both:
/// - `Callable` — a closure-style matcher, mirrors PR 3's
///   `canonical_softmax_pattern`. Maximally flexible (consumer-count
///   guards, cross-checks); less analyzable.
/// - `Declarative` — a recursive `PatternTree` that the rule engine
///   compiles to a matcher. More analyzable; auto-generation from the
///   registry entry is straightforward.
///
/// Step 3 ports SoftmaxLastDim's existing matcher to `Callable`. The
/// declarative form (see Q1 in `docs/fused-op-registry.md`) is wired up
/// in step 4 once a second op exercises the abstraction.
#[expect(
    clippy::large_enum_variant,
    reason = "matcher pattern enum; boxing the large variant is deferred to the same follow-up as PatternNode (docs/gaps.md GAP-215) -- a cross-cutting change to construction/match sites, not a lint-commit change"
)]
pub enum SubgraphPattern {
    Declarative(PatternTree),
    Callable(fn(&Graph, NodeId) -> Option<PatternMatch>),
}

/// A declarative subgraph pattern: the §3 grammar root plus the
/// [`FusedOpParams`] to stamp on the matched fused node. The rule engine
/// ([`crate::opt::FusionRule`] / `PatternKind::Declarative`) walks `root`
/// against the graph (via [`crate::jit::match_region`]) and emits
/// `Op::Fused(id, params)`.
#[derive(Debug, Clone)]
pub struct PatternTree {
    /// The region/pattern root (the subgraph SINK, fkc-fusion-patterns §3a.1).
    pub root: crate::jit::PatternNode,
    /// The per-instance params stamped on the emitted fused node. A
    /// parameterless declarative op carries its unit variant; scalar-param ops
    /// recover their payload via the `extract:` layer (a follow-up) — for now
    /// the template lives here.
    pub params: FusedOpParams,
}

/// Result of a successful pattern match. Carries the bindings — the
/// concrete `NodeId`s that the pattern's variables matched against in
/// the host graph — plus the [`FusedOpParams`] payload the fusion rule
/// should stamp onto the emitted `Op::Fused(id, params)` node.
///
/// `bindings` is index-keyed: bindings sorted by index become the
/// fused-op node's input list (`inputs[0]` ← index 0, `inputs[1]` ←
/// index 1, …). SoftmaxLastDim's match has one binding `(0, x_id)`
/// and so emits a single-input node; FusedLinear (3 inputs:
/// `[a, b, bias]`) emits three bindings indexed 0–2.
///
/// `params` is the matcher's authority on the resulting fused-op's
/// per-instance parameters. The matcher knows what variant of
/// [`FusedOpParams`] it's recognizing; carrying that decision in the
/// match result keeps `crate::opt::FusionRule::rewrite` generic
/// across all registered fused ops.
#[derive(Debug, Clone)]
pub struct PatternMatch {
    /// Variable-id → resolved NodeId. The fusion rule sorts by index
    /// and uses the resolved ids in order as the emitted node's inputs.
    pub bindings: Vec<(usize, NodeId)>,
    /// Per-instance parameters the matcher stamps onto the fused-op
    /// node it produces. Parameterless ops (SoftmaxLastDim, FusedLinear)
    /// stamp their unit variant; parameterized ops (RmsNormLastDim,
    /// FlashAttn) recover their payload from the matched subgraph.
    pub params: FusedOpParams,
}

/// The metadata-side registry. Built at process startup, frozen
/// thereafter (architecture v1.0: no runtime extensibility). Lookups
/// are by [`FusedOpId`] (O(1)) or by name.
#[derive(Default)]
pub struct FusedOpRegistry {
    entries: Vec<FusedOpEntry>,
    by_name: HashMap<&'static str, FusedOpId>,
}

impl FusedOpRegistry {
    /// Empty registry. Step 3 builds one with `with_entry(...)` calls.
    pub fn new() -> Self {
        Self::default()
    }

    /// Register an entry. Returns self for builder-style chaining.
    /// Idempotent on (id, name); a duplicate id is a programming bug
    /// (asserts in debug builds, last-writer-wins in release).
    ///
    /// Pattern-hash indexing is reserved for step 4's declarative
    /// pattern engine; the `Declarative` arm wires into it then.
    pub fn with_entry(mut self, entry: FusedOpEntry) -> Self {
        debug_assert!(
            !self.by_name.contains_key(entry.name),
            "FusedOpRegistry: duplicate entry name {:?}",
            entry.name
        );
        debug_assert!(
            entry.id.0 as usize == self.entries.len() + 1 || self.entry(entry.id).is_none(),
            "FusedOpRegistry: id {:?} already populated",
            entry.id
        );
        self.by_name.insert(entry.name, entry.id);
        // Grow `entries` so direct id-indexing works. Slot 0 stays
        // reserved (FusedOpId::UNASSIGNED).
        let slot = entry.id.0 as usize;
        if self.entries.len() <= slot {
            self.entries.reserve(slot + 1 - self.entries.len());
            // Step 1 doesn't need to fill placeholder slots — production
            // paths register every id starting from 1 contiguously, so
            // the Vec stays dense. If a gap arises, the lookup helper
            // returns `None` for missing ids.
        }
        // Insert at the matching index when contiguous; otherwise
        // append. Step 3's first registration is FusedOpId(1), so the
        // simple append-with-slot-0-empty path is what we'll exercise.
        if self.entries.is_empty() {
            // Reserve slot 0 with a unit placeholder. Direct indexing is
            // the hot path; we accept one wasted slot.
            self.entries.push(placeholder_entry());
        }
        if slot < self.entries.len() {
            self.entries[slot] = entry;
        } else {
            // Pad if necessary, then push.
            while self.entries.len() < slot {
                self.entries.push(placeholder_entry());
            }
            self.entries.push(entry);
        }
        self
    }

    /// Look up an entry by id. Returns `None` for `FusedOpId::UNASSIGNED`
    /// or any unregistered id.
    pub fn entry(&self, id: FusedOpId) -> Option<&FusedOpEntry> {
        let slot = id.0 as usize;
        if slot == 0 || slot >= self.entries.len() {
            return None;
        }
        let e = &self.entries[slot];
        if e.id == FusedOpId::UNASSIGNED {
            None
        } else {
            Some(e)
        }
    }

    /// Look up an entry by name. Returns the [`FusedOpId`] when present.
    pub fn id_for_name(&self, name: &str) -> Option<FusedOpId> {
        self.by_name.get(name).copied()
    }

    /// Iterate over every registered (non-placeholder) entry. Used by
    /// [`crate::opt::RuleRegistry::default_rules`] to auto-generate
    /// one [`crate::opt::LoweringRule`] + [`crate::opt::FusionRule`]
    /// pair per fused op without naming each entry by hand.
    pub fn entries_iter(&self) -> impl Iterator<Item = &FusedOpEntry> {
        self.entries
            .iter()
            .filter(|e| e.id != FusedOpId::UNASSIGNED)
    }

    /// Number of registered (non-placeholder) entries.
    pub fn len(&self) -> usize {
        self.entries
            .iter()
            .filter(|e| e.id != FusedOpId::UNASSIGNED)
            .count()
    }

    /// Whether the registry has any registered entries.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// Hash of a primitive-subgraph pattern. Reserved for step 4's
/// declarative pattern engine; the hashing function is filled in alongside
/// `PatternTree`. (The static catalog's `by_pattern_hash` index placeholder
/// was removed unread; re-add it with that engine.)
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct PatternHash(pub u64);

/// Internal placeholder used to keep id-indexed `entries` dense. Never
/// Decompose the fused node `id` by re-emitting a **data recipe** — the
/// static-registry bridge onto the one region re-emitter (Increment C
/// slice 1, T5 / design D6). A migrated entry's `decompose` body becomes:
///
/// ```ignore
/// fn recipe() -> &'static PatternNode              // OnceLock-built data
/// fn scalars(params: &FusedOpParams) -> Option<Vec<f64>>  // per-entry projection
/// pub fn decompose(g, id, params) -> NodeId {
///     decompose_via_recipe(g, id, recipe(), scalars(params))
/// }
/// ```
///
/// The fused node's `inputs` are the recipe's binds (in bind-index order);
/// `scalars` fill the recipe's open scalar slots in pattern pre-order (the
/// per-entry projection maps the params payload — e.g. an `eps` — onto
/// them). The `FusedOpEntry.decompose` surface is UNCHANGED — imperative
/// bodies remain legal for param-conditional ops.
///
/// **Totality (G2, the recipe principle):** ANY failure returns `id` — the
/// fused node stays in place as a surfaced opaque-op gap (the lowering
/// driver records no progress), NEVER a panic:
/// * `scalars == None` — a wrong params payload for this entry;
/// * a recipe that fails `crate::runtime_fused::validate_recipe` — a
///   semantics-absent op token with no primitive re-emission (the
///   flip-withdrawal posture: unknown/non-registry tokens are surfaced
///   honest-miss declines, never accepted), non-contiguous binds, a
///   matcher-only node, or structurally-unresolvable rel attrs;
/// * a bind-count / node-input arity mismatch;
/// * a scalar-slot count mismatch;
/// * a shape-relative resolution decline at THESE input shapes (symbolic
///   extent, negative result, …);
/// * anything else — a `catch_unwind` last-resort guard converts an
///   unanticipated panic into the same fixpoint.
///
/// The first four decline BEFORE any emission (no partial nodes). An
/// emit-time decline can leave already-pushed child nodes behind as
/// unreferenced dead nodes — inert, the documented push-only-graph posture.
/// Declines are telemetered on stderr (the surfaced-gap channel until a
/// structured gap inventory exists).
/// **A fused node's recipe as portable data — the generic accessor the JIT seam
/// needs.** Returns the same [`PatternNode`] the op's `decompose` lowers,
/// *without touching the graph*.
///
/// `fuel_kernel_seam::JitRequest::region` is documented as "the recipe's
/// `decompose` (the primitive subgraph)", so a synthesizer can be asked to build
/// one kernel for a region the consumer discovered at runtime. Lowering needs
/// every primitive bound on the target device; **handing the region over needs
/// none of them** — which is the whole reason this accessor exists.
///
/// Dispatches on the fused op id. `None` means "this op does not expose its
/// recipe as data (yet)", which is a *surfaced gap*, never a crash — the same
/// G2 posture as a `decompose` fixpoint self-return. Adding an op is one arm.
///
/// Today only [`FusedOps::PAGED_ATTN`] implements it, because it is the op that
/// forced the question: no GPU backend implements the fused form at all, so the
/// JIT route is the only one that does not depend on every primitive being
/// pre-registered.
pub fn recipe_for(graph: &Graph, id: NodeId) -> Option<PatternNode> {
    let (fused_id, params) = match &graph.node(id).op {
        crate::Op::Fused(fid, p) => (*fid, p.clone()),
        _ => return None,
    };
    match fused_id {
        FusedOps::PAGED_ATTN => crate::registry::paged_attn::recipe_for(graph, id, &params),
        _ => None,
    }
}

pub fn decompose_via_recipe(
    graph: &mut Graph,
    id: NodeId,
    recipe: &PatternNode,
    scalars: Option<Vec<f64>>,
) -> NodeId {
    let Some(scalars) = scalars else {
        eprintln!(
            "decompose_via_recipe: wrong params payload for node {id:?} — \
             typed decline, node stays fused (G2 fixpoint)"
        );
        return id;
    };
    if let Err(err) = crate::runtime_fused::validate_recipe(recipe) {
        eprintln!(
            "decompose_via_recipe: recipe declined for node {id:?}: {err:?} — \
             surfaced honest miss, node stays fused (G2 fixpoint)"
        );
        return id;
    }
    let inputs = graph.node(id).inputs.clone();
    if recipe.bind_indices().len() != inputs.len() {
        eprintln!(
            "decompose_via_recipe: recipe binds {} != node inputs {} for node {id:?} — \
             typed decline, node stays fused (G2 fixpoint)",
            recipe.bind_indices().len(),
            inputs.len(),
        );
        return id;
    }
    if scalars.len() != crate::runtime_fused::count_scalar_slots(recipe) {
        eprintln!(
            "decompose_via_recipe: {} projected scalars for {} open slots on node {id:?} — \
             typed decline, node stays fused (G2 fixpoint)",
            scalars.len(),
            crate::runtime_fused::count_scalar_slots(recipe),
        );
        return id;
    }
    // The resolving emit. A rel-attr decline at these shapes is a
    // value-dependent gap (`Err` arm); `catch_unwind` is the never-panic
    // last resort for anything validation didn't anticipate (same posture as
    // `register_runtime_fused`'s hash guard) — a partial emission left by
    // either failure mode is inert dead nodes in the push-only graph.
    let emitted = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
        crate::runtime_fused::try_emit_region(graph, recipe, &inputs, &scalars)
    }));
    match emitted {
        Ok(Ok(root)) => root,
        Ok(Err(err)) => {
            eprintln!(
                "decompose_via_recipe: rel-attr resolution declined for node {id:?}: {err:?} — \
                 surfaced gap, node stays fused (G2 fixpoint)"
            );
            id
        }
        Err(_) => {
            eprintln!(
                "decompose_via_recipe: emit panicked for node {id:?} — \
                 caught (never-panic last resort), node stays fused (G2 fixpoint)"
            );
            id
        }
    }
}

/// returned by [`FusedOpRegistry::entry`].
fn placeholder_entry() -> FusedOpEntry {
    FusedOpEntry {
        destructive_input: None,
        id: FusedOpId::UNASSIGNED,
        name: "<unassigned>",
        family: FusedOpFamily::Forward,
        pattern: SubgraphPattern::Callable(|_g, _id| None),
        decompose: |_g, id, _p| id,
        backward: BackwardKind::NotDifferentiable,
        shape_rule: |_shapes, _params| Shape::from_dims(&[]),
        dtype_rule: |_dtypes, _params| DType::F32,
        output_views: None,
    }
}

/// Associated constants for the well-known fused-op ids. Rule code
/// matches `Op::Fused(FusedOps::SOFTMAX_LAST_DIM, _)` — almost as
/// ergonomic as today's `Op::SoftmaxLastDim`. Steps 3-4 fill these in
/// as each fused op migrates to the registry.
pub struct FusedOps;

impl FusedOps {
    /// SoftmaxLastDim — proof-of-concept migration target (step 3).
    pub const SOFTMAX_LAST_DIM: FusedOpId = FusedOpId(1);
    /// FusedLinear — first multi-input fused op migrated to the
    /// registry (step 4). Three inputs `[a, b, bias]`, output
    /// `(a @ b) + bias`. The CUTLASS bias-epilogue integration in the
    /// baracuda-cutlass alpha.13 plan registers here.
    pub const FUSED_LINEAR: FusedOpId = FusedOpId(2);
    /// RmsNormLastDim — `x / sqrt(mean(x²) + eps)` along the last dim.
    /// Migrated in step 4 (continued). Single input + eps param.
    pub const RMS_NORM_LAST_DIM: FusedOpId = FusedOpId(3);
    /// LayerNormLastDim — `(x - mean) / sqrt(var + eps)` along the
    /// last dim (no affine). Single input + eps param.
    pub const LAYER_NORM_LAST_DIM: FusedOpId = FusedOpId(4);
    /// Rope — rotary position embedding with caller-supplied cos/sin
    /// tables. Three inputs `[x, cos, sin]`, parameterless (the seq/
    /// head_dim shape parameters live in the input shapes).
    pub const ROPE: FusedOpId = FusedOpId(5);
    /// Conv2D — 2-D cross-correlation. Two or three inputs `[x,
    /// weight]` or `[x, weight, bias]`; carries stride / padding /
    /// groups in [`FusedOpParams::Conv2D`].
    pub const CONV2D: FusedOpId = FusedOpId(6);
    /// SoftmaxLastDimBackward — fused backward helper for the
    /// SoftmaxLastDim forward. Two inputs `[y, upstream]`,
    /// parameterless. Wired through `BackwardKind::Fused` from
    /// `SOFTMAX_LAST_DIM`'s entry.
    pub const SOFTMAX_LAST_DIM_BACKWARD: FusedOpId = FusedOpId(7);
    /// LayerNormLastDimBackward — fused backward helper for the
    /// LayerNormLastDim forward. Two inputs `[x, upstream]` + eps.
    pub const LAYER_NORM_LAST_DIM_BACKWARD: FusedOpId = FusedOpId(8);
    /// RmsNormLastDimBackward — fused backward helper for the
    /// RmsNormLastDim forward. Two inputs `[x, upstream]` + eps.
    pub const RMS_NORM_LAST_DIM_BACKWARD: FusedOpId = FusedOpId(9);
    /// ReduceMaxToBackward — fused backward helper for the
    /// `Op::ReduceMaxTo` primitive (routes upstream to argmax
    /// positions, fair-share on ties). Two inputs `[x, upstream]`,
    /// parameterless. Note this is the backward of a *primitive*
    /// (`Op::ReduceMaxTo`), not of a fused forward — there's no
    /// `BackwardKind::Fused` edge wired to this helper, autograd
    /// reaches it directly from `Op::ReduceMaxTo`'s arm.
    pub const REDUCE_MAX_TO_BACKWARD: FusedOpId = FusedOpId(10);
    /// ConvTranspose2D — fractionally-strided 2-D convolution.
    /// Two inputs `[x, weight]` (no bias); carries
    /// stride/padding/output_padding/dilation/groups in
    /// [`FusedOpParams::ConvTranspose2D`].
    pub const CONV_TRANSPOSE2D: FusedOpId = FusedOpId(11);
    /// FlashAttn — multi-head attention with FlashAttention-shaped
    /// kernel hooks. 4-5 inputs (q, k, v, optional alibi).
    pub const FLASH_ATTN: FusedOpId = FusedOpId(12);
    /// PagedAttn — paged-cache attention. 5-6 inputs.
    pub const PAGED_ATTN: FusedOpId = FusedOpId(13);
    /// QMatMul — quantized matmul `C = A @ dequant(W_Q)`. Two
    /// inputs `[a, w_q_bytes]`.
    pub const QMATMUL: FusedOpId = FusedOpId(14);
    /// PowIBackward — fused backward helper for `Op::PowI(n)`.
    /// `grad_x = n · x^(n-1) · upstream`. Two inputs `[x, upstream]`,
    /// carries `exp: i32` in [`FusedOpParams::PowIBackward`]. Replaces
    /// the autograd primitive decomposition (PowI(n-1) → MulScalar →
    /// Mul, 3 nodes) with a single launch of baracuda alpha.31's
    /// `unary_powi_backward_*` kernel. Note this is the backward of a
    /// primitive (`Op::PowI`), not of a fused forward — autograd
    /// reaches it directly from `Op::PowI`'s arm rather than via a
    /// `BackwardKind::Fused` edge.
    pub const POWI_BACKWARD: FusedOpId = FusedOpId(15);

    /// In-place affine: `x = mul·x + add`. Single input (the mutated
    /// tensor), single launch into baracuda's `affine_inplace_*`
    /// symbol on CUDA. Carries `mul: f64, add: f64` in
    /// [`FusedOpParams::InplaceAffine`]. Marked destructive on input
    /// 0 via `Op::destructive_input`, so `opt::derive_ordering` pins
    /// it to run after every non-destructive reader. Phase 1 of the
    /// in-place ops infrastructure
    /// (`docs/session-prompts/in-place-ops-infrastructure.md`);
    /// dispatch + autograd integration deferred to Phases 3 + 4.
    pub const INPLACE_AFFINE: FusedOpId = FusedOpId(16);

    /// FusedSoftmaxCrossEntropy — fused softmax + NLL over class-indexed
    /// targets. Two inputs `[logits, targets]`; carries
    /// `{ reduction, ignore_index }` in
    /// [`FusedOpParams::FusedSoftmaxCrossEntropy`]. The forward kernel
    /// avoids materializing the `[..., V]` softmax/log-softmax
    /// intermediates that the primitive composition does (`fuel-core::
    /// train::cross_entropy_with_logits`). Backward via Decompose runs
    /// the primitive chain; the in-place Liger-style fused backward is a
    /// later session (needs in-place mutation infrastructure).
    pub const FUSED_SOFTMAX_CROSS_ENTROPY: FusedOpId = FusedOpId(17);

    /// CausalConv1d — depthwise 1-D convolution + causal masking +
    /// optional fused SiLU. Three inputs `[x, weight, bias]` where
    /// `x: [batch, channels, seq + kernel - 1]`,
    /// `weight: [channels, 1, kernel]`, `bias: [channels]`. Carries
    /// `use_silu: bool` in [`FusedOpParams::CausalConv1d`]. The forward
    /// kernel matches baracuda's `causal_conv1d_*_run` contract.
    /// `BackwardKind::NotDifferentiable` for v1 (Mamba inference path
    /// today is inference-only).
    pub const CAUSAL_CONV1D: FusedOpId = FusedOpId(18);

    /// SelectiveScan — Mamba-1's selective state-space scan
    /// (forward only). Five inputs `[u, delta, a, b, c]`; one output
    /// `y: [batch, seqlen, dim]`. Carries `delta_softplus: bool` in
    /// [`FusedOpParams::SelectiveScan`]. v1 drops the optional
    /// inputs (d_skip, z, delta_bias) and the `last_state` second
    /// output — both are mechanical follow-ups when a consumer needs
    /// them. `BackwardKind::NotDifferentiable` for v1.
    pub const SELECTIVE_SCAN: FusedOpId = FusedOpId(19);

    /// SsdChunkScan — Mamba-2's State-Space Duality chunked scan
    /// (forward only). Five inputs `[x, dt, a, b, c]`; one output
    /// `y: [batch, seqlen, heads, head_dim]`. Carries
    /// `chunk_size: usize` in [`FusedOpParams::SsdChunkScan`]
    /// (parallelism granularity for the GPU path; the CPU kernel
    /// runs sequential and any positive `chunk_size` dividing
    /// `seqlen` produces the same mathematical answer).
    /// `BackwardKind::NotDifferentiable` for v1.
    pub const SSD_CHUNK_SCAN: FusedOpId = FusedOpId(20);

    /// Nf4Matmul — bitsandbytes-style 4-bit NormalFloat quantized
    /// matrix multiply. Three inputs `[activations, w_packed, absmax]`
    /// where `w_packed: [N, K/2] U8` and `absmax: [N, K/block_size] F32`.
    /// Output `[..., M, N]` matching activations' dtype (F32/F16/BF16
    /// in v1). Carries `block_size: usize` in
    /// [`FusedOpParams::Nf4Matmul`]. `BackwardKind::NotDifferentiable`
    /// (NF4 weights are frozen; the dequant-then-matmul fallback is
    /// the wrong path for the cases NF4 exists to optimize).
    pub const NF4_MATMUL: FusedOpId = FusedOpId(21);

    /// FlashAttnBackwardQ — produces dQ from `(q, k, v, do, [alibi])`.
    /// See [`fuel_ir::dispatch::OpKind::FlashAttnBackwardQ`]. Three
    /// separate FusedOp ids (Q/K/V) is the v1 design — a single
    /// multi-output op would share the recompute pass across all three
    /// gradients but needs multi-output infrastructure that doesn't
    /// exist yet. Reuses [`FusedOpParams::FlashAttnBackward`].
    pub const FLASH_ATTN_BACKWARD_Q: FusedOpId = FusedOpId(22);
    /// FlashAttnBackwardK — produces dK. See [`Self::FLASH_ATTN_BACKWARD_Q`].
    pub const FLASH_ATTN_BACKWARD_K: FusedOpId = FusedOpId(23);
    /// FlashAttnBackwardV — produces dV. See [`Self::FLASH_ATTN_BACKWARD_Q`].
    pub const FLASH_ATTN_BACKWARD_V: FusedOpId = FusedOpId(24);

    /// AwqMatmul — AWQ (mit-han-lab) asymmetric int4 weight-only
    /// matmul. Four inputs `[activations, weight, zeros, scales]`
    /// (one more than NF4's three — AWQ's zero-point is a fourth
    /// operand, not baked into the dequant like NF4's codebook) where
    /// `weight: [N, K/8] U32`, `zeros: [N, zeros_w] U32`,
    /// `scales: [N, 8 * zeros_w] F16` — the baracuda kernel layout.
    /// Output `[..., M, N]` matching activations' dtype
    /// (F32/F16/BF16 in v1). Carries `{ group_size, split_k_iters }`
    /// in [`FusedOpParams::AwqMatmul`]. `BackwardKind::Decompose`
    /// (the weight is frozen; the activation gradient flows through
    /// the decompose). Unlike NF4 (a non-linear codebook with no
    /// primitive spelling), AWQ's affine dequant decomposes cleanly
    /// into primitives — see the module doc on
    /// [`crate::registry::awq_matmul`] for the recipe.
    pub const AWQ_MATMUL: FusedOpId = FusedOpId(25);
}

/// Process-wide default registry: the union of every fused op's
/// metadata-side entry. Built once on first access via
/// [`std::sync::OnceLock`]; immutable thereafter (architecture v1.0:
/// no runtime extensibility).
///
/// Step 3 populates only SoftmaxLastDim; step 4 fills in the other
/// 12 fused ops as they migrate.
pub fn default_registry() -> &'static FusedOpRegistry {
    use std::sync::OnceLock;
    static REGISTRY: OnceLock<FusedOpRegistry> = OnceLock::new();
    REGISTRY.get_or_init(|| {
        FusedOpRegistry::new()
            .with_entry(softmax_last_dim::entry())
            .with_entry(fused_linear::entry())
            .with_entry(rms_norm_last_dim::entry())
            .with_entry(layer_norm_last_dim::entry())
            .with_entry(rope::entry())
            .with_entry(conv2d::entry())
            .with_entry(softmax_last_dim_backward::entry())
            .with_entry(layer_norm_last_dim_backward::entry())
            .with_entry(rms_norm_last_dim_backward::entry())
            .with_entry(reduce_max_to_backward::entry())
            .with_entry(conv_transpose_2d::entry())
            .with_entry(flash_attn::entry())
            .with_entry(flash_attn_backward::entry_q())
            .with_entry(flash_attn_backward::entry_k())
            .with_entry(flash_attn_backward::entry_v())
            .with_entry(paged_attn::entry())
            .with_entry(qmatmul::entry())
            .with_entry(powi_backward::entry())
            .with_entry(inplace_affine::entry())
            .with_entry(fused_softmax_cross_entropy::entry())
            .with_entry(causal_conv1d::entry())
            .with_entry(selective_scan::entry())
            .with_entry(ssd_chunk_scan::entry())
            .with_entry(nf4_matmul::entry())
            .with_entry(awq_matmul::entry())
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Smoke: id newtype is comparable + hashable.
    #[test]
    fn fused_op_id_basic() {
        let a = FusedOpId(1);
        let b = FusedOpId(1);
        let c = FusedOpId(2);
        assert_eq!(a, b);
        assert_ne!(a, c);
    }

    /// Smoke: empty registry has no entries.
    #[test]
    fn registry_empty() {
        let r = FusedOpRegistry::new();
        assert!(r.is_empty());
        assert_eq!(r.len(), 0);
        assert!(r.entry(FusedOpId(1)).is_none());
        assert!(r.entry(FusedOpId::UNASSIGNED).is_none());
        assert!(r.id_for_name("anything").is_none());
    }

    /// Smoke: register one entry, look it up by id and by name.
    #[test]
    fn registry_register_one() {
        fn dummy_decompose(_g: &mut Graph, id: NodeId, _p: &FusedOpParams) -> NodeId {
            id
        }
        fn dummy_pattern(_g: &Graph, _id: NodeId) -> Option<PatternMatch> {
            None
        }
        fn dummy_shape(_s: &[Shape], _p: &FusedOpParams) -> Shape {
            Shape::from_dims(&[1])
        }
        fn dummy_dtype(_d: &[DType], _p: &FusedOpParams) -> DType {
            DType::F32
        }

        let r = FusedOpRegistry::new().with_entry(FusedOpEntry {
            destructive_input: None,
            id: FusedOps::SOFTMAX_LAST_DIM,
            name: "SoftmaxLastDim",
            family: FusedOpFamily::Forward,
            pattern: SubgraphPattern::Callable(dummy_pattern),
            decompose: dummy_decompose,
            backward: BackwardKind::NotDifferentiable,
            shape_rule: dummy_shape,
            dtype_rule: dummy_dtype,
            output_views: None,
        });

        assert_eq!(r.len(), 1);
        assert!(!r.is_empty());
        let entry = r.entry(FusedOps::SOFTMAX_LAST_DIM).expect("registered");
        assert_eq!(entry.name, "SoftmaxLastDim");
        assert_eq!(entry.family, FusedOpFamily::Forward);
        assert_eq!(
            r.id_for_name("SoftmaxLastDim"),
            Some(FusedOps::SOFTMAX_LAST_DIM)
        );
    }

    /// FusedOpParams::key produces a stable encoding.
    #[test]
    fn fused_op_params_key_softmax() {
        let k1 = FusedOpParams::SoftmaxLastDim.key();
        let k2 = FusedOpParams::SoftmaxLastDim.key();
        assert_eq!(k1, k2);
        assert_eq!(k1.tag, 1);
        assert!(k1.bits.is_empty());
        assert!(k1.ints.is_empty());
    }
}