Arm Backend: Update element-wise nodevisitors affected by RESCALE updates

Requires: view node visitor, rescale node visitor

Change-Id: Ie04f6dc20be22a8fa67d6527eb86bf82986236d9

Author: SaoirseARM

backends/arm/operators/op_add.py

@@ -5,12 +5,11 @@
 
 # pyre-unsafe
 
-from typing import List
+from typing import Any, List
 
 import executorch.backends.arm.tosa_quant_utils as tqutils
 import executorch.backends.arm.tosa_utils as tutils
 
-import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 from executorch.backends.arm.operators.node_visitor import (
     NodeVisitor,
     register_node_visitor,
@@ -34,10 +33,13 @@ def __init__(self, *args):
     def define_node(
         self,
         node: Node,
-        tosa_graph: ts.TosaSerializer,
+        tosa_graph: Any,
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
+
+        import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
+
         # Specification (0.80) states that input and output types
         # should all be the same
         if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
@@ -58,7 +60,7 @@ def define_node(
             if len(inputs[0].shape) > len(inputs[1].shape)
             else inputs[1].dim_order
         )
-
+        scale_back = 1.0
         if inputs[0].dtype == ts.DType.INT8:
             rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
                 tosa_graph, inputs, node
@@ -90,7 +92,9 @@ def define_node(
         if output.dtype == ts.DType.INT8:
             # Scale output back to 8 bit
             # pyre-ignore
-            tqutils.insert_rescale_op_to_int8(tosa_graph, add_output, scale_back, node)  # type: ignore[possibly-undefined]
+            tqutils.insert_rescale_op_to_int8(
+                tosa_graph, add_output, scale_back, node
+            )  # type: ignore[possibly-undefined]
 
 
 @register_node_visitor
@@ -107,10 +111,13 @@ def __init__(self, *args):
     def define_node(
         self,
         node: Node,
-        tosa_graph: ts.TosaSerializer,
+        tosa_graph: Any,
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
+
+        import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
+
         # Specification (0.80) states that input and output types
         # should all be the same
         if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
@@ -130,7 +137,7 @@ def define_node(
                     f"Expected IO data type to be FP32, got {inputs[0].dtype}"
                 )
 
-            input1, input2 = tutils.reshape_for_broadcast(tosa_graph, inputs)
+            input1, input2 = inputs
 
             # MI lowering
             tosa_graph.addOperator(
@@ -139,3 +146,122 @@ def define_node(
                 [output.name],
                 None,
             )
+
+
+@register_node_visitor
+class AddVisitor_INT(NodeVisitor):
+    target = "aten.add.Tensor"
+
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-1.0+INT"),
+    ]
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+
+        import serializer.tosa_serializer as ts  # type: ignore
+
+        # Specification (1.0) states that input and output types
+        # should all be the same
+        if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
+            raise TypeError(
+                f"All IO needs to have the same data type, got input 1: "
+                f"{inputs[0].dtype}, input 2: {inputs[1].dtype} and output: "
+                f"{output.dtype}"
+            )
+        # Handle int8 (quantized) and int32
+        supported_dtypes = [ts.DType.INT8, ts.DType.INT32]
+        if inputs[0].dtype not in supported_dtypes:
+            raise TypeError(
+                f'IO data type needs to be {supported_dtypes}, got "{inputs[0].dtype}"'
+            )
+        scale_back = 1.0
+        if inputs[0].dtype == ts.DType.INT8:
+            rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
+                tosa_graph, inputs, node, self.tosa_specs
+            )
+        else:
+            # input[0].dtype == ts.DType.INT32
+            # Non quantized input, natively support by TOSA.ADD
+            rescaled_inputs = inputs
+
+        if output.dtype == ts.DType.INT8:
+            broadcasted_shape = tutils.tosa_shape(output.shape, output.dim_order)
+            add_output = tosa_graph.addIntermediate(broadcasted_shape, ts.DType.INT32)
+        else:
+            # output.dtype == ts.DType.INT32
+            add_output = output
+
+        input1, input2 = rescaled_inputs
+
+        # Do the INT32 Add
+        tosa_graph.addOperator(
+            ts.TosaOp.Op().ADD,
+            [input1.name, input2.name],
+            [add_output.name],
+            None,
+        )
+
+        if output.dtype == ts.DType.INT8:
+            # Scale output back to 8 bit
+            # pyre-ignore
+            tqutils.insert_rescale_op_to_int8(
+                tosa_graph, add_output, scale_back, node, self.tosa_specs
+            )  # type: ignore[possibly-undefined]
+
+
+@register_node_visitor
+class AddVisitor_FP(AddVisitor_INT):
+    # inheriting 'target' from INT class
+
+    tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+FP")]
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+
+        import serializer.tosa_serializer as ts  # type: ignore
+
+        # Specification (1.0) states that input and output types
+        # should all be the same
+        if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
+            raise TypeError(
+                f"All IO needs to have the same data type, got input 1: "
+                f"{inputs[0].dtype}, input 2: {inputs[1].dtype} and output: "
+                f"{output.dtype}"
+            )
+
+        if inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]:
+            # Call the inherited define_node for handling integers
+            super().define_node(node, tosa_graph, inputs, output)
+        else:
+            # FP32 Add lowering
+            if inputs[0].dtype != ts.DType.FP32:
+                raise TypeError(
+                    f"Expected IO data type to be FP32, got {inputs[0].dtype}"
+                )
+
+            input1, input2 = inputs
+
+            # FP lowering
+            tosa_graph.addOperator(
+                ts.TosaOp.Op().ADD,
+                [input1.name, input2.name],
+                [output.name],
+                None,
+            )

backends/arm/operators/op_mul.py

@@ -5,14 +5,12 @@
 
 # pyre-unsafe
 
-from typing import List
+from typing import Any, List
 
 import executorch.backends.arm.tosa_quant_utils as tqutils
 import executorch.backends.arm.tosa_utils as tutils
 import torch
 
-import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
-
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
     get_input_qparams,
 )
@@ -37,10 +35,13 @@ class MulVisitor_080_BI(NodeVisitor):
     def define_node(
         self,
         node: torch.fx.Node,
-        tosa_graph: ts.TosaSerializer,
+        tosa_graph: Any,
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
+
+        import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
+
         if (
             inputs[0].dtype != ts.DType.INT8
             or inputs[1].dtype != ts.DType.INT8
@@ -114,10 +115,13 @@ class MulVisitor_080_MI(MulVisitor_080_BI):
     def define_node(
         self,
         node: torch.fx.Node,
-        tosa_graph: ts.TosaSerializer,
+        tosa_graph: Any,
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
+
+        import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
+
         if inputs[0].dtype == ts.DType.INT8:
             return super().define_node(node, tosa_graph, inputs, output)
 
@@ -128,3 +132,100 @@ def define_node(
         tosa_graph.addOperator(
             ts.TosaOp.Op().MUL, [input1.name, input2.name], [output.name], attr
         )
+
+
+@register_node_visitor
+class MulVisitor_INT(NodeVisitor):
+    target = "aten.mul.Tensor"
+
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-1.0+INT"),
+    ]
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+
+        import serializer.tosa_serializer as ts  # type: ignore
+
+        if (
+            inputs[0].dtype != ts.DType.INT8
+            or inputs[1].dtype != ts.DType.INT8
+            or output.dtype != ts.DType.INT8
+        ):
+            raise ValueError(
+                f"Inputs and output for {self.target} need to be INT8, got "
+                f"{inputs[0].dtype=}, {inputs[1].dtype=} and {output.dtype=}"
+            )
+
+        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,
+            [1.0],
+            tosa_spec=self.tosa_specs,
+        )
+        input_B_rescaled = tqutils.build_rescale_to_int32(
+            tosa_graph,
+            input_B,
+            input_B_qargs.zp,
+            [1.0],
+            tosa_spec=self.tosa_specs,
+        )
+
+        output_shape = tutils.tosa_shape(output.shape, output.dim_order)
+        mul_output = tosa_graph.addIntermediate(output_shape, ts.DType.INT32)
+
+        # Do the INT32 Mul
+        tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{node.name}_shift")
+        tosa_graph.addOperator(
+            ts.TosaOp.Op().MUL,
+            [input_A_rescaled.name, input_B_rescaled.name, f"{node.name}_shift"],
+            [mul_output.name],
+        )
+        output_scale = input_A_qargs.scale * input_B_qargs.scale
+        tqutils.insert_rescale_op_to_int8(
+            tosa_graph, mul_output, output_scale, node, self.tosa_specs
+        )
+
+
+@register_node_visitor
+class MulVisitor_FP(MulVisitor_INT):
+    # inheriting 'target' from INT class
+
+    tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+FP")]
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+
+        import serializer.tosa_serializer as ts  # type: ignore
+
+        if inputs[0].dtype == ts.DType.INT8:
+            return super().define_node(node, tosa_graph, inputs, output)
+
+        input1, input2 = inputs
+
+        tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{node.name}_shift")
+        tosa_graph.addOperator(
+            ts.TosaOp.Op().MUL,
+            [input1.name, input2.name, f"{node.name}_shift"],
+            [output.name],
+        )