Arm backend: Guard argmax and argmin int32 propagation

[mansnils]

Previously, once an argmax or argmin index fit int32, the pass sent every consumer through the int32 path. Later arithmetic could then overflow even though the raw index itself was safe.

Keep unsafe direct consumers on the original int64 value. When a safe int32 prefix reaches an unsafe consumer, insert an int64 boundary. This allows preceding safe operations to remain delegatable while keeping calculations correct and model output dtypes unchanged.

Cover safe and overflowing paths in both ATen and Edge graphs.

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 7105b90 with merge base cbcdfc1 ():
💚 Looks good so far! There are no failures yet. 💚

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[zingo]

Oh no, Ill try a re-run

FAILED backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py::test_clip_text_with_projection_tosa_FP - RuntimeError: Expected to find "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default" but did not find it
Searched string:
    lowered_module_0 = self.lowered_module_0
    executorch_call_delegate = torch.ops.higher_order.executorch_call_delegate(lowered_module_0);  lowered_module_0 = None
    dim_order_ops__to_dim_order_copy_default_2 = executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default(input_ids, dtype = torch.int32);  input_ids = None
From CHECK-COUNT-2: executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default
FAILED backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py::test_clip_text_with_projection_tosa_INT - RuntimeError: Expected to find "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default" but did not find it
Searched string:

def forward(self, input_ids):
    dim_order_ops__to_dim_order_copy_default = executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default(input_ids, dtype = torch.int32, dim_order = [0, 1]);  input_ids = None
    lowered_module_0 = self.lowered_module_0
    executorch_call_delegate = torch.ops.higher_order.executorch_call_delegate(lowered_module_0, dim_order_ops__to_dim_order_copy_default);  lowered_module_0 = dim_order_ops__to_dim_order_copy_default = None
From CHECK-COUNT-2: executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default
= 2 failed, 189 passed, 18 skipped, 7 xfailed, 903 warnings, 4 rerun in 3255.68s (0:54:15) =