Arm backend: Fix int64 argmax delegation rejection (#20765)

TOSA 1.0 ARGMAX returns int32 indices, while aten.argmax returns int64.
If the graph consumes the raw aten.argmax result on CPU fallback,
delegating ARGMAX changes the dtype and can fail with an Int-vs-Long
mismatch.

Keep delegation enabled only when every user immediately casts the
argmax result to int32:

  aten.argmax
    -> _to_dim_order_copy(dtype=int32)
    -> users

Otherwise reject the node during support checking so the
PyTorch-compatible int64 result is preserved.

Fixes the flow test:
  test_argmax_dim[arm_tosa_fp]

Signed-off-by: Zingo Andersen <Zingo.Andersen@arm.com>

Author: zingo

backends/arm/operator_support/tosa_supported_operators.py

@@ -9,7 +9,6 @@
 
 """
 
-
 import operator
 import typing
 from typing import final, Optional, Sequence, Type
@@ -343,6 +342,7 @@ def _negative_checks(
     additional_checks: Optional[Sequence[OperatorSupportBase]],
 ) -> list[OperatorSupportBase]:
     checks: list[OperatorSupportBase] = [RankCheck(reporter, MAX_RANK)]
+    checks.append(CheckKnownUnsupportedTOSASemantics(reporter))
 
     if not tosa_spec.support_extension("int64"):
         checks.append(CheckInt64InputsAndOutputs(exported_program, reporter, tosa_spec))
@@ -372,6 +372,52 @@ def _negative_checks(
     return checks
 
 
+class CheckKnownUnsupportedTOSASemantics(OperatorSupportBase):
+    """Reject ops whose TOSA lowering is known to differ from PyTorch."""
+
+    def __init__(self, reporter: WhyNoPartitionReporter):
+        self.reporter = reporter
+
+    def _argmax_all_users_cast_to_int32(self, node: fx.Node) -> bool:
+        # TOSA ARGMAX produces int32 indices. This is only a faithful lowering
+        # when every user immediately narrows aten.argmax's int64 result:
+        #
+        #   aten.argmax -> _to_dim_order_copy(dtype=int32) -> users
+        cast_ops = (
+            torch.ops.dim_order_ops._to_dim_order_copy.default,
+            exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+        )
+        # A live argmax should normally have users. Treat a no-user node as not
+        # proving the int64 result is intentionally narrowed to int32.
+        if not node.users:
+            return False
+        for user in node.users:
+            if user.target not in cast_ops:
+                return False
+            if user.kwargs.get("dtype") != torch.int32:
+                return False
+        return True
+
+    def _check_argmax(self, node: fx.Node) -> bool:
+        if self._argmax_all_users_cast_to_int32(node):
+            return True
+        self.reporter.report_reject(
+            node, "TOSA ARGMAX produces int32 but aten.argmax returns int64."
+        )
+        return False
+
+    def is_node_supported(
+        self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node
+    ) -> bool:
+        if node.target in (
+            torch.ops.aten.argmax.default,
+            exir_ops.edge.aten.argmax.default,
+        ):
+            return self._check_argmax(node)
+
+        return True
+
+
 def tosa_support_factory(
     tosa_spec: TosaSpecification,
     exported_program: ExportedProgram,

backends/arm/test/misc/test_tosa_operator_support.py

@@ -0,0 +1,89 @@
+# Copyright 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.
+
+import torch
+from executorch.backends.arm.operator_support.tosa_supported_operators import (
+    CheckKnownUnsupportedTOSASemantics,
+)
+from executorch.exir.backend.utils import WhyNoPartitionReporter
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+
+def _fake_tensor(shape: tuple[int, ...], dtype: torch.dtype = torch.float32):
+    with FakeTensorMode() as mode:
+        return mode.from_tensor(torch.empty(shape, dtype=dtype))
+
+
+def _placeholder(graph: torch.fx.Graph, name: str, shape, dtype=torch.float32):
+    node = graph.placeholder(name)
+    node.meta["val"] = _fake_tensor(shape, dtype)
+    return node
+
+
+def _checker() -> CheckKnownUnsupportedTOSASemantics:
+    return CheckKnownUnsupportedTOSASemantics(WhyNoPartitionReporter())
+
+
+def test_rejects_argmax_without_int32_cast_user() -> None:
+    graph = torch.fx.Graph()
+    x = _placeholder(graph, "x", (3, 4))
+    node = graph.call_function(exir_ops.edge.aten.argmax.default, (x, 1, False))
+    node.meta["val"] = _fake_tensor((3,), torch.int64)
+
+    assert not _checker().is_node_supported({}, node)
+
+
+def test_accepts_argmax_with_int32_cast_user() -> None:
+    graph = torch.fx.Graph()
+    x = _placeholder(graph, "x", (3, 4))
+    node = graph.call_function(exir_ops.edge.aten.argmax.default, (x, 1, False))
+    node.meta["val"] = _fake_tensor((3,), torch.int64)
+    cast = graph.call_function(
+        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+        (node,),
+        {"dtype": torch.int32},
+    )
+    cast.meta["val"] = _fake_tensor((3,), torch.int32)
+
+    assert _checker().is_node_supported({}, node)
+
+
+def test_accepts_argmax_with_all_users_casting_to_int32() -> None:
+    graph = torch.fx.Graph()
+    x = _placeholder(graph, "x", (3, 4))
+    node = graph.call_function(exir_ops.edge.aten.argmax.default, (x, 1, False))
+    node.meta["val"] = _fake_tensor((3,), torch.int64)
+    cast = graph.call_function(
+        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+        (node,),
+        {"dtype": torch.int32},
+    )
+    cast.meta["val"] = _fake_tensor((3,), torch.int32)
+    cast_2 = graph.call_function(
+        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+        (node,),
+        {"dtype": torch.int32},
+    )
+    cast_2.meta["val"] = _fake_tensor((3,), torch.int32)
+
+    assert _checker().is_node_supported({}, node)
+
+
+def test_rejects_argmax_with_mixed_int32_cast_and_raw_user() -> None:
+    graph = torch.fx.Graph()
+    x = _placeholder(graph, "x", (3, 4))
+    node = graph.call_function(exir_ops.edge.aten.argmax.default, (x, 1, False))
+    node.meta["val"] = _fake_tensor((3,), torch.int64)
+    cast = graph.call_function(
+        exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+        (node,),
+        {"dtype": torch.int32},
+    )
+    cast.meta["val"] = _fake_tensor((3,), torch.int32)
+    raw_user = graph.call_function(torch.ops.aten.clone.default, (node,))
+    raw_user.meta["val"] = _fake_tensor((3,), torch.int64)
+
+    assert not _checker().is_node_supported({}, node)