Only advance branch quantizes that share qparams

backends/cadence/aot/reorder_ops.py

@@ -206,13 +206,15 @@ def advance_quantize_op(self, graph_module: torch.fx.GraphModule):
                 graph.erase_node(quant_node)
                 continue
 
-            # Otherwise we have the quant descendents. Cluster them into sets
-            # that have the same scale, zero_point, and dtype. We use quant_dict
-            # for the clustering
-            quant_dict: DefaultDict[Tuple, int] = defaultdict(int)
-            for quant_node in descendent_quant_ops:
-                quant_dict[quant_node.args[1:]] += 1
-            rep_args = sorted(quant_dict.keys(), key=lambda x: x[1]).pop()
+            # Otherwise we have multiple quant descendents. Requantizing them from
+            # one hoisted quantize is only lossless if they all share qparams, so
+            # skip when they differ.
+            quant_args = {
+                tuple(quant_node.args[1:]) for quant_node in descendent_quant_ops
+            }
+            if len(quant_args) > 1:
+                continue
+            (rep_args,) = quant_args
 
             # Create a new quant node that dominates all the nodes in
             # descendent_quant_ops. Make sure that the insertion is after

backends/cadence/aot/tests/test_reorder_ops_passes.py

@@ -8,7 +8,7 @@
 
 import copy
 import unittest
-from typing import cast
+from typing import cast, Tuple
 
 import executorch.backends.cadence.aot.ops_registrations  # noqa
 import torch
@@ -232,7 +232,7 @@ def test_advance_branched_quantize(self) -> None:
         )
         quantized_decomposed_quantize_per_tensor_default_3 = builder.call_operator(
             op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
-            args=(aten_slice_copy_tensor_3, 0.2, 4, 0, 255, torch.uint8),
+            args=(aten_slice_copy_tensor_3, 0.1, 10, 0, 255, torch.uint8),
         )
         builder.output(
             [
@@ -278,16 +278,74 @@ def test_advance_branched_quantize(self) -> None:
         )
         p = FuseQuantDequantToRequantizePass()
         graph_module = cast(PassResult, p(graph_module)).graph_module
-        # We expect 3 dequant/quant pairs to be removed because they have matching params,
-        # leaving a single dequant/quant pair that is then merged into a requantize op
+        # All the branches share qparams, so every dequant/quant pair has matching
+        # params and is removed, leaving no requantize ops
         self.assertEqual(
             count_node(
                 graph_module,
                 exir_ops.edge.cadence.requantize.per_tensor,
             ),
-            1,
+            0,
         )
 
+    def _advance_two_branch_quantize(
+        self,
+        first_qparams: Tuple[float, int],
+        second_qparams: Tuple[float, int],
+    ) -> torch.fx.GraphModule:
+        builder = GraphBuilder()
+        x = builder.placeholder("x", torch.randn(24, dtype=torch.float32))
+        quantized = []
+        for start, (scale, zero_point) in (
+            (0, first_qparams),
+            (12, second_qparams),
+        ):
+            sliced = builder.call_operator(
+                op=exir_ops.edge.aten.slice_copy.Tensor,
+                args=(x, 0, start, start + 12),
+            )
+            quantized.append(
+                builder.call_operator(
+                    op=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+                    args=(sliced, scale, zero_point, 0, 255, torch.uint8),
+                )
+            )
+        builder.output(quantized)
+        graph_module = cast(
+            PassResult,
+            AdvanceQuantizeOpAboveDefInBranchPass()(builder.get_graph_module()),
+        ).graph_module
+        graph_module.graph.eliminate_dead_code()
+        return graph_module
+
+    def test_does_not_advance_quantize_when_qparams_differ(self) -> None:
+        # Requantizing from one hoisted quantize can clip or coarsen a branch, so
+        # the branches are left alone unless they already share qparams
+        for first_qparams, second_qparams in (
+            ((0.1, 10), (2.0, 4)),
+            ((0.1, 10), (0.1, 100)),
+            ((0.1, 10), (0.11, 10)),
+        ):
+            with self.subTest(first=first_qparams, second=second_qparams):
+                graph_module = self._advance_two_branch_quantize(
+                    first_qparams, second_qparams
+                )
+
+                self.assertEqual(
+                    count_node(
+                        graph_module,
+                        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+                    ),
+                    2,
+                )
+                self.assertEqual(
+                    count_node(
+                        graph_module,
+                        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+                    ),
+                    0,
+                )
+
     def test_does_not_advance_quantize_for_half_size_slice(self) -> None:
         builder = GraphBuilder()
         x = builder.placeholder("x", torch.randn(4, 60, 1, 1))