Fix torch.split fails in to_edge with alias annotations (#18700)

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`

<details>
<summary>Before fix</summary>

```
==================== 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)
```

</details>

<details>
<summary>After fix</summary>

```
==================== 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)
```

</details>

<details>
<summary>Unit test output</summary>

```
$ 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
```

</details>

This PR was authored with the assistance of Claude.

---------

Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>

Author: Lidang-Jiang

exir/program/_program.py

@@ -1081,7 +1081,7 @@ def _sanity_check_graph_for_non_decomp_ops(
                     logging.warning(warning_str)
 
 
-def _remove_invalid_ops_for_not_decompose(
+def _remove_invalid_ops_for_not_decompose(  # noqa: C901
     preserve_ops: List[torch._ops.OpOverload],
 ) -> List[torch._ops.OpOverload]:
     _logged_warnings = set()
@@ -1124,6 +1124,16 @@ def keep(op):
                 )
                 return False
 
+        # Fallback: torchgen does not detect alias annotations on ops
+        # returning lists of aliased tensors (e.g. split.Tensor returns
+        # Tensor(a)[]). Check op._schema.returns directly.
+        for ret in schema.returns:
+            if ret.alias_info is not None:
+                log_warning(
+                    f"Op {op} was requested for preservation by partitioner.  This request is ignored because it aliases output."
+                )
+                return False
+
         # Explicit block list of ops that don't work if asked for
         # preservation
         if op in [

exir/tests/test_passes.py

@@ -940,6 +940,43 @@ def body(i, h, h_accum):
             torch.allclose(prog.exported_program().module()(inp), model(inp))
         )
 
+    def test_remove_invalid_ops_filters_aliased_list_returns(self) -> None:
+        """Verify _remove_invalid_ops_for_not_decompose filters ops that return
+        aliased tensor lists (e.g. split, chunk) even when torchgen's
+        aliased_return_names() fails to detect them. Regression test for
+        https://github.com/pytorch/executorch/issues/11723
+        """
+        from executorch.exir.program._program import (
+            _remove_invalid_ops_for_not_decompose,
+        )
+
+        # These ops return Tensor(a)[] — a list of aliased views.
+        # torchgen's aliased_return_names() misses the alias annotation on
+        # list returns, so the fallback check on op._schema.returns is needed.
+        aliased_list_ops = [
+            torch.ops.aten.split.Tensor,
+            torch.ops.aten.chunk.default,
+            torch.ops.aten.tensor_split.sections,
+            torch.ops.aten.split_with_sizes.default,
+        ]
+        for op in aliased_list_ops:
+            result = _remove_invalid_ops_for_not_decompose([op])
+            self.assertNotIn(
+                op,
+                result,
+                f"{op} should be filtered out because it returns aliased tensors",
+            )
+
+        # Non-aliased ops should be preserved.
+        preserved_ops = [torch.ops.aten.linear.default]
+        for op in preserved_ops:
+            result = _remove_invalid_ops_for_not_decompose([op])
+            self.assertIn(
+                op,
+                result,
+                f"{op} should be preserved because it has no aliased returns",
+            )
+
     def test_convert_symb_ops(self) -> None:
         class Foo(torch.nn.Module):
             def forward(self, x: torch.Tensor) -> torch.Tensor: