From 587c055fb8390db3f7ab9a33e6d1bfc7d971aab7 Mon Sep 17 00:00:00 2001 From: Bryson Date: Tue, 27 Oct 2026 17:17:23 +0000 Subject: [PATCH 1/8] Arm backend: List the operators EXT-INT64 covers --- .../tosa_profile_supported_op_lists.py | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) 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..d2622b9988e 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,53 @@ operator.floordiv, } +# Operators that carry an int64 operand when the int64 extension is enabled. +# EXT-INT64 leaves out CLAMP (clamp, relu, hardtanh), MATMUL (mm, bmm), POW, +# CEIL, FLOOR, LOGICAL_* and ARGMAX, so operators lowering to one of those are +# absent, as are operators reaching one through a decomposition like cumsum. +# 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.eq.Scalar, + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.expand_copy.default, + exir_ops.edge.aten.flip.default, + exir_ops.edge.aten.full_like.default, + exir_ops.edge.aten.ge.Scalar, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.gt.Scalar, + exir_ops.edge.aten.gt.Tensor, + exir_ops.edge.aten.le.Scalar, + exir_ops.edge.aten.le.Tensor, + exir_ops.edge.aten.lt.Scalar, + exir_ops.edge.aten.lt.Tensor, + 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.permute_copy.default, + exir_ops.edge.aten.repeat.default, + exir_ops.edge.aten.select_copy.int, + exir_ops.edge.aten.slice_copy.Tensor, + 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, + exir_ops.edge.aten.where.self, + exir_ops.edge.aten.zeros_like.default, +} + __all__ = [ "TOSA_EXT_CONTROL_FLOW_SupportList", + "TOSA_EXT_INT64_SupportList", "TOSA_EXT_MXFP_SupportList", "TOSA_EXT_SHAPE_SupportList", "TOSA_PRO_INT_SupportList", From 95543953e85855ed5c70b27415111a841cd5c5ae Mon Sep 17 00:00:00 2001 From: Bryson Date: Tue, 27 Oct 2026 17:17:24 +0000 Subject: [PATCH 2/8] Arm backend: Delegate int64 only to the operators EXT-INT64 covers --- .../operator_support/slice_copy_support.py | 5 +++- .../tosa_supported_operators.py | 29 +++++++++++++++---- .../arm/operator_support/where_support.py | 5 +++- backends/arm/tosa/specification.py | 10 +++++++ 4 files changed, 41 insertions(+), 8 deletions(-) diff --git a/backends/arm/operator_support/slice_copy_support.py b/backends/arm/operator_support/slice_copy_support.py index 2b0343bafa5..852d0116273 100644 --- a/backends/arm/operator_support/slice_copy_support.py +++ b/backends/arm/operator_support/slice_copy_support.py @@ -22,6 +22,7 @@ SupportedTOSAOperatorCheck, ) from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.specification import keeps_int64 from executorch.exir.dialects._ops import ops as exir_ops @@ -81,7 +82,9 @@ def is_node_tosa_supported( values_dtype = node.args[0].meta["val"].dtype # type: ignore[union-attr] - SUPPORTED_INT_DTYPES = (torch.int8, torch.int16, torch.int32) + SUPPORTED_INT_DTYPES = (torch.int8, torch.int16, torch.int32) + ( + (torch.int64,) if keeps_int64(tosa_spec) else () + ) SUPPORTED_FLOAT_DTYPES = ( torch.float16, torch.float32, diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 97d8cc7e2aa..12177f33430 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, @@ -75,6 +76,7 @@ ) from executorch.backends.arm.tosa.specification import ( get_context_shape_env, + keeps_int64, TosaSpecification, TosaSpecMapping, ) @@ -643,13 +645,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 +1207,16 @@ def is_node_supported( return True +def _int64_operands_agree(node: fx.Node) -> bool: + """Whether no operand is int64 beside a narrower integer, bool aside.""" + dtypes = { + get_first_fake_tensor(arg).dtype + for arg in node.args + if isinstance(arg, fx.Node) and "val" in arg.meta + } - {torch.bool} + return torch.int64 not in dtypes or dtypes == {torch.int64} + + class CheckInt64InputsAndOutputs(OperatorSupportBase): """Reject general int64 tensors while allowing safe exceptions. @@ -1431,6 +1442,12 @@ def is_node_supported( """Return True when int64 use is absent or safe per exceptions.""" if is_submodule_node(node): return True + if ( + keeps_int64(self.tosa_spec) + 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/operator_support/where_support.py b/backends/arm/operator_support/where_support.py index ac6c936b1fb..3ee4d611d30 100644 --- a/backends/arm/operator_support/where_support.py +++ b/backends/arm/operator_support/where_support.py @@ -13,6 +13,7 @@ SupportedTOSAOperatorCheck, ) from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.specification import keeps_int64 from executorch.exir.dialects._ops import ops as exir_ops @@ -46,7 +47,9 @@ def is_node_tosa_supported( supported_float = [torch.bool, torch.float16, torch.float32] if tosa_spec.support_extension("bf16"): supported_float.append(torch.bfloat16) - supported_integer = (torch.bool, torch.int8, torch.int16, torch.int32) + supported_integer = [torch.bool, torch.int8, torch.int16, torch.int32] + if keeps_int64(tosa_spec): + supported_integer.append(torch.int64) if tosa_spec.support_float(): if x_dtype in supported_float and y_dtype in supported_float: diff --git a/backends/arm/tosa/specification.py b/backends/arm/tosa/specification.py index 24924ff2e94..744774dc428 100644 --- a/backends/arm/tosa/specification.py +++ b/backends/arm/tosa/specification.py @@ -496,6 +496,16 @@ def get_context_shape_env() -> ShapeEnv: raise RuntimeError("Function must be executed within a TosaLoweringContext") +def keeps_int64(tosa_spec: "TosaSpecification | None" = None) -> bool: + """Whether int64 is kept 64 bits wide; only INT has int64 lowerings.""" + if tosa_spec is None: + try: + tosa_spec = get_context_spec() + except RuntimeError: + return False + return tosa_spec.support_integer() and tosa_spec.support_extension("int64") + + def tosa_spec_in_set(spec: TosaSpecification, specs: Set[TosaSpecification]) -> bool: """Check if a specification matches any in a set, considering base specs. From e3420bfcece068ebe3d13663520d4fd00af872af Mon Sep 17 00:00:00 2001 From: Bryson Date: Thu, 29 Oct 2026 18:22:10 +0000 Subject: [PATCH 3/8] Arm backend: Test int64 partitioning under the int64 extension --- backends/arm/test/misc/test_int64.py | 40 ++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/backends/arm/test/misc/test_int64.py b/backends/arm/test/misc/test_int64.py index 46a97fff1df..1aaeab61795 100644 --- a/backends/arm/test/misc/test_int64.py +++ b/backends/arm/test/misc/test_int64.py @@ -110,3 +110,43 @@ 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]) + +# 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), + "where_int64": (torch.where, (_COND, _WIDE, _WIDE_2), 1), + "where_x_narrow": (torch.where, (_COND, _NARROW, _WIDE), 0), + "where_y_narrow": (torch.where, (_COND, _WIDE, _NARROW), 0), + "cumsum": (lambda x: torch.cumsum(x, 0), (_WIDE,), 0), +} + + +@common.parametrize("test_data", int64_extension_partition_suite) +def test_int64_extension_partition_tosa_INT(test_data: Tuple): + op, inputs, delegates = test_data + ( + ArmTester( + 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}) + ) From c00ea1fcfdc915fc0b9d81fd456688013513e5b7 Mon Sep 17 00:00:00 2001 From: Bryson Date: Thu, 29 Oct 2026 20:09:47 +0000 Subject: [PATCH 4/8] Arm backend: Check every int64 operand of a node, including list operands --- backends/arm/operator_support/tosa_supported_operators.py | 4 ++-- backends/arm/test/misc/test_int64.py | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 12177f33430..6b6f8d4a16d 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -1211,8 +1211,8 @@ def _int64_operands_agree(node: fx.Node) -> bool: """Whether no operand is int64 beside a narrower integer, bool aside.""" dtypes = { get_first_fake_tensor(arg).dtype - for arg in node.args - if isinstance(arg, fx.Node) and "val" in arg.meta + for arg in node.all_input_nodes + if "val" in arg.meta } - {torch.bool} return torch.int64 not in dtypes or dtypes == {torch.int64} diff --git a/backends/arm/test/misc/test_int64.py b/backends/arm/test/misc/test_int64.py index 1aaeab61795..d1bbe1f17f9 100644 --- a/backends/arm/test/misc/test_int64.py +++ b/backends/arm/test/misc/test_int64.py @@ -134,6 +134,9 @@ def forward(self, *args): "where_x_narrow": (torch.where, (_COND, _NARROW, _WIDE), 0), "where_y_narrow": (torch.where, (_COND, _WIDE, _NARROW), 0), "cumsum": (lambda x: torch.cumsum(x, 0), (_WIDE,), 0), + "cat_narrow": (lambda x, y: torch.cat([x, y]), (_WIDE, _NARROW), 0), + "where_x_bool": (torch.where, (_COND, _COND, _WIDE), 0), + "where_y_bool": (torch.where, (_COND, _WIDE, _COND), 0), } From bae8fa3be8a7f3b2fa264fa455e65db5c34607ac Mon Sep 17 00:00:00 2001 From: Bryson Date: Thu, 29 Oct 2026 20:21:28 +0000 Subject: [PATCH 5/8] Arm backend: Reject mixed int64 where values and skip symbolic operands --- .../arm/operator_support/tosa_supported_operators.py | 4 ++-- backends/arm/operator_support/where_support.py | 7 +++++++ backends/arm/test/misc/test_int64.py | 5 +++-- backends/arm/test/misc/test_tosa_operator_support.py | 11 +++++++++++ 4 files changed, 23 insertions(+), 4 deletions(-) diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 6b6f8d4a16d..184083d1235 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -1210,9 +1210,9 @@ def is_node_supported( def _int64_operands_agree(node: fx.Node) -> bool: """Whether no operand is int64 beside a narrower integer, bool aside.""" dtypes = { - get_first_fake_tensor(arg).dtype + arg.meta["val"].dtype for arg in node.all_input_nodes - if "val" in arg.meta + if isinstance(arg.meta.get("val"), torch.Tensor) } - {torch.bool} return torch.int64 not in dtypes or dtypes == {torch.int64} diff --git a/backends/arm/operator_support/where_support.py b/backends/arm/operator_support/where_support.py index 3ee4d611d30..809600e8835 100644 --- a/backends/arm/operator_support/where_support.py +++ b/backends/arm/operator_support/where_support.py @@ -44,6 +44,13 @@ def is_node_tosa_supported( return False x_dtype, y_dtype = x.meta["val"].dtype, y.meta["val"].dtype + if torch.int64 in (x_dtype, y_dtype) and x_dtype != y_dtype: + self.reporter.report_reject( + node, + f"Mixed {x_dtype} and {y_dtype} values in {node.target} are not supported", + ) + return False + supported_float = [torch.bool, torch.float16, torch.float32] if tosa_spec.support_extension("bf16"): supported_float.append(torch.bfloat16) diff --git a/backends/arm/test/misc/test_int64.py b/backends/arm/test/misc/test_int64.py index d1bbe1f17f9..24dfdd96151 100644 --- a/backends/arm/test/misc/test_int64.py +++ b/backends/arm/test/misc/test_int64.py @@ -125,6 +125,7 @@ def forward(self, *args): _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]) # A listed operator on matched int64 operands delegates; an unlisted operator or # a pair mixing widths falls back instead. @@ -135,8 +136,8 @@ def forward(self, *args): "where_y_narrow": (torch.where, (_COND, _WIDE, _NARROW), 0), "cumsum": (lambda x: torch.cumsum(x, 0), (_WIDE,), 0), "cat_narrow": (lambda x, y: torch.cat([x, y]), (_WIDE, _NARROW), 0), - "where_x_bool": (torch.where, (_COND, _COND, _WIDE), 0), - "where_y_bool": (torch.where, (_COND, _WIDE, _COND), 0), + "where_x_bool": (torch.where, (_COND, _FLAGS, _WIDE), 0), + "where_y_bool": (torch.where, (_COND, _WIDE, _FLAGS), 0), } 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) From 4c8ae8276f8e5fae8559b52087181889aa27a80d Mon Sep 17 00:00:00 2001 From: Bryson Date: Thu, 29 Oct 2026 20:55:17 +0000 Subject: [PATCH 6/8] Arm backend: Keep bool operands in the int64 check, list pad, inline the extension check --- backends/arm/operator_support/slice_copy_support.py | 7 +++---- .../tosa_profile_supported_op_lists.py | 2 ++ .../operator_support/tosa_supported_operators.py | 13 ++++++++----- backends/arm/operator_support/where_support.py | 3 +-- backends/arm/test/misc/test_int64.py | 3 +++ backends/arm/tosa/specification.py | 10 ---------- 6 files changed, 17 insertions(+), 21 deletions(-) diff --git a/backends/arm/operator_support/slice_copy_support.py b/backends/arm/operator_support/slice_copy_support.py index 852d0116273..77d21fb9344 100644 --- a/backends/arm/operator_support/slice_copy_support.py +++ b/backends/arm/operator_support/slice_copy_support.py @@ -22,7 +22,6 @@ SupportedTOSAOperatorCheck, ) from executorch.backends.arm.tosa import TosaSpecification -from executorch.backends.arm.tosa.specification import keeps_int64 from executorch.exir.dialects._ops import ops as exir_ops @@ -82,9 +81,9 @@ def is_node_tosa_supported( values_dtype = node.args[0].meta["val"].dtype # type: ignore[union-attr] - SUPPORTED_INT_DTYPES = (torch.int8, torch.int16, torch.int32) + ( - (torch.int64,) if keeps_int64(tosa_spec) else () - ) + SUPPORTED_INT_DTYPES = (torch.int8, torch.int16, torch.int32) + if tosa_spec.support_extension("int64"): + SUPPORTED_INT_DTYPES += (torch.int64,) SUPPORTED_FLOAT_DTYPES = ( torch.float16, torch.float32, 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 d2622b9988e..075d8449290 100644 --- a/backends/arm/operator_support/tosa_profile_supported_op_lists.py +++ b/backends/arm/operator_support/tosa_profile_supported_op_lists.py @@ -376,6 +376,7 @@ 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.eq.Scalar, exir_ops.edge.aten.eq.Tensor, exir_ops.edge.aten.expand_copy.default, @@ -393,6 +394,7 @@ 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, diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 184083d1235..585e0d863ea 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -76,7 +76,6 @@ ) from executorch.backends.arm.tosa.specification import ( get_context_shape_env, - keeps_int64, TosaSpecification, TosaSpecMapping, ) @@ -1208,12 +1207,15 @@ def is_node_supported( def _int64_operands_agree(node: fx.Node) -> bool: - """Whether no operand is int64 beside a narrower integer, bool aside.""" + """Whether int64 operands appear only beside other int64 operands.""" + operands = node.all_input_nodes + if node.target == exir_ops.edge.aten.where.self: + operands = [arg for arg in node.args[1:] if isinstance(arg, fx.Node)] dtypes = { arg.meta["val"].dtype - for arg in node.all_input_nodes + for arg in operands if isinstance(arg.meta.get("val"), torch.Tensor) - } - {torch.bool} + } return torch.int64 not in dtypes or dtypes == {torch.int64} @@ -1443,7 +1445,8 @@ def is_node_supported( if is_submodule_node(node): return True if ( - keeps_int64(self.tosa_spec) + 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) ): diff --git a/backends/arm/operator_support/where_support.py b/backends/arm/operator_support/where_support.py index 809600e8835..bbb3fb74c83 100644 --- a/backends/arm/operator_support/where_support.py +++ b/backends/arm/operator_support/where_support.py @@ -13,7 +13,6 @@ SupportedTOSAOperatorCheck, ) from executorch.backends.arm.tosa import TosaSpecification -from executorch.backends.arm.tosa.specification import keeps_int64 from executorch.exir.dialects._ops import ops as exir_ops @@ -55,7 +54,7 @@ def is_node_tosa_supported( if tosa_spec.support_extension("bf16"): supported_float.append(torch.bfloat16) supported_integer = [torch.bool, torch.int8, torch.int16, torch.int32] - if keeps_int64(tosa_spec): + if tosa_spec.support_extension("int64"): supported_integer.append(torch.int64) if tosa_spec.support_float(): diff --git a/backends/arm/test/misc/test_int64.py b/backends/arm/test/misc/test_int64.py index 24dfdd96151..b48f1830479 100644 --- a/backends/arm/test/misc/test_int64.py +++ b/backends/arm/test/misc/test_int64.py @@ -134,10 +134,13 @@ def forward(self, *args): "where_int64": (torch.where, (_COND, _WIDE, _WIDE_2), 1), "where_x_narrow": (torch.where, (_COND, _NARROW, _WIDE), 0), "where_y_narrow": (torch.where, (_COND, _WIDE, _NARROW), 0), + "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), "where_x_bool": (torch.where, (_COND, _FLAGS, _WIDE), 0), "where_y_bool": (torch.where, (_COND, _WIDE, _FLAGS), 0), + "eq_bool": (lambda x, y: x == y, (_WIDE, _FLAGS), 0), + "cat_bool": (lambda x, y: torch.cat([x, y]), (_FLAGS, _WIDE), 0), } diff --git a/backends/arm/tosa/specification.py b/backends/arm/tosa/specification.py index 744774dc428..24924ff2e94 100644 --- a/backends/arm/tosa/specification.py +++ b/backends/arm/tosa/specification.py @@ -496,16 +496,6 @@ def get_context_shape_env() -> ShapeEnv: raise RuntimeError("Function must be executed within a TosaLoweringContext") -def keeps_int64(tosa_spec: "TosaSpecification | None" = None) -> bool: - """Whether int64 is kept 64 bits wide; only INT has int64 lowerings.""" - if tosa_spec is None: - try: - tosa_spec = get_context_spec() - except RuntimeError: - return False - return tosa_spec.support_integer() and tosa_spec.support_extension("int64") - - def tosa_spec_in_set(spec: TosaSpecification, specs: Set[TosaSpecification]) -> bool: """Check if a specification matches any in a set, considering base specs. From 553a0b560eee993f8ddb24e7af4b71b9d2046621 Mon Sep 17 00:00:00 2001 From: Bryson Date: Thu, 29 Oct 2026 21:17:07 +0000 Subject: [PATCH 7/8] Arm backend: Keep the slice int64 dtype a conditional expression --- backends/arm/operator_support/slice_copy_support.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/backends/arm/operator_support/slice_copy_support.py b/backends/arm/operator_support/slice_copy_support.py index 77d21fb9344..90cd8059a97 100644 --- a/backends/arm/operator_support/slice_copy_support.py +++ b/backends/arm/operator_support/slice_copy_support.py @@ -81,9 +81,9 @@ def is_node_tosa_supported( values_dtype = node.args[0].meta["val"].dtype # type: ignore[union-attr] - SUPPORTED_INT_DTYPES = (torch.int8, torch.int16, torch.int32) - if tosa_spec.support_extension("int64"): - SUPPORTED_INT_DTYPES += (torch.int64,) + SUPPORTED_INT_DTYPES = (torch.int8, torch.int16, torch.int32) + ( + (torch.int64,) if tosa_spec.support_extension("int64") else () + ) SUPPORTED_FLOAT_DTYPES = ( torch.float16, torch.float32, From 6b44cef9da2d6fcf6bef9b549493a8fb3585b14c Mon Sep 17 00:00:00 2001 From: Bryson Date: Mon, 2 Nov 2026 18:20:49 +0000 Subject: [PATCH 8/8] Arm backend: Trim the int64 list to the operators that lower today --- .../operator_support/slice_copy_support.py | 4 +--- .../tosa_profile_supported_op_lists.py | 22 +++-------------- .../tosa_supported_operators.py | 5 +--- .../arm/operator_support/where_support.py | 11 +-------- backends/arm/test/misc/test_int64.py | 24 ++++++++++++++----- 5 files changed, 24 insertions(+), 42 deletions(-) diff --git a/backends/arm/operator_support/slice_copy_support.py b/backends/arm/operator_support/slice_copy_support.py index 90cd8059a97..2b0343bafa5 100644 --- a/backends/arm/operator_support/slice_copy_support.py +++ b/backends/arm/operator_support/slice_copy_support.py @@ -81,9 +81,7 @@ def is_node_tosa_supported( values_dtype = node.args[0].meta["val"].dtype # type: ignore[union-attr] - SUPPORTED_INT_DTYPES = (torch.int8, torch.int16, torch.int32) + ( - (torch.int64,) if tosa_spec.support_extension("int64") else () - ) + SUPPORTED_INT_DTYPES = (torch.int8, torch.int16, torch.int32) SUPPORTED_FLOAT_DTYPES = ( torch.float16, torch.float32, 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 075d8449290..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,11 +362,9 @@ operator.floordiv, } -# Operators that carry an int64 operand when the int64 extension is enabled. -# EXT-INT64 leaves out CLAMP (clamp, relu, hardtanh), MATMUL (mm, bmm), POW, -# CEIL, FLOOR, LOGICAL_* and ARGMAX, so operators lowering to one of those are -# absent, as are operators reaching one through a decomposition like cumsum. -# An absent operator narrows to int32 the way it does without the extension. +# 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, @@ -377,19 +375,8 @@ 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.eq.Scalar, - exir_ops.edge.aten.eq.Tensor, exir_ops.edge.aten.expand_copy.default, exir_ops.edge.aten.flip.default, - exir_ops.edge.aten.full_like.default, - exir_ops.edge.aten.ge.Scalar, - exir_ops.edge.aten.ge.Tensor, - exir_ops.edge.aten.gt.Scalar, - exir_ops.edge.aten.gt.Tensor, - exir_ops.edge.aten.le.Scalar, - exir_ops.edge.aten.le.Tensor, - exir_ops.edge.aten.lt.Scalar, - exir_ops.edge.aten.lt.Tensor, exir_ops.edge.aten.maximum.default, exir_ops.edge.aten.minimum.default, exir_ops.edge.aten.mul.Tensor, @@ -398,13 +385,10 @@ exir_ops.edge.aten.permute_copy.default, exir_ops.edge.aten.repeat.default, exir_ops.edge.aten.select_copy.int, - exir_ops.edge.aten.slice_copy.Tensor, 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, - exir_ops.edge.aten.where.self, - exir_ops.edge.aten.zeros_like.default, } diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 585e0d863ea..233a4d2f6a9 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -1208,12 +1208,9 @@ def is_node_supported( def _int64_operands_agree(node: fx.Node) -> bool: """Whether int64 operands appear only beside other int64 operands.""" - operands = node.all_input_nodes - if node.target == exir_ops.edge.aten.where.self: - operands = [arg for arg in node.args[1:] if isinstance(arg, fx.Node)] dtypes = { arg.meta["val"].dtype - for arg in operands + 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} diff --git a/backends/arm/operator_support/where_support.py b/backends/arm/operator_support/where_support.py index bbb3fb74c83..ac6c936b1fb 100644 --- a/backends/arm/operator_support/where_support.py +++ b/backends/arm/operator_support/where_support.py @@ -43,19 +43,10 @@ def is_node_tosa_supported( return False x_dtype, y_dtype = x.meta["val"].dtype, y.meta["val"].dtype - if torch.int64 in (x_dtype, y_dtype) and x_dtype != y_dtype: - self.reporter.report_reject( - node, - f"Mixed {x_dtype} and {y_dtype} values in {node.target} are not supported", - ) - return False - supported_float = [torch.bool, torch.float16, torch.float32] if tosa_spec.support_extension("bf16"): supported_float.append(torch.bfloat16) - supported_integer = [torch.bool, torch.int8, torch.int16, torch.int32] - if tosa_spec.support_extension("int64"): - supported_integer.append(torch.int64) + supported_integer = (torch.bool, torch.int8, torch.int16, torch.int32) if tosa_spec.support_float(): if x_dtype in supported_float and y_dtype in supported_float: diff --git a/backends/arm/test/misc/test_int64.py b/backends/arm/test/misc/test_int64.py index b48f1830479..e0150df2d74 100644 --- a/backends/arm/test/misc/test_int64.py +++ b/backends/arm/test/misc/test_int64.py @@ -127,20 +127,32 @@ def forward(self, *args): _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), - "where_int64": (torch.where, (_COND, _WIDE, _WIDE_2), 1), - "where_x_narrow": (torch.where, (_COND, _NARROW, _WIDE), 0), - "where_y_narrow": (torch.where, (_COND, _WIDE, _NARROW), 0), "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), - "where_x_bool": (torch.where, (_COND, _FLAGS, _WIDE), 0), - "where_y_bool": (torch.where, (_COND, _WIDE, _FLAGS), 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), } @@ -149,7 +161,7 @@ def test_int64_extension_partition_tosa_INT(test_data: Tuple): op, inputs, delegates = test_data ( ArmTester( - Int64Op(op), + op if isinstance(op, torch.nn.Module) else Int64Op(op), inputs, common.get_tosa_compile_spec("TOSA-1.1+INT+int64"), )