Arm backend: Decompose SDPA safe_softmax on U55

Make the U55 annotation and lowering pipelines handle SDPA
aten._safe_softmax consistently by decomposing it instead of leaving it
in the graph.

Previously, transform_for_annotation_pipeline still used
skip_safe_softmax for U55, which left aten._safe_softmax in annotated
SDPA graphs and prevented delegation.

Add a regression test that verifies U55 SDPA graphs no longer contain
aten._safe_softmax after the annotation pipeline runs.

Also warn when aten._safe_softmax is decomposed as regular softmax in
the annotation pipeline, since this is only semantics-preserving when
no row is fully masked at runtime.

Remove the unstable softmax decomposition path and its remaining
references now that the Arm backend uses the stable decomposition path.
Update the related pipeline-config and softmax tests accordingly.

Signed-off-by: per.held@arm.com
Change-Id: I7a5147d5492974ead52ea92326352f7f4407bd67

Author: perheld

backends/arm/_passes/__init__.py

@@ -85,7 +85,6 @@
 from .decompose_sinh_pass import DecomposeSinhPass  # noqa
 from .decompose_slice_scatter_pass import DecomposeSliceScatterPass  # noqa
 from .decompose_softmax_pass import DecomposeSoftmaxPass  # noqa
-from .decompose_softmax_unstable_pass import DecomposeSoftmaxUnstablePass  # noqa
 from .decompose_sqrt_pass import DecomposeSqrtPass  # noqa
 from .decompose_strided_slice_copy_pass import DecomposeStridedSliceCopyPass  # noqa
 from .decompose_sum_pass import DecomposeSumPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -86,7 +86,6 @@
     DecomposeSinhPass,
     DecomposeSliceScatterPass,
     DecomposeSoftmaxPass,
-    DecomposeSoftmaxUnstablePass,
     DecomposeSqrtPass,
     DecomposeStridedSliceCopyPass,
     DecomposeSumPass,
@@ -196,12 +195,8 @@ def configure_skip_passes(
 
         match config.softmax:
             case SoftmaxDecompositionConfig.MASKED:
-                skip_set.add(DecomposeSoftmaxUnstablePass)
-            case SoftmaxDecompositionConfig.UNSTABLE:
-                skip_set.add(DecomposeSoftmaxPass)
-                skip_set.add(DecomposeMaskedFillPass)
+                pass
             case SoftmaxDecompositionConfig.STABLE:
-                skip_set.add(DecomposeSoftmaxUnstablePass)
                 skip_set.add(DecomposeMaskedFillPass)
 
         if config.fuse_duplicate_users is FuseDuplicateUsersConfig.DISABLED:
@@ -461,9 +456,7 @@ def _tosa_pipeline(
                 ConvertMmToBmmPass(),
                 DecomposeGluPass(),
                 DecomposeDivPass(),
-                # _safe_softmax results in a ReduceMax
-                # which is not currently supported by TOSA in U55
-                DecomposeSoftmaxPass(skip_safe_softmax=self.tosa_spec.is_U55_subset),
+                DecomposeSoftmaxPass(),
                 ConvertMinMaxPass(),
                 DecomposeAnyPass(),
                 DecomposeAdaptiveAvgPool2dPass(),
@@ -593,9 +586,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 DecomposeSqrtPass(tfa_pass=True),
                 DecomposeAdaptiveAvgPool2dPass(tfa_pass=True),
                 DecomposeAvgPool2dPass(tfa_pass=True),
-                DecomposeSoftmaxUnstablePass(tfa_pass=True),
                 DecomposeSoftmaxPass(
-                    skip_safe_softmax=self.tosa_spec.is_U55_subset,
                     tfa_pass=True,
                 ),
                 ConvertMinMaxPass(tfa_pass=True),

backends/arm/_passes/decompose_softmax_pass.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 from typing import Set, Type
 
 import torch
@@ -25,6 +26,8 @@
 )
 log_softmax = (torch.ops.aten.log_softmax.int, exir_ops.edge.aten._log_softmax.default)
 
+logger = logging.getLogger(__name__)
+
 
 def _get_logsoftmax_ops(op) -> tuple:
     """Returns the (log_op, sub_op, amax_op, expo_op, sum_op, reciprocal_op),
@@ -78,6 +81,7 @@ class DecomposeSoftmaxPass(ArmPass):
     def __init__(self, skip_safe_softmax: bool = False, **kwargs):
         super().__init__(**kwargs)
         self._skip_safe_softmax = skip_safe_softmax
+        self._warned_safe_softmax = False
 
     def call_operator(self, op, args, kwargs, meta):
         if op not in torch_softmax + edge_softmax or not self.allowed_to_transform(
@@ -88,6 +92,18 @@ def call_operator(self, op, args, kwargs, meta):
         if self._skip_safe_softmax and op == torch.ops.aten._safe_softmax.default:
             return super().call_operator(op, args, kwargs, meta)
 
+        if (
+            self.is_tfa_pass
+            and op == torch.ops.aten._safe_softmax.default
+            and not self._warned_safe_softmax
+        ):
+            logger.warning(
+                "aten._safe_softmax is being decomposed as regular softmax in "
+                "the annotation pipeline; this is only semantics-preserving "
+                "when no row is fully masked at runtime."
+            )
+            self._warned_safe_softmax = True
+
         log_op, sub_op, max_op, exp_op, sum_op, reciprocal_op, mul_op = (
             _get_logsoftmax_ops(op)
         )

backends/arm/_passes/decompose_softmax_unstable_pass.py

@@ -16,7 +16,10 @@
 from dataclasses import dataclass, field
 from enum import Enum
 
-from executorch.backends.arm.common.pipeline_config import ArmPassPipelineConfig
+from executorch.backends.arm.common.pipeline_config import (
+    ArmPassPipelineConfig,
+    SoftmaxDecompositionConfig,
+)
 from executorch.backends.arm.tosa import TosaSpecification
 from executorch.exir._warnings import deprecated
 
@@ -250,7 +253,10 @@ def set_pass_pipeline_config(self, config: ArmPassPipelineConfig) -> None:
     def _create_default_pipeline_config(self) -> ArmPassPipelineConfig:
         config = ArmPassPipelineConfig()
         if self.tosa_spec.is_U55_subset:
-            config.disable_masked_softmax()
+            # Keep U55 on STABLE instead of the generic MASKED default:
+            # MASKED also enables masked_fill decomposition, which lowers to
+            # where/full_like and is not a good default fit for U55.
+            config.softmax = SoftmaxDecompositionConfig.STABLE
         return config
 
     def _get_intermediate_path(self) -> str | None:

backends/arm/common/arm_compile_spec.py

@@ -8,12 +8,9 @@
 from enum import auto, Enum
 from typing import Any
 
-from executorch.exir._warnings import deprecated
-
 
 class SoftmaxDecompositionConfig(Enum):
     MASKED = auto()  # Stable softmax + masked fill decomposition
-    UNSTABLE = auto()  # Unstable softmax, no masked fill decomposition
     STABLE = auto()  # Stable softmax, no masked fill decomposition
 
 
@@ -27,18 +24,6 @@ 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.STABLE
-
     def disable_fuse_duplicate_users(self) -> None:
         self.fuse_duplicate_users = FuseDuplicateUsersConfig.DISABLED
 

backends/arm/common/pipeline_config.py

@@ -1,12 +1,11 @@
-# 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.
 
 from executorch.backends.arm._passes import (
     DecomposeMaskedFillPass,
     DecomposeSoftmaxPass,
-    DecomposeSoftmaxUnstablePass,
     FuseDuplicateUsersPass,
 )
 from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager
@@ -25,7 +24,7 @@ def test_pipeline_config_override_outside_compile_spec_no_target():
     default_manager = ArmPassManager(compile_spec)
     default_skip_passes = default_manager._skip_pass_types
     assert FuseDuplicateUsersPass not in default_skip_passes
-    assert DecomposeSoftmaxUnstablePass in default_skip_passes
+    assert DecomposeSoftmaxPass not in default_skip_passes
 
     override_compile_spec = TosaCompileSpec(
         TosaSpecification.create_from_string("TOSA-1.00+INT")
@@ -37,10 +36,10 @@ def test_pipeline_config_override_outside_compile_spec_no_target():
     skip_passes = override_manager._skip_pass_types
 
     assert FuseDuplicateUsersPass in skip_passes
-    assert DecomposeSoftmaxUnstablePass in skip_passes
+    assert DecomposeSoftmaxPass not in skip_passes
 
 
-def test_softmax_config_masked():
+def test_softmax_config_masked_no_target():
     """Test MASKED config: stable softmax, masked fill decomposition enabled."""
     compile_spec = TosaCompileSpec(
         TosaSpecification.create_from_string("TOSA-1.00+INT")
@@ -50,31 +49,13 @@ def test_softmax_config_masked():
     manager = ArmPassManager(compile_spec)
     skip_passes = manager._skip_pass_types
 
-    # MASKED: skip unstable softmax, use stable softmax
-    assert DecomposeSoftmaxUnstablePass in skip_passes
+    # MASKED: use stable softmax
     assert DecomposeSoftmaxPass not in skip_passes
     # MASKED: masked fill decomposition is enabled (not skipped)
     assert DecomposeMaskedFillPass not in skip_passes
 
 
-def test_softmax_config_unstable():
-    """Test UNSTABLE config: unstable softmax, no masked fill decomposition."""
-    compile_spec = TosaCompileSpec(
-        TosaSpecification.create_from_string("TOSA-1.00+INT")
-    )
-    config = ArmPassPipelineConfig(softmax=SoftmaxDecompositionConfig.UNSTABLE)
-    compile_spec.set_pass_pipeline_config(config)
-    manager = ArmPassManager(compile_spec)
-    skip_passes = manager._skip_pass_types
-
-    # UNSTABLE: skip stable softmax, use unstable softmax
-    assert DecomposeSoftmaxPass in skip_passes
-    assert DecomposeSoftmaxUnstablePass not in skip_passes
-    # UNSTABLE: masked fill decomposition is disabled (skipped)
-    assert DecomposeMaskedFillPass in skip_passes
-
-
-def test_softmax_config_stable():
+def test_softmax_config_stable_no_target():
     """Test STABLE config: stable softmax, no masked fill decomposition."""
     compile_spec = TosaCompileSpec(
         TosaSpecification.create_from_string("TOSA-1.00+INT")
@@ -84,8 +65,7 @@ def test_softmax_config_stable():
     manager = ArmPassManager(compile_spec)
     skip_passes = manager._skip_pass_types
 
-    # STABLE: skip unstable softmax, use stable softmax
-    assert DecomposeSoftmaxUnstablePass in skip_passes
+    # STABLE: use stable softmax
     assert DecomposeSoftmaxPass not in skip_passes
     # STABLE: masked fill decomposition is disabled (skipped)
     assert DecomposeMaskedFillPass in skip_passes

backends/arm/test/misc/test_pass_pipeline_config.py

@@ -8,6 +8,8 @@
 
 import torch
 
+from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager
+from executorch.backends.arm.ethosu import EthosUCompileSpec
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
@@ -16,6 +18,7 @@
     TosaPipelineINT,
     VgfPipeline,
 )
+from torch.export import export
 
 
 class SDPA(torch.nn.Module):
@@ -106,13 +109,6 @@ def test_sdpa_vgf_quant(test_case: test_case_t):
 
 @common.parametrize("test_case", test_suite)
 def test_sdpa_u55_INT(test_case: test_case_t):
-    """Verify SDPA compiles on U55.
-
-    _safe_softmax from SDPA is skipped by DecomposeSoftmaxPass
-    (skip_safe_softmax=True for U55) and runs on CPU, avoiding REDUCE_MAX which
-    fails Vela compilation.
-
-    """
     model, test_input = test_case()
     pipeline = EthosU55PipelineINT[input_t](model, test_input, [], [])
     pipeline.pop_stage("check.quant_nodes")
@@ -122,14 +118,34 @@ def test_sdpa_u55_INT(test_case: test_case_t):
 
 
 @common.parametrize("test_case", test_suite)
-@common.XfailIfNoCorstone320
-def test_sdpa_u85_INT(test_case: test_case_t):
-    """Verify SDPA compiles on U85.
+def test_sdpa_u55_INT_annotation_pipeline_decomposes_safe_softmax(
+    test_case: test_case_t,
+):
+    """Verify the U55 annotation pipeline decomposes SDPA _safe_softmax.
 
-    _safe_softmax is decomposed with stable softmax (including amax/REDUCE_MAX)
-    which is supported on U85.
+    U55 now matches U85 and VGF: the annotation pipeline lowers
+    _safe_softmax to the stable softmax primitive sequence instead of leaving
+    it in the graph for partitioning.
 
     """
+    model, test_input = test_case()
+    exported_program = export(model, test_input)
+    graph_module = ArmPassManager(
+        EthosUCompileSpec("ethos-u55-128")
+    ).transform_for_annotation_pipeline(exported_program.graph_module)
+
+    softmax_targets = {
+        str(node.target)
+        for node in graph_module.graph.nodes
+        if node.op == "call_function" and "softmax" in str(node.target)
+    }
+
+    assert "aten._safe_softmax.default" not in softmax_targets
+
+
+@common.parametrize("test_case", test_suite)
+@common.XfailIfNoCorstone320
+def test_sdpa_u85_INT(test_case: test_case_t):
     model, test_input = test_case()
     pipeline = EthosU85PipelineINT[input_t](model, test_input, [], [])
     pipeline.pop_stage("check.quant_nodes")