Arm backend: preserve Q/DQ for partially quant ops (#21589)

The Q/DQ folding pass incorrectly removes quantization boundaries when
only one of an operator's peer tensor inputs is quantized. This can make
later lowering treat quantized and floating-point inputs as compatible
quantized values.

Preserve add and sub by default. For mixed TOSA profiles and VGF, use an
explicit operator allowlist while excluding index operands.

Signed-off-by: Sangwon Ha <sangwon.ha@arm.com>

Author: FabulousSuperDude

backends/arm/_passes/arm_pass_manager.py

@@ -508,7 +508,13 @@ def _tosa_pipeline(
         # Fold Q/DQ nodes, insert INT8/INT32 rescales, decompose quantization nodes.
         self.add_passes(
             [
-                FoldAndAnnotateQParamsPass(exported_program),
+                FoldAndAnnotateQParamsPass(
+                    exported_program,
+                    preserve_partial_binary_tensor_qdq=(
+                        self.tosa_spec.support_float()
+                        or self.compile_spec._get_output_format() == "vgf"
+                    ),
+                ),
                 # Both hardtanh and relu are normalized to clamp by
                 # ConvertToClampPass; after q/dq folding above, adjacent clamps
                 # (e.g. from HardTanh+ReLU) are directly connected and can be

backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py

@@ -6,7 +6,7 @@
 
 import copy
 
-from typing import cast, Optional, Set, Type
+from typing import cast, ClassVar, Optional, Set, Type
 
 import torch
 from executorch.backends.arm._passes import ArmPass
@@ -127,12 +127,36 @@ class FoldAndAnnotateQParamsPass(ArmPass):
         InsertTableOpsPass,
         RemoveNoopPass,
     }
+    _default_partial_binary_qdq_targets: ClassVar[tuple[object, ...]] = (
+        exir_ops.edge.aten.add.Tensor,
+        exir_ops.edge.aten.sub.Tensor,
+    )
+    _mixed_profile_partial_binary_qdq_targets: ClassVar[tuple[object, ...]] = (
+        *_default_partial_binary_qdq_targets,
+        exir_ops.edge.aten.mul.Tensor,
+        exir_ops.edge.aten.div.Tensor,
+        exir_ops.edge.aten.minimum.default,
+        exir_ops.edge.aten.maximum.default,
+        exir_ops.edge.aten.mm.default,
+        exir_ops.edge.aten.bmm.default,
+        exir_ops.edge.aten.eq.Tensor,
+        exir_ops.edge.aten.ge.Tensor,
+        exir_ops.edge.aten.gt.Tensor,
+        exir_ops.edge.aten.le.Tensor,
+        exir_ops.edge.aten.lt.Tensor,
+        exir_ops.edge.aten.grid_sampler_2d.default,
+    )
 
     def __init__(
-        self, exported_program: Optional[ExportedProgram] = None, *args, **kwargs
+        self,
+        exported_program: Optional[ExportedProgram] = None,
+        *args,
+        preserve_partial_binary_tensor_qdq: bool = False,
+        **kwargs,
     ) -> None:
         super().__init__(*args, **kwargs)
         self.exported_program = exported_program
+        self.preserve_partial_binary_tensor_qdq = preserve_partial_binary_tensor_qdq
 
     def _extract_input_params(
         self, arg_list: list[Node]
@@ -177,9 +201,13 @@ def _annotate_input_params(
         index: int,
         input_qparams: QuantArgs,
         nodes_to_remove: set[Node],
+        remove_qdq: bool = True,
     ) -> None:
         node.meta["input_qparams"][index] = input_qparams
 
+        if not remove_qdq:
+            return
+
         for dq in nodes_to_remove:
             if dq.target not in DQ_OPS:
                 raise RuntimeError(f"Expected one of {DQ_OPS} dq_op, got {dq.target}")
@@ -201,6 +229,35 @@ def fold_and_annotate_arg(
             graph_module, node, i, input_qparams, nodes_to_remove
         )
 
+    def _extract_arg_input_params(
+        self, arg: object
+    ) -> tuple[Optional[QuantArgs], set[Node]]:
+        if isinstance(arg, (list, tuple)):
+            return self._extract_input_params(list(arg))
+        if isinstance(arg, Node):
+            return self._extract_input_params([arg])
+        return None, set()
+
+    def _has_partial_binary_tensor_qdq_inputs(
+        self, node: Node, input_qparams: dict[int, QuantArgs]
+    ) -> bool:
+        targets = (
+            self._mixed_profile_partial_binary_qdq_targets
+            if self.preserve_partial_binary_tensor_qdq
+            else self._default_partial_binary_qdq_targets
+        )
+        if node.target not in targets:
+            return False
+
+        lhs_idx, rhs_idx = 0, 1
+        if lhs_idx >= len(node.args) or rhs_idx >= len(node.args):
+            return False
+        if not isinstance(node.args[lhs_idx], Node) or not isinstance(
+            node.args[rhs_idx], Node
+        ):
+            return False
+        return (lhs_idx in input_qparams) != (rhs_idx in input_qparams)
+
     def _handle_control_flow_node(self, node: Node, graph_module: GraphModule):
         """Fold outmost quant nodes inside submodule.
 
@@ -327,12 +384,12 @@ def _correct_output_dtype(node: torch.fx.Node):
     def call(self, graph_module: GraphModule) -> PassResult:  # noqa: C901
 
         # Loop over the graph nodes and find any node in the 'targeted_ops' list.
-        modified = False
+        graph_modified = False
+        metadata_modified = False
         for n in graph_module.graph.nodes:
             n = cast(Node, n)
             if not FoldAndAnnotateQParamsPass.is_foldable(n):
                 continue
-            modified = True
 
             # Make sure we haven't already set qparams meta information on the node
             if "input_qparams" in n.meta:
@@ -346,17 +403,32 @@ def call(self, graph_module: GraphModule) -> PassResult:  # noqa: C901
                     "output_qparams should not have been set at this point"
                 )
 
+            input_qparams: dict[int, QuantArgs] = {}
+            input_nodes_to_remove: dict[int, set[Node]] = {}
+            for i, arg in enumerate(n.args):
+                qparams, nodes_to_remove = self._extract_arg_input_params(arg)
+                if qparams is not None:
+                    input_qparams[i] = qparams
+                    input_nodes_to_remove[i] = nodes_to_remove
+
+            preserve_qdq = self._has_partial_binary_tensor_qdq_inputs(n, input_qparams)
+            graph_modified = graph_modified or not preserve_qdq
+            metadata_modified = True
+
             # for the inputs and outputs search the graph for quantization info and
             # store the information in a dict with order of the _tensor_ inputs as key,
             # ignoring any other arguments to the target node.
             n.meta["input_qparams"] = {}
             n.meta["output_qparams"] = {}
-            for i, arg in enumerate(n.args):
-                if isinstance(arg, (list, tuple)):
-                    self.fold_and_annotate_arg(graph_module, n, arg, i)  # type: ignore
-
-                elif isinstance(arg, Node):
-                    self.fold_and_annotate_arg(graph_module, n, [arg], i)
+            for i, qparams in input_qparams.items():
+                self._annotate_input_params(
+                    graph_module,
+                    n,
+                    i,
+                    qparams,
+                    input_nodes_to_remove[i],
+                    remove_qdq=not preserve_qdq,
+                )
 
             # Copy the users, since we are modifying it.
             users_copy = copy.copy(n.users)
@@ -369,8 +441,9 @@ def call(self, graph_module: GraphModule) -> PassResult:  # noqa: C901
                     user.target, user.args
                 )
 
-                user.replace_all_uses_with(n)
-                graph_module.graph.erase_node(user)
+                if not preserve_qdq:
+                    user.replace_all_uses_with(n)
+                    graph_module.graph.erase_node(user)
 
             # Some op(s) contain a "dtype" key in their node kwargs. Set this
             # to the type of output qparams.
@@ -383,10 +456,10 @@ def call(self, graph_module: GraphModule) -> PassResult:  # noqa: C901
                 self._handle_control_flow_node(n, graph_module)
 
         # retrace the graph to update the fake tensor types
-        if modified:
+        if graph_modified:
             graph_module = super().call(graph_module).graph_module
 
-        return PassResult(graph_module, modified)
+        return PassResult(graph_module, metadata_modified or graph_modified)
 
 
 class QuantizeClampArgumentsPass(ArmPass):

backends/arm/test/passes/test_fold_qdq_pass.py

@@ -3,16 +3,35 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import ClassVar, Dict, Tuple
+from typing import Callable, ClassVar, Dict, Tuple
 
+import pytest
 import torch
 from executorch.backends.arm._passes import FoldAndAnnotateQParamsPass
+from executorch.backends.arm.common.annotation_meta import ArmAnnotationInfo
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+from executorch.exir.dialects._ops import ops as exir_ops
 
 
 input_t = Tuple[torch.Tensor, torch.Tensor]  # Input x, y
 
+_MIXED_PROFILE_PARTIAL_BINARY_TENSOR_TARGETS = (
+    exir_ops.edge.aten.add.Tensor,
+    exir_ops.edge.aten.sub.Tensor,
+    exir_ops.edge.aten.mul.Tensor,
+    exir_ops.edge.aten.div.Tensor,
+    exir_ops.edge.aten.minimum.default,
+    exir_ops.edge.aten.maximum.default,
+    exir_ops.edge.aten.mm.default,
+    exir_ops.edge.aten.bmm.default,
+    exir_ops.edge.aten.eq.Tensor,
+    exir_ops.edge.aten.ge.Tensor,
+    exir_ops.edge.aten.gt.Tensor,
+    exir_ops.edge.aten.le.Tensor,
+    exir_ops.edge.aten.lt.Tensor,
+)
+
 
 class SimpleQuantizeModel(torch.nn.Module):
     test_data: ClassVar[Dict[str, input_t]] = {
@@ -49,3 +68,126 @@ def test_fold_and_annotate_q_params_tosa_INT(test_data: input_t) -> None:
     )
     pipeline.pop_stage(-1)  # Do not compare output
     pipeline.run()
+
+
+@pytest.mark.parametrize(
+    "binary_target",
+    (exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.sub.Tensor),
+)
+def test_fold_qdq_preserves_default_partial_binary_qdq(
+    binary_target: Callable[..., object],
+) -> None:
+    _check_fold_qdq_preserves_partial_binary_qdq(binary_target)
+
+
+@pytest.mark.parametrize(
+    "binary_target",
+    _MIXED_PROFILE_PARTIAL_BINARY_TENSOR_TARGETS,
+)
+def test_fold_qdq_preserves_mixed_profile_partial_binary_tensor_qdq(
+    binary_target: Callable[..., object],
+) -> None:
+    _check_fold_qdq_preserves_partial_binary_qdq(
+        binary_target, preserve_partial_binary_tensor_qdq=True
+    )
+
+
+def test_fold_qdq_preserves_mixed_profile_partial_grid_sampler_qdq() -> None:
+    _check_fold_qdq_preserves_partial_binary_qdq(
+        exir_ops.edge.aten.grid_sampler_2d.default,
+        preserve_partial_binary_tensor_qdq=True,
+        extra_args=(0, 0, False),
+    )
+
+
+def test_fold_qdq_mixed_profile_allowlist_has_test_coverage() -> None:
+    tested_targets = {
+        *_MIXED_PROFILE_PARTIAL_BINARY_TENSOR_TARGETS,
+        exir_ops.edge.aten.grid_sampler_2d.default,
+    }
+
+    assert tested_targets == set(
+        FoldAndAnnotateQParamsPass._mixed_profile_partial_binary_qdq_targets  # noqa: SLF001
+    )
+
+
+def test_fold_qdq_folds_default_partial_mul_qdq() -> None:
+    _, mul, _, _ = _partial_binary_qdq_graph(exir_ops.edge.aten.mul.Tensor)
+
+    assert not FoldAndAnnotateQParamsPass()._has_partial_binary_tensor_qdq_inputs(  # noqa: SLF001
+        mul, {0: object()}  # type: ignore[dict-item]
+    )
+
+
+@pytest.mark.parametrize(
+    "target",
+    (
+        exir_ops.edge.aten.index_select.default,
+        exir_ops.edge.aten.gather.default,
+    ),
+)
+def test_fold_qdq_does_not_treat_index_as_binary_operand(
+    target: Callable[..., object],
+) -> None:
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    index = graph.placeholder("index")
+    node = graph.call_function(target, (x, 0, index))
+
+    assert not FoldAndAnnotateQParamsPass(
+        preserve_partial_binary_tensor_qdq=True
+    )._has_partial_binary_tensor_qdq_inputs(  # noqa: SLF001
+        node, {0: object()}  # type: ignore[dict-item]
+    )
+
+
+def _check_fold_qdq_preserves_partial_binary_qdq(
+    binary_target: Callable[..., object],
+    preserve_partial_binary_tensor_qdq: bool = False,
+    extra_args: tuple[int | bool, ...] = (),
+) -> None:
+    graph_module, add, _, y = _partial_binary_qdq_graph(binary_target, extra_args)
+    x_dq = add.args[0]
+    add_q = next(iter(add.users))
+
+    FoldAndAnnotateQParamsPass(
+        preserve_partial_binary_tensor_qdq=preserve_partial_binary_tensor_qdq
+    )(graph_module)
+
+    assert set(add.meta["input_qparams"]) == {0}
+    assert set(add.meta["output_qparams"]) == {0}
+    assert add.args == (x_dq, y, *extra_args)
+    assert add_q in add.users
+
+
+def _partial_binary_qdq_graph(
+    binary_target: Callable[..., object],
+    extra_args: tuple[int | bool, ...] = (),
+) -> tuple[torch.fx.GraphModule, torch.fx.Node, torch.fx.Node, torch.fx.Node]:
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    y = graph.placeholder("y")
+    x_q = graph.call_function(
+        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+        (x, 0.5, 0, -128, 127, torch.int8),
+    )
+    x_dq = graph.call_function(
+        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+        (x_q, 0.5, 0, -128, 127, torch.int8),
+    )
+    add = graph.call_function(binary_target, (x_dq, y, *extra_args))
+    add.meta["custom"] = {
+        ArmAnnotationInfo.CUSTOM_META_KEY: ArmAnnotationInfo(quantized=True)
+    }
+    add_q = graph.call_function(
+        exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+        (add, 0.5, 0, -128, 127, torch.int8),
+    )
+    out = graph.call_function(
+        exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+        (add_q, 0.5, 0, -128, 127, torch.int8),
+    )
+    graph.output(out)
+    graph_module = torch.fx.GraphModule(torch.nn.Module(), graph)
+
+    return graph_module, add, x_q, y