Arm backend: Update rescale to handle more dtypes

Update op_rescale to handle other dtype conversion than int8 <-> int32
for TOSA 1.0.

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

Author: oscarandersson8218

backends/arm/operators/op_rescale.py

@@ -16,8 +16,8 @@
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
 )
-from executorch.backends.arm.tosa_mapping import TosaArg
-from executorch.backends.arm.tosa_quant_utils import create_const_ops_for_rescale
+from executorch.backends.arm.tosa_mapping import map_dtype, TosaArg
+from executorch.backends.arm.tosa_quant_utils import build_rescale
 
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from torch.fx import Node
@@ -98,53 +98,29 @@ def define_node(
 
         validate_num_inputs(self.target, inputs, 5)
 
-        input_dtype = node.all_input_nodes[0].meta["val"].dtype
+        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 != torch.int8 and input_zp != 0:
+        if input_dtype != map_dtype(torch.int8, self.tosa_spec) and input_zp != 0:
             raise ValueError(
                 f"If input dtype is not int8, input_zp must be 0. Got input_dtype{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=}"
+                f"If output dtype is not int8, output_zp must be 0. Got {ts.DTypeNames[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(
+        build_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,
+            scale=[scale],
+            input_node=inputs[0],
+            output_name=output.name,
+            output_type=output.dtype,
+            input_zp=input_zp,
+            output_zp=output_zp,
             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,
         )

backends/arm/tosa_quant_utils.py

@@ -236,13 +236,21 @@ def build_rescale_v0_80(
 # For TOSA spec v1.0 RESCALE operator requires multipler, shifts, input_zp and output_zp to be
 # const inputs. Create constant operators from the data already initialized.
 def create_const_ops_for_rescale(
-    tosa_fb, input_dtype, input_name, multipliers, shifts, input_zp, output_zp, ts
+    tosa_fb,
+    scale_32,
+    input_dtype,
+    input_name,
+    multipliers,
+    shifts,
+    input_zp,
+    output_zp,
+    output_dtype,
+    ts,
 ):
-    output_dtype = ts.DType.INT32 if input_dtype == ts.DType.INT8 else ts.DType.INT8
 
     multipliers = tosa_fb.addConst(
         (len(multipliers),),
-        ts.DType.INT32,
+        ts.DType.INT32 if scale_32 else ts.DType.INT16,
         multipliers,
         name=input_name + "_multipliers",
     )
@@ -275,20 +283,24 @@ def build_rescale(
 
     input_name = input_node.name
 
-    multipliers, shifts = compute_multiplier_and_shift(scale, 32)
+    scaleWidth = 32
+    is_scale32 = True
+    multipliers, shifts = compute_multiplier_and_shift(scale, scaleWidth)
     rescale_inputs = create_const_ops_for_rescale(
         tosa_fb,
+        is_scale32,
         input_node.dtype,
         input_name,
         multipliers,
         shifts,
         input_zp,
         output_zp,
+        output_type,
         ts,
     )
     attr_rescale = ts.TosaSerializerAttribute()
     attr_rescale.RescaleAttribute(
-        scale32=True,
+        scale32=is_scale32,
         rounding_mode=rounding_mode,
         per_channel=per_channel,
         input_unsigned=False,