Arm backend: Make int64 representable under EXT-INT64 (1/3) by abm-77 · Pull Request #24128 · pytorch/executorch

abm-77 · 2026-10-23T17:56:52Z

Summary

Makes int64 representable in the TOSA backend when a spec declares EXT-INT64. keeps_int64() gates it on the INT profile plus the extension, TosaArg accepts the dtype, the cast table is restricted to the bool and i32 pairs the extension defines, and the dialect validators and node visitors accept int64 for the operators it covers.

This is inert on its own. The narrowing passes still convert int64 to int32 before anything reaches these paths, so behavior is unchanged for every existing spec string. Part 1 of 3.

Test Plan

44 single-op models under TOSA-1.0+INT: identical to main.

This PR was authored with Claude Code.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

pytorch-bot · 2026-10-23T17:56:56Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24128

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures, 14 Unrelated Failures

As of commit e5e75fc with merge base 0295293 ():

NEW FAILURES - The following jobs have failed:

pull / test-arm-cortex-m-size-test (bare_metal) / linux-job (gh)
RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
pull / test-llama-runner-linux (fp32, xnnpack+custom+quantize_kv, mt-l-arm64g4-16-62, executorch-ubuntu-... / linux-job (gh)
RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
pull / unittest / macos / macos-job (gh)
.ci/scripts/tests/test_link_check_diff_selection.py::TestLinkCheckDiffSelection::test_lint_xrefs_still_catches_a_reference_the_branch_adds
trunk / test-models-macos-coreml (resnet50) / macos-job (gh)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

trunk / test-models-macos-coreml (dl3) / macos-job (gh) (similar failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (efficient_sam) / macos-job (gh) (similar failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (emformer_transcribe) / macos-job (gh) (similar failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (ic4) / macos-job (gh) (similar failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (mv2) / macos-job (gh) (similar failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (mv3) / macos-job (gh) (similar failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (vit) / macos-job (gh) (similar failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (w2l) / macos-job (gh) (similar failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

Apple / build-demo-ios / macos-job (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
trunk / test-coreml-delegate / macos-job (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (edsr) / macos-job (gh) (trunk failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (emformer_join) / macos-job (gh) (trunk failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (ic3) / macos-job (gh) (trunk failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-models-macos-coreml (mobilebert) / macos-job (gh) (trunk failure)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

meta-cla · 2026-10-23T17:56:58Z

Hi @abm-77!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

linux-foundation-easycla · 2026-10-23T17:57:00Z

The committers listed above are authorized under a signed CLA.

✅ login: abm-77 / name: Bryson (012f37f, 3e71137, 7930bc9, 92efc39, a387f1b, b8098b2, b866d99, bf72310, e2f8d86, ef7b404)

pytorch-bot · 2026-10-23T17:57:10Z

~~Workflows were awaiting approval.~~ CI has now been triggered for the ciflow labels on this PR.

abm-77 · 2026-10-23T18:03:27Z

@pytorchbot label "release notes: arm"

Copilot

🟡 Changes recommended

Unresolved extension-handling gaps still cause unwanted narrowing, lowering failures, and CPU fallback for supported int64 operations.

4 open issues

What changed in this PR

This PR addresses #24095 by making the Arm backend preserve int64 tensors when the TOSA specification declares the int64 extension.

Changes:

Gate int64 narrowing passes on extension support.
Enable INT64 mapping and extension-aware validation across dialect operations and node visitors.
Add lowering cases for int64 arithmetic and scalar comparison.

File	Description
backends/arm/tosa/specification.py	Add a context-based extension check.
backends/arm/tosa/mapping.py	Map INT64 and validate extension support.
backends/arm/tosa/dialect/ops/unary_elementwise.py	Allow extension-gated int64 unary operations.
backends/arm/tosa/dialect/ops/ternary_elementwise.py	Allow int64 SELECT values.
backends/arm/tosa/dialect/ops/comparison.py	Allow int64 comparisons.
backends/arm/tosa/dialect/ops/binary_elementwise.py	Allow int64 arithmetic and min/max validation.
backends/arm/test/misc/test_int64.py	Add int64-extension lowering cases.
backends/arm/operators/simple_node_visitor.py	Support extension-gated input dtype lists.
backends/arm/operators/operator_validation_utils.py	Include INT64 in supported layout dtypes.
backends/arm/operators/op_where.py	Accept extension-gated INT64 values.
backends/arm/operators/op_tosa_sub.py	Enable INT64 subtraction.
backends/arm/operators/op_tosa_reduce_sum.py	Enable INT64 sum visitor validation.
backends/arm/operators/op_tosa_reduce_min.py	Enable INT64 minimum reduction validation.
backends/arm/operators/op_tosa_reduce_max.py	Enable INT64 maximum reduction validation.
backends/arm/operators/op_tosa_neg.py	Enable INT64 negation.
backends/arm/operators/op_tosa_mul.py	Enable INT64 multiplication validation.
backends/arm/operators/op_tosa_minimum.py	Enable INT64 elementwise minimum.
backends/arm/operators/op_tosa_maximum.py	Enable INT64 elementwise maximum.
backends/arm/operators/op_tosa_identity.py	Enable INT64 identity.
backends/arm/operators/op_tosa_gt.py	Enable INT64 greater-than operands.
backends/arm/operators/op_tosa_ge.py	Enable INT64 greater-or-equal operands.
backends/arm/operators/op_tosa_eq.py	Enable INT64 equality operands.
backends/arm/operators/op_tosa_bitwise_not.py	Enable INT64 bitwise NOT.
backends/arm/operators/op_tosa_add.py	Enable INT64 addition.
backends/arm/operators/op_tosa_abs.py	Enable INT64 absolute value.
backends/arm/operators/op_flip.py	Enable INT64 flip validation.
backends/arm/_passes/insert_int32_casts_after_int64_placeholders.py	Skip placeholder narrowing with the extension.
backends/arm/_passes/convert_int64_output_ops_to_int32.py	Guard output narrowing with an extension check.
backends/arm/_passes/convert_int64_const_ops_to_int32.py	Skip constant-operation narrowing with the extension.
backends/arm/_passes/cast_int64_pass.py	Skip buffer narrowing with the extension.

_{Review effort: Balanced}

nil-is-all · 2026-10-23T18:18:20Z

Thanks for filing the bug report and the PR, @abm-77. As we wait on Arm maintainers to review, could you sign the CLA?

Copilot

🟡 Changes recommended

Scalar precision loss and inconsistent partitioner checks can cause incorrect results or lowering failures, while the new tests do not verify int64 preservation.

4 open issues

3 previously missed issues

In code that hasn't changed since last review

Reject unsupported int64 bitwise operators during partitioning

backends/arm/_passes/insert_int32_casts_after_int64_placeholders.py:120

Skipping placeholder narrowing lets unsupported int64 operators reach lowering. Under TOSA-1.1+INT+int64, BoolBitwiseSupported accepts int64 x | y, but the BITWISE_OR dialect validator and visitor only accept int8/int16/int32. Lowering then raises instead of leaving the operation on the portable path; BITWISE_XOR has the same mismatch. Keep the narrowing skip, but reject these unsupported dtype/operator combinations during partitioning, or implement int64 lowering where the TOSA specification permits it.

Assert delegation and preservation of INT64 operands

backends/arm/test/misc/test_int64.py:158

These cases only check that conversion completes; they do not assert delegation or preserved INT64 operands. They can therefore pass if the operators fall back to CPU or are silently narrowed. The small inputs would also hide narrowing in an output comparison. Assert that the expected operators delegate and retain INT64 operands, including the inputs to GREATER whose output is BOOL. Alternatively, assert delegation and compare execution results using values outside the int32 range.

Allow extension-enabled int64 comparison inputs in negative checks

backends/arm/tosa/dialect/ops/comparison.py:44

The dialect now accepts int64 comparisons under TOSA-1.1+FP+int64, but _floating_profile_negative_checks() still installs CheckFPComparisonInputs whenever the INT profile is absent. That check rejects int64 inputs, so eq/gt/ge with int64 placeholders remain undelegated under this valid spec. Pass the specification to the negative check and allow extension-enabled int64 comparison inputs while preserving rejection for specs without the extension.

3 resolved since last review

_{Review effort: Balanced}

Copilot

🟡 Changes recommended

Retained int64 OR/XOR operands can fail during lowering, and the new tests do not verify delegation or int64 preservation.

4 open issues

1 previously missed issue

In code that hasn't changed since last review

Add delegation, operator, and INT64 preservation assertions

backends/arm/test/misc/test_int64.py:158

This chain checks only that lowering completes. It still passes if every operation falls back to CPU or if tensors are narrowed to int32; it neither asserts delegation nor executes or inspects the lowered computation. Add delegate-count and no-leftover-operator assertions, then verify INT64 tensor dtypes in the delegated graph or compare exact outputs with width-sensitive inputs outside the int32 range. The current small inputs would not expose narrowing even if output comparison were added.

1 resolved since last review

Allow extension-supported INT64 tensors through layout lowering

_{Review effort: Balanced}

Copilot

🟡 Changes recommended

Extension-enabled paths still have partitioning and lowering gaps, and the new tests do not verify delegation or retained precision.

4 open issues

Prevent unsupported int64 shifts from raising · New
Allow FP+int64 reduction validation · New
Enable FP+int64 comparison delegation · New
Assert delegation and int64 precision in tests · New

4 resolved since last review

_{Review effort: Balanced}

Copilot

🟡 Changes recommended

Newly accepted mixed-width WHERE operands can still narrow int64 values to int32 and lose precision.

4 open issues

1 resolved since last review

Assert delegation and int64 precision in tests

_{Review effort: Balanced}

abm-77 · 2026-10-24T01:51:44Z

Drafting while I work to refactor the design.

YufengShi-dudu

Thanks for your contributions! I've left a few comments on the PR. Personally， I'd prefer to split this into several smaller PRs if feasible:

General int64 lowering support
Avoid unnecessary int64->int32 narrowing under the INT+EXT-INT64 profile
Handle other corner cases in int64 lowering
An independent fix in backends/transforms/replace_scalar_with_tensor.py
FYI, we also have some internal INT64-related patches in progress. I believe @oscarandersson8218 and @Erik-Lundell can provide more context on this work.

YufengShi-dudu · 2026-10-27T12:00:55Z

                DecomposeUnfoldToGatherPass(use_slice=self.tosa_spec.is_U55_subset),
                DecomposeEmbeddingPass(),
-                CastInt64BuffersToInt32Pass(exported_program),
+                CastInt64BuffersToInt32Pass(exported_program, tosa_spec=self.tosa_spec),


Maybe use get_context_spec to get the tosa_profile in the passes directly instead of passing the tosa_spec to the pass. @oscarandersson8218 @Erik-Lundell Have you reached an agreement to use which style?

I switched to get_context_spec(). The three passes no longer take a spec, and arm_pass_manager.py is back to upstream.

Our conclusion was that it's currently inconsistent but that the preferred way is to use the context spec.

YufengShi-dudu · 2026-10-27T12:04:24Z

        raise RuntimeError("Function must be executed within a TosaLoweringContext")


+def keeps_int64(tosa_spec: "TosaSpecification") -> bool:


Please re-evaluate whether keeps_int64 is still needed after the other cleanups are done.

I think it is required. keeps_int64() is support_integer() and support_extension("int64"), so the INT profile has to be declared. Dropping the support_integer() half means TOSA-1.1+FP+int64 keeps int64 even though it declares only FP, which introduces new errors under the FP profile.

abm-77 · 2026-10-27T17:32:44Z

Thanks for the review! I pulled the backends/transforms change out, so this is now entirely within backends/arm, and I will send that fix separately.

I will reformat this as a stacked PR in the order you listed.

abm-77 · 2026-10-27T18:31:59Z

I have cut this PR down to the first part, which makes int64 representable under EXT-INT64. It is inert on its own, since the narrowing passes still convert int64 to int32 before anything reaches it.

The other two parts are ready on branches, and I will open them as PRs once this one lands:

arm-int64-delegation adds TOSA_EXT_INT64_SupportList and the partitioner gate
arm-int64-narrowing makes the narrowing passes consult the spec and adds the tests

I tried to put these up as a GitHub stack, but a pull request base has to be a branch in this repository and mine are on a fork, so the chaining is rejected. Is there a way you normally handle stacked PRs from forks, or would you rather I open all three now, either against main or on top of each other?

Copilot

🟡 Changes recommended

The new int64 paths need direct regression coverage, and the partitioning impact needs reconciliation with the stated behavior-preserving scope.

4 open issues

_{Review effort: Balanced}

…tion

Copilot

🟢 Approval recommended

No unresolved blocking issues remain in these preparatory changes, which leave existing int64 narrowing behavior unchanged.

No open issues 🎉

4 resolved since last review

_{Review effort: Balanced}

YufengShi-dudu · 2026-10-28T13:01:19Z


        validate_num_inputs(self.target, inputs, 2)
        validate_same_dtype(self.target, [inputs[0], output])
+        if self.tosa_spec.support_extension("int64"):


Could we move this up to Line 50 to keep all the supported dtype checks together?

YufengShi-dudu · 2026-10-28T13:05:29Z


        validate_num_inputs(self.target, inputs, 1)
        validate_same_dtype(self.target, [*inputs, output])
+        if self.tosa_spec.support_extension("int64"):


YufengShi-dudu · 2026-10-28T13:05:49Z

        # Not first input, which is condition tensor.
        validate_same_dtype(self.target, inputs[1:])
        validate_valid_dtype(self.target, inputs[0], ts.DType.BOOL)
+        if self.tosa_spec.support_extension("int64"):


YufengShi-dudu · 2026-10-28T13:28:45Z

            attr_method="AbsAttribute",
            num_inputs=1,
            input_dtypes=[ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
+            extension_dtypes={"int64": [ts.DType.INT64]},


TosaArg.__validate() in backends/arm/tosa/mapping.py checks whether the required extension is enabled when constructing a TosaArg. Could we add ts.DType.INT64 directly to input_dtypes, as we already do for ts.DType.BF16, to follow the existing pattern?

YufengShi-dudu · 2026-10-28T13:29:20Z

With ts.DType.INT64 added directly to the relevant visitors’ input_dtypes, we can leave this file unchanged.

YufengShi-dudu · 2026-10-28T13:30:20Z

        raise RuntimeError("Function must be executed within a TosaLoweringContext")


+def keeps_int64(tosa_spec: "TosaSpecification | None" = None) -> bool:


Not used in this PR.

YufengShi-dudu · 2026-10-28T13:42:50Z



+def keeps_int64(tosa_spec: "TosaSpecification | None" = None) -> bool:
+    """Whether int64 is kept 64 bits wide; only INT has int64 lowerings."""


EXT-INT64 doesn't require PRO-INT. PRO-FP + EXT-INT64 is valid. For example, GATHER with FP32 input and INT64 indices is supported under this combination.
This might become clearer once EXT-INT64 is included in an officially released version of the TOSA specification.

YufengShi-dudu · 2026-10-28T13:50:15Z


+def accepts_int64(dtype: torch.dtype, op: str) -> bool:
+    """Whether `dtype` is an int64 the active specification carries."""
+    if dtype != torch.int64:


Could we move the dtype == torch.int64 check to the call sites and rename this helper to validate_int64_support(op)? This would keep dtype branching consistent with the surrounding code while sharing the extension check.

abm-77 requested review from digantdesai and rascani as code owners October 23, 2026 17:56

Copilot AI balanced review requested due to automatic review settings October 23, 2026 17:56

github-actions Bot added ciflow/trunk module: arm Issues related to arm backend labels Oct 23, 2026

Copilot started reviewing on behalf of abm-77 October 23, 2026 17:57

View session

pytorch-bot Bot added the release notes: arm Changes to the ARM backend delegate label Oct 23, 2026