NXP backend: Add tests for quantizer pattern order invariance

Author: StrycekSimon

backends/nxp/tests/models.py

@@ -236,3 +236,16 @@ def __init__(self):
     def forward(self, x):
         x = self.conv(x)
         return torch.permute(x, [0, 2, 1, 3])
+
+
+class Conv2dReLUMaxPoolModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(3, 64, 2, bias=False)
+        self.relu = torch.nn.ReLU()
+        self.pool = torch.nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.relu(x)
+        return self.pool(x)

backends/nxp/tests/test_quantizer.py

@@ -5,6 +5,8 @@
 
 # Tests for NeutronQuantizer.
 
+from copy import deepcopy
+
 import executorch.backends.nxp.tests.models as models
 import torch
 from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
@@ -271,3 +273,68 @@ def test_quantizer_conv2d_permute():
     assert nodes[7].name == "dequantize_per_tensor_default_2"
     assert nodes[8].name == "permute"
     assert nodes[9].name == "quantize_per_tensor_default_3"
+
+
+def test_multiple_shared_spec_ops_in_row():
+    """
+    This test demonstrates that having two operators in a row, both relying on quantizers
+    with SharedSpecPattern, does not break the quantization process.
+    """
+    model = models.Conv2dReLUMaxPoolModule()
+    model.eval()
+
+    example_input = (torch.ones(1, 3, 64, 64),)
+    quantizer = NeutronQuantizer()
+    graph_module = torch.export.export_for_training(
+        model, example_input, strict=True
+    ).module()
+
+    # noinspection PyTypeChecker
+    m = prepare_pt2e(graph_module, quantizer)
+    m(*example_input)
+    m = convert_pt2e(m)
+
+    # Dry run
+    m(*example_input)
+
+    nodes = list(m.graph.nodes)
+
+    assert len(nodes) == 15
+    assert nodes[-5].name == "dequantize_per_tensor_default_3"
+    assert nodes[-4].name == "max_pool2d"
+    assert nodes[-3].name == "quantize_per_tensor_default_4"
+
+
+def test_quantizers_order_invariance():
+    """
+    This test demonstrates that the order of quantizers in NeutronQuantizer
+    does not affect the resulting graph.
+    """
+    model = models.Conv2dReLUModule()
+    model.eval()
+
+    example_input = (torch.ones(1, 4, 64, 64),)
+    quantizer = NeutronQuantizer()
+
+    graph_module = torch.export.export_for_training(
+        model, example_input, strict=True
+    ).module()
+
+    m = prepare_pt2e(deepcopy(graph_module), quantizer)
+    m(*example_input)
+    m = convert_pt2e(m)
+
+    quantizer.quantizers = quantizer.quantizers[::-1]
+    m_reversed = prepare_pt2e(graph_module, quantizer)
+    m_reversed(*example_input)
+    m_reversed = convert_pt2e(m)
+
+    # Dry run
+    m(*example_input)
+    m_reversed(*example_input)
+
+    nodes = list(m.graph.nodes)
+    nodes_reversed = list(m.graph.nodes)
+
+    assert len(nodes) == len(nodes_reversed)
+    assert all(n == n_reversed for n, n_reversed in zip(nodes, nodes_reversed))