Arm backend: Add FP16 support to operators pt.7 (#17330)

Add FP16 support for operators:
 - sin
 - sub
 - tanh
 - where

Update op tests to cover the new datatype.

Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com>

Author: martinlsm

backends/arm/operators/op_sin.py

@@ -42,7 +42,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [*inputs, output],
-            [ts.DType.FP32, ts.DType.BF16],
+            [ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
             self.tosa_spec,
         )
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_sub.py

@@ -37,7 +37,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [*inputs, output],
-            [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16],
+            [ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
             self.tosa_spec,
         )
 

backends/arm/operators/op_tanh.py

@@ -44,7 +44,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [*inputs, output],
-            [ts.DType.FP32, ts.DType.BF16],
+            [ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
             self.tosa_spec,
         )
         attr = ts.TosaSerializerAttribute()

backends/arm/test/ops/test_sin.py

@@ -30,6 +30,10 @@
     "ramp": torch.arange(-16, 16, 0.2),
 }
 
+test_data_suite_fp16 = {
+    "rand_fp16": torch.rand(10, 10, dtype=torch.float16),
+}
+
 test_data_suite_bf16 = {
     "rand_bf16": torch.rand(3, 3, dtype=torch.bfloat16),
 }
@@ -41,7 +45,9 @@ def forward(self, x: torch.Tensor):
         return torch.sin(x)
 
 
-@common.parametrize("test_data", test_data_suite | test_data_suite_bf16)
+@common.parametrize(
+    "test_data", test_data_suite | test_data_suite_fp16 | test_data_suite_bf16
+)
 def test_sin_tosa_FP(test_data: Tuple):
     pipeline = TosaPipelineFP[input_t1](
         Sin(),
@@ -88,7 +94,7 @@ def test_sin_u85_INT(test_data: Tuple):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_suite)
+@common.parametrize("test_data", test_data_suite | test_data_suite_fp16)
 @common.SkipIfNoModelConverter
 def test_sin_vgf_no_quant(test_data: Tuple):
     pipeline = VgfPipeline[input_t1](

backends/arm/test/ops/test_sub.py

@@ -34,6 +34,14 @@
     "zeros": lambda: (torch.zeros(10),),
 }
 
+sub_test_data_fp16 = {
+    "rand_2D_fp16": lambda: (torch.rand(4, 4, dtype=torch.float16),),
+}
+
+sub_test_data_bf16 = {
+    "rand_2D_bf16": lambda: (torch.rand(4, 4, dtype=torch.bfloat16),),
+}
+
 # Two-input subtraction (x - y)
 sub2_test_data = {
     "rand_2D_4x4": lambda: (torch.rand(4, 4), torch.rand(4, 4)),
@@ -52,8 +60,11 @@
     "rand_3d_Scalar": lambda: (torch.rand(1, 6, 2), 1),
 }
 
-sub_test_data_bf16 = {
-    "rand_2D_bf16": lambda: (torch.rand(4, 4, dtype=torch.bfloat16),),
+sub2_test_data_fp16 = {
+    "rand_2D_pair_fp16": lambda: (
+        torch.rand(2, 3, dtype=torch.float16),
+        torch.rand(2, 3, dtype=torch.float16),
+    ),
 }
 
 sub2_test_data_bf16 = {
@@ -101,7 +112,9 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
 input_t2 = Tuple[torch.Tensor, torch.Tensor]  # Input x, y
 
 
-@common.parametrize("test_data", sub_test_data | sub_test_data_bf16)
+@common.parametrize(
+    "test_data", sub_test_data | sub_test_data_fp16 | sub_test_data_bf16
+)
 def test_sub_tensor_tosa_FP(test_data):
     """Test Subtraction (TOSA FP)"""
     pipeline = TosaPipelineFP[input_t1](
@@ -114,7 +127,9 @@ def test_sub_tensor_tosa_FP(test_data):
     pipeline.run()
 
 
-@common.parametrize("test_data", sub2_test_data | sub2_test_data_bf16)
+@common.parametrize(
+    "test_data", sub2_test_data | sub2_test_data_fp16 | sub2_test_data_bf16
+)
 def test_sub_tensor_tosa_FP_2(test_data: Tuple[torch.Tensor, torch.Tensor]):
     """Test Two-Operand Subtraction (TOSA FP)"""
     pipeline = TosaPipelineFP[input_t2](
@@ -123,7 +138,9 @@ def test_sub_tensor_tosa_FP_2(test_data: Tuple[torch.Tensor, torch.Tensor]):
     pipeline.run()
 
 
-@common.parametrize("test_data", sub_tan_test_data | sub2_test_data_bf16)
+@common.parametrize(
+    "test_data", sub_tan_test_data | sub2_test_data_fp16 | sub2_test_data_bf16
+)
 def test_sub_tensor_tosa_FP_alpha(test_data: Tuple[torch.Tensor, torch.Tensor]):
     """Test Two-Operand Subtraction with alpha (TOSA FP)"""
     pipeline = TosaPipelineFP[input_t2](
@@ -217,7 +234,7 @@ def test_sub_tensor_u85_INT(test_data: Tuple[torch.Tensor, torch.Tensor]):
     pipeline.run()
 
 
-@common.parametrize("test_data", sub_test_data)
+@common.parametrize("test_data", sub_test_data | sub_test_data_fp16)
 @common.SkipIfNoModelConverter
 def test_sub_tensor_vgf_no_quant(test_data: Tuple[torch.Tensor]):
     """Test Subtraction (VGF FP)"""
@@ -231,7 +248,7 @@ def test_sub_tensor_vgf_no_quant(test_data: Tuple[torch.Tensor]):
     pipeline.run()
 
 
-@common.parametrize("test_data", sub2_test_data)
+@common.parametrize("test_data", sub2_test_data | sub2_test_data_fp16)
 @common.SkipIfNoModelConverter
 def test_sub_tensor_vgf_no_quant_2(test_data: Tuple[torch.Tensor, torch.Tensor]):
     """Test Two-Operand Subtraction (VGF FP)"""

backends/arm/test/ops/test_tanh.py

@@ -30,6 +30,10 @@
     "ramp": lambda: torch.arange(-16, 16, 0.2),
 }
 
+test_data_suite_fp16 = {
+    "rand_fp16": lambda: torch.rand(4, 4, 2, 2, 2, dtype=torch.float16) - 0.5,
+}
+
 test_data_suite_bf16 = {
     "rand_bf16": lambda: torch.rand(4, 4, 2, 2, 2, dtype=torch.bfloat16) - 0.5,
 }
@@ -44,7 +48,9 @@ def forward(self, x):
         return self.tanh(x)
 
 
-@common.parametrize("test_data", test_data_suite | test_data_suite_bf16)
+@common.parametrize(
+    "test_data", test_data_suite | test_data_suite_fp16 | test_data_suite_bf16
+)
 def test_tanh_tosa_FP(test_data: Tuple):
     pipeline = TosaPipelineFP[input_t1](
         Tanh(),
@@ -91,7 +97,7 @@ def test_tanh_u85_INT(test_data: Tuple):
     pipeline.run()
 
 
-@common.parametrize("test_data", test_data_suite)
+@common.parametrize("test_data", test_data_suite | test_data_suite_fp16)
 @common.SkipIfNoModelConverter
 def test_tanh_vgf_no_quant(test_data: Tuple):
     pipeline = VgfPipeline[input_t1](

backends/arm/test/ops/test_where.py

@@ -44,7 +44,7 @@ def get_inputs(self):
                     self.shape,
                     dtype=self.dtype[i],
                 )
-            elif self.dtype[i] in [torch.float32, torch.bfloat16]:
+            elif self.dtype[i] in [torch.float16, torch.float32, torch.bfloat16]:
                 inputs[i] = torch.randn(*self.shape).to(self.dtype[i])
             elif self.dtype[i] is torch.bool:
                 inputs[i] = torch.randint(0, 1, self.shape, dtype=torch.bool)
@@ -114,6 +114,12 @@ def scalar_condition(input: torch.Tensor):
     tensor_condition,
 )
 
+float16_tensor_cond = Where(
+    1,
+    torch.float16,
+    tensor_condition,
+)
+
 float32_tensor_cond_tuple_dtype = Where(
     1,
     (torch.float32, torch.int8),
@@ -175,6 +181,7 @@ def scalar_condition(input: torch.Tensor):
 test_modules_FP = {
     **test_modules_common,
     "float32_tensor_cond_tuple_dtype_bool": lambda: float32_tensor_cond_tuple_dtype_bool,
+    "float16_tensor_cond": lambda: float16_tensor_cond,
 }
 
 test_modules_FP_bf16 = {