Update

[ghstack-poisoned]

Author: swolchok

.lintrunner.toml

@@ -76,6 +76,7 @@ exclude_patterns = [
     'examples/demo-apps/apple_ios/**',
     'examples/demo-apps/react-native/rnllama/ios/**',
     'extension/apple/**',
+    'extension/llm/apple/**',
     # File contains @generated
     'extension/llm/custom_ops/spinquant/fast_hadamard_transform_special.h',
     'extension/llm/custom_ops/spinquant/test/fast_hadamard_transform_special_unstrided_cpu.h',

backends/arm/_passes/__init__.py

@@ -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

backends/arm/_passes/arm_pass_manager.py

@@ -25,6 +25,7 @@
     ConvertSplitToSlicePass,
     ConvertSqueezesToViewPass,
     ConvertToClampPass,
+    DecomposeAcoshPass,
     DecomposeAtanPass,
     DecomposeAvgPool2d,
     DecomposeBatchNormNoStatsPass,
@@ -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())

backends/arm/_passes/decompose_acosh_pass.py

@@ -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

backends/arm/_passes/decompose_sqrt_pass.py

@@ -4,6 +4,8 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-unsafe
+from typing import Tuple, Union
+
 import torch
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
@@ -15,7 +17,7 @@
 )
 
 
-def get_sqrt_decomposition(op) -> tuple:
+def get_sqrt_decomposition(op) -> Union[Tuple, torch._ops.OpOverload]:
     # TODO : "MLETORCH-863 : Replace current sqrt -> pow.Tensor_Scalar workaround with pow.Tensor_Tensor"
     if op in edge_sqrt_ops:
         return exir_ops.edge.aten.pow.Tensor_Scalar

backends/arm/_passes/insert_table_ops.py

@@ -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

backends/arm/_passes/match_arg_ranks_pass.py

@@ -51,6 +51,7 @@ def __init__(self, exported_program):
         exir_ops.edge.aten.gt.Tensor,
         exir_ops.edge.aten.ge.Tensor,
         exir_ops.edge.aten.lt.Tensor,
+        exir_ops.edge.aten.le.Tensor,
         exir_ops.edge.aten.pow.Tensor_Tensor,
         exir_ops.edge.aten.where.self,
     ]

backends/arm/_passes/replace_scalar_with_tensor_pass.py

@@ -6,7 +6,7 @@
 # pyre-unsafe
 
 
-from typing import Dict
+from typing import Dict, Union
 
 import torch
 from executorch.backends.transforms.replace_scalar_with_tensor import (
@@ -18,7 +18,10 @@
 
 
 # Operators that are included for both TOSA profiles
-_common_ops: Dict[EdgeOpOverload, EdgeOpOverload] = {
+_common_ops: Dict[
+    Union[EdgeOpOverload, torch._ops.OpOverload],
+    Union[EdgeOpOverload, torch._ops.OpOverload],
+] = {
     exir_ops.edge.aten.add.Scalar: exir_ops.edge.aten.add.Tensor,
     exir_ops.edge.aten.sub.Scalar: exir_ops.edge.aten.sub.Tensor,
     exir_ops.edge.aten.mul.Scalar: exir_ops.edge.aten.mul.Tensor,
@@ -29,6 +32,7 @@
     exir_ops.edge.aten.gt.Scalar: exir_ops.edge.aten.gt.Tensor,
     exir_ops.edge.aten.ge.Scalar: exir_ops.edge.aten.ge.Tensor,
     exir_ops.edge.aten.lt.Scalar: exir_ops.edge.aten.lt.Tensor,
+    exir_ops.edge.aten.le.Scalar: exir_ops.edge.aten.le.Tensor,
     exir_ops.edge.aten.ne.Scalar: exir_ops.edge.aten.ne.Tensor,
     torch.ops.aten.add.Scalar: torch.ops.aten.add.Tensor,
     torch.ops.aten.sub.Scalar: torch.ops.aten.sub.Tensor,
@@ -40,6 +44,7 @@
     torch.ops.aten.gt.Scalar: torch.ops.aten.gt.Tensor,
     torch.ops.aten.ge.Scalar: torch.ops.aten.ge.Tensor,
     torch.ops.aten.lt.Scalar: torch.ops.aten.lt.Tensor,
+    torch.ops.aten.le.Scalar: torch.ops.aten.le.Tensor,
     torch.ops.aten.ne.Scalar: torch.ops.aten.ne.Tensor,
 }
 

backends/arm/operator_support/ethos_u55_support.py

@@ -138,6 +138,7 @@ class EthosU55NotSupported(OperatorSupportBase):
         exir_ops.edge.aten.gt.Tensor,
         exir_ops.edge.aten.gt.Scalar,
         exir_ops.edge.aten.le.Tensor,
+        exir_ops.edge.aten.le.Scalar,
         exir_ops.edge.aten.lt.Tensor,
         exir_ops.edge.aten.lt.Scalar,
         exir_ops.edge.aten.ne.Tensor,
@@ -174,6 +175,69 @@ def is_node_supported(
 shape_t = list[int]
 
 
+class EthosU55ViewCheck(OperatorSupportBase):
+
+    def __init__(self, reporter: WhyNoPartitionReporter):
+        super().__init__()
+        self.reporter = reporter
+
+    def axes_product(self, nhwc_shape: shape_t) -> int:
+        product = 1
+        for axes in nhwc_shape:
+            product *= axes
+        return product
+
+    # TODO: Extend this check to comply with u55 restrictions
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+        """
+        Check whether a given view node is supported on U55.
+
+        Currently only checks dtypes and product of axes.
+
+        It is not the view operator itself that is not supported on U55. In order for the
+        view operator to be compatible with the channels-last format of TosaBackend,
+        transposes may need to be inserted before and after the view op. If that happens
+        and that transpose operator does not adhere to the limitations then it will
+        result in the following error:
+
+            CPU performance estimation for "Transpose" not implemented.
+            ...
+            CPU operations are not supported for GraphAPI input
+
+        Args:
+            node: The FX node representing the view_copy operator.
+
+        Returns:
+            False if the operator is not support and True if it is supported.
+        """
+        if not node.target == exir_ops.edge.aten.view_copy.default:
+            return True
+
+        shape = list(get_first_fake_tensor(node).shape)
+        dtype = _try_determine_dtype(node)
+        permutation = list(typing.cast(list[int], node.args[1]))
+
+        rank = len(shape)
+        if rank > 4:
+            if dtype == torch.int32:
+                self.reporter.report_reject(
+                    node, f"No support for {permutation=} in int32."
+                )
+                return False
+
+        if dtype in (torch.int8, torch.int16):
+            if self.axes_product(shape) > 65536:
+                self.reporter.report_reject(
+                    node,
+                    f"No support for {shape=}, {dtype=}. Product of axes must be <65536",
+                )
+                return False
+
+        return True
+
+
 class EthosU55TransposeCheck(OperatorSupportBase):
 
     def __init__(self, reporter: WhyNoPartitionReporter):

backends/arm/operator_support/tosa_supported_operators.py

@@ -23,6 +23,7 @@
     EthosU55DtypeSupport,
     EthosU55NotSupported,
     EthosU55TransposeCheck,
+    EthosU55ViewCheck,
 )
 from executorch.backends.arm.tosa_quant_utils import dq_ops, q_ops
 from executorch.backends.arm.tosa_specification import TosaSpecification
@@ -133,6 +134,7 @@ def tosa_support_factory(
         negative_checks.append(EthosU55NotSupported(reporter))
         negative_checks.append(EthosU55DtypeSupport(reporter))
         negative_checks.append(EthosU55TransposeCheck(reporter))
+        negative_checks.append(EthosU55ViewCheck(reporter))
 
     return chain(
         reporter.wrap_check(
@@ -187,6 +189,7 @@ def is_node_supported(
             exir_ops.edge.aten.gt.Tensor,
             exir_ops.edge.aten.gt.Scalar,
             exir_ops.edge.aten.le.Tensor,
+            exir_ops.edge.aten.le.Scalar,
             exir_ops.edge.aten.lt.Tensor,
             exir_ops.edge.aten.lt.Scalar,
             exir_ops.edge.aten.mul.Tensor,
@@ -245,6 +248,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