NXP backend: Add test for Linear+BatchNorm fusing

Author: StrycekSimon

backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py

@@ -0,0 +1,188 @@
+# Copyright 2026 NXP
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import executorch.backends.nxp.tests.models as models
+import pytest
+import torch
+from executorch.backends.nxp.aten_passes.add_simulated_linear_bn_fusion_qat_pass import (
+    AddSimulatedLinearBatchNormFusionQATPass,
+)
+from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import (
+    FuseBatchNormWithLinearPass,
+)
+from executorch.backends.nxp.aten_passes.remove_simulated_linear_bn_fusion_qat_pass import (
+    RemoveSimulatedLinearBatchNormFusionQATPass,
+)
+
+from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
+from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec
+from torch.export import export
+from torch.fx import Node
+from torchao.quantization.pt2e.prepare import _is_activation_post_process_node
+from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e
+
+
+@pytest.mark.parametrize("input_shape", [(1, 5, 5), (1, 5, 5, 5)])
+@pytest.mark.parametrize("linear_bias", [True, False])
+def test_add_simulated_linear_bn_fusing(input_shape, linear_bias):
+    random_input = torch.randn(*input_shape)
+    model = models.LinearBNModule(
+        input_shape=input_shape,
+        out_features=5,
+        linear_bias=linear_bias,
+    )
+    model.train()
+    raw_output = model(random_input)
+
+    exported_model = export(model, (random_input,), strict=True)
+    prepared_model = prepare_qat_pt2e(
+        exported_model.module(), NeutronQuantizer(neutron_target_spec, is_qat=True)
+    )
+    prepared_model = AddSimulatedLinearBatchNormFusionQATPass()(
+        prepared_model
+    ).graph_module
+
+    graph_nodes = list(prepared_model.graph.nodes)
+    named_modules = dict(prepared_model.named_modules(remove_duplicate=False))
+    fake_quantize_output = prepared_model(random_input)
+
+    expected_number_of_nodes = 23 if linear_bias else 18
+    linear_node = next(
+        (
+            n
+            for n in graph_nodes
+            if hasattr(n, "target") and n.target == torch.ops.aten.linear.default
+        ),
+        None,
+    )
+
+    assert len(graph_nodes) == expected_number_of_nodes
+
+    # Assert Linear weight being quantized and "normalized"
+    assert linear_node is not None
+    assert all(
+        _is_activation_post_process_node(n, named_modules) for n in linear_node.args
+    )
+    assert linear_node.args[1].args[0].target == torch.ops.aten.mul.Tensor
+
+    # Assert BatchNorm input being "denormalized"
+    assert graph_nodes[-3].target == torch.ops.aten.batch_norm.default
+    if linear_bias:
+        assert graph_nodes[-3].args[0].target == torch.ops.aten.add.Tensor
+        add_arg_targets = (
+            n.target for n in graph_nodes[-3].args[0].args if hasattr(n, "target")
+        )
+        assert torch.ops.aten.div.Tensor in add_arg_targets
+    else:
+        assert graph_nodes[-3].args[0].target == torch.ops.aten.div.Tensor
+
+    assert raw_output.shape == fake_quantize_output.shape
+
+
+@pytest.mark.parametrize("input_shape", [(1, 5, 5), (1, 5, 5, 5)])
+@pytest.mark.parametrize("linear_bias", [True, False])
+def test_full_linear_bn_fusing(input_shape, linear_bias):
+    # TODO: Add pass for quantizing bias node when Linear has bias=False
+    if not linear_bias:
+        pytest.skip(
+            "Linear with bias=False is not yet supported."
+            "The graph currently produces Linear layer without quantized bias which is incorrect."
+        )
+
+    random_input = torch.randn(*input_shape)
+    model = models.LinearBNModule(
+        input_shape=input_shape,
+        out_features=5,
+        linear_bias=linear_bias,
+    )
+    model.train()
+    raw_output = model(random_input)
+
+    exported_model = export(model, (random_input,), strict=True)
+    prepared_model = prepare_qat_pt2e(
+        exported_model.module(), NeutronQuantizer(neutron_target_spec, is_qat=True)
+    )
+
+    prepared_model = AddSimulatedLinearBatchNormFusionQATPass()(
+        prepared_model
+    ).graph_module
+    prepared_model(random_input)
+    prepared_model = RemoveSimulatedLinearBatchNormFusionQATPass()(
+        prepared_model
+    ).graph_module
+    prepared_model = FuseBatchNormWithLinearPass()(prepared_model).graph_module
+    converted_model = convert_pt2e(prepared_model)
+
+    quantized_output = converted_model(random_input)
+    graph_nodes = list(converted_model.graph.nodes)
+    linear_node = graph_nodes[-4]
+
+    def _is_bn(node_: Node) -> bool:
+        return (
+            hasattr(node_, "target")
+            and node_.target == torch.ops.aten.batch_norm.default
+        )
+
+    assert len(graph_nodes) == 11
+
+    assert not any(_is_bn(node) for node in graph_nodes)
+
+    # Assert linear inputs being quantized
+    assert linear_node.target == torch.ops.aten.linear.default
+    assert (
+        linear_node.args[0].target
+        == torch.ops.quantized_decomposed.dequantize_per_tensor.default
+    )
+    assert (
+        linear_node.args[1].target
+        == torch.ops.quantized_decomposed.dequantize_per_tensor.default
+    )
+
+    # Assert linear outputs being quantized
+    assert len(linear_node.users) == 1
+    assert (
+        list(linear_node.users.keys())[0].target
+        == torch.ops.quantized_decomposed.quantize_per_tensor.default
+    )
+
+    assert raw_output.shape == quantized_output.shape
+
+
+@pytest.mark.parametrize("input_shape", [(1, 5, 5), (1, 5, 5, 5)])
+@pytest.mark.parametrize("linear_bias", [True, False])
+@pytest.mark.parametrize("bn_eps", [1e-5, 1e-6])
+def test_input_output_graph_equivalence(input_shape, linear_bias, bn_eps):
+    # TODO: Add pass for quantizing bias node when Linear has bias=False
+    if not linear_bias:
+        pytest.skip(
+            "Linear with bias=False is not yet supported."
+            "The graph currently produces Linear layer without quantized bias which is incorrect."
+        )
+
+    random_input = torch.randn(*input_shape)
+    model = models.LinearBNModule(
+        input_shape=input_shape,
+        out_features=5,
+        linear_bias=linear_bias,
+        bn_eps=bn_eps,
+    )
+    model.eval()
+
+    original_model = export(model, (random_input,), strict=True).module()
+
+    processed_model = export(model, (random_input,), strict=True).module()
+    processed_model = AddSimulatedLinearBatchNormFusionQATPass()(
+        processed_model
+    ).graph_module
+
+    assert list(processed_model.graph.nodes)[8].args[1] == bn_eps
+
+    processed_model = RemoveSimulatedLinearBatchNormFusionQATPass()(
+        processed_model
+    ).graph_module
+
+    assert list(processed_model.graph.nodes)[-2].args[7] == bn_eps
+    assert torch.equal(original_model(random_input), processed_model(random_input))
+    assert len(original_model.graph.nodes) == len(processed_model.graph.nodes)

backends/nxp/tests/models.py

@@ -481,6 +481,42 @@ def forward(self, x):
         return self.bn(x)
 
 
+class LinearBNModule(torch.nn.Module):
+    def __init__(
+        self,
+        input_shape: tuple[int],
+        out_features: int,
+        linear_bias: bool,
+        bn_eps: float = 1e-5,
+        act: nn.Module | None = None,
+    ):
+        super().__init__()
+
+        self.linear = torch.nn.Linear(
+            in_features=input_shape[-1], out_features=out_features, bias=linear_bias
+        )
+
+        num_dims = len(input_shape)
+        if num_dims == 3:
+            self.bn = torch.nn.BatchNorm1d(out_features, eps=bn_eps)
+        elif num_dims == 4:
+            self.bn = torch.nn.BatchNorm2d(out_features, eps=bn_eps)
+        elif num_dims == 5:
+            self.bn = torch.nn.BatchNorm3d(out_features, eps=bn_eps)
+        else:
+            raise ValueError(
+                f"Unsupported dimension {len(input_shape)} of the input_shape "
+                + f"({input_shape}). Only 3, 4 and 5 are supported."
+            )
+
+        self.act = act
+
+    def forward(self, x):
+        x = self.linear(x)
+        x = self.bn(x)
+        return self.act(x) if self.act is not None else x
+
+
 class MulTensorModule(torch.nn.Module):
     def __init__(self):
         super().__init__()