Arm backend: Support for sin and cos for TOSA 1.0

SIN and COS were introduced in TOSA 1.0. This patch adds support for
both. It also adds unittests for both ops.

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

Author: oscarandersson8218

backends/arm/_passes/insert_table_ops.py

@@ -48,6 +48,8 @@ class TableOps:
         exir_ops.edge.aten.reciprocal.default: torch.reciprocal,
         exir_ops.edge.aten.rsqrt.default: torch.rsqrt,
         exir_ops.edge.aten.sigmoid.default: torch.sigmoid,
+        exir_ops.edge.aten.cos.default: torch.cos,
+        exir_ops.edge.aten.sin.default: torch.sin,
         exir_ops.edge.aten.tanh.default: torch.tanh,
         exir_ops.edge.aten.hardsigmoid.default: torch.nn.functional.hardsigmoid,
         exir_ops.edge.aten.hardswish.default: torch.nn.functional.hardswish,

backends/arm/operator_support/__init__.py

@@ -12,6 +12,7 @@
     pool_2d_support,
     reduce_sum_support,
     right_shift_support,
+    sin_cos_support,
     slice_copy_support,
     to_copy_support,
     tosa_supported_operators,

backends/arm/operator_support/sin_cos_support.py

@@ -0,0 +1,32 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+
+import torch.fx as fx
+from executorch.backends.arm.operator_support.tosa_supported_operators import (
+    register_tosa_support_check,
+    SupportedTOSAOperatorCheck,
+)
+from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+@register_tosa_support_check
+class SinCosSupported(SupportedTOSAOperatorCheck):
+    targets = [
+        exir_ops.edge.aten.cos.default,
+        exir_ops.edge.aten.sin.default,
+    ]
+
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-0.80+BI"),
+        TosaSpecification.create_from_string("TOSA-1.0+INT"),
+        TosaSpecification.create_from_string("TOSA-1.0+FP"),
+    ]
+
+    def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
+        return True

backends/arm/operators/__init__.py

@@ -18,6 +18,7 @@
     op_clamp,
     op_constant_pad_nd,
     op_conv2d,
+    op_cos,
     op_eq,
     op_erf,
     op_exp,
@@ -38,6 +39,7 @@
     op_rshift_tensor,
     op_rsqrt,
     op_sigmoid,
+    op_sin,
     op_slice,
     op_sub,
     op_sum,

backends/arm/operators/op_cos.py

@@ -0,0 +1,46 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+from typing import List
+
+import serializer.tosa_serializer as ts  # type: ignore
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_specification import TosaSpecification
+from torch.fx import Node
+
+
+@register_node_visitor
+class CosVisitor(NodeVisitor):
+    target = "aten.cos.default"
+
+    # INT case should be handled by op_table
+    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: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        if len(node.all_input_nodes) != 1:
+            raise ValueError(
+                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
+            )
+        if inputs[0].dtype != ts.DType.FP32 or output.dtype != ts.DType.FP32:
+            raise ValueError(
+                f"Input and output for {self.target} need to be FP32, got input_dtype: "
+                f"{inputs[0].dtype} and output_dtype: {output.dtype}"
+            )
+
+        tosa_graph.addOperator(ts.TosaOp.Op().COS, [inputs[0].name], [output.name])

backends/arm/operators/op_sin.py

@@ -0,0 +1,46 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+from typing import List
+
+import serializer.tosa_serializer as ts  # type: ignore
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_specification import TosaSpecification
+from torch.fx import Node
+
+
+@register_node_visitor
+class SinVisitor(NodeVisitor):
+    target = "aten.sin.default"
+
+    # INT case should be handled by op_table
+    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: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        if len(node.all_input_nodes) != 1:
+            raise ValueError(
+                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
+            )
+        if inputs[0].dtype != ts.DType.FP32 or output.dtype != ts.DType.FP32:
+            raise ValueError(
+                f"Input and output for {self.target} need to be FP32, got input_dtype: "
+                f"{inputs[0].dtype} and output_dtype: {output.dtype}"
+            )
+
+        tosa_graph.addOperator(ts.TosaOp.Op().SIN, [inputs[0].name], [output.name])

backends/arm/quantizer/quantization_annotator.py

@@ -171,6 +171,8 @@ def _match_pattern(
     torch.ops.aten.reciprocal.default,
     torch.ops.aten.rsqrt.default,
     torch.ops.aten.sigmoid.default,
+    torch.ops.aten.cos.default,
+    torch.ops.aten.sin.default,
     torch.ops.aten.tanh.default,
     torch.ops.aten.sum.dim_IntList,
     torch.ops.aten.hardsigmoid.default,

backends/arm/test/misc/test_multiple_delegates.py

@@ -20,7 +20,7 @@ def get_inputs(self):
 
         def forward(self, x: torch.Tensor, y: torch.Tensor):
             z = x + y
-            s = torch.sin(z)
+            s = torch.tan(z)
             return s * z
 
     def test_tosa_MI(self):

backends/arm/test/ops/test_cos.py

@@ -0,0 +1,83 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+
+from executorch.backends.arm.test import common, conftest
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineBI,
+    EthosU85PipelineBI,
+    TosaPipelineBI,
+    TosaPipelineMI,
+)
+
+aten_op = "torch.ops.aten.cos.default"
+input_t1 = Tuple[torch.Tensor]  # Input x
+
+test_data_suite = {
+    # (test_name, test_data)
+    "zeros": torch.zeros(10, 10, 10, 10),
+    "ones": torch.ones(10, 10, 10),
+    "rand": torch.rand(10, 10) - 0.5,
+    "randn_pos": torch.randn(10) + 10,
+    "randn_neg": torch.randn(10) - 10,
+    "ramp": torch.arange(-16, 16, 0.2),
+}
+
+
+class Cos(torch.nn.Module):
+
+    def forward(self, x: torch.Tensor):
+        return torch.cos(x)
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_cos_tosa_MI(test_data: Tuple):
+    pipeline = TosaPipelineMI[input_t1](
+        Cos(),
+        (test_data,),
+        aten_op,
+        exir_op=[],
+    )
+    if conftest.get_option("tosa_version") == "1.0":
+        pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_cos_tosa_BI(test_data: Tuple):
+    pipeline = TosaPipelineBI[input_t1](
+        Cos(),
+        (test_data,),
+        aten_op,
+        exir_op=[],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_cos_tosa_u55_BI(test_data: Tuple):
+    pipeline = EthosU55PipelineBI[input_t1](
+        Cos(),
+        (test_data,),
+        aten_op,
+        exir_ops=[],
+        run_on_fvp=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_cos_tosa_u85_BI(test_data: Tuple):
+    pipeline = EthosU85PipelineBI[input_t1](
+        Cos(),
+        (test_data,),
+        aten_op,
+        exir_ops=[],
+        run_on_fvp=False,
+    )
+    pipeline.run()

backends/arm/test/ops/test_sin.py

@@ -0,0 +1,83 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+
+from executorch.backends.arm.test import common, conftest
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineBI,
+    EthosU85PipelineBI,
+    TosaPipelineBI,
+    TosaPipelineMI,
+)
+
+aten_op = "torch.ops.aten.sin.default"
+input_t1 = Tuple[torch.Tensor]  # Input x
+
+test_data_suite = {
+    # (test_name, test_data)
+    "zeros": torch.zeros(10, 10, 10, 10),
+    "ones": torch.ones(10, 10, 10),
+    "rand": torch.rand(10, 10) - 0.5,
+    "randn_pos": torch.randn(10) + 10,
+    "randn_neg": torch.randn(10) - 10,
+    "ramp": torch.arange(-16, 16, 0.2),
+}
+
+
+class Sin(torch.nn.Module):
+
+    def forward(self, x: torch.Tensor):
+        return torch.sin(x)
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_sin_tosa_MI(test_data: Tuple):
+    pipeline = TosaPipelineMI[input_t1](
+        Sin(),
+        (test_data,),
+        aten_op,
+        exir_op=[],
+    )
+    if conftest.get_option("tosa_version") == "1.0":
+        pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_sin_tosa_BI(test_data: Tuple):
+    pipeline = TosaPipelineBI[input_t1](
+        Sin(),
+        (test_data,),
+        aten_op,
+        exir_op=[],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_sin_tosa_u55_BI(test_data: Tuple):
+    pipeline = EthosU55PipelineBI[input_t1](
+        Sin(),
+        (test_data,),
+        aten_op,
+        exir_ops=[],
+        run_on_fvp=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_sin_tosa_u85_BI(test_data: Tuple):
+    pipeline = EthosU85PipelineBI[input_t1](
+        Sin(),
+        (test_data,),
+        aten_op,
+        exir_ops=[],
+        run_on_fvp=False,
+    )
+    pipeline.run()