NXP backend: Improve quantization annotation process (#11908)

### Summary
Addresses multiple subsequent nodes with SharedSpecPattern problem. 



---------

Co-authored-by: Roman Janik <[email protected]>
Co-authored-by: Simon Strycek <[email protected]>

Author: 3 people

backends/nxp/quantizer/neutron_quantizer.py

@@ -10,7 +10,6 @@
 from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import (
     NeutronAtenPassManager,
 )
-
 from executorch.backends.nxp.quantizer.patterns import (
     AddmmPattern,
     AvgPoolPattern,
@@ -24,6 +23,7 @@
     ReluInPlacePattern,
     ReluPattern,
     ReshapePattern,
+    SharedSpecPattern,
     SoftMaxPattern,
     ViewPattern,
 )
@@ -204,9 +204,36 @@ def __init__(self):
                 NeutronAtenQuantizer(ViewPattern(), static_qconfig),
             ]
         )
+        # Mapping ops defined in quantizer partition types to its quantizer
+        self.op_to_quantizer = {
+            pt: q for q in self.quantizers for pt in q.pattern.partition_types()
+        }
+        # Mapping ops to the quantizer application state
+        self.op_to_applied_quantizer = {
+            pt: False for q in self.quantizers for pt in q.pattern.partition_types()
+        }
 
     def transform_for_annotation(
         self, model: torch.fx.GraphModule
     ) -> torch.fx.GraphModule:
         pass_runner = NeutronAtenPassManager()
         return pass_runner(model).graph_module
+
+    def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
+        nodes = list(model.graph.nodes)
+        for node in nodes:
+            if (
+                node.target not in self.op_to_quantizer
+                or self.op_to_applied_quantizer[node.target]
+            ):
+                continue
+            else:
+                quantizer = self.op_to_quantizer[node.target]
+                quantizer.annotate(model)
+                if not isinstance(quantizer.pattern, SharedSpecPattern):
+                    self.op_to_applied_quantizer[node.target] = True
+
+        return model
+
+    def validate(self, model: torch.fx.GraphModule) -> None:
+        return super().validate(model)

backends/nxp/tests/models.py

@@ -239,3 +239,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,73 @@ 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"
+
+    # Assert that post-ReLU quantize and pre-MaxPool dequantize has same specs
+    assert nodes[-6].args[1:] == nodes[-5].args[1:]
+    # Assert that post-Conv quantize and pre-ReLU dequantize has same specs
+    assert nodes[6].args[1:] == nodes[7].args[1:]
+
+
+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))