Arm backend: Add atanh decomposition pass (#12483)

Decomposes atanh into other operators/lookup table for MI/BI case.

Signed-off-by: Teo Bergkvist <[email protected]>

Author: tbergkvist

backends/arm/_passes/__init__.py

@@ -26,6 +26,7 @@
 from .decompose_adaptive_avg_pool2d_pass import DecomposeAdaptiveAvgPool2dPass  # noqa
 from .decompose_asin_pass import DecomposeAsinPass  # noqa
 from .decompose_atan_pass import DecomposeAtanPass  # noqa
+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_cosine_similarity_pass import DecomposeCosineSimilarityPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -30,6 +30,7 @@
     DecomposeAcoshPass,
     DecomposeAdaptiveAvgPool2dPass,
     DecomposeAsinPass,
+    DecomposeAtanhPass,
     DecomposeAtanPass,
     DecomposeAvgPool2d,
     DecomposeBatchNormNoStatsPass,
@@ -163,6 +164,7 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(DecomposeAsinPass())
         self.add_pass(DecomposeSqrtPass())
         self.add_pass(DecomposeAtanPass())
+        self.add_pass(DecomposeAtanhPass())
         self.add_pass(ConvertIntPowToMuls())
         self.add_pass(CastBoolToInt8Pass())
         self.add_pass(DecomposeSinhPass())

backends/arm/_passes/decompose_atanh_pass.py

@@ -0,0 +1,62 @@
+# 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
+
+
+edge_atanh = exir_ops.edge.aten.atanh.default  # MI case
+
+
+def _get_atanh_ops(op):
+    """Return the primitive ops required.."""
+    if op is not edge_atanh:
+        raise RuntimeError(f"Can't decompose atanh for op {op}")
+    return (
+        exir_ops.edge.aten.mul.Tensor,
+        exir_ops.edge.aten.mul.Scalar,
+        exir_ops.edge.aten.add.Scalar,
+        exir_ops.edge.aten.reciprocal.default,
+        exir_ops.edge.aten.log.default,
+        exir_ops.edge.aten.neg.default,
+    )
+
+
+class DecomposeAtanhPass(ArmPass):
+    """
+    Decomposes the atanh operator into primitive ops.
+    atanh(x) = 0.5 * log((1 + x) / (1 - x))
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op is not edge_atanh:
+            return super().call_operator(op, args, kwargs, meta, updated=False)
+
+        ops = _get_atanh_ops(op)
+        (
+            op_mul_tensor,
+            op_mul_scalar,
+            op_add_scalar,
+            op_reciprocal,
+            op_log,
+            op_neg,
+        ) = ops
+
+        x = args[0]
+
+        nom = super().call_operator(op_add_scalar, (x, 1.0), {}, meta, updated=True)
+
+        neg_x = super().call_operator(op_neg, (x,), {}, meta, updated=True)
+        denom = super().call_operator(
+            op_add_scalar, (neg_x, 1.0), {}, meta, updated=True
+        )
+        recip = super().call_operator(op_reciprocal, (denom,), {}, meta, updated=True)
+
+        log_input = super().call_operator(
+            op_mul_tensor, (nom, recip), {}, meta, updated=True
+        )
+        log = super().call_operator(op_log, (log_input,), {}, meta, updated=True)
+
+        return super().call_operator(op_mul_scalar, (log, 0.5), {}, meta, updated=True)

backends/arm/_passes/insert_table_ops.py

@@ -52,6 +52,7 @@ class TableOps:
         exir_ops.edge.aten.sin.default: torch.sin,
         exir_ops.edge.aten.tanh.default: torch.tanh,
         exir_ops.edge.aten.atan.default: torch.atan,
+        exir_ops.edge.aten.atanh.default: torch.atanh,
         exir_ops.edge.aten.hardsigmoid.default: torch.nn.functional.hardsigmoid,
         exir_ops.edge.aten.hardswish.default: torch.nn.functional.hardswish,
         exir_ops.edge.aten.sinh.default: torch.sinh,

backends/arm/operator_support/tosa_supported_operators.py

@@ -252,6 +252,7 @@ def is_node_supported(
             exir_ops.edge.aten._adaptive_avg_pool2d.default,
             exir_ops.edge.aten.sign.default,
             exir_ops.edge.aten.asin.default,
+            exir_ops.edge.aten.atanh.default,
         ]
 
         return supported

backends/arm/quantizer/quantization_annotator.py

@@ -218,6 +218,7 @@ def _match_pattern(
     torch.ops.aten.acosh.default,
     torch.ops.aten.sign.default,
     torch.ops.aten.asin.default,
+    torch.ops.aten.atanh.default,
 ]
 
 _one_to_one_shared_input_qspec = [

backends/arm/test/ops/test_atanh.py

@@ -0,0 +1,85 @@
+# 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 (
+    EthosU55PipelineBI,
+    EthosU85PipelineBI,
+    TosaPipelineBI,
+    TosaPipelineMI,
+)
+
+aten_op = "torch.ops.aten.atanh.default"
+exir_op = "executorch_exir_dialects_edge__ops_aten__atanh_default"
+
+
+input_t1 = Tuple[torch.Tensor]
+
+
+test_data_suite = {
+    "zeros": torch.zeros(1, 10, 10, 10),
+    "zeros_alt_shape": torch.zeros(1, 10, 3, 5),
+    "ones": torch.ones(10, 10, 10),
+    "rand": torch.rand(10, 10) - 0.5,
+    "rand_alt_shape": torch.rand(1, 10, 3, 5) - 0.5,
+    "ramp": torch.arange(-1, 1, 0.2),
+    "near_bounds": torch.tensor([-0.999999, -0.999, -0.9, 0.9, 0.999, 0.999999]),
+    "on_bounds": torch.tensor([-1.0, 1.0]),
+}
+
+
+class Atanh(torch.nn.Module):
+    def forward(self, x: torch.Tensor):
+        return torch.atanh(x)
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_atanh_tosa_MI(test_data: Tuple):
+    pipeline = TosaPipelineMI[input_t1](
+        Atanh(),
+        (test_data,),
+        aten_op=aten_op,
+        exir_op=exir_op,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_atanh_tosa_BI(test_data: Tuple):
+    pipeline = TosaPipelineBI[input_t1](
+        Atanh(),
+        (test_data,),
+        aten_op=aten_op,
+        exir_op=exir_op,
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone300
+@common.parametrize("test_data", test_data_suite)
+def test_atanh_u55_BI(test_data: Tuple):
+    pipeline = EthosU55PipelineBI[input_t1](
+        Atanh(),
+        (test_data,),
+        aten_ops=aten_op,
+        exir_ops=exir_op,
+    )
+    pipeline.run()
+
+
+@common.XfailIfNoCorstone320
+@common.parametrize("test_data", test_data_suite)
+def test_atanh_u85_BI(test_data: Tuple):
+    pipeline = EthosU85PipelineBI[input_t1](
+        Atanh(),
+        (test_data,),
+        aten_ops=aten_op,
+        exir_ops=exir_op,
+    )
+    pipeline.run()