NXP backend: Add test for Linear+BatchNorm fusing

Author: StrycekSimon

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

@@ -0,0 +1,175 @@
+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.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)
+    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
+    # "activation_post_process" is not a standard operator so we need to check by name.
+    assert all(n.target.startswith("activation_post_process") 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
+    for data in (random_input,):
+        prepared_model(*data)
+    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])
+def test_input_output_graph_equivalence(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.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
+    processed_model = RemoveSimulatedLinearBatchNormFusionQATPass()(
+        processed_model
+    ).graph_module
+
+    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,41 @@ 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,
+        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)
+        bn_features = input_shape[1]
+        if num_dims == 3:
+            self.bn = torch.nn.BatchNorm1d(bn_features)
+        elif num_dims == 4:
+            self.bn = torch.nn.BatchNorm2d(bn_features)
+        elif num_dims == 5:
+            self.bn = torch.nn.BatchNorm3d(bn_features)
+        else:
+            raise ValueError(
+                f"Unknown input_dim: {len(input_shape)}, supported values are 1, 2 or 3."
+            )
+
+        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__()