Fix torch.split fails in to_edge with alias annotations

[Lidang-Jiang]

Fixes #11723

Summary

torch.split fails with RuntimeError: Found a custom (non-ATen) operator whose output has alias annotations when used with to_edge_transform_and_lower and a partitioner that requests op preservation.

Root cause: _remove_invalid_ops_for_not_decompose relies on torchgen's aliased_return_names() to detect ops with aliased returns. However, for ops returning lists of aliased tensors (e.g., split.Tensor returns Tensor(a)[]), aliased_return_names() returns [None], failing to detect the alias annotation. This lets split.Tensor pass through into the EDGE_DO_NOT_DECOMP namespace, where functionalization fails.

Fix: Add a fallback check using op._schema.returns directly, which correctly reports alias_info on list return types. This also fixes the same latent issue for chunk.default and tensor_split.sections.

Test plan

• Added test_remove_invalid_ops_filters_aliased_list_returns regression test
• Run: pytest exir/tests/test_passes.py::TestPasses::test_remove_invalid_ops_filters_aliased_list_returns -xvs
• Verified existing split-related test still passes: test_to_out_variant_singleon_tensor_list
• Verified existing broken ops test still passes: test_compile_fix_broken_ops

Before fix

==================== BEFORE FIX ====================
RESULT: FAILED
RuntimeError: Found a custom (non-ATen) operator whose output has alias annotations: EDGE_DO_NOT_DECOMP::split.Tensor(Tensor(a -> *) self, SymInt split_size, int dim=0) -> Tensor(a)[]. We only support functionalizing operators whose outputs do not have alias annotations (e.g. 'Tensor(a)' is a Tensor with an alias annotation whereas 'Tensor' is a Tensor without. The '(a)' is the alias annotation). The alias annotation specifies that the output Tensor shares storage with an input that has the same annotation. Please check if (1) the output needs to be an output (if not, don't return it), (2) if the output doesn't share storage with any inputs, then delete the alias annotation. (3) if the output indeed shares storage with an input, then add a .clone() before returning it to prevent storage sharing and then delete the alias annotation. Otherwise, please file an issue on GitHub.

While executing %split : [num_users=3] = call_function[target=torch.ops.EDGE_DO_NOT_DECOMP.split.Tensor](args = (%x, 2), kwargs = {})
Original traceback:
None
Use tlparse to see full graph. (https://github.com/pytorch/tlparse?tab=readme-ov-file#tlparse-parse-structured-pt2-logs)

After fix

==================== AFTER FIX ====================
WARNING:root:Op aten.split.Tensor was requested for preservation by partitioner.  This request is ignored because it aliases output.

Test 1: to_edge (no partitioner)
RESULT: SUCCESS - outputs match

Test 2: to_edge_transform_and_lower with split.Tensor preservation
RESULT: SUCCESS - split.Tensor correctly filtered from EDGE_DO_NOT_DECOMP
         (AttributeError from dummy partitioner partition(), not from split bug)

Test 3: _remove_invalid_ops_for_not_decompose filter check
  aten::split.Tensor                            -> FILTERED (correct)
  aten::chunk                                   -> FILTERED (correct)
  aten::tensor_split.sections                   -> FILTERED (correct)

Unit test output

$ pytest exir/tests/test_passes.py::TestPasses::test_remove_invalid_ops_filters_aliased_list_returns -xvs

============================= test session starts ==============================
platform linux -- Python 3.12.12, pytest-8.4.2
collected 1 item

exir/tests/test_passes.py::TestPasses::test_remove_invalid_ops_filters_aliased_list_returns PASSED

============================== 1 passed in 6.83s ===============================

$ pytest exir/tests/test_passes.py::TestPasses::test_to_out_variant_singleon_tensor_list -xvs
PASSED

$ pytest exir/tests/test_passes.py::TestPasses::test_compile_fix_broken_ops -xvs
PASSED

This PR was authored with the assistance of Claude.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (2 Unrelated Failures), 5 Unclassified Failures

As of commit c0191a9 with merge base 6020c29 ():

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

• Cadence Build & Test / cpu-test / test-ops / test-ops (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  examples/cadence/operators/test_g3_ops.py::ATenOpTestCases::test_g3_native_layer_norm_out_21
• pull / test-qnn-buck-build-linux / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command docker exec -t a25aa6dd7406f65269d9b14c770219612d4232deb19593cc76e2548c40fd3cd4 /exec failed with exit code 3
• pull / test-qnn-delegate-linux / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command docker exec -t 2855dbc8a7bfc57ad25c90a08694e13ee413ea05d2f0f4e7da17704377f65e95 /exec failed with exit code 2
• pull / test-qnn-passes-linux / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  backends/qualcomm/tests/test_passes.py::TestPasses::test_mha_to_sha
• pull / test-qnn-python-imports-linux / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command docker exec -t 1cb91b83aa45a1a183a75dfa4bba8878b2a407a8da60906f0d1e0a0f778acbec /exec failed with exit code 2

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

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

• pull / unittest / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[Lidang-Jiang]

@pytorchbot label "release notes: exir"

[GregoryComer]

@JacobSzwejbka Do you know what the intent is for functionalization with ops like split or chunk that return a list of aliased tensors? Do we expect that they should be functionalized? This seems like a maybe a bug in functionalization if they're not.

[JacobSzwejbka]

It should be converted to split_copy or further decomposed. Seems like the point of this change is that delegates are requesting to keep split from being decomposed and then the verifier complains that aliasing is in the graph. We should figure out how to let backends preserve aliasing ops (so they dont have to go defunctionalize them on their own) @GregoryComer

Functionalization happens at the same time as decomposition. The changes in this diff are before that code runs.

[GregoryComer]

Thanks. I think we can merge this as is. We can solve the larger question of preserving aliasing ops separately.

[GregoryComer]

@Lidang-Jiang Can you resolve the lint failure? You can just add a suppression for the function too complex. We can go ahead and merge after that. Thanks!

[Lidang-Jiang]

@GregoryComer Good catch, fixed in c0191a9. I added a function-scoped # noqa: C901 suppression for _remove_invalid_ops_for_not_decompose, and local lintrunner --skip MYPY passes on the changed files.

[GregoryComer]

Thanks. I re-triggered CI. There is a temporary code freeze at the moment, but I should be able to merge in the next few days.

[Lidang-Jiang]

Thanks for rerunning CI and for the update.