Arm backend: Update for missing operators for int16x8 (#15521)

### Summary
Updates to operators for int16: 

- avg_pool2d, clamp, constant_pad_nd, eq, ge, gt, le, lt, max_pool2d,
upsample_bilinear, upsample_nearest2d

### Test plan
Unit tests added for affected operators 


cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

---------

Signed-off-by: Saoirse Stewart <[email protected]>

Author: SaoirseARM

backends/arm/_passes/rewrite_upsample.py

@@ -11,6 +11,7 @@
     create_node,
     get_first_fake_tensor,
 )
+from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
 from executorch.backends.arm.tosa.utils import get_resize_parameters
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
@@ -52,7 +53,9 @@ def call(self, graph_module):
                 node.replace_all_uses_with(tosa_resize_node)
                 graph_module.graph.erase_node(node)
             input_dtype = get_first_fake_tensor(x).dtype
-            if input_dtype == torch.int8 and resize_mode == "bilinear":
+            if (
+                input_dtype == torch.int8 or input_dtype == torch.int16
+            ) and resize_mode == "bilinear":
                 input_size = get_first_fake_tensor(x).shape
                 input_size_xy = input_size[2:]
                 output_size = get_first_fake_tensor(node).shape
@@ -71,6 +74,11 @@ def call(self, graph_module):
                         exir_ops.backend.tosa.RESCALE.default,
                     )
                     tosa_resize_node.replace_all_uses_with(rescale_node)
+                    if input_dtype == torch.int16:
+                        tosa_resize_node.meta[TosaSpecialDtype.meta_key()] = (
+                            TosaSpecialDtype.INT48
+                        )
+
                     rescale_node.args = (
                         tosa_resize_node,
                         output_dtype,

backends/arm/operators/op_avg_pool2d.py

@@ -118,11 +118,11 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.FP32],
             output.tosa_spec,
         )
 
-        if inputs[0].dtype == ts.DType.INT8:
+        if inputs[0].dtype == ts.DType.INT8 or inputs[0].dtype == ts.DType.INT16:
             accumulator_type = ts.DType.INT32
             input_qargs = get_input_qparams(node)
             input_zp = input_qargs[0].get_zp_per_tensor()

backends/arm/operators/op_clamp.py

@@ -74,6 +74,8 @@ def _to_bytes(self, value: int | float, dtype: torch.dtype) -> bytes:
             return np.frombuffer(np.float16(value).tobytes(), dtype=np.uint8).tolist()
         elif dtype == torch.int8:
             return np.frombuffer(np.int8(value).tobytes(), dtype=np.uint8).tolist()
+        elif dtype == torch.int16:
+            return np.frombuffer(np.int16(value).tobytes(), dtype=np.uint8).tolist()
         else:
             raise ValueError(f"Unsupported dtype for to_bytes: {dtype}")
 
@@ -89,7 +91,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.FP16, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.FP16, ts.DType.FP32],
             output.tosa_spec,
         )
 

backends/arm/operators/op_constant_pad_nd.py

@@ -50,6 +50,7 @@ def define_node(
             [inputs[0], output],
             [
                 ts.DType.INT8,
+                ts.DType.INT16,
                 ts.DType.INT32,
                 ts.DType.FP32,
                 ts.DType.BOOL,
@@ -62,6 +63,11 @@ def define_node(
             qargs = input_qparams[0]
             pad_const_val = qargs.quantize_value(inputs[2].number).item()
             pad_const_dtype = ts.DType.INT8
+        elif inputs[0].dtype == ts.DType.INT16:
+            input_qparams = get_input_qparams(node)
+            qargs = input_qparams[0]
+            pad_const_val = qargs.quantize_value(inputs[2].number).item()
+            pad_const_dtype = ts.DType.INT16
         else:
             pad_const_val = inputs[2].number
             pad_const_dtype = inputs[0].dtype

backends/arm/operators/op_eq.py

@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)

backends/arm/operators/op_ge.py

@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)

backends/arm/operators/op_gt.py

@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)

backends/arm/operators/op_le.py

@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)

backends/arm/operators/op_lt.py

@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)

backends/arm/operators/op_max_pool2d.py

@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.FP32],
             output.tosa_spec,
         )