From 71b18832cca1de008f4e714aa1ddf3499f1bc5a0 Mon Sep 17 00:00:00 2001 From: Reza Sajadiany Date: Fri, 23 Oct 2026 15:05:00 -0700 Subject: [PATCH] Only advance branch quantizes that share qparams Summary: `AdvanceQuantizeOpAboveDefInBranchPass` hoists one quantize above a def whose users are all `slice`/`select` -> quantize branches, then rewrites each original quantize as `dequantize(shared) -> quantize(own qparams)`. The shared quantize used the qparams with the largest zero-point, without checking that they can represent the other branches. The pattern that breaks, seen in an end-to-end model compile: one tensor feeds several slice/select branches that each have their own quantize, and the branches have very different value ranges. Some branches hold unit-range values (roughly [-1, 1.7]) and another holds values up to about 100, because the tensor packs normalized values together with large-magnitude ones. The pass hoists a quantize with the narrow qparams onto the tensor, so the wide branch reads values already clipped to [-1, 1.7] before its own quantize runs. The information is gone, and neither the graph nor the existing structural test shows it. With int8 asymmetric activations this quantized the input at about 0.1 dB SQNR and left the model outputs at 1-2 dB SQNR, one of them at -0.3 dB (noise). Skipping the hoist moved pure-quantization SQNR to 17-19 dB on two of those outputs and 43 dB on the third. The pass now hoists only when every branch's quantize has identical qparams, in which case it behaves as before. When they differ it skips, and each branch keeps its own quantize on the original float def. Requantizing from one shared quantize either clips a branch with a wider range or coarsens one with a narrower range, so no choice of qparams is lossless for branches that differ; sharing qparams to get the hoist is a quantization decision. The old code counted qparams in `quant_dict` and never used the counts, then picked with `sorted(..., key=lambda x: x[1]).pop()`, which is the largest zero-point; that selection is removed. This is a behavior change: hoists across differing qparams are no longer done, which costs some data volume for those values. Downstream compile flows that use this pass see small structural changes where it no longer hoists across differing qparams: the expected node and delegate counts in five graph-structure tests change by a few nodes (one of them by a delegate). Those expectation updates are included in this diff so it passes on its own. Output numerics of the four other downstream models I measured are identical before and after. Differential Revision: D127215700 --- backends/cadence/aot/reorder_ops.py | 16 +++-- .../aot/tests/test_reorder_ops_passes.py | 68 +++++++++++++++++-- 2 files changed, 72 insertions(+), 12 deletions(-) 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))