From 85c948a570b2ab9238bedf2fd3930f686c56b6d9 Mon Sep 17 00:00:00 2001 From: Prashant Rawat Date: Mon, 2 Nov 2026 14:09:47 -0800 Subject: [PATCH] Keep the shape operand of view_as and expand_as out of shared qparams (#24153) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: `SharedQspecQuantizer` gives each shared-qparams cluster one set of qparams for every float input and output of its ops. `view_as(x, other)` is in that list, and the cluster was built from all of its inputs, so `other` was pulled in even though it contributes only its shape. `x` and `other` were then quantized with one scale, set by whichever has the wider range: the data loses precision, or the shape operand is clipped. For example, in `(x * y).view_as(other) + other` with `other = z + z`, `x * y` spans about ±2.5 and `other` about ±51. `x * y` was quantized at `other`'s scale of 0.349, about 18x coarser than the 0.0195 it gets when the same reshape is written `reshape(4, 8)`, and the output error against float rose from 0.258 to 0.307. `expand_as` has the same shape-only second operand but was not in the shared-qparams list at all, so every `expand_as` broadcast stayed in float behind a dequantize and a requantize. Shape-as ops (`view_as`, `expand_as`) now share qparams only through their data operand, `args[0]`. The cluster walk does not cross the edge to the shape operand in either direction, and that operand gets no shared input qspec, so it keeps the qparams of its own producer. `expand_as` joins the shared-qparams ops, so like `expand` it is quantized to int8 when its neighbours are, and lowers on Cortex-M to an int8 `expand_copy`. Authored with Claude Code. Reviewed By: rascani Differential Revision: D127380012 --- backends/arm/quantizer/arm_quantizer_utils.py | 25 ++++++++++++- backends/arm/test/misc/test_shared_qspecs.py | 35 +++++++++++++++++++ backends/arm/test/models/test_swin2sr_arm.py | 3 +- .../cortex_m/test/misc/test_portable_int8.py | 6 ++++ 4 files changed, 66 insertions(+), 3 deletions(-) diff --git a/backends/arm/quantizer/arm_quantizer_utils.py b/backends/arm/quantizer/arm_quantizer_utils.py index 58f953d8795..41b798d4347 100644 --- a/backends/arm/quantizer/arm_quantizer_utils.py +++ b/backends/arm/quantizer/arm_quantizer_utils.py @@ -411,6 +411,20 @@ def validate(self, model: torch.fx.GraphModule) -> bool: # type: ignore[overrid return True +# Ops whose second operand contributes only its shape, never its values, so it +# must not be tied into the shared qparams of the first. +_SHAPE_AS_OPS = { + torch.ops.aten.expand_as.default, + torch.ops.aten.view_as.default, +} + + +# Whether ``node`` shares qparams with its input ``input_node``: always for an +# ordinary shared op, but for a shape-as op only through its data operand. +def _shares_qparams_with(node: Node, input_node: Node) -> bool: + return node.target not in _SHAPE_AS_OPS or input_node is node.args[0] + + class SharedQspecQuantizer(Quantizer, QuantizerReporterUser): """Assures that specific ops share quantization parameters on all inputs/outputs. @@ -442,6 +456,7 @@ class SharedQspecQuantizer(Quantizer, QuantizerReporterUser): torch.ops.aten.repeat_interleave.self_int, torch.ops.aten.expand_copy.default, torch.ops.aten.expand.default, + torch.ops.aten.expand_as.default, torch.ops.aten.select.int, torch.ops.aten.select_copy.int, torch.ops.aten.slice.Tensor, @@ -543,7 +558,11 @@ def _is_annotated(self, node: Node) -> bool: return Q_ANNOTATION_KEY in node.meta def _get_input_nodes_with_float_output(self, node: Node) -> list[Node]: - return [n for n in node.all_input_nodes if has_float_output(n)] + return [ + n + for n in node.all_input_nodes + if has_float_output(n) and _shares_qparams_with(node, n) + ] def _get_user_nodes_with_float_input(self, node: Node) -> list[Node]: return [n for n in node.users.keys() if has_float_output(node)] @@ -640,6 +659,8 @@ def _get_shared_clique( shared_nodes.add(node) for input_node in node.all_input_nodes: + if not _shares_qparams_with(node, input_node): + continue self._maybe_enqueue_shared_node(input_node, shared_nodes, bfs_queue) self._append_output_qspec(input_node, adjacent_qspecs) touches_quantized_io |= self._is_quantized_io_boundary(input_node) @@ -648,6 +669,8 @@ def _get_shared_clique( ) for output_node in node.users.keys(): + if not _shares_qparams_with(output_node, node): + continue self._maybe_enqueue_shared_node(output_node, shared_nodes, bfs_queue) self._append_input_qspec(output_node, node, adjacent_qspecs) touches_quantized_io |= self._is_quantized_io_boundary(output_node) diff --git a/backends/arm/test/misc/test_shared_qspecs.py b/backends/arm/test/misc/test_shared_qspecs.py index 7287aefb8af..6174fe267e4 100644 --- a/backends/arm/test/misc/test_shared_qspecs.py +++ b/backends/arm/test/misc/test_shared_qspecs.py @@ -16,6 +16,8 @@ from executorch.backends.arm.test.tester.test_pipeline import QuantizationPipeline from executorch.backends.arm.tosa import TosaSpecification from executorch.backends.cortex_m.test.tester import McuTestCase, ramp_tensor +from torch.export import export +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e _QUANT_CONFIG_INT8 = get_symmetric_quantization_config() @@ -734,3 +736,36 @@ def test_silu_does_not_share_input_qspec(): ) pipeline.run() _check_quant_params(pipeline, model.quant_params) + + +class ShapeAsOperand(torch.nn.Module): + """``op(x + x, other)`` reads only the shape of ``other``, so its qparams + must stay ``other``'s own rather than join those of ``x + x``. + """ + + def __init__(self, op: str) -> None: + super().__init__() + self.op = op + + def forward(self, x, y): + other = y + y + return getattr(x + x, self.op)(other) + other + + +@parametrize( + "test_data", + {"view_as": ("view_as", (2, 16)), "expand_as": ("expand_as", (1, 8))}, +) +def test_shape_operand_does_not_share_qspec(test_data): + op, x_shape = test_data + inputs = (ramp_tensor(-1, 1, x_shape), ramp_tensor(-50, 50, (4, 8))) + model = ShapeAsOperand(op).eval() + prepared = prepare_pt2e(export(model, inputs).module(), _get_quantizer()) + prepared(*inputs) + quantized = convert_pt2e(prepared) + + (node,) = quantized.graph.find_nodes( + op="call_function", target=getattr(torch.ops.aten, op).default + ) + data_dequantize, other_dequantize = node.args + assert data_dequantize.args[1] != other_dequantize.args[1] diff --git a/backends/arm/test/models/test_swin2sr_arm.py b/backends/arm/test/models/test_swin2sr_arm.py index d863c7016a4..596ada74636 100644 --- a/backends/arm/test/models/test_swin2sr_arm.py +++ b/backends/arm/test/models/test_swin2sr_arm.py @@ -29,7 +29,7 @@ ] # TODO/MLETORCH-2163: Investigate Swin2SR delegation gaps around index/view -# in FP and Q/DQ, clamp, and expand_copy in INT. +# in FP and Q/DQ and clamp in INT. swin2sr_fp_lowered_outer_graph_ops = { "torch.ops.higher_order.executorch_call_delegate": 2, "executorch_exir_dialects_edge__ops_aten_index_Tensor": 2, @@ -38,7 +38,6 @@ swin2sr_int_lowered_outer_graph_ops = { "torch.ops.higher_order.executorch_call_delegate": 3, "executorch_exir_dialects_edge__ops_aten_clamp_default": 4, - "executorch_exir_dialects_edge__ops_aten_expand_copy_default": 4, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 5, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 5, } diff --git a/backends/cortex_m/test/misc/test_portable_int8.py b/backends/cortex_m/test/misc/test_portable_int8.py index ce0aeeb1de6..3c2b125f21a 100644 --- a/backends/cortex_m/test/misc/test_portable_int8.py +++ b/backends/cortex_m/test/misc/test_portable_int8.py @@ -221,6 +221,12 @@ def _quantize_and_export( (torch.randn(1, 3, 4, 5), torch.randn(1, 3, 4, 5)), None, ), + "expand_as": OpCase( + torch.ops.aten.expand_as.default, + _build_module(lambda x, y: torch.ops.aten.expand_as.default(x, y)), + (torch.randn(1, 3, 4, 5), torch.randn(2, 3, 4, 5)), + None, + ), "select": OpCase( torch.ops.aten.select.int, _build_module(lambda x, y: torch.ops.aten.select.int(x, 1, 0)),