Merge branch 'main' into issue/9131

Author: zhenyan-zhang-meta

.ci/scripts/unittest-macos.sh

@@ -35,7 +35,7 @@ if [[ "$BUILD_TOOL" == "cmake" ]]; then
     .ci/scripts/unittest-macos-cmake.sh
 elif [[ "$BUILD_TOOL" == "buck2" ]]; then
     .ci/scripts/unittest-buck2.sh
-    .ci/scripts/unittest-macos-buck2.sh
+    # .ci/scripts/unittest-macos-buck2.sh
 else
     echo "Unknown build tool $BUILD_TOOL"
     exit 1

backends/arm/_passes/arm_pass_manager.py

@@ -44,8 +44,9 @@
 from executorch.backends.arm._passes.decompose_select import (  # type: ignore[import-not-found]
     DecomposeSelectPass,
 )
-from executorch.backends.arm._passes.decompose_softmaxes_pass import (
-    DecomposeSoftmaxesPass,
+from executorch.backends.arm._passes.decompose_softmax_pass import DecomposeSoftmaxPass
+from executorch.backends.arm._passes.decompose_softmax_unstable_pass import (
+    DecomposeSoftmaxUnstablePass,
 )
 from executorch.backends.arm._passes.decompose_var_pass import DecomposeVarPass
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
@@ -81,7 +82,7 @@
 from executorch.backends.arm._passes.unsqueeze_scalar_placeholders_pass import (
     UnsqueezeScalarPlaceholdersPass,
 )
-from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.arm.tosa_specification import Tosa_0_80, TosaSpecification
 from executorch.backends.transforms.fuse_view_copy import FuseViewCopyTransform
 
 from executorch.backends.transforms.replace_scalar_with_tensor import (
@@ -155,7 +156,7 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(DecomposeMeanDimPass())
         self.add_pass(ConvertMeanDimToAveragePoolPass())
         self.add_pass(DecomposeDivPass())
-        self.add_pass(DecomposeSoftmaxesPass())
+        self.add_pass(DecomposeSoftmaxPass())
         self.add_pass(ConvertFullLikeToFullPass())
         self.add_pass(ConvertToClampPass())
         self.add_pass(ConvertMinMaxPass())
@@ -204,6 +205,12 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         self.add_pass(DecomposeVarPass())
         self.add_pass(DecomposeMeanDimPass())
         self.add_pass(DecomposeDivPass())
-        self.add_pass(DecomposeSoftmaxesPass())
+
+        if isinstance(self.tosa_spec, Tosa_0_80) and self.tosa_spec.is_U55_subset:
+            # Numerically stable softmax uses amax which is not supported on Ethos-U55
+            self.add_pass(DecomposeSoftmaxUnstablePass())
+        else:
+            self.add_pass(DecomposeSoftmaxPass())
+
         self.add_pass(ConvertMinMaxPass())
         return self._transform(graph_module)

backends/arm/_passes/decompose_softmax_pass.py

@@ -0,0 +1,77 @@
+# 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.
+
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+# For BI case
+torch_softmax = (torch.ops.aten.softmax.int, torch.ops.aten.log_softmax.int)
+# For MI case
+edge_softmax = (
+    exir_ops.edge.aten._softmax.default,
+    exir_ops.edge.aten._log_softmax.default,
+)
+log_softmax = (torch.ops.aten.log_softmax.int, exir_ops.edge.aten._log_softmax.default)
+
+
+def _get_logsoftmax_ops(op) -> tuple:
+    """
+    Returns the (log_op, sub_op, amax_op, expo_op, sum_op, reciprocal_op), where the ops depends on if
+    the softmax op is an aten or edge op.
+    """
+    if op in edge_softmax:
+        return (
+            exir_ops.edge.aten.log.default,
+            exir_ops.edge.aten.sub.Tensor,
+            exir_ops.edge.aten.amax.default,
+            exir_ops.edge.aten.exp.default,
+            exir_ops.edge.aten.sum.dim_IntList,
+            exir_ops.edge.aten.reciprocal.default,
+            exir_ops.edge.aten.mul.Tensor,
+        )
+    if op in torch_softmax:
+        return (
+            torch.ops.aten.log.default,
+            torch.ops.aten.sub.Tensor,
+            torch.ops.aten.amax.default,
+            torch.ops.aten.exp.default,
+            torch.ops.aten.sum.dim_IntList,
+            torch.ops.aten.reciprocal.default,
+            torch.ops.aten.mul.Tensor,
+        )
+    raise RuntimeError(f"Can't get logsoftmax decomposition ops for op {op}")
+
+
+class DecomposeSoftmaxPass(ExportPass):
+    """
+    This pass decomposes log_softmax or softmax into more primitive ops.
+    Example:
+        %op1 = amax(x)
+        %op2 = sub(x, %op1)
+        %op3 = exp(%op2)
+        %op4 = sum(%op3, dim)
+        %op5 = reciprocal(%op4)
+        %op6 = mul(%op3, %op5)
+        (in logsoftmax case: %op7 = log(%op6))
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in torch_softmax + edge_softmax:
+            return super().call_operator(op, args, kwargs, meta)
+        log_op, sub_op, max_op, exp_op, sum_op, reciprocal_op, mul_op = (
+            _get_logsoftmax_ops(op)
+        )
+        _input = args[0]
+        dim = [args[1]]
+        op1 = super().call_operator(max_op, (_input, dim, True), {}, meta)
+        op2 = super().call_operator(sub_op, (_input, op1), {}, meta)
+        op3 = super().call_operator(exp_op, (op2,), {}, meta)
+        op4 = super().call_operator(sum_op, (op3, dim, True), {}, meta)
+        op5 = super().call_operator(reciprocal_op, (op4,), {}, meta)
+        op6 = super().call_operator(mul_op, (op3, op5), {}, meta)
+        if op in log_softmax:
+            op6 = super().call_operator(log_op, (op6,), {}, meta)
+        return op6

backends/arm/_passes/decompose_softmaxes_pass.py renamed to backends/arm/_passes/decompose_softmax_unstable_pass.py

@@ -1,5 +1,4 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
-# All rights reserved.
+# Copyright 2024-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.
@@ -46,7 +45,7 @@ def get_logsoftmax_ops(op) -> tuple:
     raise RuntimeError(f"Can't get softmax decomposition ops for op {op}")
 
 
-class DecomposeSoftmaxesPass(ExportPass):
+class DecomposeSoftmaxUnstablePass(ExportPass):
     """
     This pass decomposes log softmax or softmax into more primitive ops.
 

backends/arm/operators/op_amax.py

@@ -5,6 +5,7 @@
 from typing import List
 
 import serializer.tosa_serializer as ts
+from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm.operators.node_visitor import (
     NodeVisitor,
     register_node_visitor,
@@ -31,6 +32,12 @@ def define_node(
 
         input = inputs[0]
         dim = inputs[1].number
+
+        if dim < 0:
+            tensor = get_first_fake_tensor(node)
+            rank = len(tensor.size())
+            dim = rank + dim
+
         keep_dims = inputs[2].number
         if not keep_dims:
             raise RuntimeError(

backends/arm/operators/op_amin.py

@@ -5,6 +5,7 @@
 from typing import List
 
 import serializer.tosa_serializer as ts
+from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm.operators.node_visitor import (
     NodeVisitor,
     register_node_visitor,
@@ -31,6 +32,12 @@ def define_node(
 
         input = inputs[0]
         dim = inputs[1].number
+
+        if dim < 0:
+            tensor = get_first_fake_tensor(node)
+            rank = len(tensor.size())
+            dim = rank + dim
+
         keep_dims = inputs[2].number
         if not keep_dims:
             raise RuntimeError(

backends/arm/test/common.py

@@ -1,5 +1,4 @@
 # Copyright 2024-2025 Arm Limited and/or its affiliates.
-# All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -159,13 +158,38 @@ def get_u85_compile_spec_unbuilt(
     not corstone300_installed() or not arm_executor_runner_exists("corstone-300"),
     reason="Did not find Corstone-300 FVP or executor_runner on path",
 )
-"""Skips a test if Corsone300 FVP is not installed, or if the executor runner is not built"""
+"""
+TO BE DEPRECATED - Use XfailIfNoCorstone300 instead
+Skips a test if Corsone300 FVP is not installed, or if the executor runner is not built
+"""
 
 SkipIfNoCorstone320 = pytest.mark.skipif(
     not corstone320_installed() or not arm_executor_runner_exists("corstone-320"),
     reason="Did not find Corstone-320 FVP or executor_runner on path",
 )
-"""Skips a test if Corsone320 FVP is not installed, or if the executor runner is not built."""
+"""
+TO BE DEPRECATED - Use XfailIfNoCorstone320 instead
+Skips a test if Corsone320 FVP is not installed, or if the executor runner is not built
+"""
+
+
+XfailIfNoCorstone300 = pytest.mark.xfail(
+    condition=not (
+        corstone300_installed() and arm_executor_runner_exists("corstone-300")
+    ),
+    raises=FileNotFoundError,
+    reason="Did not find Corstone-300 FVP or executor_runner on path",
+)
+"""Xfails a test if Corsone300 FVP is not installed, or if the executor runner is not built"""
+
+XfailIfNoCorstone320 = pytest.mark.xfail(
+    condition=not (
+        corstone320_installed() and arm_executor_runner_exists("corstone-320")
+    ),
+    raises=FileNotFoundError,
+    reason="Did not find Corstone-320 FVP or executor_runner on path",
+)
+"""Xfails a test if Corsone320 FVP is not installed, or if the executor runner is not built"""
 
 
 def parametrize(

backends/arm/test/ops/test_add.py

@@ -1,6 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
-# Copyright 2024-2025 Arm Limited and/or its affiliates.
 # All rights reserved.
+# Copyright 2024-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.
@@ -106,79 +106,47 @@ def test_add_i32_tosa_BI(test_data: input_t1):
 
 
 @common.parametrize("test_data", Add.test_data)
+@common.XfailIfNoCorstone300
 def test_add_u55_BI(test_data: input_t1):
-    pipeline = EthosU55PipelineBI[input_t1](
-        Add(), test_data, aten_op, exir_op, run_on_fvp=False
-    )
-    pipeline.run()
-
-
-@common.parametrize("test_data", Add.test_data)
-def test_add_u85_BI(test_data: input_t1):
-    pipeline = EthosU85PipelineBI[input_t1](
-        Add(), test_data, aten_op, exir_op, run_on_fvp=False
-    )
-    pipeline.run()
-
-
-@common.parametrize("test_data", Add.test_data)
-@common.SkipIfNoCorstone300
-def test_add_u55_BI_on_fvp(test_data: input_t1):
     pipeline = EthosU55PipelineBI[input_t1](
         Add(), test_data, aten_op, exir_op, run_on_fvp=True
     )
     pipeline.run()
 
 
 @common.parametrize("test_data", Add.test_data)
-@common.SkipIfNoCorstone320
-def test_add_u85_BI_on_fvp(test_data: input_t1):
+@common.XfailIfNoCorstone320
+def test_add_u85_BI(test_data: input_t1):
     pipeline = EthosU85PipelineBI[input_t1](
         Add(), test_data, aten_op, exir_op, run_on_fvp=True
     )
     pipeline.run()
 
 
 @common.parametrize("test_data", Add2.test_data)
-def test_add2_tosa_MI(test_data: input_t2):
+def test_add_2_tosa_MI(test_data: input_t2):
     pipeline = TosaPipelineMI[input_t2](Add2(), test_data, aten_op, exir_op)
     pipeline.run()
 
 
 @common.parametrize("test_data", Add2.test_data)
-def test_add2_tosa_BI(test_data: input_t2):
+def test_add_2_tosa_BI(test_data: input_t2):
     pipeline = TosaPipelineBI[input_t2](Add2(), test_data, aten_op, exir_op)
     pipeline.run()
 
 
 @common.parametrize("test_data", Add2.test_data)
-def test_add2_u55_BI(test_data: input_t2):
-    pipeline = EthosU55PipelineBI[input_t2](
-        Add2(), test_data, aten_op, exir_op, run_on_fvp=False
-    )
-    pipeline.run()
-
-
-@common.parametrize("test_data", Add2.test_data)
-@common.SkipIfNoCorstone300
-def test_add2_u55_BI_on_fvp(test_data: input_t2):
+@common.XfailIfNoCorstone300
+def test_add_2_u55_BI(test_data: input_t2):
     pipeline = EthosU55PipelineBI[input_t2](
         Add2(), test_data, aten_op, exir_op, run_on_fvp=True
     )
     pipeline.run()
 
 
 @common.parametrize("test_data", Add2.test_data)
-def test_add2_u85_BI(test_data: input_t2):
-    pipeline = EthosU85PipelineBI[input_t2](
-        Add2(), test_data, aten_op, exir_op, run_on_fvp=False
-    )
-    pipeline.run()
-
-
-@common.parametrize("test_data", Add2.test_data)
-@common.SkipIfNoCorstone320
-def test_add2_u85_BI_on_fvp(test_data: input_t2):
+@common.XfailIfNoCorstone320
+def test_add_2_u85_BI(test_data: input_t2):
     pipeline = EthosU85PipelineBI[input_t2](
         Add2(), test_data, aten_op, exir_op, run_on_fvp=True
     )