Skip to content

Arm backend: Add acosh decomposition pass and test #12105

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions backends/arm/_passes/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from .convert_split_to_slice import ConvertSplitToSlicePass # noqa
from .convert_squeezes_to_view import ConvertSqueezesToViewPass # noqa
from .convert_to_clamp import ConvertToClampPass # noqa
from .decompose_acosh_pass import DecomposeAcoshPass # noqa
from .decompose_atan_pass import DecomposeAtanPass # noqa
from .decompose_avg_pool2d import DecomposeAvgPool2d # noqa
from .decompose_batch_norm_no_stats import DecomposeBatchNormNoStatsPass # noqa
Expand Down
2 changes: 2 additions & 0 deletions backends/arm/_passes/arm_pass_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
ConvertSplitToSlicePass,
ConvertSqueezesToViewPass,
ConvertToClampPass,
DecomposeAcoshPass,
DecomposeAtanPass,
DecomposeAvgPool2d,
DecomposeBatchNormNoStatsPass,
Expand Down Expand Up @@ -151,6 +152,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(DecomposeSqrtPass())
self.add_pass(DecomposeAtanPass())
self.add_pass(ConvertIntPowToMuls())
Expand Down
52 changes: 52 additions & 0 deletions backends/arm/_passes/decompose_acosh_pass.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
# 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 executorch.backends.arm._passes import ArmPass
from executorch.exir.dialects._ops import ops as exir_ops

# For MI case
edge_acosh_op = exir_ops.edge.aten.acosh.default


class DecomposeAcoshPass(ArmPass):
"""
Decomposes acosh to supported TOSA-operations.
This decomposition is based on the mathematical identity:
acosh(x) = log(x + sqrt((x-1)(x+1))
"""

def call_operator(self, op, args, kwargs, meta, updated=False):

if op is not edge_acosh_op:
return super().call_operator(op, args, kwargs, meta, updated)

log_op, sqrt_op, mul_op, sub_op, add_op, add_op_scalar = (
exir_ops.edge.aten.log.default,
exir_ops.edge.aten.sqrt.default,
exir_ops.edge.aten.mul.Tensor,
exir_ops.edge.aten.sub.Scalar,
exir_ops.edge.aten.add.Tensor,
exir_ops.edge.aten.add.Scalar,
)

x = args[0]

# (x-1)(x+1)
sub = super().call_operator(sub_op, (x, 1.0), {}, meta, True)
add = super().call_operator(add_op_scalar, (x, 1.0), {}, meta, True)
mul = super().call_operator(mul_op, (sub, add), {}, meta, True)

# sqrt((x-1)(x+1))
sqrt = super().call_operator(sqrt_op, (mul,), {}, meta, True)

# x + sqrt((x-1)(x+1))
add = super().call_operator(add_op, (x, sqrt), {}, meta, True)

# out = ln(x + sqrt((x-1)(x+1))
out = super().call_operator(log_op, (add,), {}, meta, True)

return out
1 change: 1 addition & 0 deletions backends/arm/_passes/insert_table_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ class TableOps:
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,
exir_ops.edge.aten.acosh.default: torch.acosh,
}

# Targets that must be treated explicitly
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,7 @@ def is_node_supported(
exir_ops.edge.aten.alias_copy.default,
exir_ops.edge.aten.sinh.default,
exir_ops.edge.aten.atan.default,
exir_ops.edge.aten.acosh.default,
]

return supported
Expand Down
1 change: 1 addition & 0 deletions backends/arm/quantizer/quantization_annotator.py
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,7 @@ def _match_pattern(
torch.ops.aten.gelu.default,
torch.ops.aten.sinh.default,
torch.ops.aten.atan.default,
torch.ops.aten.acosh.default,
]

_one_to_one_shared_input_qspec = [
Expand Down
114 changes: 114 additions & 0 deletions backends/arm/test/ops/test_acosh.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
# 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 pytest

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.acosh.default"


test_data_suite = {
# Valid input cases
"ones": lambda: torch.ones(1, 7, 10, 12),
"just_above_one": lambda: torch.tensor([1.0001, 1.01, 1.1, 2.0]),
"rand_valid": lambda: torch.rand(10, 10) * 10 + 1, # [1, 11)
"ramp_valid": lambda: torch.linspace(1.0, 20.0, steps=160),
"large": lambda: torch.tensor([10.0, 100.0, 1000.0, 1e6]),
"mixed_valid": lambda: torch.tensor([1.0, 2.0, 10.0, 100.0]),
}

test_data_suite_xfails = {
# Invalid input cases (should return nan or error)
"zeros": lambda: torch.zeros(1, 5, 3, 2),
"neg_ones": lambda: -torch.ones(10, 10, 10),
"rand_invalid": lambda: torch.rand(10, 10), # [0, 1)
"ramp_invalid": lambda: torch.linspace(-10.0, 0.99, steps=160),
"near_zero": lambda: torch.tensor([-1e-6, 0.0, 1e-6]),
"large_negative": lambda: torch.tensor([-100.0, -10.0, 0.0]),
}


class Acosh(torch.nn.Module):

def forward(self, x: torch.Tensor):
return torch.acosh(x)


@common.parametrize("test_data", test_data_suite)
def test_acosh_tosa_MI(test_data: Tuple):
pipeline = TosaPipelineMI[input_t](
Acosh(),
(test_data(),),
aten_op,
exir_op=[],
)
pipeline.run()


@common.parametrize("test_data", test_data_suite)
def test_acosh_tosa_BI(test_data: Tuple):
pipeline = TosaPipelineBI[input_t](
Acosh(),
(test_data(),),
aten_op=[],
)
pipeline.run()


@common.parametrize("test_data", test_data_suite)
@common.XfailIfNoCorstone300
def test_acosh_u55_BI(test_data: Tuple):
pipeline = EthosU55PipelineBI[input_t](
Acosh(),
(test_data(),),
aten_ops=[],
)
pipeline.run()


@common.parametrize("test_data", test_data_suite_xfails)
@pytest.mark.xfail(reason="Invalid inputs are currently not handled")
def test_acosh_u55_BI_xfail(test_data: Tuple):
pipeline = EthosU55PipelineBI[input_t](
Acosh(),
(test_data(),),
aten_ops=[],
run_on_fvp=False,
)
pipeline.run()


@common.parametrize("test_data", test_data_suite)
@common.XfailIfNoCorstone320
def test_acosh_u85_BI(test_data: Tuple):
pipeline = EthosU85PipelineBI[input_t](
Acosh(),
(test_data(),),
aten_ops=[],
)
pipeline.run()


@common.parametrize("test_data", test_data_suite_xfails)
@pytest.mark.xfail(reason="Invalid inputs are currently not handled")
def test_acosh_u85_BI_xfail(test_data: Tuple):
pipeline = EthosU55PipelineBI[input_t](
Acosh(),
(test_data(),),
aten_ops=[],
run_on_fvp=False,
)
pipeline.run()
Loading