[abm-77]

Summary

Adds comparisons, where, slice_copy, full_like and zeros_like back to TOSA_PRO_INT_EXT_INT64_SupportList. Each lowers to EQUAL, GREATER, GREATER_EQUAL, SELECT, SLICE or CONST, which EXT-INT64 defines for i64.

• CastIntComparisonInputsPass compares int64 directly instead of rejecting it.
• WhereSupported and SliceCopySupported accept int64, and where only when every operand casts to int64.
• CastInt64BuffersToInt32Pass keeps an int64 buffer when every user computes in int64, returns it, or compares it with int64, int32 or bool operands, and narrows the rest as before. This covers the buffer cases deferred from Arm backend: Delegate int64 only to the operators EXT-INT64 covers (2/3) #24306.
• Casts from int32 and bool to int64 delegate instead of staying on the CPU.

I removed two rows in test_comparison_dtype_pass.py that expected int64 comparisons to be rejected.

This PR covers PRO-INT only. With FP enabled, int64 comparisons fall back to the CPU as before.

The quantized flow still narrows int64 to int32 during annotation. Part 4 will fix that.

Part 3 of 4 for #24095, following #24306.

Test Plan

The int64 partitioning tests compare outputs against the reference model with values outside the int32 range. Added unit tests for both passes. Arm ops, passes and misc tests on the reference model: same results as main apart from the new tests. 44 single-op models: identical to main under TOSA-1.0+INT and TOSA-1.1+INT+int64. A sweep of 1,018 int64 cases against main (inputs, buffers and constants, in and out of int32 range, across 38 ops and 7 operand dtypes) showed no regressions or lost delegations.

This PR was authored with Claude Code.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani @nil-is-all @YufengShi-dudu @Erik-Lundell

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 16 Awaiting Approval

As of commit 41dae1d with merge base d1f597e ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test CPU Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/trunk

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

[abm-77]

@pytorchbot label "release notes: arm"

[Copilot]

🟡 Changes recommended

Large-integer scalar comparisons lose precision and standalone static-cache pipelines crash without a lowering context, alongside mixed-dtype overflow and comparison-support regressions.

5 open issues

• Provide lowering context for static-cache pass calls · New
• Avoid float-rounding integer scalars in comparisons · New
• Preserve int64 comparison buffers using promoted operand dtype · New
• Allow int64 comparisons in FP-enabled TOSA profiles · New
• Honor explicit dtype overrides in full_like conversion · New

What changed in this PR

Extends the Arm backend’s EXT-INT64 support to comparisons, selection, slicing, and constants, continuing the work in #24095 and #24306. Quantized-flow narrowing remains deferred to part 4.

Changes:

• Enables additional int64 operators and comparison dtype promotion.
• Preserves int64 buffers when their consumers compute in int64.
• Adds pass tests and delegated-output checks using values outside the int32 range.

File	Description
backends/​arm/​test/​passes/​test_comparison_dtype_pass.py	Removes obsolete int64 rejection cases.
backends/​arm/​test/​passes/​test_cast_int64_pass.py	Tests buffer preservation and narrowing.
backends/​arm/​test/​passes/​test_cast_int_comparison_inputs_pass.py	Tests int64 comparison casts.
backends/​arm/​test/​misc/​test_int64.py	Expands delegation cases and checks outputs.
backends/​arm/​operator_support/​where_support.py	Allows int64 selection operands with the extension.
backends/​arm/​operator_support/​tosa_profile_supported_op_lists.py	Expands the EXT-INT64 operator allowlist.
backends/​arm/​operator_support/​slice_copy_support.py	Allows int64 slicing with the extension.
backends/​arm/​_passes/​cast_int64_pass.py	Adds consumer-aware int64 buffer preservation.
backends/​arm/​_passes/​cast_int_comparison_inputs_pass.py	Enables int64 comparison dtype selection and promotion.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Int64 lowering can incorrectly narrow buffers and round large integer scalars, causing range-check failures, output dtype loss, and incorrect comparison results.

7 open issues

• Preserve int64 dtype for directly returned buffers · New
• Avoid float-rounding integer scalars in comparisons
• Provide lowering context for static-cache pass calls
• Allow int64 comparisons in combined FP and INT profiles · New
• Honor explicit dtype overrides in full_like conversion
• Allow int64 comparisons in FP-enabled TOSA profiles
• Preserve int64 comparison buffers using promoted operand dtype

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The preservation check misses scalar comparison overloads, causing valid wide-int64 buffer comparisons to raise int32-range errors during narrowing even with EXT-INT64 enabled.

4 open issues

• Provide lowering context for static-cache pass calls
• Allow int64 comparisons in combined FP and INT profiles
• Honor explicit dtype overrides in full_like conversion
• Allow int64 comparisons in FP-enabled TOSA profiles

1 previously missed issue

In code that hasn't changed since last review

Preserve wide int64 buffers in scalar comparison overloads

backends/​arm/​_passes/​cast_int64_pass.py:91

target_ops contains only .Tensor comparisons, but this buffer pass runs before scalar-to-tensor replacement (arm_pass_manager.py:642–660). With EXT-INT64 enabled, a registered int64 buffer containing 2**40 used in torch.where(self.wide > 7, x, y) fails this check and raises an int32-range error during narrowing. Include the supported .Scalar comparison overloads in the preservation check and add a wide-buffer scalar-comparison regression.

3 resolved since last review

• Preserve int64 dtype for directly returned buffers
• Avoid float-rounding integer scalars in comparisons
• Preserve int64 comparison buffers using promoted operand dtype

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Human assessment is needed because dtype promotion, rank normalization, and shared int64 buffer narrowing interact in ways that can change results or cause lowering failures.

No open issues 🎉

4 resolved since last review

• Provide lowering context for static-cache pass calls
• Allow int64 comparisons in combined FP and INT profiles
• Honor explicit dtype overrides in full_like conversion
• Allow int64 comparisons in FP-enabled TOSA profiles

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

The int64-buffer preservation logic causes accepted multi-buffer indexing, mixed-width comparisons, and comparison/cast graphs to fail during lowering.

3 open issues

• Int64 comparison buffers break int8/int16 input casting · New
• Explicit int64-to-int32 casts cause buffer range errors · New
• Multi-buffer indexing preserves unsupported int64 indices · New

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Human assessment is needed because mixed-rank dtype promotion and fractional-scalar conversion can silently produce incorrect results in newly accepted operations.

No open issues 🎉

3 resolved since last review

• Multi-buffer indexing preserves unsupported int64 indices
• Explicit int64-to-int32 casts cause buffer range errors
• Int64 comparison buffers break int8/int16 input casting

_{Review effort: Balanced}

[YufengShi-dudu]

The change looks good to me overall. Could you also check these corner cases under TOSA-1.1+INT+int64?

1. x < 0.5, with int64 input x=[-1,0,1].

2. torch.where(cond,x,y), with cond=[True,False,True], int32 input x=[1,2,3], and zero-dimensional int64 input y=2**40.

3. torch.where(cond,x,self.value), with the same cond and x, and self.value registered as a zero-dimensional int64 buffer containing
   2**40+7.

[YufengShi-dudu]

Hi @zingo, I'll be off for the next two weeks. I suggest adding Erik and Oscar as reviewers for this patch.

[abm-77]

@YufengShi-dudu you were right, all three were broken. They delegated and then crashed or gave the wrong answer, whereas main just left them on the CPU.

I was only checking operand dtypes, but torch doesn't always compute in int64 when one operand is int64. 0-d tensors do not promote, so where runs in int32, and x < 0.5 runs in float. In the new changes, int64 only stays on the delegate if the op actually computes in int64. All three stay on the CPU and are in test_int64.py.

I also ran 0-d tensors and float scalars through the same op/dtype sweep as before and nothing else came up.

[Copilot]

🟢 Approval recommended

No active review issues block the Arm backend’s expanded TOSA PRO-INT int64 support.

No open issues 🎉

_{Review effort: Balanced}