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)),