Arm backend: Add Amax-support for Ethos-U55 (#17372)

Additionally start deprecation of the unstable softmax decomp for u55 to
be replaced by the now fully supported stable one.

cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

Signed-off-by: Adrian Lundell <adrian.lundell@arm.com>

Author: AdrianLundell

backends/arm/common/pipeline_config.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-2026 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.
@@ -8,6 +8,8 @@
 from enum import auto, Enum
 from typing import Any
 
+from executorch.exir._warnings import deprecated
+
 
 class SoftmaxDecompositionConfig(Enum):
     MASKED = auto()
@@ -24,7 +26,16 @@ class ArmPassPipelineConfig:
     softmax: SoftmaxDecompositionConfig = SoftmaxDecompositionConfig.MASKED
     fuse_duplicate_users: FuseDuplicateUsersConfig = FuseDuplicateUsersConfig.ENABLED
 
+    @deprecated(
+        "The stable softmax decomposition is now supported by all arm targets and will be made default in a future release. Overwrite the default config using `compile_spec.set_pass_pipeline_config(ArmPassPipelineConfig())` to use the stable algorithm and avoid this error."
+    )
     def disable_masked_softmax(self) -> None:
+        """
+            .. warning::
+
+        The stable softmax decomposition is now supported by all arm targets and will be made default in a future release. Overwrite the default config using `compile_spec.set_pass_pipeline_config(ArmPassPipelineConfig())` to use the stable algorithm and avoid this error."
+        """
+
         self.softmax = SoftmaxDecompositionConfig.UNSTABLE
 
     def disable_fuse_duplicate_users(self) -> None:

backends/arm/operator_support/ethos_u55_support.py

@@ -86,7 +86,7 @@ def __init__(self, reporter: WhyNoPartitionReporter):
         exir_ops.edge.aten.permute_copy.default,
     ]
 
-    target_ops_i8 = tuple(TableOps.included_ops())
+    target_ops_i8_i16 = (*TableOps.included_ops(), exir_ops.edge.aten.amax.default)
 
     def is_node_supported(  # noqa: C901
         self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
@@ -117,7 +117,7 @@ def is_node_supported(  # noqa: C901
                 )
                 return False
 
-        if node.target in self.target_ops_i8:
+        if node.target in self.target_ops_i8_i16:
             if dtype not in (torch.int8, torch.int16):
                 self.reporter.report_reject(
                     node, f"Unsupported dtype {dtype} (Supports i8, i16)."
@@ -187,7 +187,6 @@ class EthosU55NotSupported(OperatorSupportBase):
         exir_ops.edge.aten.logical_or.default,
         exir_ops.edge.aten.logical_xor.default,
         exir_ops.edge.aten.logical_not.default,
-        exir_ops.edge.aten.amax.default,  # REDUCE_MAX
         exir_ops.edge.aten.amin.default,  # REDUCE_MIN
         exir_ops.edge.aten.conv3d.default,  # CONV3D
         exir_ops.edge.aten.conv3d.padding,  # CONV3D (deprecated alias)

backends/arm/test/ops/test_amax.py

@@ -10,6 +10,7 @@
 import torch
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
     EthosU85PipelineINT,
     OpNotSupportedPipeline,
     TosaPipelineFP,
@@ -136,8 +137,19 @@ def test_amax_tosa_INT(test_data: Amax.input_t):
     pipeline.run()
 
 
+@common.parametrize("test_data", Amax.test_data)
+def test_amax_u55_INT(test_data: Amax.input_t):
+    data, dim, keep_dims = test_data()
+    pipeline = EthosU55PipelineINT[Amax.input_t](
+        Amax(dim, keep_dims),
+        data,
+        Amax.aten_op,
+    )
+    pipeline.run()
+
+
 def test_amax_u55_INT_not_delegated():
-    data, dim, keep_dims = Amax.test_data["rank_4_all_dim"]()
+    data, dim, keep_dims = ((torch.ones([2, 2], dtype=torch.int32),), 1, False)
     pipeline = OpNotSupportedPipeline[Amax.input_t](
         Amax(dim, keep_dims),
         data,

backends/arm/test/ops/test_softmax.py

@@ -1,13 +1,14 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
-# Copyright 2024-2025 Arm Limited and/or its affiliates.
+# Copyright 2024-2026 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.common.pipeline_config import ArmPassPipelineConfig
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
@@ -70,7 +71,28 @@ def test_softmax_u55_INT(test_data):
         data,
         [],
     )
+    pipeline.add_stage_after(
+        "quantize", pipeline.tester.check_not, [aten_op, "torch.ops.aten.amax.default"]
+    )
+    pipeline.change_args("run_method_and_compare_outputs", qtol=1)
+    pipeline.run()
+
+
+@common.parametrize("test_data", Softmax.test_data)
+@common.XfailIfNoCorstone300
+def test_softmax_u55_INT_stable(test_data):
+    data, dim = test_data()
+    pipeline = EthosU55PipelineINT[input_t1](
+        Softmax(dim),
+        data,
+        [],
+    )
+    # Override ArmPassPipelineConfig to disable the DecomposeSoftmaxUnstablePass
+    pipeline.tester.compile_spec.set_pass_pipeline_config(ArmPassPipelineConfig())
     pipeline.add_stage_after("quantize", pipeline.tester.check_not, [aten_op])
+    pipeline.add_stage_after(
+        "quantize", pipeline.tester.check, ["torch.ops.aten.amax.default"]
+    )
     pipeline.change_args("run_method_and_compare_outputs", qtol=1)
     pipeline.run()