Complete q/dq folding coverage

Add support for q/dq folding for the remaining supported ops in Arm
backend.

Signed-off-by: Oscar Andersson <[email protected]>
Change-Id: I9012b4a501ce018c9771c729706be3b031a5c7ae

Author: oscarandersson8218

backends/arm/_passes/annotate_decomposed_matmul.py

@@ -0,0 +1,91 @@
+# Copyright 2024 Arm Limited and/or its 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 itertools
+from typing import Any, Dict, List
+
+import torch
+from executorch.backends.arm._passes.arm_pass_utils import create_node
+from executorch.backends.arm.tosa_quant_utils import dq_op, q_op
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+from torch.fx import GraphModule
+from torch.fx.passes.utils.source_matcher_utils import (
+    get_source_partitions,
+    SourcePartition,
+)
+
+
+class AnnotateDecomposedMatmulPass(ExportPass):
+    """
+    torch.matmul can be decomposed in many ways, for instance:
+    dq -> matmul -> q can become
+    dq -> repeat -> view -> bmm -> view -> dq which makes quantization folding
+    difficult. This helper function find all matmul partitions and annotate its
+    matmul-op (can be mm or bmm).
+    """
+
+    def call(self, graph_module: GraphModule):
+        matmul_partitions: Dict[Any, List[SourcePartition]] = get_source_partitions(
+            graph_module.graph,
+            [
+                torch.matmul,
+            ],
+            None,
+        )
+        matmul_partitions = list(
+            itertools.chain.from_iterable(matmul_partitions.values())
+        )
+        matmul_targets = {
+            exir_ops.edge.aten.mm.default,
+            exir_ops.edge.aten.bmm.default,
+        }
+        for partition in matmul_partitions:
+            quantized_input = all(
+                input_node.target == dq_op for input_node in partition.input_nodes
+            )
+            matmul_node = [
+                node for node in partition.nodes if node.target in matmul_targets
+            ][0]
+            if quantized_input:
+                matmul_args = matmul_node.all_input_nodes
+                for i in range(len(matmul_args)):
+                    input_node = partition.input_nodes[i]
+                    matmul_input_node = matmul_args[i]
+                    # Remove partition input dq-node
+                    input_node.replace_all_uses_with(input_node.args[0])
+                    graph_module.graph.erase_node(input_node)
+                    input_node_qargs = input_node.args[1:]
+                    with graph_module.graph.inserting_before(matmul_node):
+                        # Create new dq-node before matmul
+                        dq_node = create_node(
+                            graph=graph_module.graph,
+                            op_target=dq_op,
+                        )
+                        dq_node.args = (matmul_input_node, *input_node_qargs)
+                        matmul_node.replace_input_with(matmul_input_node, dq_node)
+
+            partition_output = list(partition.output_nodes[0].users)[0]
+            quantized_output = partition_output.target == q_op
+            if quantized_output:
+                output_node_qargs = partition_output.args[1:]
+                with graph_module.graph.inserting_after(matmul_node):
+                    # Create q-node after matmul
+                    q_node = create_node(
+                        graph=graph_module.graph,
+                        op_target=q_op,
+                    )
+                    matmul_node.replace_all_uses_with(q_node)
+                    q_node.args = (matmul_node, *output_node_qargs)
+                # Remove partition output q-node
+                partition_output.replace_all_uses_with(partition_output.args[0])
+                graph_module.graph.erase_node(partition_output)
+
+        # retrace the graph to update the fake tensor types
+        graph_module = super().call(graph_module).graph_module
+
+        graph_module.recompile()
+        return PassResult(graph_module, True)

backends/arm/_passes/arm_pass_manager.py

@@ -11,6 +11,9 @@
 from executorch.backends.arm._passes.annotate_channels_last_dim_order_pass import (
     AnnotateChannelsLastDimOrder,
 )
+from executorch.backends.arm._passes.annotate_decomposed_matmul import (
+    AnnotateDecomposedMatmulPass,
+)
 from executorch.backends.arm._passes.cast_int64_pass import CastInt64ToInt32Pass
 from executorch.backends.arm._passes.conv1d_unsqueeze_pass import Conv1dUnsqueezePass
 from executorch.backends.arm._passes.convert_expand_copy_to_repeat import (
@@ -69,51 +72,64 @@ def transform_to_backend_pipeline(
         self, exported_program: ExportedProgram, compile_spec: list[CompileSpec]
     ):
         """Apply passes before transforming program to backend"""
-        self.add_pass(CastInt64ToInt32Pass(exported_program))
+        self.add_pass(DecomposeLinearPass())
         self.add_pass(RemoveGetItemPass())
-        self.add_pass(UnsqueezeScalarPlaceholdersPass(exported_program))
-        self.add_pass(SizeAdjustConv2DPass())
-        self.add_pass(RemoveClonePass())
-        self.add_pass(ConvertExpandCopyToRepeatPass())
         self.add_pass(DecomposeLayerNormPass())
-        self.add_pass(UnsqueezeBeforeRepeatPass())
         self.add_pass(DecomposeVarPass())
         self.add_pass(ConvertMeanDimToAveragePool())
         self.add_pass(DecomposeMeanDimPass())
-        self.add_pass(DecomposeLinearPass())
+        self.add_pass(ConvertSplitToSlicePass())
+        # TODO MLETORCH-558
+        self.add_pass(AnnotateDecomposedMatmulPass())
         self.add_pass(QuantizeFullArgument())
         self.add_pass(
             FoldAndAnnotateQParamsPass(
                 [
+                    exir_ops.edge.aten.minimum.default,
+                    exir_ops.edge.aten.maximum.default,
                     exir_ops.edge.aten.add.Tensor,
                     exir_ops.edge.aten.avg_pool2d.default,
+                    exir_ops.edge.aten.bmm.default,
                     exir_ops.edge.aten.cat.default,
                     exir_ops.edge.aten.convolution.default,
+                    exir_ops.edge.aten.clone.default,
                     exir_ops.edge.aten.exp.default,
+                    exir_ops.edge.aten.expand_copy.default,
                     exir_ops.edge.aten.full.default,
                     exir_ops.edge.aten.hardtanh.default,
                     exir_ops.edge.aten.log.default,
                     exir_ops.edge.aten.max_pool2d.default,
-                    exir_ops.edge.aten.maximum.default,
-                    exir_ops.edge.aten.minimum.default,
+                    exir_ops.edge.aten.mm.default,
                     exir_ops.edge.aten.mul.Tensor,
+                    exir_ops.edge.aten.permute_copy.default,
                     exir_ops.edge.aten.reciprocal.default,
                     exir_ops.edge.aten.relu.default,
+                    exir_ops.edge.aten.repeat.default,
                     exir_ops.edge.aten.rsqrt.default,
                     exir_ops.edge.aten.select_copy.int,
                     exir_ops.edge.aten.sigmoid.default,
+                    exir_ops.edge.aten.slice_copy.Tensor,
+                    exir_ops.edge.aten.squeeze_copy.dims,
                     exir_ops.edge.aten.sub.Tensor,
                     exir_ops.edge.aten.sum.dim_IntList,
                     exir_ops.edge.aten.tanh.default,
+                    exir_ops.edge.aten.unsqueeze_copy.default,
+                    exir_ops.edge.aten.upsample_nearest2d.vec,
+                    exir_ops.edge.aten.view_copy.default,
                 ]
             )
         )
         self.add_pass(RetraceFoldedDtypesPass())
         self.add_pass(InsertTableOpsPass(exported_program))
-        self.add_pass(KeepDimsFalseToSqueezePass())
+        self.add_pass(ConvertExpandCopyToRepeatPass())
+        self.add_pass(UnsqueezeBeforeRepeatPass())
+        self.add_pass(CastInt64ToInt32Pass(exported_program))
+        self.add_pass(UnsqueezeScalarPlaceholdersPass(exported_program))
+        self.add_pass(SizeAdjustConv2DPass())
+        self.add_pass(RemoveClonePass())
         self.add_pass(MatchArgRanksPass(exported_program))
         self.add_pass(DecomposeDivPass())
-        self.add_pass(ConvertSplitToSlicePass())
+        self.add_pass(KeepDimsFalseToSqueezePass())
         self.add_pass(Conv1dUnsqueezePass(exported_program))
         self.add_pass(DecomposeSoftmaxesPass())
         for spec in compile_spec:

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -12,10 +12,8 @@
 from executorch.backends.arm._passes.arm_pass_utils import (
     create_node,
     get_param_tensor,
-    insert_q_dq_pair,
     is_param_node,
 )
-from executorch.backends.arm.tosa_quant_utils import dq_op, q_op
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
@@ -27,10 +25,8 @@ class Conv1dUnsqueezePass(ExportPass):
     supports 2d and 3d convolution. This is done by modifying the graph to do the
     following:
     1) unsqueeze the convolution's input from 3d to 4d
-    2) if the input to unsqueeze is quantized, insert q/dq-pair after unsqueeze
-    3) perform a conv2d (with a modified version of the original conv1d args)
-    4) squeeze the output back down to 3d.
-    5) if all users of squeeze are quantized, insert q/dq-pair before squeeze
+    2) perform a conv2d (with a modified version of the original conv1d args)
+    3) squeeze the output back down to 3d.
     """
 
     def __init__(self, exported_program: ExportedProgram) -> None:
@@ -94,8 +90,6 @@ def call(self, graph_module: torch.fx.GraphModule):
                         continue
 
                     kernel_node = node.args[1]
-                    if kernel_node.target == dq_op:
-                        kernel_node = kernel_node.args[0]
 
                     if not is_param_node(self.exported_program, kernel_node):
                         raise AssertionError(
@@ -131,11 +125,6 @@ def call(self, graph_module: torch.fx.GraphModule):
                         )
                         node.replace_input_with(input_node, unsqueeze_before)
 
-                    # If Quantized we must insert unsqueeze --> q --> dq --> node
-                    if input_node.target == dq_op:
-                        q_params = input_node.args[1:]
-                        insert_q_dq_pair(graph, unsqueeze_before, q_params)
-
                     with graph.inserting_after(node):
                         squeeze_after = create_node(
                             graph,
@@ -151,13 +140,6 @@ def call(self, graph_module: torch.fx.GraphModule):
                         for user in original_users:
                             user.replace_input_with(node, squeeze_after)
 
-                        # If quantized, insert conv2d --> q --> dq --> squeeze
-                    if all(
-                        original_user.target == q_op for original_user in original_users
-                    ):
-                        q_params = original_users[0].args[1:]
-                        insert_q_dq_pair(graph, node, q_params)
-
         graph_module.recompile()
         # Since we are overriding "call", we need to call the parent's "call"
         # to retrace the graph and regenerate metadata

backends/arm/_passes/size_adjust_conv2d_pass.py

@@ -9,7 +9,6 @@
 from typing import cast, Optional
 
 import torch.fx
-from executorch.backends.arm.tosa_quant_utils import is_node_quantized
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch._ops import OpOverload
@@ -113,14 +112,7 @@ def call(self, graph_module: torch.fx.GraphModule):
                     slice_node = graph.create_node(
                         "call_function", self.slice_op, (last_node,) + args
                     )
-                    if is_node_quantized(last_node):
-                        q_params = last_node.args[1:]
-                        dq_node = insert_q_dq_pair(
-                            graph_module.graph, slice_node, q_params
-                        )
-                        last_node = dq_node
-                    else:
-                        last_node = slice_node
+                    last_node = slice_node
                 conv_node.replace_input_with(cast(torch.fx.Node, input_node), last_node)
                 modified_graph = True
 

backends/arm/operators/op_bmm.py

@@ -8,18 +8,17 @@
 from typing import List
 
 import serializer.tosa_serializer as ts
-import torch.fx
+import torch
+from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
+    get_input_qparams,
+    get_output_qparams,
+)
 from executorch.backends.arm.operators.node_visitor import (
     NodeVisitor,
     register_node_visitor,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
-from executorch.backends.arm.tosa_quant_utils import (
-    build_rescale,
-    get_quant_arg_downstream,
-    get_quant_arg_upstream,
-)
-from executorch.backends.arm.tosa_utils import get_two_inputs
+from executorch.backends.arm.tosa_quant_utils import build_rescale
 from serializer.tosa_serializer import TosaOp
 
 
@@ -38,42 +37,45 @@ def define_node(
         output: TosaArg,
         is_quant_node: bool,
     ) -> None:
-        input0, input1 = get_two_inputs(node)
 
+        assert inputs[0].dtype == inputs[1].dtype, "Both inputs must be of same type"
+        assert inputs[0].dtype in [
+            ts.DType.INT8,
+            ts.DType.FP32,
+        ], "Only int8 and float32 supported"
         # aten.bmm maps directly to MATMUL
         # NOTE: For now, only INT8 & FP32 is supported
 
         # For INT8, we need to get the zero points and add an intermediate tensor
         # for a later rescale.
-        if is_quant_node:
-            input0_q_params = get_quant_arg_upstream(input0)
-            input1_q_params = get_quant_arg_upstream(input1)
-            input0_zp = input0_q_params.zp
-            input1_zp = input1_q_params.zp
+
+        if inputs[0].dtype == ts.DType.INT8:
+            input_qparams = get_input_qparams(node)
+            input0_zp = input_qparams[0].zp
+            input1_zp = input_qparams[1].zp
             bmm_result = tosa_graph.addIntermediate(output.shape, ts.DType.INT32)
             bmm_output_name = bmm_result.name
         else:
-            input0_zp, input1_zp = 0, 0
             bmm_output_name = output.name
+            input0_zp, input1_zp = 0, 0
 
         # Add the MATMUL to the TOSA graph.
         attr = ts.TosaSerializerAttribute()
         attr.MatMulAttribute(A_zp=input0_zp, B_zp=input1_zp)
 
         tosa_graph.addOperator(
             TosaOp.Op().MATMUL,
-            [input0.name, input1.name],
+            [inputs[0].name, inputs[1].name],
             [bmm_output_name],
             attr,
         )
 
         # As INT8 accumulates into INT32, we need to rescale it back to INT8
-        if is_quant_node:
-            output_q_params = get_quant_arg_downstream(list(node.users)[0])
-
+        if output.dtype == ts.DType.INT8:
+            output_qparams = get_output_qparams(node)[0]
             final_output_scale = (
-                input0_q_params.scale * input1_q_params.scale
-            ) / output_q_params.scale
+                input_qparams[0].scale * input_qparams[1].scale
+            ) / output_qparams.scale
 
             build_rescale(
                 tosa_fb=tosa_graph,
@@ -84,6 +86,6 @@ def define_node(
                 output_type=ts.DType.INT8,
                 output_shape=bmm_result.shape,
                 input_zp=0,
-                output_zp=output_q_params.zp,
+                output_zp=output_qparams.zp,
                 is_double_round=False,
             )