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

backends/arm/operator_support/tosa_profile_supported_op_lists.py

@@ -362,9 +362,39 @@
     operator.floordiv,
 }
 
+# Operators that delegate with an int64 operand when the int64 extension is
+# enabled. Each lowers to TOSA operators that EXT-INT64 defines for i64. An
+# absent operator narrows to int32 the way it does without the extension.
+TOSA_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_INT64_SupportList",
     "TOSA_EXT_MXFP_SupportList",
     "TOSA_EXT_SHAPE_SupportList",
     "TOSA_PRO_INT_SupportList",

backends/arm/operator_support/tosa_supported_operators.py

@@ -67,6 +67,7 @@
 )
 from executorch.backends.arm.operator_support.tosa_profile_supported_op_lists import (
     TOSA_EXT_CONTROL_FLOW_SupportList,
+    TOSA_EXT_INT64_SupportList,
     TOSA_EXT_MXFP_SupportList,
     TOSA_EXT_SHAPE_SupportList,
     TOSA_PRO_FP_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 int64 operands appear only beside other int64 operands."""
+    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}
+
+
 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_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,62 @@ 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)
+_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 on matched int64 operands delegates; an unlisted operator or
+# a pair mixing widths 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), 0),
+    "eq_bool": (lambda x, y: x == y, (_WIDE, _FLAGS), 0),
+    "cat_bool": (lambda x, y: torch.cat([x, y]), (_FLAGS, _WIDE), 0),
+    "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)