Increase q/dq folding coverage

Add support for q/dq folding of more operators such as hardtanh,
maxpool2d, mul, relu, select, sub, to_copy.

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

Author: oscarandersson8218

backends/arm/_passes/arm_pass_manager.py

@@ -79,12 +79,6 @@ def transform_to_backend_pipeline(
         self.add_pass(DecomposeVarPass())
         self.add_pass(ConvertMeanDimToAveragePool())
         self.add_pass(DecomposeMeanDimPass())
-        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())
         self.add_pass(DecomposeLinearPass())
         self.add_pass(QuantizeFullArgument())
         self.add_pass(
@@ -96,17 +90,29 @@ def transform_to_backend_pipeline(
                     exir_ops.edge.aten.convolution.default,
                     exir_ops.edge.aten.exp.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.mul.Tensor,
                     exir_ops.edge.aten.reciprocal.default,
+                    exir_ops.edge.aten.relu.default,
                     exir_ops.edge.aten.rsqrt.default,
+                    exir_ops.edge.aten.select_copy.int,
                     exir_ops.edge.aten.sigmoid.default,
+                    exir_ops.edge.aten.sub.Tensor,
                     exir_ops.edge.aten.tanh.default,
                 ]
             )
         )
         self.add_pass(InsertTableOpsPass(exported_program))
+        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())
         for spec in compile_spec:
             if spec.key == "permute_memory_format":
                 memory_format = spec.value.decode()

backends/arm/operators/op_div.py

@@ -8,16 +8,16 @@
 
 import serializer.tosa_serializer as ts
 import torch
+from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
+    get_input_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 (
-    get_quant_arg_upstream,
-    quantize_value,
-)
+from executorch.backends.arm.tosa_quant_utils import quantize_value
 from serializer.tosa_serializer import TosaOp
 
 
@@ -38,9 +38,10 @@ def define_node(
     ) -> None:
         attr = ts.TosaSerializerAttribute()
 
-        if is_quant_node:
+        if inputs[0].dtype == ts.DType.INT8:
             # Get quant parameters
-            qargs = get_quant_arg_upstream(node.all_input_nodes[0])
+            input_qparams = get_input_qparams(node)
+            qargs = input_qparams[0]
             # Convert to quantized representation
             clamp_min_qs = quantize_value(inputs[1].number, qargs)
             clamp_max_qs = quantize_value(inputs[2].number, qargs)

backends/arm/operators/op_hardtanh.py

@@ -8,16 +8,15 @@
 
 import serializer.tosa_serializer as ts
 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 (
-    get_quant_arg_downstream,
-    get_quant_arg_upstream,
-)
-
 from serializer.tosa_serializer import TosaOp
 
 
@@ -46,19 +45,18 @@ def define_node(
         except IndexError:
             padding = [0, 0, 0, 0]
 
-        accumulator_type = input_tensor.dtype
-
-        if is_quant_node:
-            # Accumulator type always is int8 when input tensor is an integer type.
-            accumulator_type = ts.DType.INT8
+        accumulator_type = output.dtype
 
         # Initilize zero point to zero.
         input_zp = 0
-        output_zp = 0
+        if inputs[0].dtype == ts.DType.INT8:
+            input_qparams = get_input_qparams(node)
+            input_zp = input_qparams[0].zp
 
-        if is_quant_node:
-            input_zp = get_quant_arg_upstream(node.all_input_nodes[0]).zp
-            output_zp = get_quant_arg_downstream(list(node.users)[0]).zp
+        output_zp = 0
+        if output.dtype == ts.DType.INT8:
+            output_qparams = get_output_qparams(node)
+            output_zp = output_qparams[0].zp
 
         attr = ts.TosaSerializerAttribute()
         attr.PoolAttribute(

backends/arm/operators/op_max_pool2d.py

@@ -5,26 +5,34 @@
 
 # pyre-unsafe
 
-from typing import cast, List
+from typing import List
 
 import executorch.backends.arm.tosa_quant_utils as tqutils
 import executorch.backends.arm.tosa_utils as tutils
 
 import serializer.tosa_serializer as ts
 import torch
+from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
+    get_input_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_specification import TosaSpecification
 from serializer.tosa_serializer import TosaOp
 
 
 @register_node_visitor
-class MulVisitor(NodeVisitor):
+class MulVisitor_080_BI(NodeVisitor):
     target = "aten.mul.Tensor"
 
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-0.80.0+BI"),
+    ]
+
     def define_node(
         self,
         node: torch.fx.Node,
@@ -33,57 +41,68 @@ def define_node(
         output: TosaArg,
         is_quant_node: bool,
     ) -> None:
+        assert inputs[0].dtype == inputs[1].dtype == output.dtype == ts.DType.INT8
+        input_A = inputs[0]
+        input_B = inputs[1]
+        input_qparams = get_input_qparams(node)
+        input_A_qargs = input_qparams[0]
+        input_B_qargs = input_qparams[1]
+        input_A.shape = tutils.tosa_shape(input_A.shape, input_A.dim_order)
+        input_B.shape = tutils.tosa_shape(input_B.shape, input_B.dim_order)
+
+        # Rescale inputs to INT32 with zp=0
+        input_A_rescaled = tqutils.build_rescale_to_int32(
+            tosa_graph,
+            input_A,
+            input_A_qargs.zp,
+            rescale_scale=1.0,
+        )
+        input_B_rescaled = tqutils.build_rescale_to_int32(
+            tosa_graph,
+            input_B,
+            input_B_qargs.zp,
+            rescale_scale=1.0,
+        )
+
+        output_shape = tutils.tosa_shape(output.shape, output.dim_order)
+        mul_output = tosa_graph.addIntermediate(output_shape, ts.DType.INT32)
+
+        # Do the INT32 Mul
+        attr = ts.TosaSerializerAttribute()
+        attr.MulAttribute(shift=0)
+        tosa_graph.addOperator(
+            TosaOp.Op().MUL,
+            [
+                input_A_rescaled.name,
+                input_B_rescaled.name,
+            ],
+            [mul_output.name],
+            attr,
+        )
+        output_scale = input_A_qargs.scale * input_B_qargs.scale
+        tqutils.insert_rescale_op_to_int8(tosa_graph, mul_output, output_scale, node)
+
+
+@register_node_visitor
+class MulVisitor_080_MI(MulVisitor_080_BI):
+    # inheriting 'target' from BI class
+
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-0.80.0+MI"),
+    ]
 
-        if is_quant_node:
-            input_A = inputs[0]
-            input_B = inputs[1]
-            input_A_qargs = tqutils.get_quant_arg_upstream(
-                cast(torch.fx.Node, node.args[0])
-            )
-            input_B_qargs = tqutils.get_quant_arg_upstream(
-                cast(torch.fx.Node, node.args[1])
-            )
-
-            input_A.shape = tutils.tosa_shape(input_A.shape, input_A.dim_order)
-            input_B.shape = tutils.tosa_shape(input_B.shape, input_B.dim_order)
-            output_shape = tutils.tosa_shape(output.shape, output.dim_order)
-
-            # Rescale inputs to INT32 with zp=0
-            input_A_rescaled = tqutils.build_rescale_to_int32(
-                tosa_graph,
-                input_A,
-                input_A_qargs.zp,
-                rescale_scale=1.0,
-            )
-            input_B_rescaled = tqutils.build_rescale_to_int32(
-                tosa_graph,
-                input_B,
-                input_B_qargs.zp,
-                rescale_scale=1.0,
-            )
-
-            mul_output = tosa_graph.addIntermediate(output_shape, ts.DType.INT32)
-
-            # Do the INT32 Mul
-            attr = ts.TosaSerializerAttribute()
-            attr.MulAttribute(shift=0)
-            tosa_graph.addOperator(
-                TosaOp.Op().MUL,
-                [
-                    input_A_rescaled.name,
-                    input_B_rescaled.name,
-                ],
-                [mul_output.name],
-                attr,
-            )
-
-            tqutils.rescale_node_back_to_int8(
-                node, mul_output, input_A_qargs.scale * input_B_qargs.scale, tosa_graph
-            )
-
-        else:
-            attr = ts.TosaSerializerAttribute()
-            attr.MulAttribute(shift=0)
-            tosa_graph.addOperator(
-                TosaOp.Op().MUL, [inputs[0].name, inputs[1].name], [output.name], attr
-            )
+    def define_node(
+        self,
+        node: torch.fx.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)
+        attr = ts.TosaSerializerAttribute()
+        attr.MulAttribute(shift=0)
+        tosa_graph.addOperator(
+            TosaOp.Op().MUL, [inputs[0].name, inputs[1].name], [output.name], attr
+        )

backends/arm/operators/op_mul.py

@@ -8,6 +8,9 @@
 import executorch.backends.arm.tosa_quant_utils as tqutils
 import serializer.tosa_serializer as ts
 import torch.fx
+from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
+    get_output_qparams,
+)
 from executorch.backends.arm.operators.node_visitor import (
     NodeVisitor,
     register_node_visitor,
@@ -37,10 +40,10 @@ def define_node(
         clamp_max_fp = 0.0
         clamp_min_qs = 0
         clamp_max_qs = 0
-        if is_quant_node:
-            out_qargs = tqutils.get_quant_arg_downstream(list(node.users)[0])
-            clamp_min_qs = tqutils.quantize_value(0, out_qargs)
-            clamp_max_qs = tqutils.quantize_value(float("inf"), out_qargs)
+        if inputs[0].dtype == ts.DType.INT8:
+            out_qargs = get_output_qparams(node)
+            clamp_min_qs = tqutils.quantize_value(0, out_qargs[0])
+            clamp_max_qs = tqutils.quantize_value(float("inf"), out_qargs[0])
 
         else:
             clamp_min_fp = 0

backends/arm/operators/op_relu.py

@@ -50,9 +50,7 @@ def define_node(
         expanded_shape = tuple(1 if i == dim else shape[i] for i in range(rank))
         expanded_shape = tosa_shape(expanded_shape, input_node.dim_order)
 
-        output_reshaped = tosa_graph.addIntermediate(
-            expanded_shape, ts.DType.INT8 if is_quant_node else output.dtype
-        )
+        output_reshaped = tosa_graph.addIntermediate(expanded_shape, output.dtype)
 
         attr_slice = ts.TosaSerializerAttribute()