nxpu_backend_tflite/

lib.rs

//! TFLite/LiteRT backend emitter for NxPU.
//!
//! Emits TFLite FlatBuffer (`.tflite`) models from NxPU IR.
//! Targets: MediaTek APU (NeuroPilot), Google Edge TPU, Arm Ethos NPU (Vela).

use nxpu_analysis::{analyze, fusion};
use nxpu_backend_core::{
    Backend, BackendError, BackendOptions, BackendOutput, Diagnostic, DiagnosticLevel,
    OutputContent, OutputFile,
};
use nxpu_ir::Module;

mod lower;
mod schema;

/// TFLite backend that compiles NxPU IR into `.tflite` FlatBuffer files.
#[derive(Debug)]
pub struct TfLiteBackend;

/// The extent to write for dimensions the kernel leaves symbolic.
///
/// `None` means the caller did not say. One is the smallest value that
/// produces a loadable model, which is the right default for a compiler that
/// genuinely does not know the size: a WGSL kernel over `array<f32>` carries
/// no length. It is also useless for measuring anything, so a caller who wants
/// a model worth timing has to name a size.
///
/// TFLite has a place for a genuinely unknown dimension — `shape_signature` —
/// but `shape` itself must be concrete: a negative value there makes
/// `BytesRequired` overflow and no interpreter can be constructed, on any
/// device, delegate or not.
fn symbolic_extent(opts: &BackendOptions) -> i32 {
    opts.symbolic_extent
        .map(|e| e.max(1))
        .unwrap_or(1)
        .min(i32::MAX as u32) as i32
}

impl Backend for TfLiteBackend {
    fn name(&self) -> &str {
        "TFLite"
    }

    fn targets(&self) -> &[&str] {
        &["tflite", "litert"]
    }

    fn compile(
        &self,
        module: &Module,
        opts: &BackendOptions,
    ) -> Result<BackendOutput, BackendError> {
        if module.entry_points.is_empty() {
            return Err(BackendError::Other("no entry points in module".into()));
        }

        // 1. Classify all entry points.
        let mut patterns = Vec::new();
        for (i, ep) in module.entry_points.iter().enumerate() {
            let pattern = analyze::classify_entry_point(module, i).map_err(|e| {
                BackendError::Unsupported(format!("entry point '{}': {e}", ep.name))
            })?;
            if let analyze::KernelPattern::Unknown { reason } = &pattern {
                return Err(BackendError::Unsupported(format!(
                    "entry point '{}': unrecognized pattern: {reason}",
                    ep.name
                )));
            }
            patterns.push(pattern);
        }

        // 2. Fuse adjacent patterns.
        let fused = fusion::fuse_patterns(patterns);

        // 3. Lower each fused pattern.
        let extent = symbolic_extent(opts);
        let mut files = Vec::new();
        let mut diagnostics = Vec::new();

        for (fp, ep_idx) in &fused {
            let ep_name = if *ep_idx < module.entry_points.len() {
                &module.entry_points[*ep_idx].name
            } else {
                "fused"
            };

            let summary = match fp {
                fusion::FusedPattern::Single(p) => pattern_summary(p).to_string(),
                fusion::FusedPattern::ConvBatchNorm { .. } => "Conv+BatchNorm (fused)".into(),
                fusion::FusedPattern::WithActivation {
                    base, activation, ..
                } => {
                    let base_name = match base.as_ref() {
                        fusion::FusedPattern::Single(p) => pattern_summary(p).into_owned(),
                        fusion::FusedPattern::ConvBatchNorm { .. } => "Conv+BatchNorm".to_string(),
                        fusion::FusedPattern::MatMulBias { .. } => "Gemm".to_string(),
                        _ => "fused".to_string(),
                    };
                    format!("{base_name}+{activation:?}")
                }
                fusion::FusedPattern::MatMulBias { .. } => "Gemm (fused)".into(),
            };

            diagnostics.push(Diagnostic {
                level: DiagnosticLevel::Info,
                message: format!("entry point '{ep_name}': classified as {summary}"),
            });

            // Emit warnings for unsupported attention features in the TFLite backend.
            emit_attention_diagnostics(fp, ep_name, &mut diagnostics);
            emit_runtime_filter_diagnostic(fp, ep_name, opts, &mut diagnostics);

            // A caller who did not say gets the smallest valid extent, which loads
            // but measures nothing; `--symbolic-dim` is how you ask for a size
            // worth timing.
            let bytes = lower::build_fused_model(
                fp,
                lower::Lowering {
                    extent,
                    constants: &opts.constant_tensors,
                },
            )?;

            let filename = if fused.len() == 1 {
                "output.tflite".into()
            } else {
                format!("{ep_name}.tflite")
            };

            files.push(OutputFile {
                name: filename,
                content: OutputContent::Binary(bytes),
            });
        }

        // 4. Emit quantization parameters as a companion JSON file.
        if !opts.quantization_params.is_empty() || !opts.per_channel_params.is_empty() {
            let mut json = String::from("{\n");

            if !opts.quantization_params.is_empty() {
                json.push_str("  \"quantization_params\": [\n");
                for (i, qp) in opts.quantization_params.iter().enumerate() {
                    if i > 0 {
                        json.push_str(",\n");
                    }
                    json.push_str(&format!(
                        "    {{\"name\": \"{}\", \"scale\": {}, \"zero_point\": {}}}",
                        qp.name, qp.scale, qp.zero_point
                    ));
                }
                json.push_str("\n  ]");
                if !opts.per_channel_params.is_empty() {
                    json.push(',');
                }
                json.push('\n');
            }

            if !opts.per_channel_params.is_empty() {
                json.push_str("  \"per_channel_params\": [\n");
                for (i, pcp) in opts.per_channel_params.iter().enumerate() {
                    if i > 0 {
                        json.push_str(",\n");
                    }
                    let scales: Vec<String> = pcp.scales.iter().map(|s| format!("{s}")).collect();
                    let zero_points: Vec<String> =
                        pcp.zero_points.iter().map(|z| format!("{z}")).collect();
                    json.push_str(&format!(
                        "    {{\"name\": \"{}\", \"scales\": [{}], \"zero_points\": [{}], \"channel_axis\": {}}}",
                        pcp.name,
                        scales.join(", "),
                        zero_points.join(", "),
                        pcp.channel_axis
                    ));
                }
                json.push_str("\n  ]\n");
            }

            json.push_str("}\n");
            files.push(OutputFile {
                name: "quant_params.json".into(),
                content: OutputContent::Text(json),
            });
        }

        Ok(BackendOutput { files, diagnostics })
    }
}

/// Emit diagnostic warnings when the TFLite backend encounters attention patterns
/// with features it does not yet support (multi-head splitting, causal masking).
/// Say when a convolution will not be accelerated, and why.
///
/// An NNAPI driver requires a convolution's filter to be a compile-time
/// constant. A MediaTek MT6899 accelerates a convolution whose filter is one
/// and refuses the identical convolution, at the same shapes, whose filter is
/// a graph input — and reports the refusal as `unsupported`, which reads as a
/// statement about CONV_2D and is a statement about this model. Four device
/// sweeps and a byte-level diff went into learning that from the outside.
///
/// So it is said here instead. Not an error: TFLite's own GPU delegate takes
/// the convolution either way.
///
/// The GPU is not an equivalent answer, though, and the warning says so. It is
/// the engine the display composites on, so occupying it contends with
/// everything drawing on screen. Timed on the same phone with the thermal
/// conditions reported stable throughout, it is faster than the NPU in the
/// middle on 8 of the 12 operators both accept and twice as spread at p99 --
/// 2.04x its own median against the NPU's 1.27x, worst case 6.23x. Reaching
/// the NPU is what this flag is for.
fn emit_runtime_filter_diagnostic(
    fp: &fusion::FusedPattern,
    ep_name: &str,
    opts: &BackendOptions,
    diagnostics: &mut Vec<Diagnostic>,
) {
    let analyze::KernelPattern::Conv2D { weight, .. } = fp.primary_pattern() else {
        return;
    };
    if opts.constant_tensors.iter().any(|c| c.name == weight.name) {
        return;
    }
    diagnostics.push(Diagnostic {
        level: DiagnosticLevel::Warning,
        message: format!(
            "entry point '{ep_name}': `{}` is a graph input, so an NNAPI driver will \
             refuse this convolution and it will run on the CPU. TFLite's GPU delegate \
             takes it as it is, at the cost of occupying the engine the display \
             composites on. To reach an NPU, supply the contents with \
             --weights <DIR>/{}.bin",
            weight.name, weight.name
        ),
    });
}

fn emit_attention_diagnostics(
    fp: &fusion::FusedPattern,
    ep_name: &str,
    diagnostics: &mut Vec<Diagnostic>,
) {
    // Extract the inner KernelPattern from the fused pattern.
    let inner = match fp {
        fusion::FusedPattern::Single(p) => Some(p),
        fusion::FusedPattern::WithActivation { base, .. } => match base.as_ref() {
            fusion::FusedPattern::Single(p) => Some(p),
            _ => None,
        },
        _ => None,
    };

    if let Some(analyze::KernelPattern::Attention {
        num_heads, causal, ..
    }) = inner
    {
        if *num_heads > 1 {
            diagnostics.push(Diagnostic {
                level: DiagnosticLevel::Warning,
                message: format!(
                    "entry point '{ep_name}': TFLite backend does not support multi-head attention \
                     (num_heads={num_heads}); output will be single-head SDPA"
                ),
            });
        }
        if *causal {
            diagnostics.push(Diagnostic {
                level: DiagnosticLevel::Warning,
                message: format!(
                    "entry point '{ep_name}': TFLite backend does not support causal masking; \
                     output will be unmasked SDPA"
                ),
            });
        }
    }
}

/// Name the operators a pattern lowers to, for the `classified as …`
/// diagnostic.
///
/// Returns an owned string because a chain has no fixed name: it reports the
/// operator sequence it emits — `Mul+Add`, `Cast+Mul+Mul` — since a category
/// label like "chain" would say nothing about what the kernel computes.
fn pattern_summary(pattern: &analyze::KernelPattern) -> std::borrow::Cow<'static, str> {
    use std::borrow::Cow;
    Cow::Borrowed(match pattern {
        analyze::KernelPattern::ElementWiseChain { cast, steps, .. } => {
            return Cow::Owned(analyze::chain_summary(*cast, steps));
        }
        analyze::KernelPattern::MatMul { .. } => "BATCH_MATMUL",
        // Named by the whole graph rather than by its last operator: there is
        // no quantized matmul operator here, and reporting "BATCH_MATMUL"
        // would describe one node out of five and hide the dequantization
        // that makes the answer right.
        analyze::KernelPattern::QuantizedMatMul { bias, .. } => {
            if bias.is_some() {
                "TRANSPOSE+CAST+BATCH_MATMUL+MUL+ADD (int8 weights, per-channel scale)"
            } else {
                "TRANSPOSE+CAST+BATCH_MATMUL+MUL (int8 weights, per-channel scale)"
            }
        }
        analyze::KernelPattern::ElementWise { op, .. } => op.op_name(),
        analyze::KernelPattern::Conv2D { .. } => "CONV_2D",
        analyze::KernelPattern::Pool { kind, .. } => kind.op_name(),
        analyze::KernelPattern::Activation { op, .. } => op.op_name(),
        analyze::KernelPattern::Reduce { op, .. } => op.op_name(),
        analyze::KernelPattern::Transpose { .. } => "TRANSPOSE",
        analyze::KernelPattern::Reshape { .. } => "RESHAPE",
        // The kind matters: reporting every normalization as BatchNormalization
        // hid a LayerNorm being classified correctly and described wrongly,
        // which is worse than either mistake alone.
        analyze::KernelPattern::Normalization { norm_type, .. } => match norm_type {
            analyze::NormType::Batch => "BatchNormalization",
            analyze::NormType::Layer => "LayerNormalization",
        },
        analyze::KernelPattern::Concat { .. } => "CONCATENATION",
        analyze::KernelPattern::Split { .. } => "SPLIT",
        analyze::KernelPattern::Attention { .. } => "Attention",
        analyze::KernelPattern::Gather { .. } => "GATHER",
        analyze::KernelPattern::Scatter { .. } => "SCATTER_ND",
        analyze::KernelPattern::Unknown { .. } => "Unknown",
    })
}

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

    #[test]
    fn unspecified_is_the_smallest_loadable_extent() {
        assert_eq!(symbolic_extent(&BackendOptions::default()), 1);
    }

    #[test]
    fn zero_is_raised_to_one() {
        // A zero extent produces a zero-byte tensor, which loads and then
        // fails on the first attempt to hand it any data.
        let opts = BackendOptions {
            symbolic_extent: Some(0),
            ..Default::default()
        };
        assert_eq!(symbolic_extent(&opts), 1);
    }

    #[test]
    fn a_requested_extent_is_used() {
        let opts = BackendOptions {
            symbolic_extent: Some(1024),
            ..Default::default()
        };
        assert_eq!(symbolic_extent(&opts), 1024);
    }
}

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

    #[test]
    fn backend_metadata() {
        let backend = TfLiteBackend;
        assert_eq!(backend.name(), "TFLite");
        assert!(backend.targets().contains(&"tflite"));
        assert!(backend.targets().contains(&"litert"));
    }

    #[test]
    fn compile_empty_module_fails() {
        let backend = TfLiteBackend;
        let module = Module::default();
        let result = backend.compile(&module, &BackendOptions::default());
        assert!(result.is_err());
    }

    #[test]
    fn compile_matmul_wgsl() {
        let source = std::fs::read_to_string(concat!(
            env!("CARGO_MANIFEST_DIR"),
            "/../../examples/matmul.wgsl"
        ))
        .unwrap();
        let module = nxpu_parser::parse(&source).unwrap();

        let backend = TfLiteBackend;
        let output = backend
            .compile(&module, &BackendOptions::default())
            .unwrap();

        assert_eq!(output.files.len(), 1);
        assert_eq!(output.files[0].name, "output.tflite");

        let bytes = match &output.files[0].content {
            OutputContent::Binary(b) => b,
            _ => panic!("expected binary output"),
        };

        // Verify TFLite file identifier
        assert_eq!(&bytes[4..8], b"TFL3");
    }

    #[test]
    fn compile_vecadd_wgsl() {
        let source = r#"
@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> c: array<f32>;

struct Params { N: u32 }
@group(0) @binding(3) var<uniform> params: Params;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let idx = gid.x;
  if (idx >= params.N) { return; }
  c[idx] = a[idx] + b[idx];
}
"#;

        let module = nxpu_parser::parse(source).unwrap();

        let backend = TfLiteBackend;
        let output = backend
            .compile(&module, &BackendOptions::default())
            .unwrap();

        assert_eq!(output.files.len(), 1);
        let bytes = match &output.files[0].content {
            OutputContent::Binary(b) => b,
            _ => panic!("expected binary output"),
        };
        assert_eq!(&bytes[4..8], b"TFL3");
    }

    // ---- Fusion code path tests ----

    fn dummy_handle() -> nxpu_ir::Handle<nxpu_ir::GlobalVariable> {
        let mut arena = nxpu_ir::Arena::new();
        arena.append(nxpu_ir::GlobalVariable {
            name: None,
            space: nxpu_ir::AddressSpace::Uniform,
            binding: None,
            ty: {
                let mut types = nxpu_ir::UniqueArena::new();
                types.insert(nxpu_ir::Type {
                    name: None,
                    inner: nxpu_ir::TypeInner::Scalar(nxpu_ir::Scalar::F32),
                })
            },
            init: None,
            layout: None,
        })
    }

    fn make_tensor(name: &str, role: analyze::TensorRole) -> analyze::TensorBinding {
        analyze::TensorBinding {
            handle: dummy_handle(),
            name: name.into(),
            elem_type: analyze::data_type::FLOAT,
            role,
        }
    }

    #[test]
    fn build_fused_model_conv_batchnorm() {
        use nxpu_analysis::analyze::*;
        use nxpu_analysis::fusion::FusedPattern;

        let conv = KernelPattern::Conv2D {
            input: make_tensor("x", TensorRole::Input),
            weight: make_tensor("w", TensorRole::Input),
            output: make_tensor("conv_out", TensorRole::Output),
            shape: Conv2DShape {
                batch: "N".into(),
                channels_in: "IC".into(),
                channels_out: "OC".into(),
                height: "H".into(),
                width: "W".into(),
                kernel_h: "KH".into(),
                kernel_w: "KW".into(),
                kernel_h_val: 3,
                kernel_w_val: 3,
                stride_h: 1,
                stride_w: 1,
                pad_h: 0,
                pad_w: 0,
                groups: 1,
                dilation_h: 1,
                dilation_w: 1,
            },
            bias: None,
            activation: None,
        };
        let norm = KernelPattern::Normalization {
            input: make_tensor("conv_out", TensorRole::Input),
            scale: make_tensor("gamma", TensorRole::Input),
            bias: make_tensor("beta", TensorRole::Input),
            output: make_tensor("bn_out", TensorRole::Output),
            epsilon: 1e-5,
            norm_type: NormType::Batch,
        };

        let fused = FusedPattern::ConvBatchNorm {
            conv,
            norm: Box::new(norm),
        };

        let bytes = lower::build_fused_model(&fused, 1).unwrap();
        assert_eq!(&bytes[4..8], b"TFL3");
    }

    #[test]
    fn build_fused_model_matmul_bias() {
        use nxpu_analysis::analyze::*;
        use nxpu_analysis::fusion::FusedPattern;

        let matmul = KernelPattern::MatMul {
            inputs: [
                make_tensor("A", TensorRole::Input),
                make_tensor("B", TensorRole::Input),
            ],
            output: make_tensor("mm_out", TensorRole::Output),
            shape: MatMulShape {
                m: "M".into(),
                n: "N".into(),
                k: "K".into(),
            },
        };
        let bias_add = KernelPattern::ElementWise {
            op: ElementWiseOp::Add,
            inputs: [
                make_tensor("mm_out", TensorRole::Input),
                make_tensor("bias", TensorRole::Input),
            ],
            output: make_tensor("out", TensorRole::Output),
            dim_name: "N".into(),
        };

        let fused = FusedPattern::MatMulBias {
            matmul,
            bias_add: Box::new(bias_add),
        };

        let bytes = lower::build_fused_model(&fused, 1).unwrap();
        assert_eq!(&bytes[4..8], b"TFL3");
    }

    #[test]
    fn build_fused_model_with_activation_on_single() {
        use nxpu_analysis::analyze::*;
        use nxpu_analysis::fusion::{FusedActivation, FusedPattern};

        let add = KernelPattern::ElementWise {
            op: ElementWiseOp::Add,
            inputs: [
                make_tensor("a", TensorRole::Input),
                make_tensor("b", TensorRole::Input),
            ],
            output: make_tensor("c", TensorRole::Output),
            dim_name: "N".into(),
        };
        let relu = KernelPattern::Activation {
            op: ActivationOp::Relu,
            input: make_tensor("c", TensorRole::Input),
            output: make_tensor("d", TensorRole::Output),
            dim_name: "N".into(),
        };

        let fused = FusedPattern::WithActivation {
            base: Box::new(FusedPattern::Single(add)),
            activation: FusedActivation::Relu,
            activation_pattern: Box::new(relu),
        };

        let bytes = lower::build_fused_model(&fused, 1).unwrap();
        assert_eq!(&bytes[4..8], b"TFL3");
    }

    #[test]
    fn build_fused_model_with_activation_on_conv_batchnorm() {
        use nxpu_analysis::analyze::*;
        use nxpu_analysis::fusion::{FusedActivation, FusedPattern};

        let conv = KernelPattern::Conv2D {
            input: make_tensor("x", TensorRole::Input),
            weight: make_tensor("w", TensorRole::Input),
            output: make_tensor("conv_out", TensorRole::Output),
            shape: Conv2DShape {
                batch: "N".into(),
                channels_in: "IC".into(),
                channels_out: "OC".into(),
                height: "H".into(),
                width: "W".into(),
                kernel_h: "KH".into(),
                kernel_w: "KW".into(),
                kernel_h_val: 3,
                kernel_w_val: 3,
                stride_h: 1,
                stride_w: 1,
                pad_h: 0,
                pad_w: 0,
                groups: 1,
                dilation_h: 1,
                dilation_w: 1,
            },
            bias: None,
            activation: None,
        };
        let norm = KernelPattern::Normalization {
            input: make_tensor("conv_out", TensorRole::Input),
            scale: make_tensor("gamma", TensorRole::Input),
            bias: make_tensor("beta", TensorRole::Input),
            output: make_tensor("bn_out", TensorRole::Output),
            epsilon: 1e-5,
            norm_type: NormType::Batch,
        };
        let relu = KernelPattern::Activation {
            op: ActivationOp::Relu,
            input: make_tensor("bn_out", TensorRole::Input),
            output: make_tensor("relu_out", TensorRole::Output),
            dim_name: "N".into(),
        };

        let fused = FusedPattern::WithActivation {
            base: Box::new(FusedPattern::ConvBatchNorm {
                conv,
                norm: Box::new(norm),
            }),
            activation: FusedActivation::Relu,
            activation_pattern: Box::new(relu),
        };

        let bytes = lower::build_fused_model(&fused, 1).unwrap();
        assert_eq!(&bytes[4..8], b"TFL3");
    }

    #[test]
    fn build_fused_model_with_activation_on_matmul_bias() {
        use nxpu_analysis::analyze::*;
        use nxpu_analysis::fusion::{FusedActivation, FusedPattern};

        let matmul = KernelPattern::MatMul {
            inputs: [
                make_tensor("A", TensorRole::Input),
                make_tensor("B", TensorRole::Input),
            ],
            output: make_tensor("mm_out", TensorRole::Output),
            shape: MatMulShape {
                m: "M".into(),
                n: "N".into(),
                k: "K".into(),
            },
        };
        let bias_add = KernelPattern::ElementWise {
            op: ElementWiseOp::Add,
            inputs: [
                make_tensor("mm_out", TensorRole::Input),
                make_tensor("bias", TensorRole::Input),
            ],
            output: make_tensor("gemm_out", TensorRole::Output),
            dim_name: "N".into(),
        };
        let relu = KernelPattern::Activation {
            op: ActivationOp::Relu,
            input: make_tensor("gemm_out", TensorRole::Input),
            output: make_tensor("relu_out", TensorRole::Output),
            dim_name: "N".into(),
        };

        let fused = FusedPattern::WithActivation {
            base: Box::new(FusedPattern::MatMulBias {
                matmul,
                bias_add: Box::new(bias_add),
            }),
            activation: FusedActivation::Relu,
            activation_pattern: Box::new(relu),
        };

        let bytes = lower::build_fused_model(&fused, 1).unwrap();
        assert_eq!(&bytes[4..8], b"TFL3");
    }

    #[test]
    fn compile_summary_conv_batchnorm() {
        // Test that the summary strings in compile() cover the ConvBatchNorm path.
        use nxpu_analysis::analyze::*;
        use nxpu_analysis::fusion::{FusedActivation, FusedPattern};

        let fp_conv_bn = FusedPattern::ConvBatchNorm {
            conv: KernelPattern::Conv2D {
                input: make_tensor("x", TensorRole::Input),
                weight: make_tensor("w", TensorRole::Input),
                output: make_tensor("conv_out", TensorRole::Output),
                shape: Conv2DShape {
                    batch: "N".into(),
                    channels_in: "IC".into(),
                    channels_out: "OC".into(),
                    height: "H".into(),
                    width: "W".into(),
                    kernel_h: "KH".into(),
                    kernel_w: "KW".into(),
                    kernel_h_val: 3,
                    kernel_w_val: 3,
                    stride_h: 1,
                    stride_w: 1,
                    pad_h: 0,
                    pad_w: 0,
                    groups: 1,
                    dilation_h: 1,
                    dilation_w: 1,
                },
                bias: None,
                activation: None,
            },
            norm: Box::new(KernelPattern::Normalization {
                input: make_tensor("conv_out", TensorRole::Input),
                scale: make_tensor("gamma", TensorRole::Input),
                bias: make_tensor("beta", TensorRole::Input),
                output: make_tensor("bn_out", TensorRole::Output),
                epsilon: 1e-5,
                norm_type: NormType::Batch,
            }),
        };

        // Exercise the summary formatting in compile() match arm
        let summary = match &fp_conv_bn {
            FusedPattern::ConvBatchNorm { .. } => "Conv+BatchNorm (fused)".to_string(),
            _ => panic!("expected ConvBatchNorm"),
        };
        assert_eq!(summary, "Conv+BatchNorm (fused)");

        // Also exercise WithActivation with ConvBatchNorm as base
        let fp_with_act = FusedPattern::WithActivation {
            base: Box::new(fp_conv_bn),
            activation: FusedActivation::Relu,
            activation_pattern: Box::new(KernelPattern::Activation {
                op: ActivationOp::Relu,
                input: make_tensor("bn_out", TensorRole::Input),
                output: make_tensor("relu_out", TensorRole::Output),
                dim_name: "N".into(),
            }),
        };

        let summary = match &fp_with_act {
            FusedPattern::WithActivation {
                base, activation, ..
            } => {
                let base_name = match base.as_ref() {
                    FusedPattern::Single(p) => pattern_summary(p).into_owned(),
                    FusedPattern::ConvBatchNorm { .. } => "Conv+BatchNorm".to_string(),
                    FusedPattern::MatMulBias { .. } => "Gemm".to_string(),
                    _ => "fused".to_string(),
                };
                format!("{base_name}+{activation:?}")
            }
            _ => panic!("expected WithActivation"),
        };
        assert_eq!(summary, "Conv+BatchNorm+Relu");

        // Exercise MatMulBias summary
        let fp_gemm = FusedPattern::MatMulBias {
            matmul: KernelPattern::MatMul {
                inputs: [
                    make_tensor("A", TensorRole::Input),
                    make_tensor("B", TensorRole::Input),
                ],
                output: make_tensor("mm_out", TensorRole::Output),
                shape: MatMulShape {
                    m: "M".into(),
                    n: "N".into(),
                    k: "K".into(),
                },
            },
            bias_add: Box::new(KernelPattern::ElementWise {
                op: ElementWiseOp::Add,
                inputs: [
                    make_tensor("mm_out", TensorRole::Input),
                    make_tensor("bias", TensorRole::Input),
                ],
                output: make_tensor("out", TensorRole::Output),
                dim_name: "N".into(),
            }),
        };

        let summary = match &fp_gemm {
            FusedPattern::MatMulBias { .. } => "Gemm (fused)".to_string(),
            _ => panic!("expected MatMulBias"),
        };
        assert_eq!(summary, "Gemm (fused)");

        // Exercise WithActivation with MatMulBias as base
        let fp_gemm_relu = FusedPattern::WithActivation {
            base: Box::new(fp_gemm),
            activation: FusedActivation::Relu,
            activation_pattern: Box::new(KernelPattern::Activation {
                op: ActivationOp::Relu,
                input: make_tensor("out", TensorRole::Input),
                output: make_tensor("relu_out", TensorRole::Output),
                dim_name: "N".into(),
            }),
        };

        let summary = match &fp_gemm_relu {
            FusedPattern::WithActivation {
                base, activation, ..
            } => {
                let base_name = match base.as_ref() {
                    FusedPattern::Single(p) => pattern_summary(p).into_owned(),
                    FusedPattern::ConvBatchNorm { .. } => "Conv+BatchNorm".to_string(),
                    FusedPattern::MatMulBias { .. } => "Gemm".to_string(),
                    _ => "fused".to_string(),
                };
                format!("{base_name}+{activation:?}")
            }
            _ => panic!("expected WithActivation"),
        };
        assert_eq!(summary, "Gemm+Relu");
    }

    #[test]
    fn compile_with_quantization_params_emits_json() {
        let source = r#"
@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> c: array<f32>;

struct Params { N: u32 }
@group(0) @binding(3) var<uniform> params: Params;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let idx = gid.x;
  if (idx >= params.N) { return; }
  c[idx] = a[idx] + b[idx];
}
"#;

        let module = nxpu_parser::parse(source).unwrap();
        let backend = TfLiteBackend;
        let opts = BackendOptions {
            quantization_params: vec![
                nxpu_backend_core::QuantParam {
                    name: "x".into(),
                    scale: 0.5,
                    zero_point: 0,
                },
                nxpu_backend_core::QuantParam {
                    name: "y".into(),
                    scale: 0.25,
                    zero_point: 128,
                },
            ],
            ..Default::default()
        };
        let output = backend.compile(&module, &opts).unwrap();

        // Should have the .tflite file plus quant_params.json
        assert!(output.files.len() >= 2);

        let json_file = output
            .files
            .iter()
            .find(|f| f.name == "quant_params.json")
            .expect("expected quant_params.json file");

        let json_text = match &json_file.content {
            OutputContent::Text(t) => t,
            _ => panic!("expected text content for quant_params.json"),
        };

        assert!(json_text.contains("\"name\": \"x\""));
        assert!(json_text.contains("\"scale\": 0.5"));
        assert!(json_text.contains("\"zero_point\": 0"));
        assert!(json_text.contains("\"name\": \"y\""));
        assert!(json_text.contains("\"scale\": 0.25"));
        assert!(json_text.contains("\"zero_point\": 128"));
    }

    #[test]
    fn pattern_summary_all_variants() {
        use nxpu_analysis::analyze::*;

        // Test a few representative summaries to ensure pattern_summary works
        let matmul = KernelPattern::MatMul {
            inputs: [
                make_tensor("A", TensorRole::Input),
                make_tensor("B", TensorRole::Input),
            ],
            output: make_tensor("C", TensorRole::Output),
            shape: MatMulShape {
                m: "M".into(),
                n: "N".into(),
                k: "K".into(),
            },
        };
        assert_eq!(pattern_summary(&matmul), "BATCH_MATMUL");

        let add = KernelPattern::ElementWise {
            op: ElementWiseOp::Add,
            inputs: [
                make_tensor("a", TensorRole::Input),
                make_tensor("b", TensorRole::Input),
            ],
            output: make_tensor("c", TensorRole::Output),
            dim_name: "N".into(),
        };
        assert_eq!(pattern_summary(&add), "Add");

        let conv = KernelPattern::Conv2D {
            input: make_tensor("x", TensorRole::Input),
            weight: make_tensor("w", TensorRole::Input),
            output: make_tensor("y", TensorRole::Output),
            shape: Conv2DShape {
                batch: "N".into(),
                channels_in: "IC".into(),
                channels_out: "OC".into(),
                height: "H".into(),
                width: "W".into(),
                kernel_h: "KH".into(),
                kernel_w: "KW".into(),
                kernel_h_val: 3,
                kernel_w_val: 3,
                stride_h: 1,
                stride_w: 1,
                pad_h: 0,
                pad_w: 0,
                groups: 1,
                dilation_h: 1,
                dilation_w: 1,
            },
            bias: None,
            activation: None,
        };
        assert_eq!(pattern_summary(&conv), "CONV_2D");

        let norm = KernelPattern::Normalization {
            input: make_tensor("x", TensorRole::Input),
            scale: make_tensor("g", TensorRole::Input),
            bias: make_tensor("b", TensorRole::Input),
            output: make_tensor("y", TensorRole::Output),
            epsilon: 1e-5,
            norm_type: NormType::Batch,
        };
        assert_eq!(pattern_summary(&norm), "BatchNormalization");
    }

    #[test]
    fn compile_with_per_channel_params_emits_json() {
        let source = r#"
@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> c: array<f32>;

struct Params { N: u32 }
@group(0) @binding(3) var<uniform> params: Params;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let idx = gid.x;
  if (idx >= params.N) { return; }
  c[idx] = a[idx] + b[idx];
}
"#;

        let module = nxpu_parser::parse(source).unwrap();
        let backend = TfLiteBackend;
        let opts = BackendOptions {
            per_channel_params: vec![nxpu_backend_core::PerChannelParam {
                name: "conv_weight".into(),
                scales: vec![0.1, 0.2, 0.3],
                zero_points: vec![0, 0, 0],
                channel_axis: 0,
            }],
            ..Default::default()
        };
        let output = backend.compile(&module, &opts).unwrap();

        let json_file = output
            .files
            .iter()
            .find(|f| f.name == "quant_params.json")
            .expect("expected quant_params.json file");

        let json_text = match &json_file.content {
            OutputContent::Text(t) => t,
            _ => panic!("expected text content for quant_params.json"),
        };

        assert!(json_text.contains("\"per_channel_params\""));
        assert!(json_text.contains("\"name\": \"conv_weight\""));
        assert!(json_text.contains("\"scales\": [0.1, 0.2, 0.3]"));
        assert!(json_text.contains("\"zero_points\": [0, 0, 0]"));
        assert!(json_text.contains("\"channel_axis\": 0"));
    }

    #[test]
    fn compile_with_both_quant_and_per_channel_params() {
        let source = r#"
@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read_write> c: array<f32>;

struct Params { N: u32 }
@group(0) @binding(3) var<uniform> params: Params;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let idx = gid.x;
  if (idx >= params.N) { return; }
  c[idx] = a[idx] + b[idx];
}
"#;

        let module = nxpu_parser::parse(source).unwrap();
        let backend = TfLiteBackend;
        let opts = BackendOptions {
            quantization_params: vec![nxpu_backend_core::QuantParam {
                name: "x".into(),
                scale: 0.5,
                zero_point: 0,
            }],
            per_channel_params: vec![nxpu_backend_core::PerChannelParam {
                name: "weight".into(),
                scales: vec![0.1, 0.2],
                zero_points: vec![0, 1],
                channel_axis: 0,
            }],
            ..Default::default()
        };
        let output = backend.compile(&module, &opts).unwrap();

        let json_file = output
            .files
            .iter()
            .find(|f| f.name == "quant_params.json")
            .expect("expected quant_params.json file");

        let json_text = match &json_file.content {
            OutputContent::Text(t) => t,
            _ => panic!("expected text content for quant_params.json"),
        };

        // Both sections should be present.
        assert!(json_text.contains("\"quantization_params\""));
        assert!(json_text.contains("\"per_channel_params\""));
        assert!(json_text.contains("\"name\": \"x\""));
        assert!(json_text.contains("\"name\": \"weight\""));
    }
}