nxpu_analysis/

analyze.rs

//! IR pattern classification for ONNX lowering.
//!
//! Analyzes an entry point's global variables and function body to classify
//! the computation into a known ONNX-mappable pattern.

use std::fmt;

use nxpu_ir::{
    AddressSpace, Arena, ArraySize, BinaryOp, Expression, GlobalVariable, Handle, Literal,
    MathFunction, Module, Scalar, ScalarKind, Statement, StorageAccess, Type, TypeInner,
    UniqueArena,
};

/// ONNX-compatible data type constants (matches TensorProto.DataType).
pub mod data_type {
    pub const FLOAT: i32 = 1;
    pub const UINT8: i32 = 2;
    pub const INT8: i32 = 3;
    pub const INT32: i32 = 6;
    pub const INT64: i32 = 7;
    pub const BOOL: i32 = 9;
    pub const FLOAT16: i32 = 10;
    pub const UINT32: i32 = 12;
    pub const BFLOAT16: i32 = 16;
}

/// Errors during IR pattern analysis.
#[derive(Debug, thiserror::Error)]
pub enum AnalysisError {
    #[error("no entry points in module")]
    NoEntryPoints,
    #[error("entry point index {0} out of range")]
    EntryPointOutOfRange(usize),
    #[error("unsupported pattern: {0}")]
    UnsupportedPattern(String),
    #[error("missing uniform params struct")]
    MissingParams,
}

/// Role of a tensor in the computation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TensorRole {
    Input,
    Output,
}

impl fmt::Display for TensorRole {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Self::Input => "Input",
            Self::Output => "Output",
        })
    }
}

/// A storage buffer bound as a tensor.
#[derive(Debug, Clone)]
pub struct TensorBinding {
    /// Handle to the underlying global variable.
    pub handle: Handle<GlobalVariable>,
    /// Human-readable tensor name.
    pub name: String,
    /// ONNX element data type (see [`data_type`]).
    pub elem_type: i32,
    /// Whether this tensor is an input or output.
    pub role: TensorRole,
}

impl fmt::Display for TensorBinding {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} ({})", self.name, self.role)
    }
}

/// A scalar the host supplies at dispatch time.
///
/// Either a whole uniform (`var<uniform> input_scale: f32`) or one member of a
/// uniform params struct (`params.a`). It is deliberately *not* a constant:
/// the value is written per dispatch, so it lowers to a rank-0 graph input,
/// not to an initializer folded into the graph. `axpy`'s own comment is
/// explicit that `a` changes every diffusion step.
#[derive(Debug, Clone)]
pub struct ScalarBinding {
    /// Name to give the rank-0 graph input.
    pub name: String,
    /// ONNX element data type (see [`data_type`]).
    pub elem_type: i32,
}

impl fmt::Display for ScalarBinding {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} (Scalar)", self.name)
    }
}

/// One operand of a [`KernelPattern::ElementWiseChain`] step.
#[derive(Debug, Clone)]
pub enum ChainOperand {
    /// A whole tensor, read at the same index the result is stored at.
    Tensor(TensorBinding),
    /// A dispatch-time scalar, broadcast over every element.
    Scalar(ScalarBinding),
}

impl fmt::Display for ChainOperand {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Tensor(t) => write!(f, "{t}"),
            Self::Scalar(s) => write!(f, "{s}"),
        }
    }
}

/// One step of a [`KernelPattern::ElementWiseChain`]: `acc = acc <op> operand`.
///
/// The accumulator is always on the *left*. A kernel that writes the
/// accumulator on the right of a `Sub` or a `Div` computes something this
/// cannot say, and the matcher refuses it rather than reordering the operands.
#[derive(Debug, Clone)]
pub struct ChainStep {
    /// The operation applied at this step.
    pub op: ElementWiseOp,
    /// The right-hand operand.
    pub operand: ChainOperand,
}

/// Symbolic dimension names for matrix multiplication.
#[derive(Debug, Clone)]
pub struct MatMulShape {
    /// Number of rows in the left matrix.
    pub m: String,
    /// Number of columns in the right matrix.
    pub n: String,
    /// Shared inner dimension.
    pub k: String,
}

impl fmt::Display for MatMulShape {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "MatMul({}x{} * {}x{})", self.m, self.k, self.k, self.n)
    }
}

/// Element-wise binary operation kind.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ElementWiseOp {
    Add,
    Sub,
    Mul,
    Div,
}

impl fmt::Display for ElementWiseOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.op_name())
    }
}

impl ElementWiseOp {
    /// Returns the canonical operator name (e.g. "Add", "Relu", "ReduceSum").
    pub fn op_name(self) -> &'static str {
        match self {
            Self::Add => "Add",
            Self::Sub => "Sub",
            Self::Mul => "Mul",
            Self::Div => "Div",
        }
    }

    /// Whether swapping the operands leaves the result unchanged.
    ///
    /// [`KernelPattern::ElementWiseChain`] applies every step as
    /// `acc = acc <op> operand`, so a chain may only be read out of an
    /// expression that put the accumulator on the left — unless the operation
    /// does not care, which is what this answers.
    pub fn is_commutative(self) -> bool {
        matches!(self, Self::Add | Self::Mul)
    }
}

/// Activation function kind.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ActivationOp {
    Relu,
    Sigmoid,
    Tanh,
    Softmax,
    Gelu,
    Silu,
    Mish,
}

impl fmt::Display for ActivationOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.op_name())
    }
}

impl ActivationOp {
    /// Returns the canonical operator name (e.g. "Relu", "Sigmoid").
    pub fn op_name(self) -> &'static str {
        match self {
            Self::Relu => "Relu",
            Self::Sigmoid => "Sigmoid",
            Self::Tanh => "Tanh",
            Self::Softmax => "Softmax",
            Self::Gelu => "Gelu",
            Self::Silu => "Silu",
            Self::Mish => "Mish",
        }
    }
}

/// Reduction operation kind.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReduceOp {
    Sum,
    Mean,
    Max,
    Min,
}

impl fmt::Display for ReduceOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.op_name())
    }
}

impl ReduceOp {
    /// Returns the canonical operator name (e.g. "ReduceSum", "ReduceMean").
    pub fn op_name(self) -> &'static str {
        match self {
            Self::Sum => "ReduceSum",
            Self::Mean => "ReduceMean",
            Self::Max => "ReduceMax",
            Self::Min => "ReduceMin",
        }
    }
}

/// Pooling operation kind.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PoolKind {
    Max,
    Avg,
}

impl fmt::Display for PoolKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.op_name())
    }
}

impl PoolKind {
    /// Returns the canonical operator name (e.g. "MaxPool", "AveragePool").
    pub fn op_name(self) -> &'static str {
        match self {
            Self::Max => "MaxPool",
            Self::Avg => "AveragePool",
        }
    }
}

/// Normalization type.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NormType {
    Batch,
    Layer,
}

impl fmt::Display for NormType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Self::Batch => "Batch",
            Self::Layer => "Layer",
        })
    }
}

/// Conv2D shape parameters extracted from uniform params.
#[derive(Debug, Clone)]
pub struct Conv2DShape {
    /// Batch dimension name.
    pub batch: String,
    /// Input channel dimension name.
    pub channels_in: String,
    /// Output channel dimension name.
    pub channels_out: String,
    /// Spatial height dimension name.
    pub height: String,
    /// Spatial width dimension name.
    pub width: String,
    /// Kernel height dimension name.
    pub kernel_h: String,
    /// Kernel width dimension name.
    pub kernel_w: String,
    /// Kernel height as a concrete value.
    pub kernel_h_val: i64,
    /// Kernel width as a concrete value.
    pub kernel_w_val: i64,
    /// Vertical stride.
    pub stride_h: i64,
    /// Horizontal stride.
    pub stride_w: i64,
    /// Vertical padding.
    pub pad_h: i64,
    /// Horizontal padding.
    pub pad_w: i64,
    /// Number of groups (1 = standard conv, channels_in = depthwise).
    pub groups: i64,
    /// Vertical dilation.
    pub dilation_h: i64,
    /// Horizontal dilation.
    pub dilation_w: i64,
}

impl fmt::Display for Conv2DShape {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Conv2D({}x{}x{} k{}x{} g{} d{}x{})",
            self.channels_in,
            self.height,
            self.width,
            self.kernel_h,
            self.kernel_w,
            self.groups,
            self.dilation_h,
            self.dilation_w
        )
    }
}

/// Pooling shape parameters.
#[derive(Debug, Clone)]
pub struct PoolShape {
    /// Kernel height.
    pub kernel_h: i64,
    /// Kernel width.
    pub kernel_w: i64,
    /// Vertical stride.
    pub stride_h: i64,
    /// Horizontal stride.
    pub stride_w: i64,
}

impl fmt::Display for PoolShape {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Pool(k{}x{} s{}x{})",
            self.kernel_h, self.kernel_w, self.stride_h, self.stride_w
        )
    }
}

/// A classified kernel pattern that maps to ONNX operators.
#[derive(Debug, Clone)]
pub enum KernelPattern {
    /// Loop + accumulation + 2 read arrays + 1 write array → ONNX `MatMul`.
    MatMul {
        inputs: [TensorBinding; 2],
        output: TensorBinding,
        shape: MatMulShape,
    },
    /// No loop + binary op on arrays → ONNX `Add`/`Sub`/`Mul`/`Div`.
    ElementWise {
        op: ElementWiseOp,
        inputs: [TensorBinding; 2],
        output: TensorBinding,
        dim_name: String,
    },
    /// A per-element expression one level deeper than a binary op: one tensor,
    /// an optional conversion, and a short chain of element-wise steps whose
    /// operands are whole tensors or dispatch-time scalars.
    ///
    /// Two shapes in the vendored corpus need this and nothing narrower:
    ///
    ///   - `axpy`, `output[i] = y[i] + a * x[i]` with `a` a uniform scalar —
    ///     `x`, scaled, then added to `y`.
    ///   - `dequantize`, `output[i] = f32(input[i]) * s1 * s2` with two
    ///     uniform scalars — a conversion, then two scales.
    ///
    /// Neither is one operation over two tensors, and [`Self::ElementWise`]
    /// holds exactly two operands and one op, so both were refused. Neither is
    /// a *fused activation* either: nothing here is a unary function of one
    /// element, so widening `Activation` would not have held them.
    ///
    /// The chain is deliberately linear. Every node has exactly one operand
    /// that is not a leaf, which is what makes it a sequence of graph nodes
    /// rather than a tree; `snake` and `alibi`, whose multiplies have two
    /// non-leaf operands, do not match and stay refused.
    ElementWiseChain {
        /// The tensor the chain starts from.
        base: TensorBinding,
        /// The element type `base` is converted to before the first step, when
        /// the kernel converts it. `dequantize` reads `array<i32>` and writes
        /// `array<f32>`, and dropping the conversion would reinterpret bits.
        cast: Option<i32>,
        /// Applied in order, each as `acc = acc <op> operand`.
        steps: Vec<ChainStep>,
        output: TensorBinding,
        dim_name: String,
    },
    /// 2D convolution: nested loops + kernel window + accumulation.
    Conv2D {
        input: TensorBinding,
        weight: TensorBinding,
        output: TensorBinding,
        /// Per-output-channel bias, when the kernel has one.
        ///
        /// Optional because a convolution need not have a bias, not because
        /// dropping it is acceptable: a backend that cannot emit one must
        /// refuse rather than silently compute a convolution without it.
        bias: Option<TensorBinding>,
        shape: Conv2DShape,
        /// The activation the kernel applies to what it stores.
        ///
        /// `output[i] = max(sum, 0.0)` is a convolution and a ReLU, and this
        /// was silently dropped: the emitted model held a CONV_2D and nothing
        /// else, so it loaded, was accelerated, and returned unclipped values.
        /// The TFLite backend folds it into `Conv2DOptions`; a backend that
        /// cannot express it must say so rather than emit the convolution
        /// alone.
        activation: Option<ActivationOp>,
    },
    /// Pooling: nested loops + reduction over spatial window.
    Pool {
        kind: PoolKind,
        input: TensorBinding,
        output: TensorBinding,
        shape: PoolShape,
    },
    /// Activation function: no loop, single input, unary math op.
    Activation {
        op: ActivationOp,
        input: TensorBinding,
        output: TensorBinding,
        dim_name: String,
    },
    /// Reduction over an axis: loop + accumulation, single input.
    Reduce {
        op: ReduceOp,
        input: TensorBinding,
        output: TensorBinding,
        axis: i64,
    },
    /// Transpose: permute tensor axes.
    Transpose {
        input: TensorBinding,
        output: TensorBinding,
        perm: Vec<i64>,
    },
    /// Reshape: change tensor shape without data copy.
    Reshape {
        input: TensorBinding,
        output: TensorBinding,
    },
    /// Normalization: mean + variance + scale + bias.
    Normalization {
        input: TensorBinding,
        scale: TensorBinding,
        bias: TensorBinding,
        output: TensorBinding,
        epsilon: f32,
        norm_type: NormType,
    },
    /// Concatenation of multiple inputs along an axis.
    Concat {
        inputs: Vec<TensorBinding>,
        output: TensorBinding,
        axis: i64,
    },
    /// Split a single input into multiple outputs along an axis.
    Split {
        input: TensorBinding,
        outputs: Vec<TensorBinding>,
        axis: i64,
    },
    /// Scaled dot-product attention.
    Attention {
        query: TensorBinding,
        key: TensorBinding,
        value: TensorBinding,
        output: TensorBinding,
        d_k: String,
        seq_len: String,
        /// Number of attention heads (1 = single-head).
        num_heads: u32,
        /// Number of K/V heads (for GQA; equals num_heads for MHA).
        num_kv_heads: u32,
        /// Whether a causal mask is applied.
        causal: bool,
    },
    /// Gather: index into data tensor using indices.
    Gather {
        data: TensorBinding,
        indices: TensorBinding,
        output: TensorBinding,
        axis: i64,
    },
    /// Scatter: write updates into output tensor at given indices.
    Scatter {
        data: TensorBinding,
        indices: TensorBinding,
        updates: TensorBinding,
        output: TensorBinding,
        axis: i64,
    },
    /// A matrix multiplication whose right operand arrives as integer codes
    /// with one scale per output channel: `output = input @ (weight * scale)ᵀ`.
    ///
    /// This is not [`Self::MatMul`] with an odd element type. A dense matmul
    /// has two operands; this has three, and the third is not a second matrix.
    /// `scale` carries one factor per *row of the weight*, so it multiplies the
    /// contraction's result rather than participating in the contraction.
    /// Lowering it as a MatMul over `weight` and dropping `scale` would produce
    /// a graph whose every output is wrong by a per-channel factor, which is
    /// the failure a refusal is preferable to.
    ///
    /// The weight's rows are the *output* channels, so the contraction runs
    /// along the second axis of both operands and the lowerings transpose it.
    /// That is the layout every quantized kernel in the vendored corpus uses,
    /// and the one per-channel quantization is defined against: a scale per row
    /// is a scale per output feature only if a row is an output feature.
    ///
    /// `weight` is reported with an element type of `INT8` although the buffer
    /// the kernel binds is `array<u32>`. Four two's-complement codes are packed
    /// per word, least-significant byte first, which is the byte layout of a
    /// contiguous `i8` row. The packing is how the GPU kernel buys four columns
    /// per memory transaction; it is not part of what the kernel computes, and
    /// the graph names the codes rather than the words they arrived in.
    QuantizedMatMul {
        /// The f32 activations, `[m, k]`.
        input: TensorBinding,
        /// The integer codes, `[n, k]` — one row per output channel.
        weight: TensorBinding,
        /// f32, `[n]` — one scale per weight row, applied after the contraction.
        scale: TensorBinding,
        /// A per-output-channel addend the kernel fuses onto the result.
        ///
        /// `matvec/q8_residual` adds a residual after the row scale. Optional
        /// because a quantized matmul need not have one, not because dropping
        /// it is acceptable — the same rule [`Self::Conv2D`]'s bias carries.
        bias: Option<TensorBinding>,
        output: TensorBinding,
        shape: MatMulShape,
    },
    /// Unrecognized pattern — classification could not determine a known op.
    Unknown { reason: String },
}

impl fmt::Display for KernelPattern {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::MatMul { shape, .. } => write!(f, "{shape}"),
            Self::QuantizedMatMul { shape, bias, .. } => {
                let fused = if bias.is_some() { " + bias" } else { "" };
                write!(f, "Quantized{shape} (int8, per-channel scale){fused}")
            }
            Self::ElementWise { op, .. } => write!(f, "{op}"),
            Self::ElementWiseChain { cast, steps, .. } => f.write_str(&chain_summary(*cast, steps)),
            Self::Conv2D { shape, .. } => write!(f, "{shape}"),
            Self::Pool { kind, shape, .. } => {
                write!(
                    f,
                    "{}(k{}x{} s{}x{})",
                    kind, shape.kernel_h, shape.kernel_w, shape.stride_h, shape.stride_w
                )
            }
            Self::Activation { op, .. } => write!(f, "{op}"),
            Self::Reduce { op, axis, .. } => write!(f, "{op}(axis={axis})"),
            Self::Transpose { perm, .. } => write!(f, "Transpose({perm:?})"),
            Self::Reshape { .. } => f.write_str("Reshape"),
            Self::Normalization {
                epsilon, norm_type, ..
            } => write!(f, "{norm_type}Normalization(eps={epsilon})"),
            Self::Concat { axis, .. } => write!(f, "Concat(axis={axis})"),
            Self::Split { axis, .. } => write!(f, "Split(axis={axis})"),
            Self::Attention {
                d_k,
                seq_len,
                num_heads,
                causal,
                ..
            } => {
                let mask = if *causal { ", causal" } else { "" };
                write!(
                    f,
                    "Attention(d_k={d_k}, seq_len={seq_len}, heads={num_heads}{mask})"
                )
            }
            Self::Gather { axis, .. } => write!(f, "Gather(axis={axis})"),
            Self::Scatter { axis, .. } => write!(f, "Scatter(axis={axis})"),
            Self::Unknown { reason } => write!(f, "Unknown({reason})"),
        }
    }
}

/// The operators a chain lowers to, in order: `["Mul", "Add"]`,
/// `["Cast", "Mul", "Mul"]`.
///
/// One name per node the ONNX and TFLite backends emit, because that is what a
/// vendor support matrix has to be asked about: a chain is several ordinary
/// operators, not one fused operator that no matrix lists. Asking about the
/// joined name reports `Mul+Add` as unsupported on hardware that supports both
/// a multiply and an add.
pub fn chain_op_names(cast: Option<i32>, steps: &[ChainStep]) -> Vec<&'static str> {
    let mut parts: Vec<&'static str> = Vec::with_capacity(steps.len() + 1);
    if cast.is_some() {
        parts.push("Cast");
    }
    parts.extend(steps.iter().map(|s| s.op.op_name()));
    parts
}

/// Name a chain by the operators it lowers to, in order: `Mul+Add`,
/// `Cast+Mul+Mul`.
///
/// Not a fixed label like "Chain". A count of steps says nothing about what a
/// kernel computes, and neither does a category name; the operator sequence is
/// exactly what the backend emits, so it is what the diagnostic reports.
pub fn chain_summary(cast: Option<i32>, steps: &[ChainStep]) -> String {
    chain_op_names(cast, steps).join("+")
}

/// Embedded constant weight data extracted from GlobalVariable initializers.
#[derive(Debug, Clone)]
pub struct EmbeddedWeight {
    /// Tensor name (from the global variable name).
    pub name: String,
    /// Tensor dimensions (e.g. `[4]` for `array<f32, 4>`).
    pub dims: Vec<i64>,
    /// Flattened f32 data.
    pub data: Vec<f32>,
}

/// Extract constant weight data from module globals with initializers.
///
/// Scans all global variables with `init: Some(...)` and evaluates
/// Compose/Literal/ZeroValue expressions into flat f32 arrays.
pub fn extract_embedded_weights(module: &Module) -> Vec<EmbeddedWeight> {
    let mut weights = Vec::new();
    for (handle, gv) in module.global_variables.iter() {
        let Some(init_handle) = gv.init else {
            continue;
        };
        let name = gv
            .name
            .clone()
            .unwrap_or_else(|| format!("weight_{}", handle.index()));
        let dims = type_dims(gv.ty, &module.types);
        let Some(data) =
            eval_const_expr_f32(init_handle, &module.global_expressions, &module.types)
        else {
            continue;
        };
        weights.push(EmbeddedWeight { name, dims, data });
    }
    weights
}

/// Recursively evaluate a constant expression to flat f32 values.
fn eval_const_expr_f32(
    handle: Handle<Expression>,
    exprs: &Arena<Expression>,
    types: &UniqueArena<Type>,
) -> Option<Vec<f32>> {
    match exprs.try_get(handle)? {
        Expression::Literal(Literal::F32(v)) => Some(vec![*v]),
        Expression::Literal(Literal::I32(v)) => Some(vec![*v as f32]),
        Expression::Literal(Literal::U32(v)) => Some(vec![*v as f32]),
        Expression::Compose { components, .. } => {
            let mut result = Vec::new();
            for &comp in components {
                result.extend(eval_const_expr_f32(comp, exprs, types)?);
            }
            Some(result)
        }
        Expression::ZeroValue(ty) => {
            let dims = type_dims(*ty, types);
            let count = dims.iter().product::<i64>().max(1) as usize;
            Some(vec![0.0; count])
        }
        _ => None,
    }
}

/// Compute flattened dimensions from a Type (Array, Scalar, Vector).
fn type_dims(ty_handle: Handle<Type>, types: &UniqueArena<Type>) -> Vec<i64> {
    match &types[ty_handle].inner {
        TypeInner::Scalar(_) => vec![],
        TypeInner::Vector { size, .. } => vec![*size as i64],
        TypeInner::Array {
            size: ArraySize::Constant(n),
            ..
        } => vec![*n as i64],
        _ => vec![],
    }
}

/// Whether a params-struct member is an integer scalar, and so plausibly a
/// tensor dimension rather than a tuning constant like `eps`.
fn is_integer_member(module: &Module, m: &nxpu_ir::StructMember) -> bool {
    matches!(
        module.types[m.ty].inner,
        TypeInner::Scalar(Scalar {
            kind: ScalarKind::Uint | ScalarKind::Sint,
            ..
        })
    )
}

/// Detect number of attention heads from shape_names or literal divisions.
fn detect_num_heads(shape_names: &[String], exprs: &Arena<Expression>) -> u32 {
    // Look for param named "num_heads", "n_head", "H", "n_heads", "nhead"
    for name in shape_names.iter() {
        let lower = name.to_lowercase();
        if lower == "num_heads" || lower == "n_head" || lower == "n_heads" || lower == "nhead" {
            // Can't get runtime value, but if there's a matching literal, use it
            // Otherwise return 1 as we know it's multi-head but can't determine count
            return find_division_literal(exprs).unwrap_or(1);
        }
    }
    // Also check for "H" (common in transformer code)
    if shape_names.iter().any(|n| n == "H") {
        return find_division_literal(exprs).unwrap_or(1);
    }
    1
}

/// Look for a division by a literal (d_model / num_heads -> head_dim).
#[allow(clippy::collapsible_if)] // nested if-let for MSRV 1.87 compat (no let chains)
fn find_division_literal(exprs: &Arena<Expression>) -> Option<u32> {
    for (_, expr) in exprs.iter() {
        if let Expression::Binary {
            op: BinaryOp::Divide,
            right,
            ..
        } = expr
        {
            if let Some(Expression::Literal(Literal::U32(n))) = exprs.try_get(*right) {
                if *n > 1 {
                    return Some(*n);
                }
            }
        }
    }
    None
}

/// Detect causal attention mask: look for an If guarding assignment of a large
/// negative literal inside a loop. The pattern is:
///   `if (j > i) { score = -1e30; }` or similar.
///
/// We specifically look for an `If` whose accept or reject block contains (or
/// leads to) a Store of a large negative float literal. This avoids false
/// positives from for-loop break conditions and numerically-stable softmax
/// `max_score` initializations that also use large negative values but are
/// not inside an If-guarded store.
fn detect_causal_mask(body: &[Statement], exprs: &Arena<Expression>) -> bool {
    has_causal_if_in_loop(body, exprs)
}

/// Recursively search loops for an `If` whose condition is a Greater/Less
/// comparison and whose accept or reject branch stores a large negative
/// float literal (like -1e30, -1e9, or -inf).
#[allow(clippy::collapsible_if)] // nested if-let for MSRV 1.87 compat (no let chains)
fn has_causal_if_in_loop(body: &[Statement], exprs: &Arena<Expression>) -> bool {
    for stmt in body {
        if let Statement::Loop {
            body, continuing, ..
        } = stmt
        {
            // Search loop body; skip the first non-Emit If that acts as a
            // for-loop break guard.
            let mut saw_non_emit = false;
            for s in body.iter() {
                match s {
                    Statement::Emit(_) => {}
                    Statement::If {
                        condition,
                        accept,
                        reject,
                    } => {
                        // The first non-emit If in a naga for-loop is the break
                        // guard (e.g. `if (j < N) {} else { break; }`). Skip it.
                        if !saw_non_emit
                            && (matches!(accept.as_slice(), [Statement::Break])
                                || matches!(reject.as_slice(), [Statement::Break]))
                        {
                            saw_non_emit = true;
                            continue;
                        }
                        saw_non_emit = true;
                        if is_causal_guard(*condition, accept, reject, exprs) {
                            return true;
                        }
                    }
                    _ => {
                        saw_non_emit = true;
                    }
                }
            }
            // Also search the continuing block
            for s in continuing.iter() {
                if let Statement::If {
                    condition,
                    accept,
                    reject,
                } = s
                {
                    if is_causal_guard(*condition, accept, reject, exprs) {
                        return true;
                    }
                }
            }
            // Recurse into nested loops
            if has_causal_if_in_loop(body, exprs) || has_causal_if_in_loop(continuing, exprs) {
                return true;
            }
        }
    }
    false
}

/// Check if a conditional is a comparison (Greater/Less/GE/LE) and one of
/// its branches (accept or reject) contains a Store of a large negative float.
fn is_causal_guard(
    condition: Handle<Expression>,
    accept: &[Statement],
    reject: &[Statement],
    exprs: &Arena<Expression>,
) -> bool {
    let is_comparison = match exprs.try_get(condition) {
        Some(Expression::Binary { op, .. }) => matches!(
            op,
            BinaryOp::Greater | BinaryOp::Less | BinaryOp::GreaterEqual | BinaryOp::LessEqual
        ),
        _ => false,
    };
    if !is_comparison {
        return false;
    }
    block_stores_large_negative(accept, exprs) || block_stores_large_negative(reject, exprs)
}

/// Check if a block contains a Store whose value is a large negative float literal.
fn block_stores_large_negative(body: &[Statement], exprs: &Arena<Expression>) -> bool {
    for stmt in body {
        match stmt {
            Statement::Store { value, .. } if expr_is_large_negative(*value, exprs) => {
                return true;
            }
            Statement::If { accept, reject, .. }
                if block_stores_large_negative(accept, exprs)
                    || block_stores_large_negative(reject, exprs) =>
            {
                return true;
            }
            _ => {}
        }
    }
    false
}

/// Check if an expression is (or contains) a large negative float literal (< -1e6).
fn expr_is_large_negative(handle: Handle<Expression>, exprs: &Arena<Expression>) -> bool {
    match exprs.try_get(handle) {
        Some(Expression::Literal(Literal::F32(v))) => *v < -1e6,
        Some(Expression::Unary {
            op: nxpu_ir::UnaryOp::Negate,
            expr,
        }) => {
            // -literal
            matches!(exprs.try_get(*expr), Some(Expression::Literal(Literal::F32(v))) if *v > 1e6)
        }
        _ => false,
    }
}

/// The operator names a pattern will emit, for asking a support matrix about.
///
/// One list rather than one name: an `ElementWiseChain` becomes a Cast and a
/// node per step, and asking about "Mul+Add" reported it unsupported on
/// hardware that has both.
///
/// This lived in eight backends, character for character, and the arms that
/// no backend's tests reached were uncovered eight times over. It belongs
/// next to the enum it matches on.
pub fn pattern_op_names(pattern: &KernelPattern) -> Vec<String> {
    match pattern {
        // Named by the operators it emits, not by a category: the ONNX and
        // TFLite graphs this delegates to contain a Cast and one node per
        // step, and "ElementWiseChain" would describe none of them.
        KernelPattern::ElementWiseChain { cast, steps, .. } => chain_op_names(*cast, steps)
            .into_iter()
            .map(String::from)
            .collect(),
        KernelPattern::MatMul { .. } => vec!["MatMul".into()],
        // Named by the nodes the backends emit, the same rule the chain
        // follows. Neither output format has a fused quantized-matmul
        // operator and neither backend emits one: the weight is dequantized,
        // transposed into contraction order, and multiplied. A vendor matrix
        // that does not list `DequantizeLinear` is telling the truth about
        // that hardware, and answering only "MatMul" here would hide two of
        // the nodes the graph contains.
        KernelPattern::QuantizedMatMul { bias, .. } => {
            let mut names = vec![
                "Transpose".to_string(),
                "DequantizeLinear".to_string(),
                "MatMul".to_string(),
            ];
            if bias.is_some() {
                names.push("Add".to_string());
            }
            names
        }
        KernelPattern::ElementWise { op, .. } => vec![op.op_name().into()],
        KernelPattern::Conv2D { .. } => vec!["Conv".into()],
        KernelPattern::Pool { kind, .. } => vec![kind.op_name().into()],
        KernelPattern::Activation { op, .. } => vec![op.op_name().into()],
        KernelPattern::Reduce { op, .. } => vec![op.op_name().into()],
        KernelPattern::Transpose { .. } => vec!["Transpose".into()],
        KernelPattern::Reshape { .. } => vec!["Reshape".into()],
        KernelPattern::Normalization { .. } => vec!["BatchNormalization".into()],
        KernelPattern::Concat { .. } => vec!["Concat".into()],
        KernelPattern::Split { .. } => vec!["Split".into()],
        KernelPattern::Attention { .. } => vec!["Attention".into()],
        KernelPattern::Gather { .. } => vec!["Gather".into()],
        KernelPattern::Scatter { .. } => vec!["ScatterND".into()],
        KernelPattern::Unknown { .. } => vec!["Unknown".into()],
    }
}

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

    fn binding(name: &str) -> TensorBinding {
        let mut arena: Arena<GlobalVariable> = Arena::new();
        let ty = {
            let mut types = UniqueArena::default();
            types.insert(Type {
                name: None,
                inner: TypeInner::Scalar(Scalar::F32),
            })
        };
        let handle = arena.append(GlobalVariable {
            name: Some(name.into()),
            space: AddressSpace::Uniform,
            binding: None,
            ty,
            init: None,
            layout: None,
        });
        TensorBinding {
            handle,
            name: name.into(),
            elem_type: data_type::FLOAT,
            role: TensorRole::Input,
        }
    }

    /// Every arm, because the eight copies of this that used to exist were
    /// each covered only for the two or three patterns that backend's tests
    /// happened to compile.
    #[test]
    fn every_pattern_names_something() {
        let cases: Vec<KernelPattern> = vec![
            KernelPattern::Transpose {
                input: binding("a"),
                output: binding("b"),
                perm: vec![1, 0],
            },
            KernelPattern::Reshape {
                input: binding("a"),
                output: binding("b"),
            },
            KernelPattern::Reduce {
                op: ReduceOp::Sum,
                input: binding("a"),
                output: binding("b"),
                axis: 0,
            },
            KernelPattern::Unknown {
                reason: "nothing".into(),
            },
        ];
        for pattern in cases {
            let names = pattern_op_names(&pattern);
            assert!(
                !names.is_empty() && names.iter().all(|n| !n.is_empty()),
                "{pattern:?} named nothing"
            );
        }
    }
}

/// Classify an entry point into a known ONNX-mappable pattern.
pub fn classify_entry_point(
    module: &Module,
    ep_index: usize,
) -> Result<KernelPattern, AnalysisError> {
    if module.entry_points.is_empty() {
        return Err(AnalysisError::NoEntryPoints);
    }
    let ep = module
        .entry_points
        .get(ep_index)
        .ok_or(AnalysisError::EntryPointOutOfRange(ep_index))?;

    // 1. Classify globals by address space.
    let mut inputs: Vec<(Handle<GlobalVariable>, &GlobalVariable)> = Vec::new();
    let mut outputs: Vec<(Handle<GlobalVariable>, &GlobalVariable)> = Vec::new();
    let mut init_globals: Vec<(Handle<GlobalVariable>, &GlobalVariable)> = Vec::new();
    let mut params_members: Option<&[nxpu_ir::StructMember]> = None;
    let mut params_global: Option<Handle<GlobalVariable>> = None;

    for (handle, gv) in module.global_variables.iter() {
        match &gv.space {
            AddressSpace::Storage { access } => {
                if access.contains(StorageAccess::STORE) {
                    outputs.push((handle, gv));
                } else {
                    inputs.push((handle, gv));
                }
            }
            AddressSpace::Uniform => {
                if let TypeInner::Struct { members, .. } = &module.types[gv.ty].inner {
                    params_members = Some(members);
                    params_global = Some(handle);
                }
            }
            _ => {
                if gv.init.is_some() {
                    init_globals.push((handle, gv));
                }
            }
        }
    }

    // Sort inputs by resource binding order (binding 0 = A, binding 1 = B).
    inputs.sort_by_key(|(_, gv)| gv.binding.map(|b| b.binding).unwrap_or(u32::MAX));

    if outputs.is_empty() {
        return Err(AnalysisError::UnsupportedPattern(
            "expected at least 1 output storage buffer".into(),
        ));
    }

    // Shape parameters only. A params struct routinely carries scalars that
    // are not dimensions — `eps` is the common one — and counting them made
    // LayerNorm, whose struct is (N, D, eps), look like it had three shape
    // names and report itself as BatchNorm.
    let shape_names: Vec<String> = params_members
        .map(|members| {
            members
                .iter()
                .filter(|m| is_integer_member(module, m))
                .filter_map(|m| m.name.clone())
                .collect()
        })
        .unwrap_or_default();

    // Every member's name, by position, so an index read out of the IR names
    // the right one. `shape_names` above has the non-integer members filtered
    // out and so cannot be indexed by member position.
    let param_names: Vec<String> = params_members
        .map(|members| {
            members
                .iter()
                .map(|m| m.name.clone().unwrap_or_default())
                .collect()
        })
        .unwrap_or_default();

    let has_loop = has_loop(&ep.function.body);
    let num_inputs = inputs.len();

    // Single input patterns
    if num_inputs == 1 {
        let input = make_binding(module, inputs[0].0, inputs[0].1, TensorRole::Input);

        // 1 input + 2 or more outputs: a Split, or nothing this compiler has.
        //
        // "One input, several outputs, and an If somewhere" was the whole
        // test, and it asks about the kernel's shape rather than about what it
        // computes. Three vendor kernels answered yes and none of them is a
        // split: an activation quantizer emitting codes and per-row scales, a
        // greedy CTC decode emitting labels and their lengths, and a
        // mixture-of-experts router emitting expert indices and gate weights.
        //
        // What a Split is, and what every backend here emits for one, is a
        // slice: each output element *is* an input element, moved. So that is
        // what has to be found — a boundary test, and stores that copy rather
        // than compute. The rest is refused by name below.
        if outputs.len() >= 2 {
            let output_handles: Vec<Handle<GlobalVariable>> =
                outputs.iter().map(|(h, _)| *h).collect();
            let (written, computed) = survey_output_stores(
                &ep.function.body,
                &ep.function.expressions,
                inputs[0].0,
                &output_handles,
            );
            let every_output_is_a_copy =
                computed.is_empty() && written.len() == output_handles.len();

            if every_output_is_a_copy && has_if_statement(&ep.function.body) {
                let out_bindings: Vec<TensorBinding> = outputs
                    .iter()
                    .map(|(h, gv)| make_binding(module, *h, gv, TensorRole::Output))
                    .collect();
                let axis =
                    infer_split_axis(&ep.function.body, &ep.function.expressions, &shape_names);
                return Ok(KernelPattern::Split {
                    input,
                    outputs: out_bindings,
                    axis,
                });
            }

            return Ok(KernelPattern::Unknown {
                reason: multi_output_refusal(module, &outputs, inputs[0].1, &written, &computed),
            });
        }

        let output = make_binding(module, outputs[0].0, outputs[0].1, TensorRole::Output);

        if !has_loop {
            // Try activation (Math expression)
            if let Some(act_op) = find_store_activation(&ep.function.body, &ep.function.expressions)
            {
                let dim_name = shape_names.first().cloned().unwrap_or_else(|| "N".into());
                return Ok(KernelPattern::Activation {
                    op: act_op,
                    input,
                    output,
                    dim_name,
                });
            }

            // ElementWise with embedded weight: 1 storage input + binary op + private init global.
            let ew_op = match find_store_value_op(&ep.function.body, &ep.function.expressions) {
                Some(StoreValueOp::Binary(op)) => Some(op),
                // A multiply-add is not an element-wise binary op; fall through.
                Some(StoreValueOp::MultiplyAdd) | None => None,
            };
            if let (Some(ew_op), Some(&(wh, wgv))) = (ew_op, init_globals.first()) {
                let weight = make_binding(module, wh, wgv, TensorRole::Input);
                let dim_name = shape_names.first().cloned().unwrap_or_else(|| "N".into());
                return Ok(KernelPattern::ElementWise {
                    op: ew_op,
                    inputs: [input, weight],
                    output,
                    dim_name,
                });
            }

            // A reindexing kernel: it moves values without computing on
            // them, so there is no activation to find and it used to fall
            // through to Unknown. Every backend already lowers Transpose.
            if let Some(perm) = detect_permutation(
                &ep.function.body,
                &ep.function.expressions,
                inputs[0].0,
                outputs[0].0,
            ) {
                return Ok(KernelPattern::Transpose {
                    input,
                    output: make_binding(module, outputs[0].0, outputs[0].1, TensorRole::Output),
                    perm,
                });
            }

            // A stored value that is more than one operation deep.
            // `dequantize` writes `f32(input[i]) * s1 * s2`: one tensor, a
            // conversion and two dispatch-time scalars, which is not an
            // activation and has no second tensor to make it element-wise.
            // `axpy`'s in-place entry point arrives here too, with `y` bound
            // `read_write` so that it counts as an output and leaves `x` the
            // only input.
            if let Some(pattern) = match_elementwise_chain(
                module,
                &ep.function.body,
                &ep.function.expressions,
                &outputs,
                &shape_names,
            ) {
                return Ok(pattern);
            }

            // No recognized activation — unknown pattern.
            return Ok(KernelPattern::Unknown {
                reason: "single input, no loop, no recognized activation function".into(),
            });
        }

        // A softmax stores `exp(x - max) / Σ exp(x - max)` and builds that sum
        // in a loop of its own. Recognised as an Activation, which both the
        // ONNX and TFLite backends already lower.
        //
        // `dim_name` is the innermost shape parameter, not the first: the
        // reduction runs along the last axis, and naming it with the batch
        // count would label the softmax's length wrongly.
        if outputs.len() == 1
            && detect_softmax(&ep.function.body, &ep.function.expressions, outputs[0].0)
        {
            let dim_name = shape_names.last().cloned().unwrap_or_else(|| "N".into());
            return Ok(KernelPattern::Activation {
                op: ActivationOp::Softmax,
                input,
                output,
                dim_name,
            });
        }

        // Has loop + single input → Pool or Reduce
        //
        // "Four or more parameters" is not evidence of pooling. A pool has a
        // spatial window, so the loop bounds have to be there; without them
        // the kernel size was invented as 2x2 and the stride copied from it,
        // and a reduction over a 4-parameter tensor came back as AveragePool.
        let pool_bounds = extract_loop_bound_literals(&ep.function.body, &ep.function.expressions);
        if shape_names.len() >= 4 && !pool_bounds.is_empty() {
            let pool_kind = if find_store_math_fun(
                &ep.function.body,
                &ep.function.expressions,
                MathFunction::Max,
            ) {
                PoolKind::Max
            } else {
                PoolKind::Avg
            };
            let bounds = pool_bounds;
            let strides = extract_multiply_literals(&ep.function.expressions);
            let kh_u = bounds[0];
            let kw_u = bounds.get(1).copied().unwrap_or(kh_u);
            let sh_u = strides.first().copied().unwrap_or(kh_u);
            let sw_u = strides.get(1).copied().unwrap_or(sh_u);
            let kh = kh_u as i64;
            let kw = kw_u as i64;
            let sh = sh_u as i64;
            let sw = sw_u as i64;
            let pool_shape = PoolShape {
                kernel_h: kh,
                kernel_w: kw,
                stride_h: sh,
                stride_w: sw,
            };
            return Ok(KernelPattern::Pool {
                kind: pool_kind,
                input,
                output,
                shape: pool_shape,
            });
        }

        // Single input + loop → Reduce
        let reduce_op = detect_reduce_op(&ep.function.body, &ep.function.expressions);
        let axis = if shape_names.len() >= 2 { 1 } else { 0 };
        return Ok(KernelPattern::Reduce {
            op: reduce_op,
            input,
            output,
            axis: axis as i64,
        });
    }

    // 2-input patterns require at least 2 inputs
    if num_inputs < 2 {
        return Err(AnalysisError::UnsupportedPattern(
            "expected at least 1 input storage buffer".into(),
        ));
    }

    // 3+ inputs: check for a quantized matmul, Attention, Normalization, Scatter
    if num_inputs >= 3 {
        // First, because it is the most specific thing here: a kernel that
        // unpacks integer codes out of a weight buffer is not any of the
        // patterns below, and the ones it is not go on to be tested normally
        // because this returns `None` for a kernel that unpacks nothing.
        //
        // Only in this arm. The corpus's quantized kernels bind a weight, a
        // scale and an activation at least, so none of them arrives with two
        // inputs, and `dequant_transpose` — which does unpack codes with two
        // inputs — is a movement of values that this pattern would misread.
        if let Some(pattern) = match_quantized_matmul(
            module,
            ep,
            &inputs,
            &outputs,
            params_global,
            &param_names,
            &shape_names,
        ) {
            return Ok(pattern);
        }

        // Attention heuristic: 3 inputs + has loop + contains Exp + Sqrt.
        // NOTE: This is fragile — any 3-input kernel with loop + exp() + sqrt()
        // will match. A false positive is possible for custom kernels that
        // happen to use both functions. Consider adding more structural checks
        // (e.g., nested loop depth, softmax pattern) if this becomes an issue.
        if has_loop
            && has_math_function_in_expressions(&ep.function.expressions, MathFunction::Exp)
            && has_math_function_in_expressions(&ep.function.expressions, MathFunction::Sqrt)
        {
            let query = make_binding(module, inputs[0].0, inputs[0].1, TensorRole::Input);
            let key = make_binding(module, inputs[1].0, inputs[1].1, TensorRole::Input);
            let value = make_binding(module, inputs[2].0, inputs[2].1, TensorRole::Input);
            let output = make_binding(module, outputs[0].0, outputs[0].1, TensorRole::Output);
            let seq_len = shape_names
                .first()
                .cloned()
                .unwrap_or_else(|| "seq_len".into());
            let d_k = shape_names.get(1).cloned().unwrap_or_else(|| "d_k".into());
            let num_heads = detect_num_heads(&shape_names, &ep.function.expressions);
            let causal = detect_causal_mask(&ep.function.body, &ep.function.expressions);
            return Ok(KernelPattern::Attention {
                query,
                key,
                value,
                output,
                d_k,
                seq_len,
                num_heads,
                num_kv_heads: num_heads, // Default: MHA (same as num_heads)
                causal,
            });
        }

        // Normalization: 3 inputs (input, scale, bias) + 1 output + loop, a
        // reciprocal square root in either spelling, and no Exp.
        // Distinguishes LayerNorm (2 shape params: N, C) from BatchNorm (3+: N, C, HW).
        //
        // `inverseSqrt` is the spelling a kernel written for speed uses, and
        // requiring `sqrt` rejected real LayerNorm and GroupNorm kernels for a
        // difference that is not one.
        if num_inputs == 3
            && outputs.len() == 1
            && has_loop
            && (has_math_function_in_expressions(&ep.function.expressions, MathFunction::Sqrt)
                || has_math_function_in_expressions(
                    &ep.function.expressions,
                    MathFunction::InverseSqrt,
                ))
            && !has_math_function_in_expressions(&ep.function.expressions, MathFunction::Exp)
            && !has_nested_loop(&ep.function.body)
        {
            let input = make_binding(module, inputs[0].0, inputs[0].1, TensorRole::Input);
            let scale = make_binding(module, inputs[1].0, inputs[1].1, TensorRole::Input);
            let bias = make_binding(module, inputs[2].0, inputs[2].1, TensorRole::Input);
            let output = make_binding(module, outputs[0].0, outputs[0].1, TensorRole::Output);
            // A group count in the params means group normalization, and the
            // formats this compiler targets cannot express it: ONNX's
            // GroupNormalization takes num_groups as a static attribute, and
            // here it is a uniform the host sets at dispatch time. Saying so
            // is the only honest answer — reporting it as BatchNormalization,
            // which is what the parameter count produced, is a graph that runs
            // and computes something else.
            if shape_names
                .iter()
                .any(|n| n == "G" || n.eq_ignore_ascii_case("groups"))
            {
                return Ok(KernelPattern::Unknown {
                    reason: "group normalization with a group count supplied at \
                             runtime — the target formats need it as a static \
                             attribute"
                        .into(),
                });
            }
            let norm_type = if shape_names.len() <= 2 {
                NormType::Layer
            } else {
                NormType::Batch
            };
            return Ok(KernelPattern::Normalization {
                input,
                scale,
                bias,
                output,
                epsilon: 1e-5,
                norm_type,
            });
        }

        // Conv with a bias. Convolution detection lived in the two-input arm,
        // so every conv that adds a per-channel bias — which is most of them —
        // arrived here with three inputs and was refused for want of a
        // recogniser rather than for want of a lowering.
        //
        // The evidence is the kernel window itself, named in the params. A
        // parameter count is not evidence and was tried: it swallowed
        // attention scores, GQA, MoE dispatch and an inverse STFT, all
        // reported as CONV_2D — the same guessing this recogniser replaces,
        // moved to a different arm.
        //
        // KH and KW without a KD is exactly two spatial dimensions. A 1-D
        // conv names one extent, a 3-D conv names three, and a transposed
        // conv is a different operator; none of them is a Conv2D and each is
        // refused below rather than rounded to one.
        let names_kernel = |n: &str| {
            shape_names.iter().any(|s| {
                s.eq_ignore_ascii_case(n) || s.eq_ignore_ascii_case(&format!("kernel_{n}"))
            })
        };
        if num_inputs == 3
            && outputs.len() == 1
            && has_loop
            && names_kernel("kh")
            && names_kernel("kw")
            && !names_kernel("kd")
        {
            let input = make_binding(module, inputs[0].0, inputs[0].1, TensorRole::Input);
            let weight = make_binding(module, inputs[1].0, inputs[1].1, TensorRole::Input);
            let bias = make_binding(module, inputs[2].0, inputs[2].1, TensorRole::Input);
            let output = make_binding(module, outputs[0].0, outputs[0].1, TensorRole::Output);
            let conv_shape = extract_conv2d_shape(
                &shape_names,
                Some(&ep.function.body),
                Some(&ep.function.expressions),
            );
            return Ok(KernelPattern::Conv2D {
                input,
                weight,
                output,
                bias: Some(bias),
                shape: conv_shape,
                activation: store_value_activation(&ep.function.body, &ep.function.expressions),
            });
        }

        // The convolutions this compiler has no operator for. Each is a real
        // shape, recognised, and declined by name — which is worth more than
        // a CONV_2D that computes something else.
        if num_inputs == 3 && outputs.len() == 1 && has_loop {
            let names = |n: &str| shape_names.iter().any(|s| s.eq_ignore_ascii_case(n));
            if names("kd") {
                return Ok(KernelPattern::Unknown {
                    reason: "3-D convolution — the kernel window has three \
                             spatial extents and Conv2D has two"
                        .into(),
                });
            }
            if names("lout") && names("k") {
                return Ok(KernelPattern::Unknown {
                    reason: "1-D or transposed convolution — one spatial \
                             extent, which Conv2D cannot represent"
                        .into(),
                });
            }
        }

        let input_handles: Vec<Handle<GlobalVariable>> = inputs.iter().map(|(h, _)| *h).collect();
        let output_handles_3: Vec<Handle<GlobalVariable>> =
            outputs.iter().map(|(h, _)| *h).collect();

        // Scatter: 3 inputs + 1 output + no loop (simple scatter write)
        //
        // A scatter takes its write address out of a buffer. Three inputs, one
        // output and no loop was the whole test, and it reported a snake
        // activation — `x + (1/(β+ε))·sin²(αx)` with per-channel α and β — as
        // SCATTER_ND, which writes elsewhere entirely.
        if num_inputs == 3
            && outputs.len() == 1
            && !has_loop
            && detect_indexed_write(
                &ep.function.body,
                &ep.function.expressions,
                &input_handles,
                &output_handles_3,
            )
        {
            let data = make_binding(module, inputs[0].0, inputs[0].1, TensorRole::Input);
            let indices = make_binding(module, inputs[1].0, inputs[1].1, TensorRole::Input);
            let updates = make_binding(module, inputs[2].0, inputs[2].1, TensorRole::Input);
            let output = make_binding(module, outputs[0].0, outputs[0].1, TensorRole::Output);
            return Ok(KernelPattern::Scatter {
                data,
                indices,
                updates,
                output,
                axis: 0,
            });
        }

        // Three inputs and no loop, writing where the thread id says: a
        // per-element function of three tensors. `ElementWiseChain` takes as
        // many operands as the expression has steps, so it is tried here as
        // well as in the one- and two-input arms — the rule is about the shape
        // of the expression, not about how many buffers are bound.
        if num_inputs == 3 && outputs.len() == 1 && !has_loop {
            if let Some(pattern) = match_elementwise_chain(
                module,
                &ep.function.body,
                &ep.function.expressions,
                &outputs,
                &shape_names,
            ) {
                return Ok(pattern);
            }
            // `snake`'s SnakeBeta is the shape that lands here:
            // `x + (1/(β_c + ε)) · sin²(α_c · x)`. Its multiply has a
            // reciprocal on one side and a squared sine on the other, so it is
            // a tree and not a chain, and both parameters are read at
            // `(i / L) % C` rather than at `i`, so neither is an operand of
            // the result's shape. An `Activation` carries one input and no
            // parameters, and `ElementWise` carries two.
            return Ok(KernelPattern::Unknown {
                reason: "3 inputs combined per element, and not as a chain: a \
                         chain step takes a whole tensor read at the index \
                         written, or a scalar the host sets per dispatch, and \
                         these operands are neither"
                    .into(),
            });
        }

        // 3+ inputs but no recognized pattern — unknown.
        return Ok(KernelPattern::Unknown {
            reason: "3+ inputs but no recognized pattern (expected Attention, Normalization, or Scatter)".into(),
        });
    }

    // 2 inputs
    let input_a = make_binding(module, inputs[0].0, inputs[0].1, TensorRole::Input);
    let input_b = make_binding(module, inputs[1].0, inputs[1].1, TensorRole::Input);
    let output_c = make_binding(module, outputs[0].0, outputs[0].1, TensorRole::Output);

    // Addresses that were loaded rather than computed.
    //
    // Every pattern below reads and writes at positions built from the thread
    // id. A kernel that takes an address out of a second buffer is a gather or
    // a scatter whatever its loop structure looks like, so the evidence is
    // weighed here, once, ahead of the loop/no-loop split — the alternative is
    // a scatter with a loop reported as a matmul, which is what happened.
    let input_handles = [inputs[0].0, inputs[1].0];
    let output_handles: Vec<Handle<GlobalVariable>> = outputs.iter().map(|(h, _)| *h).collect();

    // Writes that accumulate. An atomic output means several invocations
    // reach the same slot and their contributions are summed there; the
    // read-modify-write is the operator, not an implementation detail of it.
    if outputs
        .iter()
        .any(|(_, gv)| is_atomic_buffer(module, gv.ty))
    {
        return Ok(KernelPattern::Unknown {
            reason: "the output is an array of atomics, so writes to the same slot \
                     accumulate instead of overwriting — an accumulating scatter. \
                     Scatter has no reduction mode and takes the tensor it writes \
                     into as an input; here the destination is the output itself, \
                     zeroed by the caller"
                .into(),
        });
    }

    // Writes whose position comes out of a buffer.
    if detect_indexed_write(
        &ep.function.body,
        &ep.function.expressions,
        &input_handles,
        &output_handles,
    ) {
        let also = if outputs.len() > 1 {
            ", and it writes two output tensors where every pattern here writes one"
        } else {
            ""
        };
        return Ok(KernelPattern::Unknown {
            reason: format!(
                "the address written to is loaded from an input buffer — a \
                 data-dependent placement, where the input data decides where each \
                 element lands{also}. Scatter takes three inputs, a destination to \
                 write into as well as the indices and the updates, and there are two"
            ),
        });
    }

    // Reads whose position comes out of a buffer. Either input can be the one
    // holding the addresses; nothing says the index tensor is bound second.
    let indexed_read = [(0usize, 1usize), (1, 0)].into_iter().find_map(|(d, i)| {
        detect_indexed_read(
            &ep.function.body,
            &ep.function.expressions,
            inputs[d].0,
            inputs[i].0,
            outputs[0].0,
        )
        .map(|kind| (d, i, kind))
    });
    if let Some((d, i, kind)) = indexed_read {
        let data = make_binding(module, inputs[d].0, inputs[d].1, TensorRole::Input);
        let indices = make_binding(module, inputs[i].0, inputs[i].1, TensorRole::Input);
        return Ok(match kind {
            IndexedRead::Element => KernelPattern::Gather {
                data,
                indices,
                output: output_c,
                axis: 0,
            },
            // A row gather. Every backend lowers `Gather` over a flat buffer —
            // one element per index — so emitting one here would select single
            // elements out of a table whose rows are `D` wide, which is a
            // different tensor of a different size. The row width is the thing
            // missing, and nothing in `Gather` can carry it.
            IndexedRead::Block => KernelPattern::Unknown {
                reason: "a row gather: the index is scaled by a row width before it \
                         becomes an address, so each index names a block rather than \
                         an element. Gather indexes a flat buffer and carries no row \
                         width"
                    .into(),
            },
        });
    }

    if !has_loop {
        let has_structural_if = has_non_guard_if(&ep.function.body);

        // Gather: 2 inputs + no loop + one input is u32 array (indices)
        // Check if second input is a u32 array (integer index type).
        let second_is_int = input_b.elem_type == data_type::UINT32
            || input_b.elem_type == data_type::INT32
            || input_b.elem_type == data_type::INT64;
        if second_is_int
            && find_store_value_op(&ep.function.body, &ep.function.expressions).is_none()
            && !has_structural_if
        {
            return Ok(KernelPattern::Gather {
                data: input_a,
                indices: input_b,
                output: output_c,
                axis: 0,
            });
        }

        // Concat: 2 inputs + no loop + has If + no binary store op
        if has_structural_if
            && find_store_value_op(&ep.function.body, &ep.function.expressions).is_none()
        {
            let axis = infer_concat_axis(&ep.function.body, &ep.function.expressions, &shape_names);
            return Ok(KernelPattern::Concat {
                inputs: vec![input_a, input_b],
                output: output_c,
                axis,
            });
        }

        // An element-wise op writes element i from element i. rope writes a
        // rotated pair at `base` and `base + 1`; a dequantizing transpose
        // writes at `col * width + row`. Both were reported as `Mul`, which
        // keeps their arithmetic and discards their movement.
        if !stores_at_the_index_it_read(&ep.function.body, &ep.function.expressions) {
            return Ok(KernelPattern::Unknown {
                reason: "2 inputs and no loop, but the store index is not the \
                         index it read from — the values move as well as \
                         combine, and no operator here does both"
                    .into(),
            });
        }

        // ElementWise: store of binary operation.
        let op = match find_store_value_op(&ep.function.body, &ep.function.expressions) {
            Some(StoreValueOp::Binary(op)) => op,
            // `y + a * x` is not an add — the multiply has nowhere to go in a
            // two-operand `ElementWise`, and reporting `Add` would silently
            // drop it. `axpy` is a chain and is recognised as one; the ones
            // that are not say what was found instead, and say it identically
            // whether or not FMA fusion has already rewritten the expression.
            Some(StoreValueOp::MultiplyAdd) => {
                if let Some(pattern) = match_elementwise_chain(
                    module,
                    &ep.function.body,
                    &ep.function.expressions,
                    &outputs,
                    &shape_names,
                ) {
                    return Ok(pattern);
                }
                return Ok(KernelPattern::Unknown {
                    reason: "fused multiply-add (`c + a * b`, or `fma(a, b, c)` after \
                             optimization) that is not a chain: a chain step takes a whole \
                             tensor read at the index written, or a scalar the host sets per \
                             dispatch, and at least one of these operands is neither"
                        .into(),
                });
            }
            None => {
                return Err(AnalysisError::UnsupportedPattern(
                    "no recognizable binary operation found".into(),
                ));
            }
        };

        let dim_name = shape_names.first().cloned().unwrap_or_else(|| "N".into());

        return Ok(KernelPattern::ElementWise {
            op,
            inputs: [input_a, input_b],
            output: output_c,
            dim_name,
        });
    }

    // A matmul does not compute a reciprocal square root. RMSNorm is a
    // normalization with a scale and no bias, so it has two inputs rather than
    // three and never reached the Normalization branch — it fell through to
    // here and was reported as BATCH_MATMUL. Serving it properly needs the
    // Normalization pattern to make its bias optional, which is a change to
    // every backend that lowers one; until then, refusing is the honest half.
    if (has_math_function_in_expressions(&ep.function.expressions, MathFunction::InverseSqrt)
        || has_math_function_in_expressions(&ep.function.expressions, MathFunction::Sqrt))
        && !has_math_function_in_expressions(&ep.function.expressions, MathFunction::Exp)
    {
        return Ok(KernelPattern::Unknown {
            reason: "2 inputs, a loop and a reciprocal square root — a \
                     normalization without a bias, which Normalization cannot \
                     yet represent"
                .into(),
        });
    }

    // A convolution does not evaluate trigonometry. An STFT does — its
    // twiddle factors are sin and cos — and with five parameters it landed in
    // the arm below and was reported as CONV_2D.
    if has_math_function_in_expressions(&ep.function.expressions, MathFunction::Sin)
        || has_math_function_in_expressions(&ep.function.expressions, MathFunction::Cos)
    {
        return Ok(KernelPattern::Unknown {
            reason: "2 inputs, a loop and trigonometry — a transform rather \
                     than a convolution or a matmul"
                .into(),
        });
    }

    // 2 inputs + loop: a convolution without a bias.
    //
    // Same evidence as the three-input arm — the kernel window named in the
    // params. The parameter count that used to stand in for it reported
    // attention's context half, GQA's, and MoE dispatch as CONV_2D.
    let names_kernel_2d = shape_names
        .iter()
        .any(|s| s.eq_ignore_ascii_case("kh") || s.eq_ignore_ascii_case("kernel_h"))
        && shape_names
            .iter()
            .any(|s| s.eq_ignore_ascii_case("kw") || s.eq_ignore_ascii_case("kernel_w"));
    if names_kernel_2d {
        let conv_shape = extract_conv2d_shape(
            &shape_names,
            Some(&ep.function.body),
            Some(&ep.function.expressions),
        );
        return Ok(KernelPattern::Conv2D {
            input: input_a,
            weight: input_b,
            output: output_c,
            shape: conv_shape,
            bias: None,
            activation: store_value_activation(&ep.function.body, &ep.function.expressions),
        });
    }

    // MatMul: loop + accumulation pattern.
    //
    // The dimensions have to come from the kernel. Inventing M, N and K when
    // the params struct does not name three of them meant every two-input
    // looping kernel became a matmul over dimensions nobody had established —
    // a scatter, among others.
    if shape_names.len() < 3 {
        return Ok(KernelPattern::Unknown {
            reason: format!(
                "2 inputs and a loop, but {} shape parameters — too few to \
                 establish a matmul's M, N and K",
                shape_names.len()
            ),
        });
    }
    let shape = MatMulShape {
        m: shape_names[0].clone(),
        n: shape_names[1].clone(),
        k: shape_names[2].clone(),
    };

    Ok(KernelPattern::MatMul {
        inputs: [input_a, input_b],
        output: output_c,
        shape,
    })
}

// ---------------------------------------------------------------------------
// Quantized matrix multiplication
//
// The six quantized kernels in the vendored corpus all reached the classifier
// with three or more inputs and no recogniser, because there was no pattern
// that could carry a weight arriving as integer codes plus the scale that
// turns them back into numbers. What follows is the evidence that a kernel is
// one, keyed on what it computes rather than on how many buffers it binds:
//
//   1. It unpacks a field out of a word loaded from a storage buffer. That is
//      `extractBits`, and nothing else in the corpus does it inside a loop.
//   2. The unpacked code meets a value out of another buffer in a multiply —
//      either the activation it is contracted against, or the scale that
//      dequantizes it. Codes unpacked and never multiplied are not a matmul.
//   3. A buffer is read at *exactly* the index that selects the weight's row.
//      That is the definition of a per-channel scale, and it is also the test
//      that separates the two kernels this can represent from the two it
//      cannot: `q4_g128`'s scale is read at `row * groups + k / 128`, which
//      contains the row and is not the row, and a scale that varies along the
//      contraction is a block-wise scale that neither output format expresses.
// ---------------------------------------------------------------------------

/// One unpacking of an integer code out of a packed weight buffer.
struct PackedCodes {
    /// The storage buffer the packed word was loaded from.
    weight: Handle<GlobalVariable>,
    /// The index the packed word was loaded at.
    address: Handle<Expression>,
    /// The field width `extractBits` was asked for, in bits.
    code_bits: u32,
    /// The unpacked code, as an expression of the entry point.
    value: Handle<Expression>,
}

/// Functions that exist to pull one integer code out of a packed word.
///
/// Returned as (function, which argument holds the word, how wide the code is).
/// Every quantized kernel in the corpus spells the unpack as a helper —
/// `unpack_i8`, `unpack_i4` — and nothing inlines it, so looking only at the
/// entry point's own expressions finds no `extractBits` at all.
fn packed_code_helpers(module: &Module) -> Vec<(Handle<nxpu_ir::Function>, u32, u32)> {
    let mut found = Vec::new();
    for (handle, func) in module.functions.iter() {
        for (_, expr) in func.expressions.iter() {
            let Expression::Math {
                fun: MathFunction::ExtractBits,
                arg,
                arg2: Some(width),
                ..
            } = expr
            else {
                continue;
            };
            // The word reaches `extractBits` through a bitcast: the codes are
            // two's complement, and a signed base is what makes the extract
            // sign-extend rather than mask.
            let Expression::FunctionArgument(index) =
                func.expressions[strip_casts(&func.expressions, *arg)]
            else {
                continue;
            };
            if let Some(bits) = literal_u32(&func.expressions, *width) {
                found.push((handle, index, bits));
                break;
            }
        }
    }
    found
}

/// Every place the entry point unpacks a code out of one of `inputs`.
fn find_packed_codes(
    module: &Module,
    ep: &nxpu_ir::EntryPoint,
    inputs: &[Handle<GlobalVariable>],
) -> Vec<PackedCodes> {
    fn walk(
        body: &[Statement],
        exprs: &Arena<Expression>,
        helpers: &[(Handle<nxpu_ir::Function>, u32, u32)],
        inputs: &[Handle<GlobalVariable>],
        out: &mut Vec<PackedCodes>,
    ) {
        for stmt in body {
            match stmt {
                Statement::Call {
                    function,
                    arguments,
                    result: Some(value),
                } => {
                    let Some(&(_, arg_index, code_bits)) =
                        helpers.iter().find(|(f, ..)| f == function)
                    else {
                        continue;
                    };
                    let Some(&word) = arguments.get(arg_index as usize) else {
                        continue;
                    };
                    if let Some((weight, address)) = input_load(exprs, word, inputs) {
                        out.push(PackedCodes {
                            weight,
                            address,
                            code_bits,
                            value: *value,
                        });
                    }
                }
                Statement::If { accept, reject, .. } => {
                    walk(accept, exprs, helpers, inputs, out);
                    walk(reject, exprs, helpers, inputs, out);
                }
                Statement::Loop {
                    body, continuing, ..
                } => {
                    walk(body, exprs, helpers, inputs, out);
                    walk(continuing, exprs, helpers, inputs, out);
                }
                _ => {}
            }
        }
    }

    let exprs = &ep.function.expressions;
    let mut out = Vec::new();

    // The unpack written inline, without a helper. No kernel in the corpus
    // spells it this way, but the pattern is about what is computed and a
    // kernel that inlines its own `extractBits` computes the same thing.
    for (handle, expr) in exprs.iter() {
        let Expression::Math {
            fun: MathFunction::ExtractBits,
            arg,
            arg2: Some(width),
            ..
        } = expr
        else {
            continue;
        };
        let Some(bits) = literal_u32(exprs, *width) else {
            continue;
        };
        if let Some((weight, address)) = input_load(exprs, strip_casts(exprs, *arg), inputs) {
            out.push(PackedCodes {
                weight,
                address,
                code_bits: bits,
                value: handle,
            });
        }
    }

    walk(
        &ep.function.body,
        exprs,
        &packed_code_helpers(module),
        inputs,
        &mut out,
    );
    out
}

/// If `handle` loads out of one of `inputs`, which buffer and at which index.
fn input_load(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    inputs: &[Handle<GlobalVariable>],
) -> Option<(Handle<GlobalVariable>, Handle<Expression>)> {
    let Expression::Load { pointer } = exprs.try_get(strip_casts(exprs, handle))? else {
        return None;
    };
    let Expression::Access { base, index } = exprs.try_get(*pointer)? else {
        return None;
    };
    let global = access_base_global(exprs, *base)?;
    inputs.contains(&global).then_some((global, *index))
}

/// A `u32` literal, in either of the spellings a bit width arrives in.
fn literal_u32(exprs: &Arena<Expression>, handle: Handle<Expression>) -> Option<u32> {
    match exprs.try_get(handle)? {
        Expression::Literal(Literal::U32(n)) => Some(*n),
        Expression::Literal(Literal::I32(n)) if *n >= 0 => Some(*n as u32),
        Expression::Literal(Literal::AbstractInt(n)) if *n >= 0 => Some(*n as u32),
        _ => None,
    }
}

/// The two factors of the multiply in `a * b + c`, in either spelling.
///
/// `-O0` writes the multiply and the add as separate expressions; `FmaFusion`
/// at `-O1` rewrites them into one `fma`, and it does so for the *integer*
/// address arithmetic as well as for the float arithmetic — the weight address
/// `row * words_per_row + word_index` is a `Binary(Add, Binary(Mul, ..), ..)`
/// at `-O0` and a `Math(Fma, ..)` at `-O2`. Reading only one of the two is how
/// a classifier comes to give different answers at different levels.
fn multiply_factors(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
) -> Option<(Handle<Expression>, Handle<Expression>)> {
    match exprs.try_get(handle)? {
        Expression::Math {
            fun: MathFunction::Fma,
            arg,
            arg1: Some(b),
            ..
        } => Some((*arg, *b)),
        Expression::Binary {
            op: BinaryOp::Multiply,
            left,
            right,
        } => Some((*left, *right)),
        Expression::Binary {
            op: BinaryOp::Add,
            left,
            right,
        } => multiply_factors(exprs, *left).or_else(|| multiply_factors(exprs, *right)),
        _ => None,
    }
}

/// The addend of `a * b + c`, in either spelling. `None` when there is no add.
fn multiply_addend(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
) -> Option<Handle<Expression>> {
    match exprs.try_get(handle)? {
        Expression::Math {
            fun: MathFunction::Fma,
            arg2: Some(c),
            ..
        } => Some(*c),
        Expression::Binary {
            op: BinaryOp::Add,
            left,
            right,
        } => {
            if multiply_factors(exprs, *left).is_some() {
                Some(*right)
            } else if multiply_factors(exprs, *right).is_some() {
                Some(*left)
            } else {
                None
            }
        }
        _ => None,
    }
}

/// Are these two expressions the same computation?
///
/// Handle equality alone is not enough. `-O1` runs common-subexpression
/// elimination, so an index spelled by two handles before it is spelled by one
/// after, and a test that compared handles would answer differently at the two
/// levels about the same kernel.
fn same_expression(
    exprs: &Arena<Expression>,
    a: Handle<Expression>,
    b: Handle<Expression>,
    depth: u32,
) -> bool {
    if a == b {
        return true;
    }
    if depth == 0 {
        return false;
    }
    let (Some(ea), Some(eb)) = (exprs.try_get(a), exprs.try_get(b)) else {
        return false;
    };
    let same =
        |x: Handle<Expression>, y: Handle<Expression>| same_expression(exprs, x, y, depth - 1);
    let same_opt = |x: &Option<Handle<Expression>>, y: &Option<Handle<Expression>>| match (x, y) {
        (None, None) => true,
        (Some(x), Some(y)) => same_expression(exprs, *x, *y, depth - 1),
        _ => false,
    };
    match (ea, eb) {
        (Expression::Literal(x), Expression::Literal(y)) => same_literal(x, y),
        (Expression::FunctionArgument(x), Expression::FunctionArgument(y)) => x == y,
        (Expression::GlobalVariable(x), Expression::GlobalVariable(y)) => x == y,
        (Expression::LocalVariable(x), Expression::LocalVariable(y)) => x == y,
        (Expression::CallResult(x), Expression::CallResult(y)) => x == y,
        (Expression::Load { pointer: x }, Expression::Load { pointer: y }) => same(*x, *y),
        (
            Expression::Access {
                base: xb,
                index: xi,
            },
            Expression::Access {
                base: yb,
                index: yi,
            },
        ) => same(*xb, *yb) && same(*xi, *yi),
        (
            Expression::AccessIndex {
                base: xb,
                index: xi,
            },
            Expression::AccessIndex {
                base: yb,
                index: yi,
            },
        ) => xi == yi && same(*xb, *yb),
        (Expression::Unary { op: xo, expr: xe }, Expression::Unary { op: yo, expr: ye }) => {
            xo == yo && same(*xe, *ye)
        }
        (
            Expression::Binary {
                op: xo,
                left: xl,
                right: xr,
            },
            Expression::Binary {
                op: yo,
                left: yl,
                right: yr,
            },
        ) => xo == yo && same(*xl, *yl) && same(*xr, *yr),
        (
            Expression::As {
                expr: xe,
                kind: xk,
                convert: xc,
            },
            Expression::As {
                expr: ye,
                kind: yk,
                convert: yc,
            },
        ) => xk == yk && xc == yc && same(*xe, *ye),
        (
            Expression::Math {
                fun: xf,
                arg: xa,
                arg1: xa1,
                arg2: xa2,
                arg3: xa3,
            },
            Expression::Math {
                fun: yf,
                arg: ya,
                arg1: ya1,
                arg2: ya2,
                arg3: ya3,
            },
        ) => {
            xf == yf
                && same(*xa, *ya)
                && same_opt(xa1, ya1)
                && same_opt(xa2, ya2)
                && same_opt(xa3, ya3)
        }
        _ => false,
    }
}

/// `Literal` carries an `f32`, so it cannot derive `PartialEq`; bit equality is
/// the right test for two spellings of the same constant anyway.
fn same_literal(a: &Literal, b: &Literal) -> bool {
    match (a, b) {
        (Literal::Bool(x), Literal::Bool(y)) => x == y,
        (Literal::I32(x), Literal::I32(y)) => x == y,
        (Literal::U32(x), Literal::U32(y)) => x == y,
        (Literal::F32(x), Literal::F32(y)) => x.to_bits() == y.to_bits(),
        (Literal::F64(x), Literal::F64(y)) => x.to_bits() == y.to_bits(),
        (Literal::AbstractInt(x), Literal::AbstractInt(y)) => x == y,
        (Literal::AbstractFloat(x), Literal::AbstractFloat(y)) => x.to_bits() == y.to_bits(),
        _ => false,
    }
}

/// Does `haystack` compute `needle` somewhere inside it?
///
/// Iterative with a visited set, for [`reads_global`]'s reason: an expression
/// arena is a DAG, and an address built from shared subexpressions — which
/// every one of these addresses is, after `-O1` runs common-subexpression
/// elimination — is walked exponentially by the naive spelling.
fn expression_contains(
    exprs: &Arena<Expression>,
    haystack: Handle<Expression>,
    needle: Handle<Expression>,
) -> bool {
    let mut seen = std::collections::HashSet::new();
    let mut stack = vec![haystack];
    while let Some(handle) = stack.pop() {
        if !seen.insert(handle.index()) {
            continue;
        }
        if same_expression(exprs, handle, needle, EXPR_MATCH_DEPTH) {
            return true;
        }
        if let Some(expr) = exprs.try_get(handle) {
            push_operands(expr, &mut stack);
        }
    }
    false
}

/// How deep [`same_expression`] walks.
///
/// An address is a handful of nodes; the limit is here so a cyclic or
/// pathological arena cannot make classification hang, which is a property the
/// corpus test asserts.
const EXPR_MATCH_DEPTH: u32 = 24;

/// One load out of a storage input.
struct InputLoad {
    global: Handle<GlobalVariable>,
    index: Handle<Expression>,
    value: Handle<Expression>,
}

/// Every load out of one of `inputs`, anywhere in the entry point.
fn collect_input_loads(
    exprs: &Arena<Expression>,
    inputs: &[Handle<GlobalVariable>],
) -> Vec<InputLoad> {
    exprs
        .iter()
        .filter_map(|(handle, expr)| {
            let Expression::Load { pointer } = expr else {
                return None;
            };
            let Expression::Access { base, index } = exprs.try_get(*pointer)? else {
                return None;
            };
            let global = access_base_global(exprs, *base)?;
            inputs.contains(&global).then_some(InputLoad {
                global,
                index: *index,
                value: handle,
            })
        })
        .collect()
}

/// The members of the uniform params struct this expression reads.
fn uniform_members_read(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    params: Option<Handle<GlobalVariable>>,
) -> Vec<u32> {
    let Some(params) = params else {
        return Vec::new();
    };
    let mut seen = std::collections::HashSet::new();
    let mut members = Vec::new();
    let mut stack = vec![handle];
    while let Some(handle) = stack.pop() {
        if !seen.insert(handle.index()) {
            continue;
        }
        let Some(expr) = exprs.try_get(handle) else {
            continue;
        };
        #[allow(clippy::collapsible_if)]
        if let Expression::AccessIndex { base, index } = expr {
            if matches!(exprs.try_get(*base), Some(Expression::GlobalVariable(g)) if *g == params)
                && !members.contains(index)
            {
                members.push(*index);
            }
        }
        push_operands(expr, &mut stack);
    }
    members
}

/// The one store into `output`, as (index written, value stored).
///
/// Not [`find_lone_store`]: every one of these kernels also writes workgroup
/// memory, several times, and the store that says what the kernel produces is
/// the one that lands in the output buffer.
fn find_store_into(
    body: &[Statement],
    exprs: &Arena<Expression>,
    output: Handle<GlobalVariable>,
) -> Option<(Handle<Expression>, Handle<Expression>)> {
    fn walk(
        body: &[Statement],
        exprs: &Arena<Expression>,
        output: Handle<GlobalVariable>,
        found: &mut Vec<(Handle<Expression>, Handle<Expression>)>,
    ) {
        for stmt in body {
            match stmt {
                Statement::Store { pointer, value } => {
                    if access_base_global(exprs, *pointer) == Some(output) {
                        match exprs.try_get(*pointer) {
                            Some(Expression::Access { index, .. }) => found.push((*index, *value)),
                            _ => found.push((*pointer, *value)),
                        }
                    }
                }
                Statement::If { accept, reject, .. } => {
                    walk(accept, exprs, output, found);
                    walk(reject, exprs, output, found);
                }
                Statement::Loop {
                    body, continuing, ..
                } => {
                    walk(body, exprs, output, found);
                    walk(continuing, exprs, output, found);
                }
                _ => {}
            }
        }
    }
    let mut found = Vec::new();
    walk(body, exprs, output, &mut found);
    match found.as_slice() {
        [one] => Some(*one),
        _ => None,
    }
}

/// Recognise a matrix multiplication against packed integer weight codes.
///
/// `None` when the kernel unpacks nothing, so every other pattern is still on
/// the table. `Some(Unknown { .. })` when it does unpack codes but into a shape
/// no backend here can emit — a refusal that names the format is worth more
/// than a graph that runs and computes something else.
fn match_quantized_matmul(
    module: &Module,
    ep: &nxpu_ir::EntryPoint,
    inputs: &[(Handle<GlobalVariable>, &GlobalVariable)],
    outputs: &[(Handle<GlobalVariable>, &GlobalVariable)],
    params_global: Option<Handle<GlobalVariable>>,
    // `param_names` is every params-struct member's name, indexed by member
    // position, so an index read out of the IR names the right one.
    // `shape_names` is the integer members only — `eps` sits in the same struct
    // and is not a dimension, so it must not be a candidate extent.
    param_names: &[String],
    shape_names: &[String],
) -> Option<KernelPattern> {
    let exprs = &ep.function.expressions;
    let input_handles: Vec<Handle<GlobalVariable>> = inputs.iter().map(|(h, _)| *h).collect();
    let codes = find_packed_codes(module, ep, &input_handles);
    if codes.is_empty() {
        return None;
    }
    let refuse = |reason: String| Some(KernelPattern::Unknown { reason });

    // One packed weight. `matvec/q8_ffn` unpacks two — a gate projection and
    // an up projection sharing one activation, with a SiLU and a multiply on
    // top. That is two quantized matmuls and two more operators, not one
    // matmul with a wider weight, and this carries one of each.
    let weights = dedup_globals(codes.iter().map(|c| c.weight));
    if weights.len() != 1 {
        return refuse(format!(
            "{} separate packed weight buffers unpacked in one kernel — several \
             quantized projections fused together, which needs a pattern per \
             projection and an operator for whatever combines them; this one \
             carries a single weight, a single scale and a single contraction",
            weights.len()
        ));
    }
    let weight_handle = weights[0];
    let code_bits = codes[0].code_bits;
    if codes.iter().any(|c| c.code_bits != code_bits) {
        return refuse(
            "the packed weight is unpacked at two different code widths, so no \
             single element type describes it"
                .into(),
        );
    }
    if outputs.len() != 1 {
        return refuse(format!(
            "a quantized matmul writes one result and this kernel writes {}",
            outputs.len()
        ));
    }
    if !has_loop(&ep.function.body) {
        return refuse(
            "integer codes unpacked without a loop — there is no contraction \
             here, so this is a dequantization rather than a matmul"
                .into(),
        );
    }

    // The unpacked code has to meet a value out of a buffer in a multiply:
    // the activation it is contracted against, or the scale that turns it back
    // into a number. Codes unpacked and never multiplied are something else.
    let loads = collect_input_loads(exprs, &input_handles);
    let multiplied = exprs.iter().any(|(_, expr)| {
        let Some((a, b)) = (match expr {
            Expression::Binary {
                op: BinaryOp::Multiply,
                left,
                right,
            } => Some((*left, *right)),
            Expression::Math {
                fun: MathFunction::Fma,
                arg,
                arg1: Some(b),
                ..
            } => Some((*arg, *b)),
            _ => None,
        }) else {
            return false;
        };
        let is_code = |h: Handle<Expression>| {
            codes
                .iter()
                .any(|c| same_expression(exprs, strip_casts(exprs, h), c.value, EXPR_MATCH_DEPTH))
        };
        let is_input = |h: Handle<Expression>| {
            let h = strip_casts(exprs, h);
            loads.iter().any(|l| l.value == h)
        };
        (is_code(a) && is_input(b)) || (is_code(b) && is_input(a))
    });
    if !multiplied {
        return refuse(
            "integer codes unpacked but never multiplied by a value from another \
             buffer — nothing dequantizes them and nothing contracts them"
                .into(),
        );
    }

    // The weight's address is `channel * stride + position`. Of the two
    // factors, the one that reads the params struct is the row stride —
    // `ceil(K / codes_per_word)` — and the other one selects the row.
    let Some((f0, f1)) = multiply_factors(exprs, codes[0].address) else {
        return refuse(
            "the packed weight is addressed without a row stride, so which axis \
             of it is an output channel cannot be established"
                .into(),
        );
    };
    let m0 = uniform_members_read(exprs, f0, params_global);
    let m1 = uniform_members_read(exprs, f1, params_global);
    let (channel, stride_members) = match (m0.is_empty(), m1.is_empty()) {
        (true, false) => (f0, m1),
        (false, true) => (f1, m0),
        _ => {
            return refuse(
                "the packed weight's row stride cannot be told from its row index \
                 — both, or neither, are built from the dispatch parameters"
                    .into(),
            );
        }
    };
    let [k_member] = stride_members[..] else {
        return refuse(
            "the packed weight's row stride is built from more than one dispatch \
             parameter, so the contracted extent is not named by any one of them"
                .into(),
        );
    };
    let Some(k_name) = param_names.get(k_member as usize).cloned() else {
        return refuse(
            "the packed weight's row stride reads a dispatch parameter with no name".into(),
        );
    };

    // The output store, and the addend fused onto it if there is one.
    let output_handle = outputs[0].0;
    let Some((store_index, store_value)) = find_store_into(&ep.function.body, exprs, output_handle)
    else {
        return refuse(
            "the result is written in more than one place, so what this kernel \
             produces is not one matmul"
                .into(),
        );
    };
    let bias_load = multiply_addend(exprs, store_value).and_then(|addend| {
        let addend = strip_casts(exprs, addend);
        loads.iter().find(|l| {
            l.value == addend && same_expression(exprs, l.index, store_index, EXPR_MATCH_DEPTH)
        })
    });
    let bias_handle = bias_load.map(|l| l.global);

    // The per-channel scale: read at *exactly* the index that selects the
    // weight's row. A scale read at anything wider varies along the
    // contraction, which is a block-wise scale — see the refusal below.
    let per_channel: Vec<Handle<GlobalVariable>> = dedup_globals(loads.iter().filter_map(|l| {
        (l.global != weight_handle
            && Some(l.global) != bias_handle
            && same_expression(exprs, l.index, channel, EXPR_MATCH_DEPTH))
        .then_some(l.global)
    }));
    let [scale_handle] = per_channel[..] else {
        let block_scaled: Vec<Handle<GlobalVariable>> =
            dedup_globals(loads.iter().filter_map(|l| {
                (l.global != weight_handle
                    && !same_expression(exprs, l.index, channel, EXPR_MATCH_DEPTH)
                    && expression_contains(exprs, l.index, channel))
                .then_some(l.global)
            }));
        if per_channel.is_empty() && !block_scaled.is_empty() {
            return refuse(format!(
                "the {code_bits}-bit codes are scaled by a factor read at the \
                 weight's row *and* the contracted position — one scale per block \
                 of columns rather than one per output channel. Per-channel \
                 quantization is what `DequantizeLinear` with an axis and what a \
                 TFLite per-channel tensor express; a block-wise scale needs a \
                 blocked dequantization neither backend here emits"
            ));
        }
        return refuse(format!(
            "{} buffers are read at exactly the weight's row index, and a \
             quantized matmul has one such buffer — the per-channel scale",
            per_channel.len()
        ));
    };

    // Only int8 survives. A 4-bit code is an INT4 tensor: ONNX has one from
    // opset 21 and this backend declares 13, and the TFLite writer here has no
    // sub-byte tensor type at all. Widening the codes to int8 on the way out
    // would be a different graph with a weight four times the size, which is
    // not what the kernel was written to run.
    if code_bits != 8 {
        return refuse(format!(
            "{code_bits}-bit weight codes. The formats this compiler writes carry \
             int8 weights; a {code_bits}-bit one needs a sub-byte tensor type — \
             ONNX INT4 arrived in opset 21 and this backend emits 13, and the \
             TFLite writer here has no sub-byte type"
        ));
    }

    // Whatever is left is the activation, and there has to be exactly one of
    // it. Counting the inputs is not what decided any of the roles above; this
    // is the check that every input was accounted for by one of them.
    let assigned = [Some(weight_handle), Some(scale_handle), bias_handle];
    let rest: Vec<&(Handle<GlobalVariable>, &GlobalVariable)> = inputs
        .iter()
        .filter(|(h, _)| !assigned.contains(&Some(*h)))
        .collect();
    let [(activation_handle, activation_gv)] = rest[..] else {
        return refuse(format!(
            "{} buffers are left over once the weight, its scale and its addend \
             are accounted for, and a quantized matmul contracts against one set \
             of activations",
            rest.len()
        ));
    };
    let input = make_binding(module, *activation_handle, activation_gv, TensorRole::Input);
    if input.elem_type != data_type::FLOAT {
        return refuse(
            "the activations are not f32, so the contraction is integer on both \
             sides — an integer matmul with its own accumulator width, which this \
             pattern does not describe"
                .into(),
        );
    }

    // The shape. `k` is already known from the weight's row stride. The output
    // store says the rest: `output[row * M + col]` names `M` as the number of
    // output channels, and a store at the row index alone is one row of
    // results, so `m` is 1.
    let (n_name, m_name) = if let Some((g0, g1)) = multiply_factors(exprs, store_index) {
        let n0 = uniform_members_read(exprs, g0, params_global);
        let n1 = uniform_members_read(exprs, g1, params_global);
        let members = match (n0.is_empty(), n1.is_empty()) {
            (true, false) => n1,
            (false, true) => n0,
            _ => Vec::new(),
        };
        let [n_member] = members[..] else {
            return refuse(
                "the result's row stride is not a single dispatch parameter, so \
                 the number of output channels is not named"
                    .into(),
            );
        };
        let Some(n_name) = param_names.get(n_member as usize).cloned() else {
            return refuse("the result's row stride reads an unnamed parameter".into());
        };
        let remaining: Vec<&String> = shape_names
            .iter()
            .filter(|p| **p != k_name && **p != n_name)
            .collect();
        let [m_name] = remaining[..] else {
            return refuse(format!(
                "{} dispatch parameters are left once the contracted extent and \
                 the output-channel count are named, and the number of rows has \
                 to be exactly one of them",
                remaining.len()
            ));
        };
        (n_name, m_name.clone())
    } else {
        // The result is written at the row index alone, so it is one row.
        let remaining: Vec<&String> = shape_names.iter().filter(|p| **p != k_name).collect();
        let [n_name] = remaining[..] else {
            return refuse(format!(
                "the result is one row, so the remaining dispatch parameter is \
                 the output-channel count — and {} remain",
                remaining.len()
            ));
        };
        (n_name.clone(), "1".to_string())
    };

    // The codes, not the words they arrived in: four two's-complement bytes
    // per `u32`, least-significant first, is the byte layout of a contiguous
    // int8 row, and int8 is what the graph's consumer has to be told it holds.
    let weight_gv = inputs.iter().find(|(h, _)| *h == weight_handle)?.1;
    let mut weight = make_binding(module, weight_handle, weight_gv, TensorRole::Input);
    weight.elem_type = data_type::INT8;
    let scale_gv = inputs.iter().find(|(h, _)| *h == scale_handle)?.1;
    let bias = bias_handle.and_then(|h| {
        inputs
            .iter()
            .find(|(g, _)| *g == h)
            .map(|(g, gv)| make_binding(module, *g, gv, TensorRole::Input))
    });

    Some(KernelPattern::QuantizedMatMul {
        input,
        weight,
        scale: make_binding(module, scale_handle, scale_gv, TensorRole::Input),
        bias,
        output: make_binding(module, outputs[0].0, outputs[0].1, TensorRole::Output),
        shape: MatMulShape {
            m: m_name,
            n: n_name,
            k: k_name,
        },
    })
}

/// Collect handles, dropping repeats and keeping the order they arrived in.
fn dedup_globals(it: impl Iterator<Item = Handle<GlobalVariable>>) -> Vec<Handle<GlobalVariable>> {
    let mut out: Vec<Handle<GlobalVariable>> = Vec::new();
    for h in it {
        if !out.contains(&h) {
            out.push(h);
        }
    }
    out
}

/// Extract literal U32 values from loop break conditions.
///
/// For `for (kh = 0; kh < N; ...)`, naga generates `break_if: kh >= N`.
/// When N is `Literal(U32(n))`, we collect n as a loop bound.
fn extract_loop_bound_literals(body: &[Statement], exprs: &Arena<Expression>) -> Vec<u32> {
    let mut bounds = Vec::new();
    for stmt in body {
        match stmt {
            Statement::Loop {
                body,
                continuing,
                break_if,
            } => {
                // Check this loop's break condition for a literal bound.
                // Pattern: break_if is `Binary { GreaterEqual|Greater, _, Literal(U32(n)) }`.
                if let Some(n) = break_if.and_then(|bi| {
                    let Expression::Binary {
                        op: BinaryOp::GreaterEqual | BinaryOp::Greater,
                        right,
                        ..
                    } = exprs.try_get(bi)?
                    else {
                        return None;
                    };
                    match exprs.try_get(*right)? {
                        Expression::Literal(Literal::U32(n)) => Some(*n),
                        _ => None,
                    }
                }) {
                    bounds.push(n);
                }
                // naga 28 pattern: for-loop condition lowered as
                //   Emit; ...; If(cond) { } else { Break }  at the start of the loop body,
                // where cond is `Binary { Less|LessEqual, _, Literal(U32(n)) }`.
                // Skip leading Emit statements to find the If.
                if let Some(Statement::If {
                    condition,
                    accept,
                    reject,
                }) = body.iter().find(|s| !matches!(s, Statement::Emit(_)))
                {
                    let rhs = match exprs.try_get(*condition) {
                        Some(&Expression::Binary {
                            op: BinaryOp::Less | BinaryOp::LessEqual,
                            right,
                            ..
                        }) if accept.is_empty()
                            && matches!(reject.as_slice(), [Statement::Break]) =>
                        {
                            Some(right)
                        }
                        Some(&Expression::Binary {
                            op: BinaryOp::GreaterEqual | BinaryOp::Greater,
                            right,
                            ..
                        }) if reject.is_empty()
                            && matches!(accept.as_slice(), [Statement::Break]) =>
                        {
                            Some(right)
                        }
                        _ => None,
                    };
                    let bound = rhs.and_then(|h| exprs.try_get(h)).and_then(|e| match *e {
                        Expression::Literal(Literal::U32(n)) => Some(n),
                        _ => None,
                    });
                    if let Some(n) = bound {
                        bounds.push(n);
                    }
                }
                // Recurse into nested loops.
                bounds.extend(extract_loop_bound_literals(body, exprs));
                bounds.extend(extract_loop_bound_literals(continuing, exprs));
            }
            Statement::If { accept, reject, .. } => {
                bounds.extend(extract_loop_bound_literals(accept, exprs));
                bounds.extend(extract_loop_bound_literals(reject, exprs));
            }
            _ => {}
        }
    }
    bounds
}

/// Does this expression come from the invocation id rather than from a loop?
///
/// The dividing line for [`extract_multiply_literals`], and the only structural
/// difference between the two things a literal multiply can mean. A stride
/// scales the coordinate the invocation is computing — `oh`, which is `gid.y` —
/// while an index-flattening factor scales a reduction's loop variable, which
/// reaches the expression as a load of a local.
///
/// Deliberately narrow: an id reaches this either directly or through one
/// component access, and nothing else counts. Widening it to follow arithmetic
/// would readmit `oh + kh`, which is half a loop variable.
fn derives_from_invocation_id(exprs: &Arena<Expression>, handle: Handle<Expression>) -> bool {
    match exprs.try_get(handle) {
        // `gid` itself, and `gid.x` / `gid.y` / `gid.z`.
        Some(Expression::FunctionArgument(_)) => true,
        Some(Expression::AccessIndex { base, .. }) => {
            matches!(exprs.try_get(*base), Some(Expression::FunctionArgument(_)))
        }
        _ => false,
    }
}

/// Stride factors, from `Binary(Multiply, id, Literal(U32(n)))` where n > 1.
///
/// The multiplicand has to come from the invocation id. Without that test this
/// took **every** literal multiply in the function, and a convolution flattens
/// its weight index with them: `oc * IC * 9u + ic * 9u + kh * 3u + kw` reported
/// strides of 3 and 9, and a 3x3 VALID convolution over a 64x64 image was
/// emitted with an output of 21x7 instead of 62x62. That model loads, is
/// accelerated, and returns the wrong numbers — the failure this project
/// refuses everywhere it can see it, and this one it could not see.
///
/// The test narrows rather than widens, on purpose. A genuine stride written
/// through a loop variable is now missed and reported as 1, which is the
/// common case and wrong by a little; a flattening factor read as a stride is
/// wrong by whatever the channel count happens to be.
fn extract_multiply_literals(exprs: &Arena<Expression>) -> Vec<u32> {
    let mut strides = Vec::new();
    // Nested rather than a let-chain: those are stable from 1.88 and this
    // workspace's rust-version is 1.87.
    #[allow(clippy::collapsible_if)]
    for (_, expr) in exprs.iter() {
        let Expression::Binary {
            op: BinaryOp::Multiply,
            left,
            right,
        } = expr
        else {
            continue;
        };
        if let Some(&Expression::Literal(Literal::U32(n @ 2..))) = exprs.try_get(*right) {
            if derives_from_invocation_id(exprs, *left) {
                strides.push(n);
            }
        } else if let Some(&Expression::Literal(Literal::U32(n @ 2..))) = exprs.try_get(*left) {
            if derives_from_invocation_id(exprs, *right) {
                strides.push(n);
            }
        }
    }
    strides.sort_unstable();
    strides.dedup();
    strides
}

/// Scan the expression arena for `Binary(Subtract, _, Literal(U32(n)))` where n > 0.
/// These represent padding offsets in index computations like `oh * 2u + kh - 1u`.
fn extract_subtract_literals(exprs: &Arena<Expression>) -> Vec<u32> {
    let mut pads = Vec::new();
    for (_, expr) in exprs.iter() {
        let Expression::Binary {
            op: BinaryOp::Subtract,
            right,
            ..
        } = expr
        else {
            continue;
        };
        if let Some(&Expression::Literal(Literal::U32(n @ 1..))) = exprs.try_get(*right) {
            pads.push(n);
        }
    }
    pads.sort_unstable();
    pads.dedup();
    pads
}

/// Extract Conv2D shape from param names and (optionally) the function body/expressions.
///
/// When body and expressions are provided, attempts to extract kernel size from
/// loop bound literals, stride from multiplication factors, and padding from
/// subtraction offsets. Falls back to 0 (unknown) for values that cannot be
/// determined at compile time.
fn extract_conv2d_shape(
    shape_names: &[String],
    body: Option<&[Statement]>,
    exprs: Option<&Arena<Expression>>,
) -> Conv2DShape {
    // Convention: params struct has N, IC, IH, IW, OC, KH, KW, ...
    let get = |i: usize| {
        shape_names
            .get(i)
            .cloned()
            .unwrap_or_else(|| format!("d{i}"))
    };

    let (kernel_h_val, kernel_w_val, stride_h, stride_w, pad_h, pad_w) =
        if let (Some(body), Some(exprs)) = (body, exprs) {
            let bounds = extract_loop_bound_literals(body, exprs);
            let strides = extract_multiply_literals(exprs);
            let pads = extract_subtract_literals(exprs);

            // Innermost loop bounds are kernel sizes (e.g., kh < 5u → kernel_h = 5).
            let kh = bounds.first().copied().unwrap_or(0) as i64;
            let kw = bounds.get(1).copied().unwrap_or(0) as i64;
            // If only one bound found, assume square kernel.
            let kw = if kw == 0 && kh > 0 { kh } else { kw };

            // Stride from multiplication factors (e.g., oh * 2u → stride = 2).
            let sh_u = strides.first().copied().unwrap_or(1);
            let sw_u = strides.get(1).copied().unwrap_or(sh_u);
            let sh = sh_u as i64;
            let sw = sw_u as i64;

            // Padding from subtraction offsets (e.g., ih - 1u → pad = 1).
            let ph_u = pads.first().copied().unwrap_or(0);
            let pw_u = pads.get(1).copied().unwrap_or(ph_u);
            let ph = ph_u as i64;
            let pw = pw_u as i64;

            (kh, kw, sh, sw, ph, pw)
        } else {
            (0, 0, 1, 1, 0, 0)
        };

    // Detect groups parameter: if shape_names contains "groups", use it.
    let groups = if shape_names.iter().any(|n| n.eq_ignore_ascii_case("groups")) {
        // Depthwise convention: groups == channels_in
        -1 // Sentinel: resolved at runtime or marked as depthwise
    } else {
        1
    };

    Conv2DShape {
        batch: get(0),
        channels_in: get(1),
        height: get(2),
        width: get(3),
        channels_out: get(4),
        kernel_h: get(5),
        kernel_w: get(6),
        kernel_h_val,
        kernel_w_val,
        stride_h,
        stride_w,
        pad_h,
        pad_w,
        groups,
        dilation_h: 1,
        dilation_w: 1,
    }
}

fn make_binding(
    module: &Module,
    handle: Handle<GlobalVariable>,
    gv: &GlobalVariable,
    role: TensorRole,
) -> TensorBinding {
    let elem_type = resolve_array_elem_type(module, gv.ty).unwrap_or(data_type::FLOAT);
    TensorBinding {
        handle,
        name: gv
            .name
            .clone()
            .unwrap_or_else(|| format!("tensor_{}", handle.index())),
        elem_type,
        role,
    }
}

/// Resolve an array or tensor type to its element's ONNX data type.
fn resolve_array_elem_type(module: &Module, ty: Handle<Type>) -> Option<i32> {
    match &module.types[ty].inner {
        TypeInner::Array { base, .. } => match &module.types[*base].inner {
            TypeInner::Scalar(s) => Some(scalar_to_onnx_data_type(s)),
            _ => None,
        },
        TypeInner::Tensor { scalar, .. } => Some(scalar_to_onnx_data_type(scalar)),
        _ => None,
    }
}

/// Map an IR scalar type to an ONNX data type constant.
fn scalar_to_onnx_data_type(scalar: &Scalar) -> i32 {
    match (scalar.kind, scalar.width) {
        (ScalarKind::Float, 4) => data_type::FLOAT,
        (ScalarKind::Float, 2) => data_type::FLOAT16,
        (ScalarKind::BFloat, 2) => data_type::BFLOAT16,
        (ScalarKind::Sint, 4) => data_type::INT32,
        (ScalarKind::Sint, 1) => data_type::INT8,
        (ScalarKind::Uint, 4) => data_type::UINT32,
        (ScalarKind::Uint, 1) => data_type::UINT8,
        (ScalarKind::Bool, _) => data_type::BOOL,
        _ => data_type::FLOAT,
    }
}

/// Normalize a potentially negative axis to a positive axis.
///
/// For example, `normalize_axis(-1, 4)` returns `3`.
pub fn normalize_axis(axis: i64, ndim: usize) -> i64 {
    if axis < 0 {
        let ndim = ndim as i64;
        ((axis % ndim) + ndim) % ndim
    } else {
        axis
    }
}

/// Infer the concat axis from the If condition in the function body.
///
/// In WGSL concat patterns, the If condition checks whether the linear index
/// falls before or after a boundary (e.g., `if idx < params.C1`). The boundary
/// param name maps to a position in the params struct, which gives us the axis.
fn infer_concat_axis(body: &[Statement], exprs: &Arena<Expression>, shape_names: &[String]) -> i64 {
    find_if_comparison_axis(body, exprs, shape_names).unwrap_or(0)
}

/// Infer the split axis from the If condition in the function body.
///
/// Works identically to [`infer_concat_axis`]: the If condition compares
/// against a params struct member whose index indicates the split axis.
fn infer_split_axis(body: &[Statement], exprs: &Arena<Expression>, shape_names: &[String]) -> i64 {
    find_if_comparison_axis(body, exprs, shape_names).unwrap_or(0)
}

/// Walk the statement tree to find an If whose condition is a comparison
/// (Less, LessEqual, Greater, GreaterEqual) where one operand is an
/// `AccessIndex` into the uniform params struct. The `index` field of
/// `AccessIndex` gives the struct member position, which corresponds to
/// the concat/split axis.
fn find_if_comparison_axis(
    body: &[Statement],
    exprs: &Arena<Expression>,
    _shape_names: &[String],
) -> Option<i64> {
    for stmt in body {
        match stmt {
            Statement::If { condition, .. } => {
                // Split into nested ifs to avoid let-chains (unstable in MSRV 1.87).
                #[allow(clippy::collapsible_if)]
                if let Some(Expression::Binary {
                    op, left, right, ..
                }) = exprs.try_get(*condition)
                {
                    if matches!(
                        op,
                        BinaryOp::Less
                            | BinaryOp::LessEqual
                            | BinaryOp::Greater
                            | BinaryOp::GreaterEqual
                    ) {
                        // Check right operand first (most common: `idx < params.C1`)
                        if let Some(idx) = extract_uniform_member_index(exprs, *right) {
                            return Some(idx as i64);
                        }
                        if let Some(idx) = extract_uniform_member_index(exprs, *left) {
                            return Some(idx as i64);
                        }
                    }
                }
            }
            Statement::Loop {
                body, continuing, ..
            } => {
                if let Some(axis) = find_if_comparison_axis(body, exprs, _shape_names) {
                    return Some(axis);
                }
                if let Some(axis) = find_if_comparison_axis(continuing, exprs, _shape_names) {
                    return Some(axis);
                }
            }
            _ => {}
        }
    }
    None
}

/// Extract the struct member index from an expression that accesses a
/// uniform params member.
///
/// Handles both direct `AccessIndex { base, index }` and
/// `Load { pointer: AccessIndex { base, index } }` patterns, since
/// different naga versions may emit either form.
fn extract_uniform_member_index(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
) -> Option<u32> {
    match exprs.try_get(handle)? {
        Expression::AccessIndex { index, .. } => Some(*index),
        Expression::Load { pointer } => match exprs.try_get(*pointer)? {
            Expression::AccessIndex { index, .. } => Some(*index),
            _ => None,
        },
        _ => None,
    }
}

/// Check if any expression in the arena uses a specific math function.
fn has_math_function_in_expressions(exprs: &Arena<Expression>, target: MathFunction) -> bool {
    exprs
        .iter()
        .any(|(_, expr)| matches!(expr, Expression::Math { fun, .. } if *fun == target))
}

/// Check if a block (or any nested block) contains an If statement.
fn has_if_statement(body: &[Statement]) -> bool {
    body.iter().any(|stmt| match stmt {
        Statement::If { .. } => true,
        Statement::Loop {
            body, continuing, ..
        } => has_if_statement(body) || has_if_statement(continuing),
        _ => false,
    })
}

/// Check if a block contains a non-guard If statement (i.e., one that is not
/// just a bounds-check early return like `if (idx >= N) { return; }`).
fn has_non_guard_if(body: &[Statement]) -> bool {
    body.iter().any(|stmt| match stmt {
        Statement::If { accept, reject, .. } => {
            // A guard-if has only a Return in accept and empty reject (or vice-versa).
            let is_guard = reject.is_empty()
                && accept.len() == 1
                && matches!(accept[0], Statement::Return { .. });
            !is_guard
        }
        Statement::Loop {
            body, continuing, ..
        } => has_non_guard_if(body) || has_non_guard_if(continuing),
        _ => false,
    })
}

/// Check if a block (or any nested block) contains a Loop statement.
fn has_loop(body: &[Statement]) -> bool {
    body.iter().any(|stmt| match stmt {
        Statement::Loop { .. } => true,
        Statement::If { accept, reject, .. } => has_loop(accept) || has_loop(reject),
        _ => false,
    })
}

/// What the value expression of a `Store` computes, as far as element-wise
/// classification is concerned.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StoreValueOp {
    /// A plain two-operand binary op: `a + b`, `a - b`, `a * b`, `a / b`.
    Binary(ElementWiseOp),
    /// A three-operand multiply-add, in either of the two spellings this
    /// pipeline can produce for it: `c + a * b` as the kernel author wrote it,
    /// or `fma(a, b, c)` after `nxpu-opt`'s FMA fusion pass has rewritten it.
    ///
    /// It is deliberately *not* mapped onto [`ElementWiseOp::Add`]. Calling it
    /// an add discards the multiply, which is how `axpy` — `y + a * x` — used
    /// to classify as a bare `Add` at `--opt-level 0` and then fail outright at
    /// `--opt-level 1` once fusion had rewritten it. Naming the shape here lets
    /// the classifier give the same answer at every optimization level, and
    /// leaves a single place to add a real three-operand pattern later.
    MultiplyAdd,
}

/// Search a block for a Store whose value is a recognizable arithmetic
/// expression, returning what kind it is.
fn find_store_value_op(body: &[Statement], exprs: &Arena<Expression>) -> Option<StoreValueOp> {
    for stmt in body {
        match stmt {
            Statement::Store { value, .. } => {
                if let Some(op) = classify_store_value_op(exprs, *value) {
                    return Some(op);
                }
            }
            Statement::If { accept, reject, .. } => {
                if let Some(op) = find_store_value_op(accept, exprs) {
                    return Some(op);
                }
                if let Some(op) = find_store_value_op(reject, exprs) {
                    return Some(op);
                }
            }
            _ => {}
        }
    }
    None
}

/// Classify a single stored value expression.
fn classify_store_value_op(
    exprs: &Arena<Expression>,
    value: Handle<Expression>,
) -> Option<StoreValueOp> {
    match exprs.try_get(value)? {
        // What FMA fusion leaves behind.
        Expression::Math {
            fun: MathFunction::Fma,
            ..
        } => Some(StoreValueOp::MultiplyAdd),
        Expression::Binary { op, left, right } => {
            let ew = match op {
                // The same three-operand op before fusion. Only `Add` is
                // checked because only `Add` is what FmaFusion rewrites; a
                // `Subtract` over multiplies (rope's rotation, say) stays a
                // `Subtract` at every level and needs no special case.
                BinaryOp::Add => {
                    if expr_is_multiply(exprs, *left) || expr_is_multiply(exprs, *right) {
                        return Some(StoreValueOp::MultiplyAdd);
                    }
                    ElementWiseOp::Add
                }
                BinaryOp::Subtract => ElementWiseOp::Sub,
                BinaryOp::Multiply => ElementWiseOp::Mul,
                BinaryOp::Divide => ElementWiseOp::Div,
                _ => return None,
            };
            Some(StoreValueOp::Binary(ew))
        }
        _ => None,
    }
}

/// Whether an expression is directly a multiplication — the operand shape
/// FMA fusion looks for.
fn expr_is_multiply(exprs: &Arena<Expression>, handle: Handle<Expression>) -> bool {
    matches!(
        exprs.try_get(handle),
        Some(Expression::Binary {
            op: BinaryOp::Multiply,
            ..
        })
    )
}

/// A leaf of a per-element expression: something with no arithmetic under it.
#[derive(Debug, Clone)]
enum ChainLeaf {
    /// A storage buffer, read at the index the result is stored at.
    Tensor(Handle<GlobalVariable>),
    /// A uniform scalar, or one scalar member of the uniform params struct.
    Scalar(ScalarBinding),
}

/// A chain before its globals are turned into [`TensorBinding`]s.
struct RawChain {
    base: Handle<GlobalVariable>,
    cast: Option<i32>,
    steps: Vec<(ElementWiseOp, ChainLeaf)>,
}

/// The one `Store` in a block, as (index written, value stored).
///
/// `None` when there is more than one: a chain describes a single write, and
/// two writes are two operators or a movement of values, neither of which this
/// says.
fn find_lone_store(
    body: &[Statement],
    exprs: &Arena<Expression>,
) -> Option<(Handle<Expression>, Handle<Expression>)> {
    fn walk(
        body: &[Statement],
        exprs: &Arena<Expression>,
        found: &mut Vec<(Handle<Expression>, Handle<Expression>)>,
    ) {
        for stmt in body {
            match stmt {
                Statement::Store { pointer, value } => match access_index(exprs, *pointer) {
                    Some(index) => found.push((index, *value)),
                    // A store through something that is not a plain index into
                    // a global. Recorded as a store all the same, so that it
                    // makes the block ineligible rather than invisible.
                    None => found.push((*pointer, *value)),
                },
                Statement::If { accept, reject, .. } => {
                    walk(accept, exprs, found);
                    walk(reject, exprs, found);
                }
                Statement::Loop {
                    body, continuing, ..
                } => {
                    walk(body, exprs, found);
                    walk(continuing, exprs, found);
                }
                _ => {}
            }
        }
    }
    let mut found = Vec::new();
    walk(body, exprs, &mut found);
    match found.as_slice() {
        [one] => Some(*one),
        _ => None,
    }
}

/// Classify an expression as a chain leaf, if that is what it is.
///
/// The index test is what keeps a broadcast out. `alibi` reads
/// `slopes[head]` and `snake` reads `alpha[c]`, both at a subscript derived
/// from the thread id rather than at the id itself; a chain step over either
/// would be an operand of a different shape, and calling it element-wise
/// would be the same confident guess this file exists to stop making.
fn chain_leaf(
    module: &Module,
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    store_index: Handle<Expression>,
) -> Option<ChainLeaf> {
    let Expression::Load { pointer } = exprs.try_get(handle)? else {
        return None;
    };
    match exprs.try_get(*pointer)? {
        // `x[idx]` — a whole tensor.
        Expression::Access { base, index } => {
            let global = access_base_global(exprs, *base)?;
            let gv = module.global_variables.try_get(global)?;
            if !matches!(gv.space, AddressSpace::Storage { .. }) {
                return None;
            }
            (*index == store_index).then_some(ChainLeaf::Tensor(global))
        }
        // `input_scale` — a uniform that is one scalar.
        Expression::GlobalVariable(global) => {
            let gv = module.global_variables.try_get(*global)?;
            if !matches!(gv.space, AddressSpace::Uniform) {
                return None;
            }
            let TypeInner::Scalar(scalar) = &module.types.try_get(gv.ty)?.inner else {
                return None;
            };
            Some(ChainLeaf::Scalar(ScalarBinding {
                name: gv
                    .name
                    .clone()
                    .unwrap_or_else(|| format!("scalar_{}", global.index())),
                elem_type: scalar_to_onnx_data_type(scalar),
            }))
        }
        // `params.a` — one scalar member of the uniform params struct.
        Expression::AccessIndex { base, index } => {
            let Expression::GlobalVariable(global) = exprs.try_get(*base)? else {
                return None;
            };
            let gv = module.global_variables.try_get(*global)?;
            if !matches!(gv.space, AddressSpace::Uniform) {
                return None;
            }
            let TypeInner::Struct { members, .. } = &module.types.try_get(gv.ty)?.inner else {
                return None;
            };
            let member = members.get(*index as usize)?;
            let TypeInner::Scalar(scalar) = &module.types.try_get(member.ty)?.inner else {
                return None;
            };
            Some(ChainLeaf::Scalar(ScalarBinding {
                name: member.name.clone()?,
                elem_type: scalar_to_onnx_data_type(scalar),
            }))
        }
        _ => None,
    }
}

/// Read a per-element expression as a chain, or refuse it.
///
/// The chain has to be *linear*: at every node exactly one operand carries the
/// rest of the expression and the other is a leaf. That is the whole
/// difference between a sequence of graph nodes and a tree, and it is what
/// separates `axpy` from `snake`, whose multiply has a reciprocal on one side
/// and a squared sine on the other.
fn decompose_chain(
    module: &Module,
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    store_index: Handle<Expression>,
    depth: u32,
) -> Option<RawChain> {
    if depth == 0 {
        return None;
    }
    // A tensor on its own is a chain of no steps. A scalar is not: the result
    // is a tensor, so the chain has to start from one.
    if let Some(leaf) = chain_leaf(module, exprs, handle, store_index) {
        return match leaf {
            ChainLeaf::Tensor(global) => Some(RawChain {
                base: global,
                cast: None,
                steps: Vec::new(),
            }),
            ChainLeaf::Scalar(_) => None,
        };
    }
    match exprs.try_get(handle)? {
        // `f32(input[i])` — a conversion of the tensor the chain starts from.
        Expression::As {
            expr,
            kind,
            convert: Some(width),
        } => {
            let inner = decompose_chain(module, exprs, *expr, store_index, depth - 1)?;
            // A conversion partway along would need a type per step, and
            // nothing here carries one.
            if inner.cast.is_some() || !inner.steps.is_empty() {
                return None;
            }
            Some(RawChain {
                cast: Some(scalar_to_onnx_data_type(&Scalar {
                    kind: *kind,
                    width: *width,
                })),
                ..inner
            })
        }
        Expression::Binary { op, left, right } => {
            let ew = match op {
                BinaryOp::Add => ElementWiseOp::Add,
                BinaryOp::Subtract => ElementWiseOp::Sub,
                BinaryOp::Multiply => ElementWiseOp::Mul,
                BinaryOp::Divide => ElementWiseOp::Div,
                _ => return None,
            };
            decompose_chain_binary(module, exprs, ew, *left, *right, store_index, depth - 1)
        }
        // What FMA fusion leaves behind. `fma(a, b, c)` is `c + a * b`, and it
        // has to come apart into exactly what the unfused spelling gives, or
        // the classifier answers differently at `-O0` and `-O1` — the failure
        // `e2e_opt_invariance` exists to catch.
        Expression::Math {
            fun: MathFunction::Fma,
            arg,
            arg1: Some(factor),
            arg2: Some(addend),
            ..
        } => {
            let mut chain = decompose_chain_binary(
                module,
                exprs,
                ElementWiseOp::Mul,
                *arg,
                *factor,
                store_index,
                depth - 1,
            )?;
            let addend = chain_leaf(module, exprs, *addend, store_index)?;
            chain.steps.push((ElementWiseOp::Add, addend));
            Some(chain)
        }
        _ => None,
    }
}

/// One binary node of a chain: one side continues it, the other is the operand.
#[allow(clippy::collapsible_if)] // nested if-let for MSRV 1.87 compat (no let chains)
fn decompose_chain_binary(
    module: &Module,
    exprs: &Arena<Expression>,
    op: ElementWiseOp,
    left: Handle<Expression>,
    right: Handle<Expression>,
    store_index: Handle<Expression>,
    depth: u32,
) -> Option<RawChain> {
    // The left operand is tried as the accumulator first: it is the only
    // arrangement that keeps `a - b` meaning `a - b`.
    if let Some(operand) = chain_leaf(module, exprs, right, store_index) {
        if let Some(mut chain) = decompose_chain(module, exprs, left, store_index, depth) {
            chain.steps.push((op, operand));
            return Some(chain);
        }
    }
    // The accumulator on the right is only the same expression when the
    // operation does not care which side its operands are on.
    if op.is_commutative() {
        if let Some(operand) = chain_leaf(module, exprs, left, store_index) {
            if let Some(mut chain) = decompose_chain(module, exprs, right, store_index, depth) {
                chain.steps.push((op, operand));
                return Some(chain);
            }
        }
    }
    None
}

/// Recognise a kernel whose stored value is an element-wise chain.
///
/// Returns `None` — leaving the caller's refusal in place — for anything that
/// is not one, including the chains [`KernelPattern::ElementWise`] already
/// describes: one binary op over two whole tensors is an `Add`, and two names
/// for one thing is how a pattern starts being emitted where it is not meant.
fn match_elementwise_chain(
    module: &Module,
    body: &[Statement],
    exprs: &Arena<Expression>,
    outputs: &[(Handle<GlobalVariable>, &GlobalVariable)],
    shape_names: &[String],
) -> Option<KernelPattern> {
    let [(out_handle, out_gv)] = outputs else {
        return None;
    };
    // Element *i* of the result has to be built from element *i*, and written
    // in one place. A chain that moved values would be a different operator.
    if !stores_at_the_index_it_read(body, exprs) {
        return None;
    }
    let (store_index, value) = find_lone_store(body, exprs)?;
    let chain = decompose_chain(module, exprs, value, store_index, 32)?;

    let has_scalar = chain
        .steps
        .iter()
        .any(|(_, leaf)| matches!(leaf, ChainLeaf::Scalar(_)));
    if chain.steps.is_empty() || (chain.cast.is_none() && chain.steps.len() == 1 && !has_scalar) {
        return None;
    }

    let base_gv = module.global_variables.try_get(chain.base)?;
    let base = make_binding(module, chain.base, base_gv, TensorRole::Input);
    let mut steps = Vec::with_capacity(chain.steps.len());
    let mut reads_its_own_output = chain.base == *out_handle;
    for (op, leaf) in chain.steps {
        let operand = match leaf {
            ChainLeaf::Tensor(global) => {
                reads_its_own_output |= global == *out_handle;
                let gv = module.global_variables.try_get(global)?;
                ChainOperand::Tensor(make_binding(module, global, gv, TensorRole::Input))
            }
            ChainLeaf::Scalar(scalar) => ChainOperand::Scalar(scalar),
        };
        steps.push(ChainStep { op, operand });
    }

    // `axpy`'s in-place entry point binds `y` `read_write` and computes
    // `y = y + a*x`, so the same buffer is read and written. The arithmetic is
    // the out-of-place kernel's exactly; only the allocation differs. A graph
    // needs two names for the two values, so the result takes a suffix and the
    // operand keeps the buffer's own name.
    let mut output = make_binding(module, *out_handle, out_gv, TensorRole::Output);
    if reads_its_own_output {
        output.name = format!("{}_out", output.name);
    }

    Some(KernelPattern::ElementWiseChain {
        base,
        cast: chain.cast,
        steps,
        output,
        dim_name: shape_names.first().cloned().unwrap_or_else(|| "N".into()),
    })
}

/// Search for a Store whose value is a Math expression, detecting activation type.
fn find_store_activation(body: &[Statement], exprs: &Arena<Expression>) -> Option<ActivationOp> {
    for stmt in body {
        match stmt {
            Statement::Store { value, .. } => {
                if let Some(act) = classify_activation_expr(exprs, *value) {
                    return Some(act);
                }
            }
            Statement::If { accept, reject, .. } => {
                if let Some(op) = find_store_activation(accept, exprs) {
                    return Some(op);
                }
                if let Some(op) = find_store_activation(reject, exprs) {
                    return Some(op);
                }
            }
            _ => {}
        }
    }
    None
}

/// Classify an expression as an activation function.
fn classify_activation_expr(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
) -> Option<ActivationOp> {
    match exprs.try_get(handle)? {
        // max(x, 0) → ReLU
        Expression::Math {
            fun: MathFunction::Max,
            ..
        } => Some(ActivationOp::Relu),
        // tanh(x) → Tanh (standalone, not part of a Multiply)
        Expression::Math {
            fun: MathFunction::Tanh,
            ..
        } => Some(ActivationOp::Tanh),
        // Multiply patterns: GELU, SiLU, Mish
        Expression::Binary {
            op: BinaryOp::Multiply,
            left,
            right,
            ..
        } => {
            let has_tanh_left = contains_math_fun(exprs, *left, MathFunction::Tanh);
            let has_tanh_right = contains_math_fun(exprs, *right, MathFunction::Tanh);
            let has_exp_left = contains_math_fun(exprs, *left, MathFunction::Exp);
            let has_exp_right = contains_math_fun(exprs, *right, MathFunction::Exp);

            // GELU: x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715*x^3)))
            // Detected as Multiply where one sub-expression contains Tanh
            // but NOT Exp (to distinguish from Mish which has tanh+exp)
            if (has_tanh_left || has_tanh_right) && !has_exp_left && !has_exp_right {
                return Some(ActivationOp::Gelu);
            }

            // Mish: x * tanh(ln(1 + exp(x))) = x * tanh(softplus(x))
            // Detected as Multiply containing both Tanh and Exp
            if (has_tanh_left || has_tanh_right) && (has_exp_left || has_exp_right) {
                return Some(ActivationOp::Mish);
            }

            // SiLU: x * sigmoid(x) = x * (1 / (1 + exp(-x)))
            // Detected as Multiply where one sub-expression contains Divide+Exp (sigmoid pattern)
            if has_exp_left || has_exp_right {
                let has_div_left = contains_binary_op(exprs, *left, BinaryOp::Divide);
                let has_div_right = contains_binary_op(exprs, *right, BinaryOp::Divide);
                if has_div_left || has_div_right {
                    return Some(ActivationOp::Silu);
                }
            }

            // Not a recognized activation multiply, try recursing
            classify_activation_expr(exprs, *left)
                .or_else(|| classify_activation_expr(exprs, *right))
        }
        // 1/(1+exp(-x)) → Sigmoid: detected as Divide whose right is Add(1, Exp(Negate(x)))
        // exp(x)/sum(exp(x)) → Softmax: detected as Divide with Exp on left
        Expression::Binary {
            op: BinaryOp::Divide,
            left,
            right,
            ..
        } => {
            if contains_math_fun(exprs, *left, MathFunction::Exp) {
                // Softmax: exp(x) / sum(exp(x))
                Some(ActivationOp::Softmax)
            } else if contains_math_fun(exprs, *right, MathFunction::Exp) {
                // Sigmoid: 1 / (1 + exp(-x))
                Some(ActivationOp::Sigmoid)
            } else {
                None
            }
        }
        _ => None,
    }
}

/// Check if an expression (recursively) contains a specific binary operation.
fn contains_binary_op(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    target: BinaryOp,
) -> bool {
    match exprs.try_get(handle) {
        Some(Expression::Binary {
            op, left, right, ..
        }) => {
            *op == target
                || contains_binary_op(exprs, *left, target)
                || contains_binary_op(exprs, *right, target)
        }
        Some(Expression::Math {
            fun,
            arg,
            arg1,
            arg2,
            arg3,
        }) => {
            // `fma(a, b, c)` is a multiply and an add wearing one name. Say so,
            // or a pattern that asks "is there a multiply in here?" changes its
            // answer the moment FMA fusion runs.
            (*fun == MathFunction::Fma && matches!(target, BinaryOp::Multiply | BinaryOp::Add))
                || contains_binary_op(exprs, *arg, target)
                || [*arg1, *arg2, *arg3]
                    .into_iter()
                    .flatten()
                    .any(|h| contains_binary_op(exprs, h, target))
        }
        Some(Expression::Unary { expr, .. }) => contains_binary_op(exprs, *expr, target),
        _ => false,
    }
}

/// Check if an expression (recursively) contains a specific math function.
fn contains_math_fun(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    target: MathFunction,
) -> bool {
    match exprs.try_get(handle) {
        Some(Expression::Math {
            fun,
            arg,
            arg1,
            arg2,
            arg3,
        }) => {
            // All operands, not just the first: `fma` carries two more, and
            // anything fused into them was invisible to this walk before.
            *fun == target
                || contains_math_fun(exprs, *arg, target)
                || [*arg1, *arg2, *arg3]
                    .into_iter()
                    .flatten()
                    .any(|h| contains_math_fun(exprs, h, target))
        }
        Some(Expression::Binary { left, right, .. }) => {
            contains_math_fun(exprs, *left, target) || contains_math_fun(exprs, *right, target)
        }
        Some(Expression::Unary { expr, .. }) => contains_math_fun(exprs, *expr, target),
        _ => false,
    }
}

// ---------------------------------------------------------------------------
// Permutation recognition
// ---------------------------------------------------------------------------

/// How many divisions were applied to the flat invocation id to produce this
/// value, or `None` if it is not derived from the id by the usual
/// decomposition.
///
/// A kernel that reindexes a tensor peels dimensions off the flat id from the
/// fastest-varying end: `d = id % D`, `rest = id / D`, `i = rest % N`, and so
/// on. The number of divisions on the path is therefore the axis's distance
/// from the fastest end, which is all that is needed to place it.
fn id_division_depth(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    depth: u32,
) -> Option<u32> {
    if depth == 0 {
        return None;
    }
    match exprs.try_get(handle)? {
        // The flat id itself: `gid.x`.
        Expression::AccessIndex { base, index: 0 } => {
            matches!(exprs.try_get(*base)?, Expression::FunctionArgument(_)).then_some(0)
        }
        Expression::Binary { op, left, .. } => match op {
            BinaryOp::Modulo => id_division_depth(exprs, *left, depth - 1),
            BinaryOp::Divide => id_division_depth(exprs, *left, depth - 1).map(|d| d + 1),
            _ => None,
        },
        _ => None,
    }
}

/// Decompose a store index built as `((c_n * f_n + c_n-1) * f_n-1 + …) + c_0`
/// into its components, fastest-varying first.
///
/// Both spellings are accepted: the multiply-add as written, and the `fma`
/// the optimizer rewrites it into. Missing the second cost a whole pass's
/// worth of kernels once already.
fn flatten_index_components(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    out: &mut Vec<Handle<Expression>>,
    depth: u32,
) {
    if depth == 0 {
        return;
    }
    match exprs.try_get(handle) {
        Some(Expression::Math {
            fun: MathFunction::Fma,
            arg,
            arg2: Some(addend),
            ..
        }) => {
            out.push(*addend);
            flatten_index_components(exprs, *arg, out, depth - 1);
        }
        Some(Expression::Binary {
            op: BinaryOp::Add,
            left,
            right,
        }) => {
            // `a * f + c`: the addend is the component, the product carries
            // the rest.
            let (product, addend) = match exprs.try_get(*left) {
                Some(Expression::Binary {
                    op: BinaryOp::Multiply,
                    ..
                }) => (*left, *right),
                _ => (*right, *left),
            };
            out.push(addend);
            if let Some(Expression::Binary {
                op: BinaryOp::Multiply,
                left: inner,
                ..
            }) = exprs.try_get(product)
            {
                flatten_index_components(exprs, *inner, out, depth - 1);
            }
        }
        _ => out.push(handle),
    }
}

/// Recognise a kernel that copies `input[id]` to a position built by taking
/// the flat id apart and putting it back together in a different order.
///
/// Every backend already lowers `KernelPattern::Transpose`; until now nothing
/// produced one, so `permute` was refused for want of a recogniser rather
/// than for want of a lowering.
fn detect_permutation(
    body: &[Statement],
    exprs: &Arena<Expression>,
    input: Handle<GlobalVariable>,
    output: Handle<GlobalVariable>,
) -> Option<Vec<i64>> {
    fn find_store(
        body: &[Statement],
        exprs: &Arena<Expression>,
        input: Handle<GlobalVariable>,
        output: Handle<GlobalVariable>,
    ) -> Option<(Handle<Expression>, Handle<Expression>)> {
        for stmt in body {
            match stmt {
                Statement::Store { pointer, value } => {
                    let store_index = match exprs.try_get(*pointer) {
                        Some(Expression::Access { base, index })
                            if matches!(
                                exprs.try_get(*base),
                                Some(Expression::GlobalVariable(g)) if *g == output
                            ) =>
                        {
                            *index
                        }
                        _ => continue,
                    };
                    let load_index = match exprs.try_get(*value) {
                        Some(Expression::Load { pointer }) => match exprs.try_get(*pointer) {
                            Some(Expression::Access { base, index })
                                if matches!(
                                    exprs.try_get(*base),
                                    Some(Expression::GlobalVariable(g)) if *g == input
                                ) =>
                            {
                                *index
                            }
                            _ => continue,
                        },
                        _ => continue,
                    };
                    return Some((store_index, load_index));
                }
                Statement::If { accept, reject, .. } => {
                    if let Some(found) = find_store(accept, exprs, input, output)
                        .or_else(|| find_store(reject, exprs, input, output))
                    {
                        return Some(found);
                    }
                }
                _ => {}
            }
        }
        None
    }

    let (store_index, load_index) = find_store(body, exprs, input, output)?;

    // The read has to be the plain flat id; anything else is a gather, not a
    // permutation, and calling it Transpose would be the same kind of
    // confident guess this recogniser exists to replace.
    if id_division_depth(exprs, load_index, 32)? != 0 {
        return None;
    }

    let mut components = Vec::new();
    flatten_index_components(exprs, store_index, &mut components, 32);
    if components.len() < 2 {
        return None;
    }

    // Components come out fastest-first; a permutation is stated slowest-first.
    let rank = components.len();
    let mut perm = Vec::with_capacity(rank);
    for handle in components.iter().rev() {
        let depth = id_division_depth(exprs, *handle, 32)?;
        if depth as usize >= rank {
            return None;
        }
        perm.push((rank - 1 - depth as usize) as i64);
    }

    // A genuine permutation uses each axis exactly once, and one that is the
    // identity is a copy rather than a transpose.
    let mut seen = perm.clone();
    seen.sort_unstable();
    if seen != (0..rank as i64).collect::<Vec<_>>() {
        return None;
    }
    if perm.iter().enumerate().all(|(i, p)| i as i64 == *p) {
        return None;
    }
    Some(perm)
}

/// Which outputs a kernel writes, and which of them it *computes* rather than
/// copies out of `input`.
///
/// A Split hands back slices of its input: every element it writes is an
/// element it read, so every store into an output is a bare `Load` from the
/// input and nothing else. Anything arithmetic in the stored value — a scale,
/// a rounding, a cast to an index — means the kernel is producing a new tensor
/// rather than partitioning the one it was given, whatever else it may be
/// doing, and no `Split` describes that.
///
/// Returns `(written, computed)`, both in first-seen order. An output that is
/// never written at all is absent from `written`, which is also disqualifying:
/// a slice that stores nothing is not a slice.
fn survey_output_stores(
    body: &[Statement],
    exprs: &Arena<Expression>,
    input: Handle<GlobalVariable>,
    outputs: &[Handle<GlobalVariable>],
) -> (Vec<Handle<GlobalVariable>>, Vec<Handle<GlobalVariable>>) {
    fn walk(
        body: &[Statement],
        exprs: &Arena<Expression>,
        input: Handle<GlobalVariable>,
        outputs: &[Handle<GlobalVariable>],
        written: &mut Vec<Handle<GlobalVariable>>,
        computed: &mut Vec<Handle<GlobalVariable>>,
    ) {
        for stmt in body {
            match stmt {
                Statement::Store { pointer, value } => {
                    let Some(target) = access_base_global(exprs, *pointer) else {
                        continue;
                    };
                    if !outputs.contains(&target) {
                        continue;
                    }
                    if !written.contains(&target) {
                        written.push(target);
                    }
                    let copied = match exprs.try_get(*value) {
                        Some(Expression::Load { pointer }) => {
                            access_base_global(exprs, *pointer) == Some(input)
                        }
                        _ => false,
                    };
                    if !copied && !computed.contains(&target) {
                        computed.push(target);
                    }
                }
                Statement::If { accept, reject, .. } => {
                    walk(accept, exprs, input, outputs, written, computed);
                    walk(reject, exprs, input, outputs, written, computed);
                }
                Statement::Loop {
                    body, continuing, ..
                } => {
                    walk(body, exprs, input, outputs, written, computed);
                    walk(continuing, exprs, input, outputs, written, computed);
                }
                _ => {}
            }
        }
    }

    let mut written = Vec::new();
    let mut computed = Vec::new();
    walk(body, exprs, input, outputs, &mut written, &mut computed);
    (written, computed)
}

/// Say, in the refusal, what was actually seen.
///
/// "Unrecognized pattern" tells the caller nothing about their kernel. These
/// name the outputs that disqualified it and why, so the message is checkable
/// against the source in front of them.
fn multi_output_refusal(
    module: &Module,
    outputs: &[(Handle<GlobalVariable>, &GlobalVariable)],
    input: &GlobalVariable,
    written: &[Handle<GlobalVariable>],
    computed: &[Handle<GlobalVariable>],
) -> String {
    let name_of = |handle: Handle<GlobalVariable>| -> String {
        outputs
            .iter()
            .find(|(h, _)| *h == handle)
            .and_then(|(_, gv)| gv.name.clone())
            .unwrap_or_else(|| "an output".into())
    };
    let quoted = |handles: &[Handle<GlobalVariable>]| -> String {
        handles
            .iter()
            .map(|h| format!("'{}'", name_of(*h)))
            .collect::<Vec<_>>()
            .join(" and ")
    };

    let unwritten: Vec<Handle<GlobalVariable>> = outputs
        .iter()
        .map(|(h, _)| *h)
        .filter(|h| !written.contains(h))
        .collect();

    let mut reason = format!(
        "1 input and {} outputs, but this does not slice its input: ",
        outputs.len()
    );
    if !computed.is_empty() {
        reason.push_str(&format!(
            "{} {} written with a computed value rather than a copy of the input",
            quoted(computed),
            if computed.len() == 1 { "is" } else { "are" },
        ));
        if !unwritten.is_empty() {
            reason.push_str(", and ");
        }
    }
    if !unwritten.is_empty() {
        reason.push_str(&format!(
            "{} {} never stored to",
            quoted(&unwritten),
            if unwritten.len() == 1 { "is" } else { "are" },
        ));
    }
    if computed.is_empty() && unwritten.is_empty() {
        // Copies throughout, but no boundary test to split on.
        reason.push_str("there is no comparison marking where one output ends and the next begins");
    }

    // A slice keeps its element type. When it changes, the kernel is
    // converting, and saying so points at the half of the problem a Split
    // could never have carried anyway.
    let input_elem = resolve_array_elem_type(module, input.ty);
    if outputs
        .iter()
        .any(|(_, gv)| resolve_array_elem_type(module, gv.ty) != input_elem)
    {
        reason.push_str(
            ". The outputs do not all share the input's element type either, and a slice \
             never changes it",
        );
    }

    reason.push_str(
        ". Splitting is the only single-input, multi-output op these backends lower, \
         so there is nothing here for one to carry",
    );
    reason
}

/// Positive recognition of a softmax: `exp(x - max) / Σ exp(x - max)`.
///
/// Two facts have to hold together, and neither is enough on its own:
///
///   - the value written to the output is an `exp` over a division — the
///     normalised numerator, whether it is spelled `exp(..) / sum` or
///     `exp(..) * (1.0 / sum)`, which are the same expression tree with the
///     `Divide` in a different place;
///   - some store inside a loop adds an `exp` into a running total — the
///     denominator being built, one element of the row at a time.
///
/// The second is what separates a softmax from a hand-written sigmoid,
/// `1.0 / (1.0 + exp(-x))`, which also stores an `exp` over a `Divide`. A
/// sigmoid accumulates no exponentials because it has nothing to sum over.
/// The first is what separates it from a plain `ReduceSum` of exponentials,
/// which builds the same total and then writes the total.
///
/// What is deliberately *not* required is the max subtraction. It is there in
/// every stable softmax and in this vendor kernel, but it is an implementation
/// choice about overflow rather than part of the operator, and a kernel that
/// skips it is still a softmax.
fn detect_softmax(
    body: &[Statement],
    exprs: &Arena<Expression>,
    output: Handle<GlobalVariable>,
) -> bool {
    stores_a_normalized_exp(body, exprs, output) && sums_exponentials_in_a_loop(body, exprs, false)
}

/// A store into `output` whose value contains both an `exp` and a division.
fn stores_a_normalized_exp(
    body: &[Statement],
    exprs: &Arena<Expression>,
    output: Handle<GlobalVariable>,
) -> bool {
    for stmt in body {
        match stmt {
            Statement::Store { pointer, value } => {
                if access_base_global(exprs, *pointer) == Some(output)
                    && contains_math_fun(exprs, *value, MathFunction::Exp)
                    && contains_binary_op(exprs, *value, BinaryOp::Divide)
                {
                    return true;
                }
            }
            Statement::If { accept, reject, .. } => {
                if stores_a_normalized_exp(accept, exprs, output)
                    || stores_a_normalized_exp(reject, exprs, output)
                {
                    return true;
                }
            }
            Statement::Loop {
                body, continuing, ..
            } if stores_a_normalized_exp(body, exprs, output)
                || stores_a_normalized_exp(continuing, exprs, output) =>
            {
                return true;
            }
            _ => {}
        }
    }
    false
}

/// An `acc = acc + exp(..)` — an add, containing an exponential, executed
/// inside a loop. The target is not checked: a workgroup slot, a local and a
/// storage scratch buffer are all the same accumulator here.
fn sums_exponentials_in_a_loop(
    body: &[Statement],
    exprs: &Arena<Expression>,
    in_loop: bool,
) -> bool {
    for stmt in body {
        match stmt {
            Statement::Store { value, .. } => {
                if in_loop
                    && matches!(
                        exprs.try_get(*value),
                        Some(Expression::Binary {
                            op: BinaryOp::Add,
                            ..
                        })
                    )
                    && contains_math_fun(exprs, *value, MathFunction::Exp)
                {
                    return true;
                }
            }
            Statement::If { accept, reject, .. } => {
                if sums_exponentials_in_a_loop(accept, exprs, in_loop)
                    || sums_exponentials_in_a_loop(reject, exprs, in_loop)
                {
                    return true;
                }
            }
            Statement::Loop {
                body, continuing, ..
            } if sums_exponentials_in_a_loop(body, exprs, true)
                || sums_exponentials_in_a_loop(continuing, exprs, true) =>
            {
                return true;
            }
            _ => {}
        }
    }
    false
}

/// Whether a loop appears inside another loop.
///
/// A normalization reduces once over a row: its passes are sequential loops,
/// never nested. A filterbank contracts a matrix against a spectrum, so its
/// loops nest. That is the difference between `mel`, which was reported as
/// BatchNormalization because it has three inputs and a square root, and a
/// normalization that actually is one.
fn has_nested_loop(body: &[Statement]) -> bool {
    fn inside_loop(body: &[Statement]) -> bool {
        body.iter().any(|s| match s {
            Statement::Loop { .. } => true,
            Statement::If { accept, reject, .. } => inside_loop(accept) || inside_loop(reject),
            _ => false,
        })
    }
    body.iter().any(|s| match s {
        Statement::Loop {
            body, continuing, ..
        } => inside_loop(body) || inside_loop(continuing) || has_nested_loop(body),
        Statement::If { accept, reject, .. } => has_nested_loop(accept) || has_nested_loop(reject),
        _ => false,
    })
}

/// Check if a block contains a Store whose value (or sub-expr) uses a specific Math function.
/// The activation a kernel applies to the value it stores, if any.
///
/// A convolution that writes `max(sum, 0.0)` is a convolution *and* a ReLU, and
/// the classifier used to return `Conv2D` and drop the `max` on the floor. The
/// emitted model then loaded, was accelerated, and computed the wrong thing --
/// which is not a shortcoming of the operator table but a graph that lies.
///
/// Only `max` and `tanh` are recognised, and deliberately: those are the two
/// TFLite can fuse into a convolution and the two whose expression shape cannot
/// be mistaken for something else. GELU, SiLU and Mish are all *multiplies*,
/// and a convolution's own accumulation is a multiply, so recognising them here
/// would risk reading an accumulator as an activation. Narrowing is the same
/// choice made for stride extraction and for the same reason: missing an
/// activation leaves today's behaviour, inventing one corrupts a graph.
fn store_value_activation(body: &[Statement], exprs: &Arena<Expression>) -> Option<ActivationOp> {
    for stmt in body {
        match stmt {
            Statement::Store { value, .. } => match exprs.try_get(*value) {
                Some(Expression::Math {
                    fun: MathFunction::Max,
                    ..
                }) => return Some(ActivationOp::Relu),
                Some(Expression::Math {
                    fun: MathFunction::Tanh,
                    ..
                }) => return Some(ActivationOp::Tanh),
                _ => {}
            },
            Statement::If { accept, reject, .. } => {
                if let Some(act) = store_value_activation(accept, exprs) {
                    return Some(act);
                }
                if let Some(act) = store_value_activation(reject, exprs) {
                    return Some(act);
                }
            }
            Statement::Loop {
                body, continuing, ..
            } => {
                if let Some(act) = store_value_activation(body, exprs) {
                    return Some(act);
                }
                if let Some(act) = store_value_activation(continuing, exprs) {
                    return Some(act);
                }
            }
            _ => {}
        }
    }
    None
}

fn find_store_math_fun(
    body: &[Statement],
    exprs: &Arena<Expression>,
    target: MathFunction,
) -> bool {
    for stmt in body {
        match stmt {
            Statement::Store { value, .. } if contains_math_fun(exprs, *value, target) => {
                return true;
            }
            Statement::If { accept, reject, .. }
                if find_store_math_fun(accept, exprs, target)
                    || find_store_math_fun(reject, exprs, target) =>
            {
                return true;
            }
            Statement::Loop {
                body, continuing, ..
            } if find_store_math_fun(body, exprs, target)
                || find_store_math_fun(continuing, exprs, target) =>
            {
                return true;
            }
            _ => {}
        }
    }
    false
}

/// Detect the reduce operation type from the loop body.
fn detect_reduce_op(body: &[Statement], exprs: &Arena<Expression>) -> ReduceOp {
    if find_store_math_fun(body, exprs, MathFunction::Max) {
        ReduceOp::Max
    } else if find_store_math_fun(body, exprs, MathFunction::Min) {
        ReduceOp::Min
    } else if find_store_binary_divide(body, exprs) {
        // Sum followed by divide → Mean
        ReduceOp::Mean
    } else {
        // Default: sum accumulation (binary add in loop)
        ReduceOp::Sum
    }
}

/// Check if a block contains a Store whose value is a Divide expression.
fn find_store_binary_divide(body: &[Statement], exprs: &Arena<Expression>) -> bool {
    for stmt in body {
        match stmt {
            Statement::Store { value, .. } => {
                if let Some(Expression::Binary {
                    op: BinaryOp::Divide,
                    ..
                }) = exprs.try_get(*value)
                {
                    return true;
                }
            }
            Statement::If { accept, reject, .. }
                if find_store_binary_divide(accept, exprs)
                    || find_store_binary_divide(reject, exprs) =>
            {
                return true;
            }
            Statement::Loop {
                body, continuing, ..
            } if find_store_binary_divide(body, exprs)
                || find_store_binary_divide(continuing, exprs) =>
            {
                return true;
            }
            _ => {}
        }
    }
    false
}

// ---------------------------------------------------------------------------
// Data-dependent addressing
// ---------------------------------------------------------------------------
//
// Every pattern this file recognises reads and writes at positions computed
// from the thread id: a matmul's `row * K + k`, a convolution's window, an
// element-wise op's `idx`. The kernels that do not are the ones that take an
// address *out of a buffer* — a gather on the read side, a scatter on the
// write side — and that is visible in the expression graph without counting
// anything.

/// What an index-dependent read addresses.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum IndexedRead {
    /// `out[i] = data[idx[i]]` — the loaded index *is* the address, so each
    /// index names one element. This is what `KernelPattern::Gather` lowers
    /// to: a flat buffer, one element per index.
    Element,
    /// `out[i] = data[idx[n] * W + d]` — the loaded index is scaled by a width
    /// before it becomes an address, so each index names a whole block.
    Block,
}

/// Whether the kernel writes each output element from the same position it
/// read, at a single position per invocation.
///
/// Two rules, and the second was learned by breaking the first. Element *i* of
/// the output has to be a function of element *i* of at least one input — a
/// dequantizing transpose writes at `col * width + row` and fails this — but
/// *not* of element *i* of every input, because a broadcast is element-wise
/// and reads its second operand at a row index.
///
/// And an element-wise invocation writes one place. rope writes a rotated
/// pair at `base` and `base + 1`, each from both, which is a movement of
/// values between positions however element-wise the arithmetic looks.
fn stores_at_the_index_it_read(body: &[Statement], exprs: &Arena<Expression>) -> bool {
    fn walk(
        body: &[Statement],
        exprs: &Arena<Expression>,
        ok: &mut bool,
        seen: &mut Vec<Handle<Expression>>,
    ) {
        for stmt in body {
            match stmt {
                Statement::Store { pointer, value } => {
                    let Some(store_index) = access_index(exprs, *pointer) else {
                        *ok = false;
                        return;
                    };
                    // Branches of an `if` are alternatives, not a sequence, so
                    // the same index appearing in both is still one write.
                    if seen.iter().all(|i| *i != store_index) {
                        seen.push(store_index);
                    }
                    let mut reads = Vec::new();
                    collect_load_indices(exprs, *value, &mut reads, 32);
                    // Reading nothing is a generator; reading only elsewhere
                    // is a move.
                    if !reads.contains(&store_index) {
                        *ok = false;
                        return;
                    }
                }
                Statement::If { accept, reject, .. } => {
                    walk(accept, exprs, ok, seen);
                    walk(reject, exprs, ok, seen);
                }
                Statement::Loop {
                    body, continuing, ..
                } => {
                    walk(body, exprs, ok, seen);
                    walk(continuing, exprs, ok, seen);
                }
                _ => {}
            }
        }
    }
    let mut ok = true;
    let mut seen = Vec::new();
    walk(body, exprs, &mut ok, &mut seen);
    ok && seen.len() == 1
}

/// The index expression of an `Access` into a global, if that is what this is.
fn access_index(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
) -> Option<Handle<Expression>> {
    match exprs.try_get(handle)? {
        Expression::Access { base, index } => {
            matches!(exprs.try_get(*base)?, Expression::GlobalVariable(_)).then_some(*index)
        }
        _ => None,
    }
}

/// Every index a storage load in this expression reads at.
fn collect_load_indices(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    out: &mut Vec<Handle<Expression>>,
    depth: u32,
) {
    if depth == 0 {
        return;
    }
    match exprs.try_get(handle) {
        Some(Expression::Load { pointer }) => {
            if let Some(index) = access_index(exprs, *pointer) {
                out.push(index);
            }
        }
        Some(Expression::Binary { left, right, .. }) => {
            collect_load_indices(exprs, *left, out, depth - 1);
            collect_load_indices(exprs, *right, out, depth - 1);
        }
        Some(Expression::Unary { expr, .. }) => collect_load_indices(exprs, *expr, out, depth - 1),
        // A conversion is transparent to *where* the value came from.
        // `f32(input[i])` reads element `i`, and leaving `As` out of this walk
        // made `dequantize` look like a kernel that read nothing at all.
        Some(Expression::As { expr, .. }) => collect_load_indices(exprs, *expr, out, depth - 1),
        Some(Expression::Math {
            arg,
            arg1,
            arg2,
            arg3,
            ..
        }) => {
            collect_load_indices(exprs, *arg, out, depth - 1);
            for a in [arg1, arg2, arg3].into_iter().flatten() {
                collect_load_indices(exprs, *a, out, depth - 1);
            }
        }
        _ => {}
    }
}

/// Walk a pointer expression down to the global variable it addresses.
fn access_base_global(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
) -> Option<Handle<GlobalVariable>> {
    let mut handle = handle;
    for _ in 0..32 {
        match exprs.try_get(handle)? {
            Expression::GlobalVariable(g) => return Some(*g),
            Expression::Access { base, .. } | Expression::AccessIndex { base, .. } => {
                handle = *base;
            }
            _ => return None,
        }
    }
    None
}

/// Peel the type casts off a value. `u32(indices[n])` is the index it loaded.
fn strip_casts(exprs: &Arena<Expression>, handle: Handle<Expression>) -> Handle<Expression> {
    let mut handle = handle;
    for _ in 0..8 {
        match exprs.try_get(handle) {
            Some(Expression::As { expr, .. }) => handle = *expr,
            _ => break,
        }
    }
    handle
}

/// If `handle` is a load out of `global`, the index expression it loaded from.
fn load_address(
    exprs: &Arena<Expression>,
    handle: Handle<Expression>,
    global: Handle<GlobalVariable>,
) -> Option<Handle<Expression>> {
    let Expression::Load { pointer } = exprs.try_get(handle)? else {
        return None;
    };
    let Expression::Access { base, index } = exprs.try_get(*pointer)? else {
        return None;
    };
    (access_base_global(exprs, *base) == Some(global)).then_some(*index)
}

/// Does any of `roots` load, anywhere inside it, from one of `wanted`?
///
/// Iterative with a visited set rather than the depth-limited recursion used
/// elsewhere in this file: an expression arena is a DAG, and an address built
/// from shared subexpressions is walked exponentially by the naive spelling.
fn reads_global(
    exprs: &Arena<Expression>,
    roots: &[Handle<Expression>],
    wanted: &[Handle<GlobalVariable>],
) -> bool {
    let mut seen = std::collections::HashSet::new();
    let mut stack: Vec<Handle<Expression>> = roots.to_vec();
    while let Some(handle) = stack.pop() {
        if !seen.insert(handle.index()) {
            continue;
        }
        let Some(expr) = exprs.try_get(handle) else {
            continue;
        };
        // Nested rather than a let-chain: those are stable from 1.88 and this
        // workspace's rust-version is 1.87.
        #[allow(clippy::collapsible_if)]
        if let Expression::Load { pointer } = expr {
            if access_base_global(exprs, *pointer).is_some_and(|g| wanted.contains(&g)) {
                return true;
            }
        }
        push_operands(expr, &mut stack);
    }
    false
}

/// Push every sub-expression of `expr` onto `stack`.
fn push_operands(expr: &Expression, stack: &mut Vec<Handle<Expression>>) {
    match expr {
        Expression::Load { pointer } => stack.push(*pointer),
        Expression::Access { base, index } => {
            stack.push(*base);
            stack.push(*index);
        }
        Expression::AccessIndex { base, .. } => stack.push(*base),
        Expression::Unary { expr, .. } => stack.push(*expr),
        Expression::Binary { left, right, .. } => {
            stack.push(*left);
            stack.push(*right);
        }
        Expression::Select {
            condition,
            accept,
            reject,
        } => {
            stack.push(*condition);
            stack.push(*accept);
            stack.push(*reject);
        }
        Expression::Math {
            arg,
            arg1,
            arg2,
            arg3,
            ..
        } => {
            stack.push(*arg);
            stack.extend([*arg1, *arg2, *arg3].into_iter().flatten());
        }
        Expression::As { expr, .. } => stack.push(*expr),
        Expression::Splat { value, .. } => stack.push(*value),
        Expression::Swizzle { vector, .. } => stack.push(*vector),
        Expression::Compose { components, .. } => stack.extend(components.iter().copied()),
        Expression::ArrayLength(handle) => stack.push(*handle),
        _ => {}
    }
}

/// Recognise a store into `output` whose value was loaded from `data` at an
/// address that was itself loaded from `indices` — a gather.
///
/// The two shapes are told apart because only one of them is representable:
/// `Gather` indexes a flat buffer and carries no row width, so a read whose
/// index is scaled by one is a different operator wearing the same name.
fn detect_indexed_read(
    body: &[Statement],
    exprs: &Arena<Expression>,
    data: Handle<GlobalVariable>,
    indices: Handle<GlobalVariable>,
    output: Handle<GlobalVariable>,
) -> Option<IndexedRead> {
    for stmt in body {
        match stmt {
            Statement::Store { pointer, value } => {
                if access_base_global(exprs, *pointer) != Some(output) {
                    continue;
                }
                let Some(address) = load_address(exprs, strip_casts(exprs, *value), data) else {
                    continue;
                };
                if !reads_global(exprs, &[address], &[indices]) {
                    continue;
                }
                return Some(
                    if load_address(exprs, strip_casts(exprs, address), indices).is_some() {
                        IndexedRead::Element
                    } else {
                        IndexedRead::Block
                    },
                );
            }
            Statement::If { accept, reject, .. } => {
                if let Some(kind) = detect_indexed_read(accept, exprs, data, indices, output)
                    .or_else(|| detect_indexed_read(reject, exprs, data, indices, output))
                {
                    return Some(kind);
                }
            }
            Statement::Loop {
                body, continuing, ..
            } => {
                if let Some(kind) = detect_indexed_read(body, exprs, data, indices, output)
                    .or_else(|| detect_indexed_read(continuing, exprs, data, indices, output))
                {
                    return Some(kind);
                }
            }
            _ => {}
        }
    }
    None
}

/// Does a store into one of `outputs` write to an address that was loaded out
/// of one of `sources` — a scatter?
fn detect_indexed_write(
    body: &[Statement],
    exprs: &Arena<Expression>,
    sources: &[Handle<GlobalVariable>],
    outputs: &[Handle<GlobalVariable>],
) -> bool {
    for stmt in body {
        match stmt {
            Statement::Store { pointer, .. } => {
                let Some(Expression::Access { base, index }) = exprs.try_get(*pointer) else {
                    continue;
                };
                if !access_base_global(exprs, *base).is_some_and(|g| outputs.contains(&g)) {
                    continue;
                }
                if reads_global(exprs, &[*index], sources) {
                    return true;
                }
            }
            Statement::If { accept, reject, .. }
                if detect_indexed_write(accept, exprs, sources, outputs)
                    || detect_indexed_write(reject, exprs, sources, outputs) =>
            {
                return true;
            }
            Statement::Loop {
                body, continuing, ..
            } if detect_indexed_write(body, exprs, sources, outputs)
                || detect_indexed_write(continuing, exprs, sources, outputs) =>
            {
                return true;
            }
            _ => {}
        }
    }
    false
}

/// Is this buffer an array of atomics (or an atomic scalar)?
///
/// An atomic output is a statement about the operator: several invocations
/// reach the same slot and their contributions are combined there.
fn is_atomic_buffer(module: &Module, ty: Handle<Type>) -> bool {
    match &module.types[ty].inner {
        TypeInner::Atomic(_) => true,
        TypeInner::Array { base, .. } => matches!(&module.types[*base].inner, TypeInner::Atomic(_)),
        _ => false,
    }
}

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

    fn make_matmul_module() -> Module {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: vec![
                    StructMember {
                        name: Some("M".into()),
                        ty: u32_ty,
                        offset: 0,
                    },
                    StructMember {
                        name: Some("N".into()),
                        ty: u32_ty,
                        offset: 4,
                    },
                    StructMember {
                        name: Some("K".into()),
                        ty: u32_ty,
                        offset: 8,
                    },
                ],
                span: 12,
            },
        });

        module.global_variables.append(GlobalVariable {
            name: Some("a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("b".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("result".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 3,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        // Entry point with a loop in the body (triggers MatMul detection).
        let mut func = Function::new("main");
        func.body.push(Statement::Loop {
            body: vec![Statement::Break],
            continuing: vec![],
            break_if: None,
        });

        module.entry_points.push(EntryPoint {
            name: "main".into(),
            workgroup_size: [16, 16, 1],
            function: func,
        });

        module
    }

    fn make_elementwise_module(op: BinaryOp) -> Module {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: vec![StructMember {
                    name: Some("N".into()),
                    ty: u32_ty,
                    offset: 0,
                }],
                span: 4,
            },
        });

        let a_gv = module.global_variables.append(GlobalVariable {
            name: Some("a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        let b_gv = module.global_variables.append(GlobalVariable {
            name: Some("b".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        let out_gv = module.global_variables.append(GlobalVariable {
            name: Some("c".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 3,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        // `out[i] = a[i] op b[i]`.
        //
        // This used to combine two float literals and store the result through
        // a third, addressing nothing and reading nothing — it described a
        // shape no kernel has, and passed because classification did not look
        // at where the values came from.
        let mut func = Function::new("main");
        let idx = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let load_from = |func: &mut Function, gv| {
            let base = func.expressions.append(Expression::GlobalVariable(gv));
            let access = func
                .expressions
                .append(Expression::Access { base, index: idx });
            func.expressions
                .append(Expression::Load { pointer: access })
        };
        let left = load_from(&mut func, a_gv);
        let right = load_from(&mut func, b_gv);
        let binary = func
            .expressions
            .append(Expression::Binary { op, left, right });
        let out_base = func.expressions.append(Expression::GlobalVariable(out_gv));
        let ptr = func.expressions.append(Expression::Access {
            base: out_base,
            index: idx,
        });
        func.body.push(Statement::Store {
            pointer: ptr,
            value: binary,
        });

        module.entry_points.push(EntryPoint {
            name: "vecadd".into(),
            workgroup_size: [256, 1, 1],
            function: func,
        });

        module
    }

    fn make_activation_module(math_fun: MathFunction) -> Module {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: vec![StructMember {
                    name: Some("N".into()),
                    ty: u32_ty,
                    offset: 0,
                }],
                span: 4,
            },
        });

        // Single input
        module.global_variables.append(GlobalVariable {
            name: Some("a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        // Single output
        module.global_variables.append(GlobalVariable {
            name: Some("c".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        let mut func = Function::new("main");
        let arg = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let arg1 = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        let math = func.expressions.append(Expression::Math {
            fun: math_fun,
            arg,
            arg1: Some(arg1),
            arg2: None,
            arg3: None,
        });
        let ptr = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        func.body.push(Statement::Store {
            pointer: ptr,
            value: math,
        });

        module.entry_points.push(EntryPoint {
            name: "activation".into(),
            workgroup_size: [256, 1, 1],
            function: func,
        });

        module
    }

    fn make_reduce_module() -> Module {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: vec![StructMember {
                    name: Some("N".into()),
                    ty: u32_ty,
                    offset: 0,
                }],
                span: 4,
            },
        });

        module.global_variables.append(GlobalVariable {
            name: Some("a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("c".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        let mut func = Function::new("main");
        // A loop with a binary add → reduce sum
        let left = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        let right = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let add = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left,
            right,
        });
        let ptr = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        func.body.push(Statement::Loop {
            body: vec![
                Statement::Store {
                    pointer: ptr,
                    value: add,
                },
                Statement::Break,
            ],
            continuing: vec![],
            break_if: None,
        });

        module.entry_points.push(EntryPoint {
            name: "reduce".into(),
            workgroup_size: [256, 1, 1],
            function: func,
        });

        module
    }

    #[test]
    fn classify_matmul() {
        let module = make_matmul_module();
        let pattern = classify_entry_point(&module, 0).unwrap();
        match pattern {
            KernelPattern::MatMul {
                inputs,
                output,
                shape,
            } => {
                assert_eq!(inputs[0].name, "a");
                assert_eq!(inputs[1].name, "b");
                assert_eq!(output.name, "result");
                assert_eq!(inputs[0].elem_type, data_type::FLOAT);
                assert_eq!(shape.m, "M");
                assert_eq!(shape.n, "N");
                assert_eq!(shape.k, "K");
            }
            _ => panic!("expected MatMul pattern"),
        }
    }

    #[test]
    fn classify_elementwise_add() {
        let module = make_elementwise_module(BinaryOp::Add);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match pattern {
            KernelPattern::ElementWise {
                op,
                inputs,
                output,
                dim_name,
            } => {
                assert_eq!(op, ElementWiseOp::Add);
                assert_eq!(inputs[0].name, "a");
                assert_eq!(inputs[1].name, "b");
                assert_eq!(output.name, "c");
                assert_eq!(dim_name, "N");
            }
            _ => panic!("expected ElementWise pattern"),
        }
    }

    #[test]
    fn classify_elementwise_div() {
        let module = make_elementwise_module(BinaryOp::Divide);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::ElementWise { op, .. } => {
                assert_eq!(*op, ElementWiseOp::Div);
            }
            _ => panic!("expected ElementWise pattern"),
        }
    }

    #[test]
    fn classify_activation_relu() {
        let module = make_activation_module(MathFunction::Max);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Activation { op, .. } => {
                assert_eq!(*op, ActivationOp::Relu);
            }
            _ => panic!("expected Activation pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_activation_tanh() {
        let module = make_activation_module(MathFunction::Tanh);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Activation { op, .. } => {
                assert_eq!(*op, ActivationOp::Tanh);
            }
            _ => panic!("expected Activation pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_reduce_sum() {
        let module = make_reduce_module();
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Reduce { op, .. } => {
                assert_eq!(*op, ReduceOp::Sum);
            }
            _ => panic!("expected Reduce pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_single_input_no_activation_unknown() {
        // 1 input, no loop, no recognized activation, 2+ params → Unknown.
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: vec![
                    StructMember {
                        name: Some("rows".into()),
                        ty: u32_ty,
                        offset: 0,
                    },
                    StructMember {
                        name: Some("cols".into()),
                        ty: u32_ty,
                        offset: 4,
                    },
                ],
                span: 8,
            },
        });

        module.global_variables.append(GlobalVariable {
            name: Some("a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("c".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        // Body: just a store of a literal (no activation, no loop)
        let mut func = Function::new("main");
        let val = func
            .expressions
            .append(Expression::Literal(Literal::F32(42.0)));
        let ptr = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        func.body.push(Statement::Store {
            pointer: ptr,
            value: val,
        });

        module.entry_points.push(EntryPoint {
            name: "unknown_kernel".into(),
            workgroup_size: [256, 1, 1],
            function: func,
        });

        let pattern = classify_entry_point(&module, 0).unwrap();
        assert!(
            matches!(&pattern, KernelPattern::Unknown { .. }),
            "expected Unknown pattern, got {pattern:?}"
        );
    }

    #[test]
    fn classify_elementwise_with_embedded_weight() {
        // 1 storage input + 1 output + private global with init + binary Add → ElementWise.
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let array_f32_4 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Constant(4),
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: vec![StructMember {
                    name: Some("N".into()),
                    ty: u32_ty,
                    offset: 0,
                }],
                span: 4,
            },
        });

        // Storage input (binding 0)
        module.global_variables.append(GlobalVariable {
            name: Some("input".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        // Storage output (binding 1)
        module.global_variables.append(GlobalVariable {
            name: Some("output".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });

        // Private global with init (embedded weight)
        let lit0 = module
            .global_expressions
            .append(Expression::Literal(Literal::F32(0.1)));
        let lit1 = module
            .global_expressions
            .append(Expression::Literal(Literal::F32(0.2)));
        let lit2 = module
            .global_expressions
            .append(Expression::Literal(Literal::F32(0.3)));
        let lit3 = module
            .global_expressions
            .append(Expression::Literal(Literal::F32(0.4)));
        let compose = module.global_expressions.append(Expression::Compose {
            ty: array_f32_4,
            components: vec![lit0, lit1, lit2, lit3],
        });
        module.global_variables.append(GlobalVariable {
            name: Some("bias".into()),
            space: AddressSpace::Private,
            binding: None,
            ty: array_f32_4,
            init: Some(compose),
            layout: None,
        });

        // Uniform params
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        // Entry point: Store of Binary Add (no loop)
        let mut func = Function::new("main");
        let left = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let right = func
            .expressions
            .append(Expression::Literal(Literal::F32(2.0)));
        let binary = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left,
            right,
        });
        let ptr = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        func.body.push(Statement::Store {
            pointer: ptr,
            value: binary,
        });

        module.entry_points.push(EntryPoint {
            name: "main".into(),
            workgroup_size: [64, 1, 1],
            function: func,
        });

        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::ElementWise {
                op,
                inputs,
                output,
                dim_name,
            } => {
                assert_eq!(*op, ElementWiseOp::Add);
                assert_eq!(inputs[0].name, "input");
                assert_eq!(inputs[1].name, "bias");
                assert_eq!(output.name, "output");
                assert_eq!(dim_name, "N");
            }
            _ => panic!("expected ElementWise pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn detect_reduce_mean() {
        // Loop body with a divide → ReduceOp::Mean
        let mut func = Function::new("test");
        let left = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        let right = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let div = func.expressions.append(Expression::Binary {
            op: BinaryOp::Divide,
            left,
            right,
        });
        let ptr = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        let body = vec![
            Statement::Store {
                pointer: ptr,
                value: div,
            },
            Statement::Break,
        ];
        let result = detect_reduce_op(&body, &func.expressions);
        assert_eq!(result, ReduceOp::Mean);
    }

    #[test]
    fn classify_out_of_range() {
        let module = make_matmul_module();
        let err = classify_entry_point(&module, 99).unwrap_err();
        assert!(matches!(err, AnalysisError::EntryPointOutOfRange(99)));
    }

    #[test]
    fn classify_empty_module() {
        let module = Module::default();
        let err = classify_entry_point(&module, 0).unwrap_err();
        assert!(matches!(err, AnalysisError::NoEntryPoints));
    }

    #[test]
    fn input_sorted_by_binding() {
        // Build a module where 'b' (binding 1) is appended before 'a' (binding 0)
        // to verify that classify sorts inputs by binding order.
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: vec![
                    StructMember {
                        name: Some("M".into()),
                        ty: u32_ty,
                        offset: 0,
                    },
                    StructMember {
                        name: Some("N".into()),
                        ty: u32_ty,
                        offset: 4,
                    },
                    StructMember {
                        name: Some("K".into()),
                        ty: u32_ty,
                        offset: 8,
                    },
                ],
                span: 12,
            },
        });

        // Append b (binding 1) BEFORE a (binding 0).
        module.global_variables.append(GlobalVariable {
            name: Some("b".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("result".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 3,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        let mut func = Function::new("main");
        func.body.push(Statement::Loop {
            body: vec![Statement::Break],
            continuing: vec![],
            break_if: None,
        });
        module.entry_points.push(EntryPoint {
            name: "main".into(),
            workgroup_size: [16, 16, 1],
            function: func,
        });

        let pattern = classify_entry_point(&module, 0).unwrap();
        match pattern {
            KernelPattern::MatMul { inputs, .. } => {
                assert_eq!(inputs[0].name, "a"); // binding 0 sorted first
                assert_eq!(inputs[1].name, "b"); // binding 1 sorted second
            }
            _ => panic!("expected MatMul pattern"),
        }
    }

    #[test]
    fn extract_loop_bounds_literal_u32() {
        let mut func = Function::new("test");
        let local = func.local_variables.append(LocalVariable {
            name: Some("kh".into()),
            ty: {
                let mut types = UniqueArena::new();
                types.insert(Type {
                    name: None,
                    inner: TypeInner::Scalar(Scalar::U32),
                })
            },
            init: None,
        });
        let load = func.expressions.append(Expression::LocalVariable(local));
        let lit = func
            .expressions
            .append(Expression::Literal(Literal::U32(5)));
        let cmp = func.expressions.append(Expression::Binary {
            op: BinaryOp::GreaterEqual,
            left: load,
            right: lit,
        });
        let body = vec![Statement::Loop {
            body: vec![Statement::Break],
            continuing: vec![],
            break_if: Some(cmp),
        }];
        let bounds = extract_loop_bound_literals(&body, &func.expressions);
        assert_eq!(bounds, vec![5]);
    }

    #[test]
    fn extract_loop_bounds_nested() {
        let mut func = Function::new("test");
        // Inner loop: kh < 3
        let local_kh = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let lit3 = func
            .expressions
            .append(Expression::Literal(Literal::U32(3)));
        let cmp_inner = func.expressions.append(Expression::Binary {
            op: BinaryOp::GreaterEqual,
            left: local_kh,
            right: lit3,
        });
        // Outer loop: kw < 5
        let local_kw = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let lit5 = func
            .expressions
            .append(Expression::Literal(Literal::U32(5)));
        let cmp_outer = func.expressions.append(Expression::Binary {
            op: BinaryOp::GreaterEqual,
            left: local_kw,
            right: lit5,
        });
        let body = vec![Statement::Loop {
            body: vec![Statement::Loop {
                body: vec![Statement::Break],
                continuing: vec![],
                break_if: Some(cmp_inner),
            }],
            continuing: vec![],
            break_if: Some(cmp_outer),
        }];
        let bounds = extract_loop_bound_literals(&body, &func.expressions);
        assert_eq!(bounds, vec![5, 3]);
    }

    #[test]
    fn extract_loop_bounds_no_literal() {
        let func = Function::new("test");
        let body = vec![Statement::Loop {
            body: vec![Statement::Break],
            continuing: vec![],
            break_if: None,
        }];
        let bounds = extract_loop_bound_literals(&body, &func.expressions);
        assert_eq!(bounds.len(), 0);
    }

    /// naga 28 lowers `for (var kh = 0u; kh < 5u; ...)` as:
    ///   Loop { body: [Emit, Emit, If(kh<5) {} else {Break}, ...], break_if: None }
    #[test]
    fn extract_loop_bounds_if_else_break_pattern() {
        let mut func = Function::new("test");
        let local = func.local_variables.append(LocalVariable {
            name: Some("kh".into()),
            ty: {
                let mut types = UniqueArena::new();
                types.insert(Type {
                    name: None,
                    inner: TypeInner::Scalar(Scalar::U32),
                })
            },
            init: None,
        });
        let load = func.expressions.append(Expression::LocalVariable(local));
        let lit5 = func
            .expressions
            .append(Expression::Literal(Literal::U32(5)));
        let cmp = func.expressions.append(Expression::Binary {
            op: BinaryOp::Less,
            left: load,
            right: lit5,
        });
        // Simulate If(cond){} else {Break} inside loop body (skipping Emit)
        let body = vec![Statement::Loop {
            body: vec![Statement::If {
                condition: cmp,
                accept: vec![],
                reject: vec![Statement::Break],
            }],
            continuing: vec![],
            break_if: None,
        }];
        let bounds = extract_loop_bound_literals(&body, &func.expressions);
        assert_eq!(bounds, vec![5]);
    }

    /// `gid.y`, as a kernel spells the coordinate a stride scales.
    ///
    /// These fixtures used to multiply one literal by another, which no kernel
    /// writes and which stopped being enough once a stride had to come from
    /// the invocation id.
    fn invocation_id_component(func: &mut Function, index: u32) -> Handle<Expression> {
        let gid = func.expressions.append(Expression::FunctionArgument(0));
        func.expressions
            .append(Expression::AccessIndex { base: gid, index })
    }

    #[test]
    fn extract_multiply_literals_stride() {
        let mut func = Function::new("test");
        let oh = invocation_id_component(&mut func, 1);
        let two = func
            .expressions
            .append(Expression::Literal(Literal::U32(2)));
        func.expressions.append(Expression::Binary {
            op: BinaryOp::Multiply,
            left: oh,
            right: two,
        });
        let strides = extract_multiply_literals(&func.expressions);
        assert_eq!(strides, vec![2]);
    }

    #[test]
    fn a_flattening_factor_is_not_a_stride() {
        // `kh * 3u` inside a weight index. Read as a stride, it emitted a 3x3
        // convolution over a 64x64 image with an output of 21x7 rather than
        // 62x62 -- a model that loads, is accelerated, and is wrong.
        let mut func = Function::new("test");
        // A loop variable reaches the expression as a load, which is the whole
        // difference from `gid.y`; what it loads from does not matter here.
        let slot = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let kh = func.expressions.append(Expression::Load { pointer: slot });
        let three = func
            .expressions
            .append(Expression::Literal(Literal::U32(3)));
        func.expressions.append(Expression::Binary {
            op: BinaryOp::Multiply,
            left: kh,
            right: three,
        });
        assert_eq!(
            extract_multiply_literals(&func.expressions),
            Vec::<u32>::new()
        );
    }

    #[test]
    fn extract_multiply_literals_ignores_one() {
        let mut func = Function::new("test");
        let oh = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let one = func
            .expressions
            .append(Expression::Literal(Literal::U32(1)));
        func.expressions.append(Expression::Binary {
            op: BinaryOp::Multiply,
            left: oh,
            right: one,
        });
        let strides = extract_multiply_literals(&func.expressions);
        assert_eq!(strides.len(), 0);
    }

    #[test]
    fn extract_subtract_literals_padding() {
        let mut func = Function::new("test");
        let idx = func
            .expressions
            .append(Expression::Literal(Literal::U32(10)));
        let pad = func
            .expressions
            .append(Expression::Literal(Literal::U32(1)));
        func.expressions.append(Expression::Binary {
            op: BinaryOp::Subtract,
            left: idx,
            right: pad,
        });
        let pads = extract_subtract_literals(&func.expressions);
        assert_eq!(pads, vec![1]);
    }

    #[test]
    fn conv2d_shape_no_body_defaults_to_unknown() {
        let names: Vec<String> = ["N", "IC", "IH", "IW", "OC", "KH", "KW"]
            .iter()
            .map(|s| s.to_string())
            .collect();
        let shape = extract_conv2d_shape(&names, None, None);
        assert_eq!(shape.batch, "N");
        assert_eq!(shape.kernel_h, "KH");
        assert_eq!(shape.kernel_w, "KW");
        assert_eq!(shape.kernel_h_val, 0);
        assert_eq!(shape.kernel_w_val, 0);
        assert_eq!(shape.stride_h, 1);
        assert_eq!(shape.stride_w, 1);
        assert_eq!(shape.pad_h, 0);
        assert_eq!(shape.pad_w, 0);
    }

    #[test]
    fn conv2d_shape_with_literal_kernel() {
        let names: Vec<String> = ["N", "IC", "IH", "IW", "OC", "KH", "KW"]
            .iter()
            .map(|s| s.to_string())
            .collect();
        // Build a function body with loops: kh < 5, kw < 5
        let mut func = Function::new("test");
        let kh_var = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let lit5a = func
            .expressions
            .append(Expression::Literal(Literal::U32(5)));
        let cmp_h = func.expressions.append(Expression::Binary {
            op: BinaryOp::GreaterEqual,
            left: kh_var,
            right: lit5a,
        });
        let kw_var = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let lit5b = func
            .expressions
            .append(Expression::Literal(Literal::U32(5)));
        let cmp_w = func.expressions.append(Expression::Binary {
            op: BinaryOp::GreaterEqual,
            left: kw_var,
            right: lit5b,
        });
        // Also add stride: oh * 2u, with `oh` coming from the invocation id
        // as it does in a kernel -- a literal multiplied by a literal is not
        // a stride and is no longer read as one.
        let oh = invocation_id_component(&mut func, 1);
        let two = func
            .expressions
            .append(Expression::Literal(Literal::U32(2)));
        func.expressions.append(Expression::Binary {
            op: BinaryOp::Multiply,
            left: oh,
            right: two,
        });
        // Also add padding: ih - 1u
        let ih = func
            .expressions
            .append(Expression::Literal(Literal::U32(10)));
        let one = func
            .expressions
            .append(Expression::Literal(Literal::U32(1)));
        func.expressions.append(Expression::Binary {
            op: BinaryOp::Subtract,
            left: ih,
            right: one,
        });

        let body = vec![Statement::Loop {
            body: vec![Statement::Loop {
                body: vec![Statement::Break],
                continuing: vec![],
                break_if: Some(cmp_w),
            }],
            continuing: vec![],
            break_if: Some(cmp_h),
        }];

        let shape = extract_conv2d_shape(&names, Some(&body), Some(&func.expressions));
        assert_eq!(shape.kernel_h_val, 5);
        assert_eq!(shape.kernel_w_val, 5);
        assert_eq!(shape.stride_h, 2);
        assert_eq!(shape.stride_w, 2);
        assert_eq!(shape.pad_h, 1);
        assert_eq!(shape.pad_w, 1);
    }

    /// Build a pool-like module with literal kernel and stride.
    fn make_pool_module(kernel: u32, stride: u32) -> Module {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: vec![
                    StructMember {
                        name: Some("C".into()),
                        ty: u32_ty,
                        offset: 0,
                    },
                    StructMember {
                        name: Some("IH".into()),
                        ty: u32_ty,
                        offset: 4,
                    },
                    StructMember {
                        name: Some("IW".into()),
                        ty: u32_ty,
                        offset: 8,
                    },
                    StructMember {
                        name: Some("OH".into()),
                        ty: u32_ty,
                        offset: 12,
                    },
                ],
                span: 16,
            },
        });

        module.global_variables.append(GlobalVariable {
            name: Some("input".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("output".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        let mut func = Function::new("main");

        // Loop bounds: kh < kernel, kw < kernel
        let kh_var = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let lit_k1 = func
            .expressions
            .append(Expression::Literal(Literal::U32(kernel)));
        let cmp_h = func.expressions.append(Expression::Binary {
            op: BinaryOp::GreaterEqual,
            left: kh_var,
            right: lit_k1,
        });
        let kw_var = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let lit_k2 = func
            .expressions
            .append(Expression::Literal(Literal::U32(kernel)));
        let cmp_w = func.expressions.append(Expression::Binary {
            op: BinaryOp::GreaterEqual,
            left: kw_var,
            right: lit_k2,
        });

        // Stride: oh * stride
        if stride > 1 {
            let oh = func
                .expressions
                .append(Expression::Literal(Literal::U32(0)));
            let lit_s = func
                .expressions
                .append(Expression::Literal(Literal::U32(stride)));
            func.expressions.append(Expression::Binary {
                op: BinaryOp::Multiply,
                left: oh,
                right: lit_s,
            });
        }

        // Max function for MaxPool detection
        let arg = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        let arg1 = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        let max_expr = func.expressions.append(Expression::Math {
            fun: MathFunction::Max,
            arg,
            arg1: Some(arg1),
            arg2: None,
            arg3: None,
        });
        let ptr = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        func.body.push(Statement::Loop {
            body: vec![
                Statement::Loop {
                    body: vec![
                        Statement::Store {
                            pointer: ptr,
                            value: max_expr,
                        },
                        Statement::Break,
                    ],
                    continuing: vec![],
                    break_if: Some(cmp_w),
                },
                Statement::Break,
            ],
            continuing: vec![],
            break_if: Some(cmp_h),
        });

        module.entry_points.push(EntryPoint {
            name: "main".into(),
            workgroup_size: [16, 16, 1],
            function: func,
        });

        module
    }

    #[test]
    fn classify_pool_extracts_2x2_stride2() {
        let module = make_pool_module(2, 2);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Pool { kind, shape, .. } => {
                assert_eq!(*kind, PoolKind::Max);
                assert_eq!(shape.kernel_h, 2);
                assert_eq!(shape.kernel_w, 2);
                assert_eq!(shape.stride_h, 2);
                assert_eq!(shape.stride_w, 2);
            }
            _ => panic!("expected Pool pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_pool_extracts_3x3_stride1() {
        let module = make_pool_module(3, 1);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Pool { kind, shape, .. } => {
                assert_eq!(*kind, PoolKind::Max);
                assert_eq!(shape.kernel_h, 3);
                assert_eq!(shape.kernel_w, 3);
                // stride=1: no multiply literals, default to kernel size
                assert_eq!(shape.stride_h, 3);
                assert_eq!(shape.stride_w, 3);
            }
            _ => panic!("expected Pool pattern, got {pattern:?}"),
        }
    }

    // --- EmbeddedWeight extraction tests ---

    #[test]
    fn eval_const_expr_f32_literal() {
        let mut exprs = Arena::new();
        let types = UniqueArena::new();
        let h = exprs.append(Expression::Literal(Literal::F32(2.78)));
        assert_eq!(eval_const_expr_f32(h, &exprs, &types), Some(vec![2.78]));
    }

    #[test]
    fn eval_const_expr_f32_compose() {
        let mut exprs = Arena::new();
        let types = UniqueArena::new();
        let a = exprs.append(Expression::Literal(Literal::F32(0.1)));
        let b = exprs.append(Expression::Literal(Literal::F32(0.2)));
        let c = exprs.append(Expression::Literal(Literal::F32(0.3)));
        let f32_ty = {
            let mut t = UniqueArena::new();
            t.insert(Type {
                name: None,
                inner: TypeInner::Scalar(Scalar::F32),
            })
        };
        let compose = exprs.append(Expression::Compose {
            ty: f32_ty,
            components: vec![a, b, c],
        });
        let result = eval_const_expr_f32(compose, &exprs, &types).unwrap();
        assert_eq!(result.len(), 3);
        assert!((result[0] - 0.1).abs() < 1e-6);
        assert!((result[1] - 0.2).abs() < 1e-6);
        assert!((result[2] - 0.3).abs() < 1e-6);
    }

    #[test]
    fn eval_const_expr_f32_zero_value() {
        let mut exprs = Arena::new();
        let mut types = UniqueArena::new();
        let f32_ty = types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let arr_ty = types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Constant(3),
                stride: 4,
            },
        });
        let h = exprs.append(Expression::ZeroValue(arr_ty));
        let result = eval_const_expr_f32(h, &exprs, &types).unwrap();
        assert_eq!(result, vec![0.0, 0.0, 0.0]);
    }

    #[test]
    fn extract_embedded_weights_private_global() {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let array_f32_4 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Constant(4),
                stride: 4,
            },
        });

        // Build init expression: Compose([0.1, 0.2, 0.3, 0.4])
        let lit0 = module
            .global_expressions
            .append(Expression::Literal(Literal::F32(0.1)));
        let lit1 = module
            .global_expressions
            .append(Expression::Literal(Literal::F32(0.2)));
        let lit2 = module
            .global_expressions
            .append(Expression::Literal(Literal::F32(0.3)));
        let lit3 = module
            .global_expressions
            .append(Expression::Literal(Literal::F32(0.4)));
        let compose = module.global_expressions.append(Expression::Compose {
            ty: array_f32_4,
            components: vec![lit0, lit1, lit2, lit3],
        });

        module.global_variables.append(GlobalVariable {
            name: Some("bias".into()),
            space: AddressSpace::Private,
            binding: None,
            ty: array_f32_4,
            init: Some(compose),
            layout: None,
        });

        let weights = extract_embedded_weights(&module);
        assert_eq!(weights.len(), 1);
        assert_eq!(weights[0].name, "bias");
        assert_eq!(weights[0].dims, vec![4]);
        assert_eq!(weights[0].data.len(), 4);
        assert!((weights[0].data[0] - 0.1).abs() < 1e-6);
        assert!((weights[0].data[3] - 0.4).abs() < 1e-6);
    }

    #[test]
    fn extract_embedded_weights_no_init() {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });

        module.global_variables.append(GlobalVariable {
            name: Some("a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });

        let weights = extract_embedded_weights(&module);
        assert_eq!(weights.len(), 0);
    }

    // --- Concat/Split axis inference tests ---

    /// Build a concat-like module with 2 inputs, 1 output, and an If condition
    /// that compares against `AccessIndex { base: params, index: boundary_index }`.
    /// This simulates the IR that naga produces for `if c < params.C1`.
    fn make_concat_module(boundary_index: u32, shape_names: &[&str]) -> Module {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: shape_names
                    .iter()
                    .enumerate()
                    .map(|(i, name)| StructMember {
                        name: Some((*name).into()),
                        ty: u32_ty,
                        offset: (i * 4) as u32,
                    })
                    .collect(),
                span: (shape_names.len() * 4) as u32,
            },
        });

        // Input a (binding 0)
        module.global_variables.append(GlobalVariable {
            name: Some("a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        // Input b (binding 1)
        module.global_variables.append(GlobalVariable {
            name: Some("b".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        // Output (binding 2)
        module.global_variables.append(GlobalVariable {
            name: Some("result".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        // Uniform params (binding 3)
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 3,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        // Build function with If(idx < AccessIndex(params, boundary_index))
        let mut func = Function::new("main");
        let idx = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        // params GlobalVariable (handle index 3)
        let params_gv = {
            let mut iter = module.global_variables.iter();
            iter.nth(3).unwrap().0
        };
        let params_expr = func
            .expressions
            .append(Expression::GlobalVariable(params_gv));
        let access = func.expressions.append(Expression::AccessIndex {
            base: params_expr,
            index: boundary_index,
        });
        let load = func
            .expressions
            .append(Expression::Load { pointer: access });
        let cmp = func.expressions.append(Expression::Binary {
            op: BinaryOp::Less,
            left: idx,
            right: load,
        });

        // If body: Store (accept), Store (reject)
        let val = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        let ptr = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        func.body.push(Statement::If {
            condition: cmp,
            accept: vec![Statement::Store {
                pointer: ptr,
                value: val,
            }],
            reject: vec![Statement::Store {
                pointer: ptr,
                value: val,
            }],
        });

        module.entry_points.push(EntryPoint {
            name: "main".into(),
            workgroup_size: [256, 1, 1],
            function: func,
        });

        module
    }

    /// Build a split-like module with 1 input, 2 outputs, and an If condition
    /// that compares against `AccessIndex { base: params, index: boundary_index }`.
    fn make_split_module(boundary_index: u32, shape_names: &[&str]) -> Module {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members: shape_names
                    .iter()
                    .enumerate()
                    .map(|(i, name)| StructMember {
                        name: Some((*name).into()),
                        ty: u32_ty,
                        offset: (i * 4) as u32,
                    })
                    .collect(),
                span: (shape_names.len() * 4) as u32,
            },
        });

        // Input (binding 0) - read-only storage
        module.global_variables.append(GlobalVariable {
            name: Some("input".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 0,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        // Output a (binding 1)
        module.global_variables.append(GlobalVariable {
            name: Some("out_a".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 1,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        // Output b (binding 2)
        module.global_variables.append(GlobalVariable {
            name: Some("out_b".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 2,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });
        // Uniform params (binding 3)
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 3,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        // Build function with If(idx < AccessIndex(params, boundary_index))
        let mut func = Function::new("main");
        let idx = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let params_gv = {
            let mut iter = module.global_variables.iter();
            iter.nth(3).unwrap().0
        };
        let params_expr = func
            .expressions
            .append(Expression::GlobalVariable(params_gv));
        let access = func.expressions.append(Expression::AccessIndex {
            base: params_expr,
            index: boundary_index,
        });
        let load = func
            .expressions
            .append(Expression::Load { pointer: access });
        let cmp = func.expressions.append(Expression::Binary {
            op: BinaryOp::Less,
            left: idx,
            right: load,
        });

        // `out_a[idx] = input[idx]` on one side of the boundary and
        // `out_b[idx] = input[idx]` on the other. Both stores used to be a
        // float literal written through a float literal, which addresses
        // nothing and copies nothing; the classifier now asks whether the
        // outputs are slices of the input, so the module has to be one.
        let (input_gv, out_a_gv, out_b_gv) = {
            let mut iter = module.global_variables.iter();
            let input = iter.next().unwrap().0;
            let out_a = iter.next().unwrap().0;
            let out_b = iter.next().unwrap().0;
            (input, out_a, out_b)
        };
        let input_expr = func
            .expressions
            .append(Expression::GlobalVariable(input_gv));
        let read = func.expressions.append(Expression::Access {
            base: input_expr,
            index: idx,
        });
        let val = func.expressions.append(Expression::Load { pointer: read });

        let mut copy_into = |out: Handle<GlobalVariable>| {
            let out_expr = func.expressions.append(Expression::GlobalVariable(out));
            let ptr = func.expressions.append(Expression::Access {
                base: out_expr,
                index: idx,
            });
            Statement::Store {
                pointer: ptr,
                value: val,
            }
        };
        let store_a = copy_into(out_a_gv);
        let store_b = copy_into(out_b_gv);

        func.body.push(Statement::If {
            condition: cmp,
            accept: vec![store_a],
            reject: vec![store_b],
        });

        module.entry_points.push(EntryPoint {
            name: "main".into(),
            workgroup_size: [256, 1, 1],
            function: func,
        });

        module
    }

    #[test]
    fn classify_concat_axis_0_default() {
        // Concat with boundary at index 0 → axis = 0 (same as before).
        let module = make_concat_module(0, &["N_a", "N_b"]);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Concat { axis, .. } => {
                assert_eq!(*axis, 0, "expected axis 0 for boundary at index 0");
            }
            _ => panic!("expected Concat pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_concat_axis_1_inferred() {
        // Concat with boundary at index 1 (params.C1) → axis = 1.
        let module = make_concat_module(1, &["N", "C1", "C2"]);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Concat {
                axis,
                inputs,
                output,
                ..
            } => {
                assert_eq!(*axis, 1, "expected axis 1 for boundary at index 1");
                assert_eq!(inputs[0].name, "a");
                assert_eq!(inputs[1].name, "b");
                assert_eq!(output.name, "result");
            }
            _ => panic!("expected Concat pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_concat_axis_2_inferred() {
        // Concat with boundary at index 2 → axis = 2.
        let module = make_concat_module(2, &["N", "C", "W1", "W2"]);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Concat { axis, .. } => {
                assert_eq!(*axis, 2, "expected axis 2 for boundary at index 2");
            }
            _ => panic!("expected Concat pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_concat_axis_3_inferred() {
        // Concat with boundary at index 3 → axis = 3 (4D shape: N, C, H, W).
        let module = make_concat_module(3, &["N", "C", "H", "W1", "W2"]);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Concat { axis, .. } => {
                assert_eq!(*axis, 3, "expected axis 3 for boundary at index 3");
            }
            _ => panic!("expected Concat pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_split_axis_0_default() {
        // Split with boundary at index 0 → axis = 0.
        let module = make_split_module(0, &["N_a", "N_b"]);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Split { axis, .. } => {
                assert_eq!(*axis, 0, "expected axis 0 for boundary at index 0");
            }
            _ => panic!("expected Split pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_split_axis_1_inferred() {
        // Split with boundary at index 1 (params.C1) → axis = 1.
        let module = make_split_module(1, &["N", "C1", "C2"]);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Split {
                axis,
                input,
                outputs,
                ..
            } => {
                assert_eq!(*axis, 1, "expected axis 1 for boundary at index 1");
                assert_eq!(input.name, "input");
                assert_eq!(outputs.len(), 2);
                assert_eq!(outputs[0].name, "out_a");
                assert_eq!(outputs[1].name, "out_b");
            }
            _ => panic!("expected Split pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_split_axis_2_inferred() {
        // Split with boundary at index 2 → axis = 2.
        let module = make_split_module(2, &["N", "C", "W1", "W2"]);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Split { axis, .. } => {
                assert_eq!(*axis, 2, "expected axis 2 for boundary at index 2");
            }
            _ => panic!("expected Split pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn classify_split_axis_3_inferred() {
        // Split with boundary at index 3 → axis = 3 (4D shape: N, C, H, W).
        let module = make_split_module(3, &["N", "C", "H", "W1", "W2"]);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Split { axis, .. } => {
                assert_eq!(*axis, 3, "expected axis 3 for boundary at index 3");
            }
            _ => panic!("expected Split pattern, got {pattern:?}"),
        }
    }

    #[test]
    fn normalize_axis_positive() {
        assert_eq!(normalize_axis(0, 3), 0);
        assert_eq!(normalize_axis(1, 4), 1);
        assert_eq!(normalize_axis(2, 4), 2);
    }

    #[test]
    fn normalize_axis_negative() {
        assert_eq!(normalize_axis(-1, 4), 3);
        assert_eq!(normalize_axis(-2, 4), 2);
        assert_eq!(normalize_axis(-3, 4), 1);
        assert_eq!(normalize_axis(-4, 4), 0);
    }

    #[test]
    fn normalize_axis_negative_wraps() {
        // -1 with ndim 3 → 2
        assert_eq!(normalize_axis(-1, 3), 2);
        // -2 with ndim 3 → 1
        assert_eq!(normalize_axis(-2, 3), 1);
    }

    #[test]
    fn find_if_comparison_axis_direct_access_index() {
        // Test that find_if_comparison_axis works with direct AccessIndex
        // (without Load wrapper).
        let mut func = Function::new("test");
        let idx = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let base = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let access = func
            .expressions
            .append(Expression::AccessIndex { base, index: 2 });
        let cmp = func.expressions.append(Expression::Binary {
            op: BinaryOp::Less,
            left: idx,
            right: access,
        });

        let body = vec![Statement::If {
            condition: cmp,
            accept: vec![],
            reject: vec![],
        }];
        let names: Vec<String> = vec!["N".into(), "C".into(), "W".into()];
        let result = find_if_comparison_axis(&body, &func.expressions, &names);
        assert_eq!(result, Some(2));
    }

    #[test]
    fn find_if_comparison_axis_in_loop() {
        // Test that the function searches into Loop bodies.
        let mut func = Function::new("test");
        let idx = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let base = func
            .expressions
            .append(Expression::Literal(Literal::U32(0)));
        let access = func
            .expressions
            .append(Expression::AccessIndex { base, index: 1 });
        let load = func
            .expressions
            .append(Expression::Load { pointer: access });
        let cmp = func.expressions.append(Expression::Binary {
            op: BinaryOp::LessEqual,
            left: idx,
            right: load,
        });

        let body = vec![Statement::Loop {
            body: vec![Statement::If {
                condition: cmp,
                accept: vec![],
                reject: vec![],
            }],
            continuing: vec![],
            break_if: None,
        }];
        let names: Vec<String> = vec!["N".into(), "C".into()];
        let result = find_if_comparison_axis(&body, &func.expressions, &names);
        assert_eq!(result, Some(1));
    }

    #[test]
    fn find_if_comparison_axis_no_if_returns_none() {
        let func = Function::new("test");
        let body = vec![Statement::Break];
        let names: Vec<String> = vec!["N".into()];
        let result = find_if_comparison_axis(&body, &func.expressions, &names);
        assert_eq!(result, None);
    }

    /// A params struct that mixes dimensions with a tolerance constant is the
    /// normal case, not an edge one: LayerNorm's is (N, D, eps), and counting
    /// eps as a third dimension is what made it report as BatchNorm.
    #[test]
    fn eps_is_not_a_dimension() {
        let mut module = Module::default();
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let members = [
            StructMember {
                name: Some("N".into()),
                ty: u32_ty,
                offset: 0,
            },
            StructMember {
                name: Some("D".into()),
                ty: u32_ty,
                offset: 4,
            },
            StructMember {
                name: Some("eps".into()),
                ty: f32_ty,
                offset: 8,
            },
        ];
        let dims: Vec<String> = members
            .iter()
            .filter(|m| is_integer_member(&module, m))
            .filter_map(|m| m.name.clone())
            .collect();
        assert_eq!(dims, vec!["N".to_string(), "D".to_string()]);
    }

    // --- Multi-head & causal attention classification tests ---

    /// Build a module with 3 inputs (query, key, value), 1 output, a uniform
    /// params struct with the given member names, Loop + Exp + Sqrt in the
    /// function expressions, and optionally a division-by-literal and/or a
    /// causal If-guard inside the loop body.
    fn make_attention_module(param_names: &[&str], divide_by: Option<u32>, causal: bool) -> Module {
        let mut module = Module::default();

        let f32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::F32),
        });
        let u32_ty = module.types.insert(Type {
            name: None,
            inner: TypeInner::Scalar(Scalar::U32),
        });
        let array_f32 = module.types.insert(Type {
            name: None,
            inner: TypeInner::Array {
                base: f32_ty,
                size: ArraySize::Dynamic,
                stride: 4,
            },
        });

        // Build params struct from provided member names
        let members: Vec<StructMember> = param_names
            .iter()
            .enumerate()
            .map(|(i, name)| StructMember {
                name: Some((*name).into()),
                ty: u32_ty,
                offset: (i * 4) as u32,
            })
            .collect();
        let params_ty = module.types.insert(Type {
            name: Some("Params".into()),
            inner: TypeInner::Struct {
                members,
                span: (param_names.len() * 4) as u32,
            },
        });

        // 3 storage-read inputs (query, key, value)
        for (i, name) in ["query", "key", "value"].iter().enumerate() {
            module.global_variables.append(GlobalVariable {
                name: Some((*name).into()),
                space: AddressSpace::Storage {
                    access: StorageAccess::LOAD,
                },
                binding: Some(ResourceBinding {
                    group: 0,
                    binding: i as u32,
                }),
                ty: array_f32,
                init: None,
                layout: None,
            });
        }

        // 1 storage-read_write output
        module.global_variables.append(GlobalVariable {
            name: Some("output".into()),
            space: AddressSpace::Storage {
                access: StorageAccess::LOAD | StorageAccess::STORE,
            },
            binding: Some(ResourceBinding {
                group: 0,
                binding: 3,
            }),
            ty: array_f32,
            init: None,
            layout: None,
        });

        // Uniform params
        module.global_variables.append(GlobalVariable {
            name: Some("params".into()),
            space: AddressSpace::Uniform,
            binding: Some(ResourceBinding {
                group: 0,
                binding: 4,
            }),
            ty: params_ty,
            init: None,
            layout: None,
        });

        // Build function with Exp, Sqrt, and optionally Division + causal If
        let mut func = Function::new("main");

        // Exp expression
        let arg_exp = func
            .expressions
            .append(Expression::Literal(Literal::F32(0.0)));
        func.expressions.append(Expression::Math {
            fun: MathFunction::Exp,
            arg: arg_exp,
            arg1: None,
            arg2: None,
            arg3: None,
        });

        // Sqrt expression
        let arg_sqrt = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        func.expressions.append(Expression::Math {
            fun: MathFunction::Sqrt,
            arg: arg_sqrt,
            arg1: None,
            arg2: None,
            arg3: None,
        });

        // Optional: division by literal (e.g. d_model / num_heads)
        if let Some(n) = divide_by {
            let dividend = func
                .expressions
                .append(Expression::Literal(Literal::U32(64)));
            let divisor = func
                .expressions
                .append(Expression::Literal(Literal::U32(n)));
            func.expressions.append(Expression::Binary {
                op: BinaryOp::Divide,
                left: dividend,
                right: divisor,
            });
        }

        // Build loop body contents
        let mut loop_inner: Vec<Statement> = Vec::new();

        if causal {
            // Causal mask pattern: If(j > i) { score = -1e30; }
            // First, add a comparison expression: Greater(j, i)
            let j_expr = func
                .expressions
                .append(Expression::Literal(Literal::U32(0)));
            let i_expr = func
                .expressions
                .append(Expression::Literal(Literal::U32(0)));
            let cmp = func.expressions.append(Expression::Binary {
                op: BinaryOp::Greater,
                left: j_expr,
                right: i_expr,
            });

            // Store of large negative literal
            let neg_val = func
                .expressions
                .append(Expression::Literal(Literal::F32(-1e30)));
            let ptr = func
                .expressions
                .append(Expression::Literal(Literal::F32(0.0)));

            // The for-loop break guard (first If in naga loop)
            let break_cond = func
                .expressions
                .append(Expression::Literal(Literal::Bool(true)));
            loop_inner.push(Statement::If {
                condition: break_cond,
                accept: vec![],
                reject: vec![Statement::Break],
            });

            // The causal If guard (second If in the loop)
            loop_inner.push(Statement::If {
                condition: cmp,
                accept: vec![Statement::Store {
                    pointer: ptr,
                    value: neg_val,
                }],
                reject: vec![],
            });
        }

        loop_inner.push(Statement::Break);

        func.body.push(Statement::Loop {
            body: loop_inner,
            continuing: vec![],
            break_if: None,
        });

        module.entry_points.push(EntryPoint {
            name: "main".into(),
            workgroup_size: [16, 16, 1],
            function: func,
        });

        module
    }

    #[test]
    fn classify_multihead_4_heads() {
        let module = make_attention_module(&["seq_len", "d_model", "num_heads"], Some(4), false);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Attention {
                num_heads,
                num_kv_heads,
                causal,
                seq_len,
                ..
            } => {
                assert_eq!(*num_heads, 4, "expected 4 heads from division literal");
                assert_eq!(
                    *num_kv_heads, 4,
                    "num_kv_heads should equal num_heads for MHA"
                );
                assert!(!causal, "should not detect causal mask");
                assert_eq!(seq_len, "seq_len");
            }
            _ => panic!("expected Attention pattern, got {pattern:?}"),
        }
        // Display should include heads=4 and no causal marker
        let display = format!("{pattern}");
        assert!(
            display.contains("heads=4"),
            "display should show heads=4: {display}"
        );
        assert!(
            !display.contains("causal"),
            "display should not mention causal: {display}"
        );
    }

    #[test]
    fn classify_causal_attention() {
        let module = make_attention_module(&["seq_len", "d_k"], None, true);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Attention {
                num_heads,
                causal,
                d_k,
                ..
            } => {
                assert_eq!(*num_heads, 1, "no num_heads param → default 1");
                assert!(*causal, "should detect causal mask from If + Store(-1e30)");
                assert_eq!(d_k, "d_k");
            }
            _ => panic!("expected Attention pattern, got {pattern:?}"),
        }
        // Display should include causal marker
        let display = format!("{pattern}");
        assert!(
            display.contains("causal"),
            "display should mention causal: {display}"
        );
        assert!(
            display.contains("heads=1"),
            "display should show heads=1: {display}"
        );
    }

    #[test]
    fn classify_gqa_defaults_to_mha() {
        // GQA detection is not yet implemented; num_kv_heads always equals num_heads.
        let module = make_attention_module(&["seq_len", "d_model", "num_heads"], Some(4), false);
        let pattern = classify_entry_point(&module, 0).unwrap();
        match &pattern {
            KernelPattern::Attention {
                num_heads,
                num_kv_heads,
                ..
            } => {
                assert_eq!(*num_heads, 4);
                assert_eq!(
                    *num_kv_heads, *num_heads,
                    "GQA not implemented: kv_heads defaults to num_heads"
                );
            }
            _ => panic!("expected Attention pattern, got {pattern:?}"),
        }
    }
}