Arm backend: Add cosh decomposition pass and test (#13181)

Add decomposition and tests for cosh.

Signed-off-by: Emma Kujala <[email protected]>

Author: emmakujala

backends/arm/_passes/__init__.py

@@ -31,6 +31,7 @@
 from .decompose_atanh_pass import DecomposeAtanhPass  # noqa
 from .decompose_avg_pool2d import DecomposeAvgPool2d  # noqa
 from .decompose_batch_norm_no_stats import DecomposeBatchNormNoStatsPass  # noqa
+from .decompose_cosh_pass import DecomposeCoshPass  # noqa
 from .decompose_cosine_similarity_pass import DecomposeCosineSimilarityPass  # noqa
 from .decompose_div_pass import DecomposeDivPass  # noqa
 from .decompose_embedding_pass import DecomposeEmbeddingPass  # noqa  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -36,6 +36,7 @@
     DecomposeAtanPass,
     DecomposeAvgPool2d,
     DecomposeBatchNormNoStatsPass,
+    DecomposeCoshPass,
     DecomposeCosineSimilarityPass,
     DecomposeDivPass,
     DecomposeEmbeddingPass,
@@ -167,6 +168,7 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(DecomposeAcoshPass())
         self.add_pass(DecomposeAsinPass())
         self.add_pass(DecomposeAsinhPass())
+        self.add_pass(DecomposeCoshPass())
         self.add_pass(DecomposeSqrtPass())
         self.add_pass(DecomposeAtanPass())
         self.add_pass(DecomposeAtanhPass())

backends/arm/_passes/decompose_cosh_pass.py

@@ -0,0 +1,48 @@
+# 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.
+
+from executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+
+# For MI case
+edge_cosh = exir_ops.edge.aten.cosh.default
+
+
+class DecomposeCoshPass(ArmPass):
+    """
+    This pass replaces the cosh operator with a sequence of TOSA-equivalent operations that
+    compute the hyperbolic cosine using the formula:
+
+        cosh(x) = 0.5 * (e^x + e^(-x))
+
+    """
+
+    def call_operator(self, op, args, kwargs, meta, updated=False):
+        if op is not edge_cosh:
+            return super().call_operator(op, args, kwargs, meta, updated)
+
+        x = args
+
+        exp_op, mul_op, neg_op, add_op = (
+            exir_ops.edge.aten.exp.default,
+            exir_ops.edge.aten.mul.Scalar,
+            exir_ops.edge.aten.neg.default,
+            exir_ops.edge.aten.add.Tensor,
+        )
+
+        # exp1 = e^x
+        exp1 = super().call_operator(exp_op, x, {}, meta, updated=True)
+
+        # exp2 = e^(⁻x)
+        neg_x = super().call_operator(neg_op, x, {}, meta, updated=True)
+        exp2 = super().call_operator(exp_op, (neg_x,), {}, meta, updated=True)
+
+        # numer = exp1 + exp2
+        numer = super().call_operator(add_op, (exp1, exp2), {}, meta, updated=True)
+
+        # out = 0.5 * numer
+        out = super().call_operator(mul_op, (numer, 0.5), {}, meta, updated=True)
+
+        return out

backends/arm/_passes/insert_table_ops.py

@@ -59,6 +59,7 @@ class TableOps:
         exir_ops.edge.aten.acosh.default: torch.acosh,
         exir_ops.edge.aten.asin.default: torch.asin,
         exir_ops.edge.aten.asinh.default: torch.asinh,
+        exir_ops.edge.aten.cosh.default: torch.cosh,
     }
 
     # Targets that must be treated explicitly

backends/arm/operator_support/tosa_supported_operators.py

@@ -257,6 +257,7 @@ def is_node_supported(
             exir_ops.edge.aten.addmm.default,
             exir_ops.edge.aten.masked_fill.Scalar,
             exir_ops.edge.aten.asinh.default,
+            exir_ops.edge.aten.cosh.default,
         ]
 
         return supported

backends/arm/quantizer/quantization_annotator.py

@@ -287,6 +287,7 @@ def _match_pattern(
     torch.ops.aten.asin.default,
     torch.ops.aten.atanh.default,
     torch.ops.aten.asinh.default,
+    torch.ops.aten.cosh.default,
 ]
 
 _one_to_one_shared_input_qspec = [

backends/arm/test/ops/test_cosh.py

@@ -0,0 +1,107 @@
+# 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.
+from typing import Tuple
+
+import torch
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
+    TosaPipelineFP,
+    TosaPipelineINT,
+    VgfPipeline,
+)
+
+aten_op = "torch.ops.aten.cosh.default"
+exir_op = "executorch_exir_dialects_edge__ops_aten__cosh_default"
+
+input_t1 = Tuple[torch.Tensor]  # Input x
+
+test_data_suite = {
+    # (test_name, test_data)
+    "zeros": torch.zeros(10, 10, 10),
+    "zeros_4D": torch.zeros(1, 10, 32, 7),
+    "zeros_alt_shape": torch.zeros(10, 3, 5),
+    "ones": torch.ones(15, 10, 7),
+    "ones_4D": torch.ones(1, 3, 32, 16),
+    "rand": torch.rand(10, 10) - 0.5,
+    "rand_alt_shape": torch.rand(10, 3, 5) - 0.5,
+    "rand_4D": torch.rand(1, 6, 5, 7) - 0.5,
+    "randn_pos": torch.randn(10) + 10,
+    "randn_neg": torch.randn(10) - 10,
+    "ramp": torch.arange(-16, 16, 0.2),
+    "large": 100 * torch.ones(1, 1),
+    "small": 0.000001 * torch.ones(1, 1),
+    "small_rand": torch.rand(100) * 0.01,
+    "biggest": torch.tensor([700.0, 710.0, 750.0]),
+}
+
+
+class Cosh(torch.nn.Module):
+    def forward(self, x: torch.Tensor):
+        return torch.cosh(x)
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_cosh_tosa_FP(test_data: Tuple):
+    pipeline = TosaPipelineFP[input_t1](
+        Cosh(),
+        (test_data,),
+        aten_op,
+        exir_op,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_cosh_tosa_INT(test_data: Tuple):
+    pipeline = TosaPipelineINT[input_t1](
+        Cosh(), (test_data,), aten_op=aten_op, exir_op=exir_op
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone300
+@common.parametrize("test_data", test_data_suite)
+def test_cosh_u55_INT(test_data: Tuple):
+    pipeline = EthosU55PipelineINT[input_t1](
+        Cosh(), (test_data,), aten_ops=aten_op, exir_ops=exir_op
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone320
+@common.parametrize("test_data", test_data_suite)
+def test_cosh_u85_INT(test_data: Tuple):
+    pipeline = EthosU85PipelineINT[input_t1](
+        Cosh(), (test_data,), aten_ops=aten_op, exir_ops=exir_op
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.SkipIfNoModelConverter
+def test_cosh_vgf_FP(test_data: Tuple):
+    pipeline = VgfPipeline[input_t1](
+        Cosh(),
+        (test_data,),
+        [],
+        [],
+        tosa_version="TOSA-1.0+FP",
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.SkipIfNoModelConverter
+def test_cosh_vgf_INT(test_data: Tuple):
+    pipeline = VgfPipeline[input_t1](
+        Cosh(),
+        (test_data,),
+        [],
+        [],
+        tosa_version="TOSA-1.0+INT",
+    )
+    pipeline.run()