Arm backend: Support for rescale for TOSA 1.0

Signed-off-by: Per Åstrand <[email protected]>
Change-Id: I88822c93470c4ff54f0e75df2f5ff271fa6bd1e6

Author: per

backends/arm/operators/op_rescale.py

@@ -5,32 +5,35 @@
 
 # pyre-unsafe
 
-from typing import cast, List
+from typing import Any, cast, List
 
 import executorch.backends.arm.tosa_quant_utils as tosa_quant_utils
 import torch
-import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
-
-import tosa_tools.v0_80.tosa.Op as TosaOp  # type: ignore
 from executorch.backends.arm.operators.node_visitor import (
     NodeVisitor,
     register_node_visitor,
 )
 from executorch.backends.arm.tosa_mapping import map_dtype, TosaArg
+from executorch.backends.arm.tosa_quant_utils import create_const_ops_for_rescale
+
+from executorch.backends.arm.tosa_specification import TosaSpecification
 from torch.fx import Node
 
 
 @register_node_visitor
-class RescaleVisitor(NodeVisitor):
+class RescaleVisitor_0_80(NodeVisitor):
     target = "_rescale.default"
 
+    tosa_specs = NodeVisitor.tosa_specs_0_80
+
     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
 
         input_dtype = inputs[0].dtype
         output_dtype = cast(torch.dtype, node.args[1])
@@ -68,5 +71,73 @@ def define_node(
         )
 
         tosa_graph.addOperator(
-            TosaOp.Op().RESCALE, [inputs[0].name], [output.name], attr_rescale
+            ts.TosaOp.Op().RESCALE, [inputs[0].name], [output.name], attr_rescale
+        )
+
+
+@register_node_visitor
+class RescaleVisitor_INT(NodeVisitor):
+    target = "_rescale.default"
+
+    tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+INT")]
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        import serializer.tosa_serializer as ts  # type: ignore
+        from tosa.RoundingMode import RoundingMode  # type: ignore
+
+        input_dtype = inputs[0].dtype
+        output_dtype = cast(torch.dtype, node.args[1])
+        scale = cast(float, node.args[2])
+        input_zp = cast(int, node.args[3])
+        output_zp = cast(int, node.args[4])
+
+        if input_dtype != map_dtype(torch.int8) and input_zp != 0:
+            raise ValueError(
+                f"If input dtype is not int8, input_zp must be 0. Got input_dtype{ts.DTypeNames[input_dtype]}, {input_zp=}"
+            )
+        if output_dtype != torch.int8 and output_zp != 0:
+            raise ValueError(
+                f"If output dtype is not int8, output_zp must be 0. Got {output_dtype=}, {output_zp=}"
+            )
+
+        # scale32 gives higher accuracy but for a higher HW cost.
+        # For now, always go for scale32.
+        scale_32 = True
+        scale_width = 32 if scale_32 else 16
+        multipliers, shifts = tosa_quant_utils.compute_multiplier_and_shift(
+            [scale], scale_width
+        )
+
+        rescale_inputs = create_const_ops_for_rescale(
+            tosa_graph,
+            input_dtype,
+            inputs[0].name,
+            multipliers,
+            shifts,
+            input_zp,
+            output_zp,
+            ts,
+        )
+
+        attr_rescale = ts.TosaSerializerAttribute()
+
+        attr_rescale.RescaleAttribute(
+            scale32=scale_32,
+            rounding_mode=RoundingMode.SINGLE_ROUND,
+            per_channel=False,
+            input_unsigned=False,
+            output_unsigned=False,
+        )
+
+        tosa_graph.addOperator(
+            ts.TosaOp.Op().RESCALE,
+            [inputs[0].name, *rescale_inputs],
+            [output.name],
+            attr_rescale,
         )