[abm-77]

Map torch.int64 to TOSA INT64 and accept it in TosaArg when the
specification enables EXT-INT64. Extend dialect validators and node
visitors to accept int64 for operators covered by the extension.

Add dialect tests that verify int64 acceptance with the extension and
rejection without it.

Existing narrowing passes still convert int64 to int32. Preserving
int64 through those passes is deferred to a follow-up change.
Part 1 of 3.

This PR was authored with Claude Code.

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

[pytorch-bot]

🔗 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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• PRs based before #202327 fail all Linux builds: can't find build-external-packages action

⏳ No Failures, 1 Pending

As of commit 160c2b2 with merge base 4fe0253 ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-cla]

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]

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]

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

[abm-77]

@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

• Honor supplied TOSA spec for INT64 reduction paths · New
• Enable INT64 BITWISE_AND lowering through binary operator factory · New
• Allow extension-supported INT64 tensors through layout lowering · New
• Support INT64 WHERE branches under the extension gate · New

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]

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

• Preserve large integer scalar precision during materialization · New
• Allow extension-supported INT64 tensors through layout lowering
• Accept extension-enabled int64 slices independently of INT profile · New
• Allow extension-enabled int64 where comparisons independently of INT profile · New

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

• Enable INT64 BITWISE_AND lowering through binary operator factory
• Honor supplied TOSA spec for INT64 reduction paths
• Support INT64 WHERE branches under the extension gate

_{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

• Reject unsupported INT64 bitwise OR/XOR delegation · New
• Preserve large integer scalar precision during materialization
• Allow extension-enabled int64 where comparisons independently of INT profile
• Accept extension-enabled int64 slices independently of INT profile

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

• Reject unsupported INT64 bitwise OR/XOR delegation
• Preserve large integer scalar precision during materialization
• Allow extension-enabled int64 where comparisons independently of INT profile
• Accept extension-enabled int64 slices independently of INT profile

_{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

• Preserve mixed-width WHERE operands without lossy casts · New
• Enable FP+int64 comparison delegation
• Allow FP+int64 reduction validation
• Prevent unsupported int64 shifts from raising

1 resolved since last review

• Assert delegation and int64 precision in tests

_{Review effort: Balanced}

[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

• Reject unsupported mixed-width binary operators
• Reject unsupported INT64 gather and embedding inputs
• Add partitioning coverage for restricted int64 cast pairs · New
• Test int64 dialect acceptance and rejection paths · New

_{Review effort: Balanced}

[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

• Reject unsupported mixed-width binary operators
• Reject unsupported INT64 gather and embedding inputs
• Test int64 dialect acceptance and rejection paths
• Add partitioning coverage for restricted int64 cast pairs

_{Review effort: Balanced}

[YufengShi-dudu]

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.

[abm-77]

Thanks, that makes sense. keeps_int64 is out of this PR now, but the same question comes up in the second one, so I tried dropping the INT check there and ran single-op models under TOSA-1.1+FP+int64.

Data movement ops on int64 (squeeze, unsqueeze, permute, expand, select, stack, transpose, concat, flip, tile) start failing in InsertDataLayoutCastsPass. Under an FP-only spec it moves integer data by casting it to float, and _int_to_fp_map stops at int32, so int64 raises TypeError: Data type torch.int64 of operator aten.view_copy.default is not supported by TOSA-1.1+FP+int64. Integer arithmetic on int64 (add, sub, mul, maximum, minimum, abs, neg) also gets delegated, which the FP profile rejects for int32 today.

Handling PRO-FP + EXT-INT64 fully looks like it would need per-profile capability lists and per-node narrowing, so for now I'd lean toward scoping this series to PRO-INT + EXT-INT64. I'm not attached to the INT check itself, though. If you have a cleaner way to express this, or would rather structure it differently, I'm happy to rework it.

[YufengShi-dudu]

Make sense. Let's focus on PRO-INT + EXT-INT64 in the next PR.

[Copilot]

🟢 Approval recommended

No active review issues are recorded for the scoped EXT-INT64 mapping, validation, and visitor changes.

No open issues 🎉

_{Review effort: Balanced}

[Copilot]

🟢 Approval recommended

No active review issues block the scoped EXT-INT64 validation and serialization changes.

No open issues 🎉

_{Review effort: Balanced}

[YufengShi-dudu]

Nice work! I think it's good to go now. I've also updated the commit message a bit.
@zingo Could you take another look? Thanks!

[abm-77]

Thank you for the reviews! I'll have the next PR up soon. Also, I don't have permission to merge, so I think someone would have to do so on my behalf.

[zingo]

Merged, the tests on "main" after merge can be monitored on this link https://hud.pytorch.org/hud/pytorch/executorch/main/1?per_page=50&name_filter=test-arm%7Ccortex-m%7Csize&useRegexFilter=true
(only showing Arm backends related tests :) )

I do not expect and problems so this is mostly for you info. :) It should be the exact same tests that as was runned on the PR but just once more on them merged version.

Thanks again for the great work.

[zingo]

And also congratulation to the first PR in the project!!!