diff --git a/backends/arm/operator_support/tosa_profile_supported_op_lists.py b/backends/arm/operator_support/tosa_profile_supported_op_lists.py index d8f7ab54c74..90fe8458970 100644 --- a/backends/arm/operator_support/tosa_profile_supported_op_lists.py +++ b/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", diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 97d8cc7e2aa..233a4d2f6a9 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/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] diff --git a/backends/arm/test/misc/test_int64.py b/backends/arm/test/misc/test_int64.py index 46a97fff1df..e0150df2d74 100644 --- a/backends/arm/test/misc/test_int64.py +++ b/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}) + ) diff --git a/backends/arm/test/misc/test_tosa_operator_support.py b/backends/arm/test/misc/test_tosa_operator_support.py index 635566dffdf..25c40d63905 100644 --- a/backends/arm/test/misc/test_tosa_operator_support.py +++ b/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)