Arm backend: Delegate int64 only to the operators EXT-INT64 covers (2/3)

backends/arm/_passes/match_arg_dtype_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.
 
+from functools import reduce
 from typing import Set, Type
 
 import torch
@@ -33,6 +34,10 @@ class MatchArgDtypePass(ArmOpTargetedPass):
         exir_ops.edge.aten.add.Tensor,
         exir_ops.edge.aten.mul.Tensor,
         exir_ops.edge.aten.where.self,
+        exir_ops.edge.aten.bitwise_and.Tensor,
+        exir_ops.edge.aten.bitwise_or.Tensor,
+        exir_ops.edge.aten.bitwise_xor.Tensor,
+        exir_ops.edge.aten.cat.default,
     }
 
     def call(self, graph_module: torch.fx.GraphModule):
@@ -42,20 +47,27 @@ def call(self, graph_module: torch.fx.GraphModule):
         for node in graph.nodes:
             if node.op != "call_function" or node.target not in self.target_ops:
                 continue
-            args_to_cast = (
-                node.args[1:]
-                if node.target == exir_ops.edge.aten.where.self
-                else node.args
-            )
-            input_ = args_to_cast[0]
-            other_ = args_to_cast[1]
-            target_dtype = torch.result_type(input_.meta["val"], other_.meta["val"])
+            if node.target == exir_ops.edge.aten.cat.default:
+                args_to_cast = node.args[0]
+                target_dtype = reduce(
+                    torch.promote_types,
+                    (arg.meta["val"].dtype for arg in args_to_cast),
+                )
+            else:
+                args_to_cast = (
+                    node.args[1:]
+                    if node.target == exir_ops.edge.aten.where.self
+                    else node.args
+                )
+                input_ = args_to_cast[0]
+                other_ = args_to_cast[1]
+                target_dtype = torch.result_type(input_.meta["val"], other_.meta["val"])
 
             for arg in args_to_cast:
                 arg_dtype = arg.meta["val"].dtype
 
                 if arg_dtype != target_dtype:
-                    if arg_dtype == torch.bool:
+                    if arg_dtype == torch.bool and target_dtype != torch.int64:
                         # Bool is an edge case which cannot necessarily be directly
                         # converted to the target data type.
                         with graph.inserting_after(arg):

backends/arm/operator_support/tosa_profile_supported_op_lists.py

@@ -362,11 +362,42 @@
     operator.floordiv,
 }
 
+# Operators that delegate with an int64 operand under PRO-INT with EXT-INT64.
+# Each lowers to TOSA operators that EXT-INT64 defines for i64 under PRO-INT.
+# The list is not exhaustive. An absent operator narrows to int32 the way it
+# does without the extension.
+TOSA_PRO_INT_EXT_INT64_SupportList: Final[Set] = {
+    exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+    exir_ops.edge.aten.abs.default,
+    exir_ops.edge.aten.add.Tensor,
+    exir_ops.edge.aten.bitwise_and.Tensor,
+    exir_ops.edge.aten.bitwise_not.default,
+    exir_ops.edge.aten.bitwise_or.Tensor,
+    exir_ops.edge.aten.bitwise_xor.Tensor,
+    exir_ops.edge.aten.cat.default,
+    exir_ops.edge.aten.constant_pad_nd.default,
+    exir_ops.edge.aten.expand_copy.default,
+    exir_ops.edge.aten.flip.default,
+    exir_ops.edge.aten.maximum.default,
+    exir_ops.edge.aten.minimum.default,
+    exir_ops.edge.aten.mul.Tensor,
+    exir_ops.edge.aten.neg.default,
+    exir_ops.edge.aten.pad.default,
+    exir_ops.edge.aten.permute_copy.default,
+    exir_ops.edge.aten.repeat.default,
+    exir_ops.edge.aten.select_copy.int,
+    exir_ops.edge.aten.squeeze_copy.dims,
+    exir_ops.edge.aten.sub.Tensor,
+    exir_ops.edge.aten.unsqueeze_copy.default,
+    exir_ops.edge.aten.view_copy.default,
+}
+
 
 __all__ = [
     "TOSA_EXT_CONTROL_FLOW_SupportList",
     "TOSA_EXT_MXFP_SupportList",
     "TOSA_EXT_SHAPE_SupportList",
+    "TOSA_PRO_INT_EXT_INT64_SupportList",
     "TOSA_PRO_INT_SupportList",
     "TOSA_PRO_MIXED_INT_SupportList",
     "TOSA_PRO_FP_SupportList",

backends/arm/operator_support/tosa_supported_operators.py

@@ -70,6 +70,7 @@
     TOSA_EXT_MXFP_SupportList,
     TOSA_EXT_SHAPE_SupportList,
     TOSA_PRO_FP_SupportList,
+    TOSA_PRO_INT_EXT_INT64_SupportList,
     TOSA_PRO_INT_SupportList,
     TOSA_PRO_MIXED_INT_SupportList,
 )
@@ -643,13 +644,12 @@ def _negative_checks(
     if tosa_spec.support_integer() and tosa_spec.support_float():
         checks.append(CheckMixedProfileOperatorSupport(reporter))
 
-    if not tosa_spec.support_extension("int64"):
-        checks.append(
-            _AllowDynamicW8A8QParamsOr(
-                CheckInt64InputsAndOutputs(exported_program, reporter, tosa_spec),
-                tosa_spec,
-            )
+    checks.append(
+        _AllowDynamicW8A8QParamsOr(
+            CheckInt64InputsAndOutputs(exported_program, reporter, tosa_spec),
+            tosa_spec,
         )
+    )
 
     checks.append(CheckScalarReductionInputs(reporter))
 
@@ -1206,6 +1206,16 @@ def is_node_supported(
         return True
 
 
+def _int64_operands_agree(node: fx.Node) -> bool:
+    """Whether every operand beside an int64 operand casts to int64."""
+    dtypes = {
+        arg.meta["val"].dtype
+        for arg in node.all_input_nodes
+        if isinstance(arg.meta.get("val"), torch.Tensor)
+    }
+    return torch.int64 not in dtypes or dtypes <= {torch.int64, torch.int32, torch.bool}
+
+
 class CheckInt64InputsAndOutputs(OperatorSupportBase):
     """Reject general int64 tensors while allowing safe exceptions.
 
@@ -1431,6 +1441,13 @@ def is_node_supported(
         """Return True when int64 use is absent or safe per exceptions."""
         if is_submodule_node(node):
             return True
+        if (
+            self.tosa_spec.support_integer()
+            and self.tosa_spec.support_extension("int64")
+            and node.target in TOSA_PRO_INT_EXT_INT64_SupportList
+            and _int64_operands_agree(node)
+        ):
+            return True
         vals = node.meta["val"]
         tensor_list = vals if isinstance(vals, (list, tuple)) else [vals]
 

backends/arm/test/misc/test_int64.py

@@ -110,3 +110,65 @@ def test_int64_tosa_INT(test_data: Tuple):
         .to_executorch()
         .run_method_and_compare_outputs(inputs)
     )
+
+
+class Int64Op(torch.nn.Module):
+    def __init__(self, op):
+        super().__init__()
+        self.op = op
+
+    def forward(self, *args):
+        return self.op(*args)
+
+
+_WIDE = torch.tensor([2**40, -(2**38), 2**31, -1], dtype=torch.int64)
+_WIDE_2 = torch.tensor([2**39, 2**37, 1, -(2**40)], dtype=torch.int64)
+_NARROW = torch.tensor([1, 2, 3, 4], dtype=torch.int32)
+_BYTES = torch.tensor([1, 2, 3, 4], dtype=torch.int8)
+_COND = torch.tensor([True, False, True, False])
+_FLAGS = torch.tensor([False, True, True, False])
+
+
+class Int64BufferWhere(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("wide", _WIDE)
+
+    def forward(self, cond, x):
+        return torch.where(cond, x, self.wide)
+
+
+# A listed operator delegates when every operand casts to int64; an unlisted
+# operator or an operand narrower than int32 falls back instead.
+int64_extension_partition_suite = {
+    "add": (lambda x, y: x + y, (_WIDE, _WIDE_2), 1),
+    "pad": (lambda x: torch.nn.functional.pad(x, (1, 1)), (_WIDE,), 1),
+    "cumsum": (lambda x: torch.cumsum(x, 0), (_WIDE,), 0),
+    "cat_narrow": (lambda x, y: torch.cat([x, y]), (_WIDE, _NARROW), 1),
+    "or_narrow": (torch.bitwise_or, (_WIDE, _NARROW), 1),
+    "cat_byte": (lambda x, y: torch.cat([x, y]), (_WIDE, _BYTES), 0),
+    "eq_bool": (lambda x, y: x == y, (_WIDE, _FLAGS), 0),
+    "cat_bool": (lambda x, y: torch.cat([x, y]), (_FLAGS, _WIDE), 1),
+    "where": (torch.where, (_COND, _WIDE, _WIDE_2), 0),
+    "where_buffer": (Int64BufferWhere(), (_COND, _WIDE_2), 0),
+    "strided_slice": (lambda x: x[::2], (torch.arange(5, dtype=torch.int64),), 0),
+    "gt": (lambda x, y: x > y, (_WIDE, _WIDE_2), 0),
+    "gt_scalar": (lambda x: x > 7, (_WIDE,), 0),
+    "full_like": (lambda x: torch.full_like(x, 5) + x, (_WIDE,), 1),
+    "zeros_like": (lambda x: torch.zeros_like(x) + x, (_WIDE,), 1),
+}
+
+
+@common.parametrize("test_data", int64_extension_partition_suite)
+def test_int64_extension_partition_tosa_INT(test_data: Tuple):
+    op, inputs, delegates = test_data
+    (
+        ArmTester(
+            op if isinstance(op, torch.nn.Module) else Int64Op(op),
+            inputs,
+            common.get_tosa_compile_spec("TOSA-1.1+INT+int64"),
+        )
+        .export()
+        .to_edge_transform_and_lower()
+        .check_count({"torch.ops.higher_order.executorch_call_delegate": delegates})
+    )

backends/arm/test/misc/test_tosa_operator_support.py

@@ -13,6 +13,7 @@
 )
 from executorch.backends.arm.operator_support.topk_support import TopKSupported
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
+    _int64_operands_agree,
     CheckFPComparisonInputs,
     CheckKnownUnsupportedTOSASemantics,
 )
@@ -323,3 +324,13 @@ def test_topk_keyword_arguments(dim) -> None:
     )
     assert config is not None and config.k == 4
     assert reason is None
+
+
+def test_int64_operand_check_skips_symbolic_inputs() -> None:
+    graph = torch.fx.Graph()
+    x = _placeholder(graph, "x", (8,), torch.int32)
+    size = graph.placeholder("size")
+    size.meta["val"] = ShapeEnv().create_unbacked_symint()
+    node = graph.call_function(exir_ops.edge.aten.view_copy.default, (x, [size, -1]))
+
+    assert _int64_operands_agree(node)