Arm backend: Add decomposition pass and test for asin (#12241)

Add decomposition pass and test for asin

Signed-off-by: Emma Kujala <emma.kujala@arm.com>

Author: emmakujala

backends/arm/_passes/__init__.py

@@ -24,6 +24,7 @@
 from .convert_to_clamp import ConvertToClampPass  # noqa
 from .decompose_acosh_pass import DecomposeAcoshPass  # noqa
 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_avg_pool2d import DecomposeAvgPool2d  # noqa
 from .decompose_batch_norm_no_stats import DecomposeBatchNormNoStatsPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -29,6 +29,7 @@
     ConvertToClampPass,
     DecomposeAcoshPass,
     DecomposeAdaptiveAvgPool2dPass,
+    DecomposeAsinPass,
     DecomposeAtanPass,
     DecomposeAvgPool2d,
     DecomposeBatchNormNoStatsPass,
@@ -158,6 +159,7 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
     def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModule:
         self.add_pass(DecomposeRoundPass())
         self.add_pass(DecomposeAcoshPass())
+        self.add_pass(DecomposeAsinPass())
         self.add_pass(DecomposeSqrtPass())
         self.add_pass(DecomposeAtanPass())
         self.add_pass(ConvertIntPowToMuls())

backends/arm/_passes/decompose_asin_pass.py

@@ -0,0 +1,201 @@
+# 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 logging
+from math import pi
+
+import torch
+
+from executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+
+# For MI case
+edge_asin_op = (exir_ops.edge.aten.asin.default,)
+
+
+def get_asin_decomposition(op) -> tuple:
+    if op in edge_asin_op:
+        return (
+            exir_ops.edge.aten.mul.Tensor,
+            exir_ops.edge.aten.add.Tensor,
+            exir_ops.edge.aten.mul.Scalar,
+            exir_ops.edge.aten.sqrt.default,
+            exir_ops.edge.aten.abs.default,
+            exir_ops.edge.aten.sub.Scalar,
+            exir_ops.edge.aten.div.Tensor,
+            exir_ops.edge.aten.gt.Scalar,
+            exir_ops.edge.aten.lt.Scalar,
+            exir_ops.edge.aten.sub.Tensor,
+            exir_ops.edge.aten.full_like.default,
+            exir_ops.edge.aten.where.self,
+            exir_ops.edge.aten.neg.default,
+        )
+
+    raise RuntimeError(f"Can't get asin decomposition for op {op}")
+
+
+class DecomposeAsinPass(ArmPass):
+    """
+    This pass decomposes asin into a rational approximation for small values
+    and a transformed rational approximation for large values.
+    Example:
+        y = asin(x)
+    Becomes:
+        if abs(x) < 0.5:
+            y = x + P(x^2) / Q(x^2)
+        else:
+            y = π/2 - 2 * (s + s^3 * Q(z) / P(z))
+    where P and Q are polynomials defined in the function.
+    """
+
+    def _build_polynomial(
+        self, coefficients: list[float], variable: torch.Tensor, meta: dict[str, str]
+    ) -> torch.Tensor:
+        """
+        Helper function to build polynomial from coefficients and variable.
+        """
+        full_like_op, add_op, mul_op_scalar, mul_op = (
+            exir_ops.edge.aten.full_like.default,
+            exir_ops.edge.aten.add.Tensor,
+            exir_ops.edge.aten.mul.Scalar,
+            exir_ops.edge.aten.mul.Tensor,
+        )
+        result = super().call_operator(
+            full_like_op, (variable, coefficients[0]), {}, meta, True
+        )
+        for coeff in coefficients[1:]:
+            result = super().call_operator(
+                add_op,
+                (
+                    result,
+                    super().call_operator(
+                        mul_op_scalar, (variable, coeff), {}, meta, True
+                    ),
+                ),
+                {},
+                meta,
+            )
+            variable = super().call_operator(
+                mul_op, (variable, variable), {}, meta, True
+            )
+        return result
+
+    def call_operator(self, op, args, kwargs, meta):
+        logging.info(
+            f"Approximating asin. This may introduce small numerical errors. For details, see {__file__}."
+        )
+        if op not in edge_asin_op:
+            return super().call_operator(op, args, kwargs, meta)
+
+        x = args[0]
+        half = 0.5
+        one = 1.0
+        neg_half = -0.5
+        two = 2.0
+        pi_over_2 = pi / 2.0
+        zero = 0.0
+        neg_one = -1.0
+
+        (
+            mul_op,
+            add_op,
+            mul_op_scalar,
+            sqrt_op,
+            abs_op,
+            sub_op_scalar,
+            div_op,
+            gt_op,
+            lt_op,
+            sub_op,
+            full_like_op,
+            where_op,
+            neg_op,
+        ) = get_asin_decomposition(op)
+
+        # Coefficients for the rational approximation, calculated with the Minimax (Remez) method
+        p_coefficients = [
+            1.6666667163e-01,
+            -3.2556581497e-01,
+            2.0121252537e-01,
+            -4.0055535734e-02,
+            7.9153501429e-04,
+        ]
+
+        q_coefficients = [1.0, -2.4033949375e00, 2.0209457874e00, -6.8828397989e-01]
+
+        x_abs = super().call_operator(abs_op, (x,), {}, meta, True)
+
+        # Step 1: compute asin_small - rational approximation for [0,0.5]
+
+        y = super().call_operator(mul_op, (x_abs, x_abs), {}, meta, True)
+        x3 = super().call_operator(mul_op, (x_abs, y), {}, meta, True)
+
+        P = self._build_polynomial(p_coefficients, x_abs, meta)
+        Q = self._build_polynomial(q_coefficients, x_abs, meta)
+        numer = super().call_operator(mul_op, (x3, P), {}, meta, True)
+        r_small = super().call_operator(div_op, (numer, Q), {}, meta, True)
+        asin_small = super().call_operator(add_op, (x_abs, r_small), {}, meta, True)
+
+        # Step 2: Compute the transformed approximation for large values
+        # Calculate z = -0.5 * (|x| - 1)
+        tmp_ones = super().call_operator(full_like_op, (x_abs, one), {}, meta, True)
+        tmp = super().call_operator(sub_op, (x_abs, tmp_ones), {}, meta, True)
+        z = super().call_operator(mul_op_scalar, (tmp, neg_half), {}, meta, True)
+
+        # Calculate s-terms
+        s = super().call_operator(sqrt_op, (z,), {}, meta, True)
+        s2 = super().call_operator(mul_op, (s, s), {}, meta, True)
+        s3 = super().call_operator(mul_op, (s2, s), {}, meta, True)
+
+        Pz = self._build_polynomial(p_coefficients, z, meta)
+        Qz = self._build_polynomial(q_coefficients, z, meta)
+
+        numer = super().call_operator(mul_op, (s3, Pz), {}, meta, True)
+        # Calculate r_large = P(z) / Q(z)
+        r_large = super().call_operator(div_op, (numer, Qz), {}, meta, True)
+
+        # Calculate asin_large = pi/2 - 2 * (s + s^3 * Q(z) / P(z))
+        t1 = super().call_operator(add_op, (s, r_large), {}, meta, True)
+        t2 = super().call_operator(mul_op_scalar, (t1, two), {}, meta, True)
+        diff = super().call_operator(sub_op_scalar, (t2, pi_over_2), {}, meta, True)
+        tmp_neg_ones = super().call_operator(
+            full_like_op, (diff, neg_one), {}, meta, True
+        )
+        asin_large = super().call_operator(mul_op, (diff, tmp_neg_ones), {}, meta, True)
+
+        # Combine branches
+        is_large = super().call_operator(gt_op, (x_abs, half), {}, meta, True)
+        asin_unsigned = super().call_operator(
+            where_op,
+            (
+                is_large,
+                asin_large,
+                asin_small,
+            ),
+            {},
+            meta,
+            True,
+        )
+
+        # Handle x < 0
+        is_neg = super().call_operator(lt_op, (x, zero), {}, meta, True)
+        # Compute -asin_unsigned
+        negated_asin = super().call_operator(neg_op, (asin_unsigned,), {}, meta, True)
+        # Combine branches for signed asin
+        asin_signed = super().call_operator(
+            where_op,
+            (
+                is_neg,
+                negated_asin,
+                asin_unsigned,
+            ),
+            {},
+            meta,
+            True,
+        )
+
+        return asin_signed

backends/arm/_passes/insert_table_ops.py

@@ -56,6 +56,7 @@ class TableOps:
         exir_ops.edge.aten.hardswish.default: torch.nn.functional.hardswish,
         exir_ops.edge.aten.sinh.default: torch.sinh,
         exir_ops.edge.aten.acosh.default: torch.acosh,
+        exir_ops.edge.aten.asin.default: torch.asin,
     }
 
     # Targets that must be treated explicitly

backends/arm/operator_support/tosa_supported_operators.py

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

backends/arm/quantizer/quantization_annotator.py

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

backends/arm/test/ops/test_asin.py

@@ -0,0 +1,80 @@
+# 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,
+)
+
+input_t = Tuple[torch.Tensor]  # Input x
+aten_op = "torch.ops.aten.asin.default"
+
+test_data_suite = {
+    "zeros": lambda: torch.zeros(1, 5, 3, 2),  # valid: asin(0) = 0
+    "ones": lambda: torch.ones(10, 5, 15),  # edge case: asin(1) = pi/2
+    "neg_ones": lambda: -torch.ones(10, 5, 15),  # edge case: asin(-1) = -pi/2
+    "rand": lambda: (torch.rand(10, 10, 5) * 2) - 1,  # uniform random in [-1, 1]
+    "ramp": lambda: torch.linspace(-1.0, 1.0, steps=160),  # full domain coverage
+    "near_bounds": lambda: torch.tensor(
+        [-0.999, -0.9, -0.5, 0.0, 0.5, 0.9, 0.999]
+    ),  # precision edge values
+    "pos_rand": lambda: torch.rand(7, 10, 2),  # positive random values in [0, 1]
+}
+
+
+class Asin(torch.nn.Module):
+    def forward(self, x):
+        return torch.asin(x)
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_asin_tosa_MI(test_data: Tuple):
+    pipeline = TosaPipelineMI[input_t](
+        Asin(),
+        (test_data(),),
+        aten_op,
+        exir_op=[],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_asin_tosa_BI(test_data: Tuple):
+    pipeline = TosaPipelineBI[input_t](
+        Asin(),
+        (test_data(),),
+        aten_op=[],
+        exir_op=[],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.XfailIfNoCorstone300
+def test_asin_u55_BI(test_data: Tuple):
+    pipeline = EthosU55PipelineBI[input_t](
+        Asin(),
+        (test_data(),),
+        aten_ops=[],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.XfailIfNoCorstone320
+def test_asin_u85_BI(test_data: Tuple):
+    pipeline = EthosU85PipelineBI[input_t](
+        Asin(),
+        (test_data(),),
+        aten_ops=[],
+    )
+    pipeline.run()