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

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/passes/BUCK

@@ -31,6 +31,7 @@ fbcode_target(_kind = runtime.python_library,
         "activation_fusion_pass.py",
         "aten_to_cortex_m_pass.py",
         "clamp_hardswish_pass.py",
+        "compute_quantized_constant_ops_pass.py",
         "cortex_m_pass.py",
         "cortex_m_pass_manager.py",
         "decompose_hardswish_pass.py",

backends/cortex_m/passes/compute_quantized_constant_ops_pass.py

@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from executorch.backends.arm._passes import ComputeConstantOpsAOTPass
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+class ComputeQuantizedConstantOpsPass(ComputeConstantOpsAOTPass):
+    """Give the quantized constant ops their int8 values.
+
+    Folding removes the Q node after these ops but leaves them producing float,
+    so the int8 kernels they feed would read float data. A ``full`` stays a
+    runtime fill with its fill value quantized, so it adds no constant data; the
+    other ops are computed in float and stored as quantized constants. Runs
+    after ``ConvertFullLikeToFullPass``, so a ``full_like`` arrives as ``full``.
+    """
+
+    def compute_node_aot(self, node: torch.fx.Node) -> bool:
+        output_qparams = node.meta.get("output_qparams")
+        if not output_qparams:
+            return False
+        if node.target == exir_ops.edge.aten.full.default:
+            qparams = output_qparams[0]
+            fill_value = node.args[1]
+            assert isinstance(fill_value, (int, float))
+            fill = qparams.quantize_value(torch.tensor(float(fill_value)))
+            node.args = (node.args[0], int(fill.item()))
+            node.kwargs = {**node.kwargs, "dtype": qparams.dtype}
+            return True
+        return super().compute_node_aot(node)

backends/cortex_m/passes/cortex_m_pass_manager.py

@@ -10,6 +10,7 @@
 
 import torch
 from executorch.backends.arm._passes import (
+    ConvertFullLikeToFullPass,
     DeduplicateGetAttrPass,
     FoldAndAnnotateQParamsPass,
     ScalarsToAttributePass,
@@ -48,6 +49,7 @@
 from .activation_fusion_pass import ActivationFusionPass
 from .aten_to_cortex_m_pass import AtenToCortexMPass
 from .clamp_hardswish_pass import ClampHardswishPass
+from .compute_quantized_constant_ops_pass import ComputeQuantizedConstantOpsPass
 from .decompose_hardswish_pass import DecomposeHardswishPass
 from .decompose_mean_pass import DecomposeMeanPass
 from .decompose_sdpa_pass import DecomposeSDPAPass
@@ -149,6 +151,8 @@ class CortexMPassManager(ExportedProgramPassManager):
         FoldAndAnnotateQParamsPass,
         ActivationFusionPass,
         QuantizedClampActivationPass,
+        ConvertFullLikeToFullPass,
+        ComputeQuantizedConstantOpsPass,
         DecomposeHardswishPass,
         AtenToCortexMPass,
         # Retrace only after quantized ATen ops have become Cortex-M kernels.
@@ -165,6 +169,8 @@ class CortexMPassManager(ExportedProgramPassManager):
         ReplaceScalarWithTensorArgPass,
         ActivationFusionPass,
         QuantizedClampActivationPass,
+        ConvertFullLikeToFullPass,
+        ComputeQuantizedConstantOpsPass,
         DecomposeHardswishPass,
         ConvertConv1dToConv2dPass,
         CortexMReplaceOpsWithChannelsLastVariants,

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

backends/cortex_m/test/misc/test_quantization.py

@@ -6,7 +6,7 @@
 
 import torch
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
-from executorch.backends.arm.test.common import parametrize, xfail_type
+from executorch.backends.arm.test.common import parametrize
 from executorch.backends.cortex_m.test.tester import (
     CortexMTester,
     McuTestCase,
@@ -358,12 +358,8 @@ def forward(self, x):
     ),
 }
 
-xfails: dict[str, xfail_type] = {
-    "surrounded_quantized_op_constant": "Numerical error since the add is forced to have non-correct qparams.",
-}
-
 
-@parametrize("test_case", test_cases, xfails=xfails)
+@parametrize("test_case", test_cases)
 def test_shared_qspec_quantizer(test_case):
     """
     Test that ops which does not change dynamic range are able to use int8 portable kernels.

backends/cortex_m/test/test_explicit_layout_pipeline.py

@@ -546,14 +546,126 @@ def test_convs_around_a_reshape_behind_lut_activation_lower_to_cmsis_nn():
     tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
 
 
-def test_explicit_layout_rejects_float_input_to_int8_conv():
+def test_explicit_layout_lowers_conv_on_a_quantized_constant():
+    inputs = (torch.randn(1, 8, 16),)
     model = ThenConv1d(lambda x: torch.full_like(x, 0.5, dtype=torch.float32))
-    tester = CortexMTester(model.eval(), (torch.randn(1, 8, 16),))
+    tester = _run_explicit_layout_passes(CortexMTester(model.eval(), inputs))
+    program = tester.get_artifact(StageType.RUN_PASSES).exported_program()
 
-    with pytest.raises(Exception) as caught:
-        _run_explicit_layout_passes(tester)
+    assert _count(program, exir_ops.edge.aten.full_like.default) == 0
+    assert _count(program, exir_ops.edge.aten.full.default) == 1
+    assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
 
-    error = caught.value
-    while error.__cause__ is not None:
-        error = error.__cause__
-    assert "must move an int8 tensor" in str(error)
+
+@pytest.mark.parametrize("use_explicit_layout", [False, True])
+@pytest.mark.parametrize(
+    "constant,constant_op,shape",
+    [
+        (
+            lambda: torch.arange(-1.0, 1.0, 0.125),
+            exir_ops.edge.aten.arange.start_step,
+            (16,),
+        ),
+        (
+            lambda: torch.scalar_tensor(0.5),
+            torch.ops.aten.scalar_tensor.default,
+            (),
+        ),
+    ],
+    ids=["arange", "scalar_tensor"],
+)
+def test_quantized_constant_is_computed_before_quantizing(
+    use_explicit_layout, constant, constant_op, shape
+):
+    torch.manual_seed(0)
+    inputs = (torch.randn(shape),)
+    tester = CortexMTester(Apply(lambda x: x + constant()).eval(), inputs)
+    tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout)))
+    tester.export().to_edge()
+    program = _run_passes(tester, use_explicit_layout)
+
+    assert _count(program, constant_op) == 0
+    assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 1
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
+
+
+@pytest.mark.parametrize("use_explicit_layout", [False, True])
+@pytest.mark.parametrize(
+    "fill",
+    [lambda x: torch.full((1, 8, 16), 0.5), lambda x: torch.full_like(x, 0.5)],
+    ids=["full", "full_like"],
+)
+def test_quantized_fill_stays_a_runtime_int8_full(use_explicit_layout, fill):
+    torch.manual_seed(0)
+    inputs = (torch.randn(1, 8, 16),)
+    tester = CortexMTester(Apply(lambda x: x + fill(x)).eval(), inputs)
+    tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout)))
+    tester.export().to_edge()
+    program = _run_passes(tester, use_explicit_layout)
+
+    (full,) = program.graph.find_nodes(
+        op="call_function", target=exir_ops.edge.aten.full.default
+    )
+    assert full.meta["val"].dtype == torch.int8
+    assert _count(program, exir_ops.edge.aten.full_like.default) == 0
+    assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 1
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
+
+
+@pytest.mark.parametrize("use_explicit_layout", [False, True])
+def test_clamp_with_constant_tensor_bounds_quantizes_them_once(use_explicit_layout):
+    torch.manual_seed(0)
+    inputs = (torch.randn(2, 3, 4, 5),)
+    tester = CortexMTester(
+        Apply(
+            lambda x: torch.clamp(x, torch.full_like(x, -1.0), torch.full_like(x, 1.0))
+        ).eval(),
+        inputs,
+    )
+    tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout)))
+    tester.export().to_edge()
+    program = _run_passes(tester, use_explicit_layout)
+
+    (clamp,) = program.graph.find_nodes(
+        op="call_function", target=exir_ops.edge.aten.clamp.default
+    )
+    qparams = clamp.meta["output_qparams"][0]
+    expected = tuple(
+        int(qparams.quantize_value(torch.tensor(bound)).item()) for bound in (-1.0, 1.0)
+    )
+    assert tuple(clamp.args[1:]) == expected
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
+
+
+@pytest.mark.parametrize("use_explicit_layout", [False, True])
+def test_unquantized_full_like_keeps_int8_after_float_promoting_op(use_explicit_layout):
+    inputs = (torch.randn(1, 8, 16),)
+    tester = CortexMTester(
+        Apply(
+            lambda x: (
+                torch.sigmoid(x).permute(0, 2, 1),
+                torch.full_like(x, 2, dtype=torch.int32),
+            )
+        ).eval(),
+        inputs,
+    )
+    tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout)))
+    tester.export().to_edge()
+    program = _run_passes(tester, use_explicit_layout)
+
+    assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 1
+    assert _count(program, exir_ops.edge.aten.permute_copy.default) == 0
+
+
+@pytest.mark.parametrize("use_explicit_layout", [False, True])
+def test_unquantized_constant_is_not_precomputed(use_explicit_layout):
+    inputs = (torch.randn(1, 8, 16),)
+    tester = CortexMTester(
+        Apply(lambda x: x + torch.full((1, 8, 16), 0.5)).eval(), inputs
+    )
+    tester.export().to_edge()
+    program = _run_passes(tester, use_explicit_layout)
+
+    assert _count(program, exir_ops.edge.aten.full.default) == 1
+    tester.run_method_and_compare_outputs(inputs=inputs)