diff --git a/backends/cadence/aot/reorder_ops.py b/backends/cadence/aot/reorder_ops.py index 2700e1ccf28..4ee7bb193d6 100644 --- a/backends/cadence/aot/reorder_ops.py +++ b/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 diff --git a/backends/cadence/aot/tests/test_reorder_ops_passes.py b/backends/cadence/aot/tests/test_reorder_ops_passes.py index e0a9e7937a0..b7644d28525 100644 --- a/backends/cadence/aot/tests/test_reorder_ops_passes.py +++ b/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))