Add support for sum q/dq folding

oscarandersson8218 · oscarandersson8218 · commit f9a5d461149f · 2024-12-17T07:46:47.000+01:00
sum is retraced to an int64 dtype of operator after q/dq folding.
This patch adds a pass to manually force the dtype to be int8.

Signed-off-by: Oscar Andersson &lt;oscar.andersson@arm.com&gt;
Change-Id: Ifa737a398c5a878d52cd76a2392499905da085ce
diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py
@@ -32,6 +32,7 @@
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
     FoldAndAnnotateQParamsPass,
     QuantizeFullArgument,
+    RetraceFoldedDtypesPass,
 )
 from executorch.backends.arm._passes.insert_table_ops import InsertTableOpsPass
 from executorch.backends.arm._passes.keep_dims_false_to_squeeze_pass import (
@@ -102,14 +103,16 @@ def transform_to_backend_pipeline(
                     exir_ops.edge.aten.select_copy.int,
                     exir_ops.edge.aten.sigmoid.default,
                     exir_ops.edge.aten.sub.Tensor,
+                    exir_ops.edge.aten.sum.dim_IntList,
                     exir_ops.edge.aten.tanh.default,
                 ]
             )
         )
+        self.add_pass(RetraceFoldedDtypesPass())
         self.add_pass(InsertTableOpsPass(exported_program))
+        self.add_pass(KeepDimsFalseToSqueezePass())
         self.add_pass(MatchArgRanksPass(exported_program))
         self.add_pass(DecomposeDivPass())
-        self.add_pass(KeepDimsFalseToSqueezePass())
         self.add_pass(ConvertSplitToSlicePass())
         self.add_pass(Conv1dUnsqueezePass(exported_program))
         self.add_pass(DecomposeSoftmaxesPass())
diff --git a/backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py b/backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py
@@ -195,3 +195,33 @@ def call(self, graph_module: GraphModule) -> PassResult:
                 modified = True
 
         return PassResult(graph_module, modified)
+
+
+class RetraceFoldedDtypesPass(ExportPass):
+    """
+    FoldAndAnnotateQParamsPass folds dq and q nodes. When the graph is retraced
+    some operators are retraced to types that cannot be handled by TOSA. One
+    such example is sum.dim_IntList:
+        q (int8) -> dq (fp32) -> sum (fp32) -> q (int8) ...
+    After folding it becomes:
+        q (int8)              -> sum (int64) ->         ...
+    This pass changes types of ops in self.targeted_ops, such as sum, so that
+    the output type of that matches the type of the output_qparams.
+    """
+
+    targeted_ops = {
+        exir_ops.edge.aten.sum.dim_IntList,
+    }
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in self.targeted_ops:
+            return super().call_operator(op, args, kwargs, meta)
+
+        node_kwargs = kwargs.copy()
+        output_qparams = meta["output_qparams"]
+        if len(output_qparams) == 0:
+            return super().call_operator(op, args, kwargs, meta)
+
+        output_dtype = output_qparams[0].dtype
+        node_kwargs["dtype"] = output_dtype
+        return super().call_operator(op, args, node_kwargs, meta)
diff --git a/backends/arm/operators/op_sum.py b/backends/arm/operators/op_sum.py
@@ -16,14 +16,19 @@
     register_node_visitor,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_specification import TosaSpecification
 from serializer.tosa_serializer import TosaOp
 from torch.fx import Node
 
 
 @register_node_visitor
-class AddVisitor(NodeVisitor):
+class SumVisitor_080_BI(NodeVisitor):
     target = "aten.sum.dim_IntList"
 
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-0.80.0+BI"),
+    ]
+
     def __init__(self, *args):
         super().__init__(*args)
 
@@ -35,64 +40,89 @@ def define_node(
         output: TosaArg,
         is_quant_node: bool,
     ) -> None:
-        input_node = inputs[0]
-        input_shape = list(input_node.shape)
+        input_shape = list(inputs[0].shape)
         dim_list = cast(list[int], inputs[1].special)
-        dim_list = [dim % len(input_node.shape) for dim in dim_list]
+        dim_list = [dim % len(input_shape) for dim in dim_list]
         keep_dim = cast(bool, inputs[2].number if len(inputs) > 2 else False)
         assert keep_dim, "This case should be handled by InsertSqueezeAfterSumPass"
 
-        if is_quant_node:
+        # Rescale input to 32 bit
+        rescaled_inputs, scale = tqutils.insert_rescale_ops_to_int32(
+            tosa_graph,
+            [inputs[0]],
+            node,
+        )
+
+        prev_node = rescaled_inputs[0]
+        reduced_shape = input_shape
+
+        # Reduce all dims in dim_list one-by-one.
+        for dim in dim_list:
+            # When reduced, the size of the dim becomes 1.
+            reduced_shape[dim] = 1
+
+            attr = ts.TosaSerializerAttribute()
+            attr.AxisAttribute(inputs[0].dim_order.index(dim))
+
+            next_node = tosa_graph.addIntermediate(
+                tutils.tosa_shape(reduced_shape, inputs[0].dim_order),
+                dtype=ts.DType.INT32,
+            )
+
+            tosa_graph.addOperator(
+                TosaOp.Op().REDUCE_SUM, [prev_node.name], [next_node.name], attr
+            )
+
+            prev_node = next_node
+        tqutils.insert_rescale_op_to_int8(tosa_graph, prev_node, scale, node)
+
+
+@register_node_visitor
+class SumVisitor_080_MI(SumVisitor_080_BI):
+    # inheriting 'target' from BI class
+
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-0.80.0+MI"),
+    ]
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+        is_quant_node: bool,
+    ) -> None:
+        if inputs[0].dtype == ts.DType.INT8:
+            return super().define_node(node, tosa_graph, inputs, output, is_quant_node)
+        input_name = inputs[0].name
+        reduced_shape = list(inputs[0].shape)
+        dim_list = cast(list[int], inputs[1].special)
+        dim_list = [dim % len(reduced_shape) for dim in dim_list]
+        keep_dim = cast(bool, inputs[2].number if len(inputs) > 2 else False)
+        assert keep_dim, "This case should be handled by InsertSqueezeAfterSumPass"
+
+        # Reduce all dims in dim_list one-by-one.
+        for dim in dim_list:
+            # When reduced, the size of the dim becomes 1
+            reduced_shape[dim] = 1
+
+            attr = ts.TosaSerializerAttribute()
+            attr.AxisAttribute(inputs[0].dim_order.index(dim))
+
+            if dim == dim_list[-1]:
+                output_name = output.name
+            else:
+                output_name = tosa_graph.addIntermediate(
+                    tutils.tosa_shape(reduced_shape, inputs[0].dim_order),
+                    dtype=ts.DType.FP32,
+                ).name
 
-            # Rescale input to 32 bit
-            rescaled_inputs, scale = tqutils.rescale_nodes_to_int32(
-                [node.all_input_nodes[0]], tosa_graph
+            tosa_graph.addOperator(
+                TosaOp.Op().REDUCE_SUM, [input_name], [output_name], attr
             )
 
-            prev_node = rescaled_inputs[0]
-            reduced_shape = input_shape
-
-            # Reduce all dims in dim_list one-by-one.
-            for dim in dim_list:
-                # When reduced, the size of the dim becomes 1.
-                reduced_shape[dim] = 1
-
-                attr = ts.TosaSerializerAttribute()
-                attr.AxisAttribute(input_node.dim_order.index(dim))
-
-                next_node = tosa_graph.addIntermediate(
-                    tutils.tosa_shape(reduced_shape, input_node.dim_order),
-                    dtype=ts.DType.INT32,
-                )
-
-                tosa_graph.addOperator(
-                    TosaOp.Op().REDUCE_SUM, [prev_node.name], [next_node.name], attr
-                )
-
-                prev_node = next_node
-            tqutils.rescale_node_back_to_int8(node, prev_node, scale, tosa_graph)
-        else:
-            input_name = input_node.name
-            reduced_shape = input_shape
-
-            # Reduce all dims in dim_list one-by-one.
-            for dim in dim_list:
-                # When reduced, the size of the dim becomes 1
-                reduced_shape[dim] = 1
-
-                attr = ts.TosaSerializerAttribute()
-                attr.AxisAttribute(input_node.dim_order.index(dim))
-
-                if dim == dim_list[-1]:
-                    output_name = output.name
-                else:
-                    output_name = tosa_graph.addIntermediate(
-                        tutils.tosa_shape(reduced_shape, input_node.dim_order),
-                        dtype=ts.DType.FP32,
-                    ).name
-
-                tosa_graph.addOperator(
-                    TosaOp.Op().REDUCE_SUM, [input_name], [output_name], attr
-                )
-
-                input_name = output_name
+            input_name = output_name
