Keep the shape operand of view_as and expand_as out of shared qparams (#24153)

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

Author: pssrawat

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)

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]

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

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