From 4330ac28a9b7fdca50e11ee3771ede1fd3801c88 Mon Sep 17 00:00:00 2001 From: Lidang-Jiang Date: Sat, 4 Apr 2026 17:10:38 +0800 Subject: [PATCH 1/2] Fix torch.split fails in to_edge with alias annotations Fixes #11723 _remove_invalid_ops_for_not_decompose relied on torchgen's aliased_return_names() to detect ops with aliased returns, but it returns [None] for ops returning lists of aliased tensors (e.g., split.Tensor returns Tensor(a)[]). This let split.Tensor through into the EDGE_DO_NOT_DECOMP namespace where functionalization failed. 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 and tensor_split. Signed-off-by: Lidang-Jiang --- exir/program/_program.py | 10 ++++++++++ exir/tests/test_passes.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/exir/program/_program.py b/exir/program/_program.py index baacd5eaec4..6baff52e78d 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1122,6 +1122,16 @@ def keep(op): ) return False + # Fallback: torchgen may fail to detect alias annotations on ops + # returning lists of tensors (e.g. split.Tensor returns Tensor(a)[]). + # Check op._schema.returns directly as a more reliable source. + 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 [ diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index f683384f8f9..d38014421f6 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -940,6 +940,42 @@ 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, + ] + 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: From c79aca72112be2fe872794e5eed46b1f43f5878c Mon Sep 17 00:00:00 2001 From: Lidang-Jiang Date: Sat, 4 Apr 2026 17:13:39 +0800 Subject: [PATCH 2/2] Address review: fix comment wording and extend test coverage - Change 'may fail' to 'does not detect' (torchgen structurally cannot handle ListType alias annotations) - Add split_with_sizes.default to test to document overlap with blocklist Signed-off-by: Lidang-Jiang --- exir/program/_program.py | 6 +++--- exir/tests/test_passes.py | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/exir/program/_program.py b/exir/program/_program.py index 6baff52e78d..dd0c1b0e5da 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1122,9 +1122,9 @@ def keep(op): ) return False - # Fallback: torchgen may fail to detect alias annotations on ops - # returning lists of tensors (e.g. split.Tensor returns Tensor(a)[]). - # Check op._schema.returns directly as a more reliable source. + # 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( diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index d38014421f6..8a084ba491a 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -957,6 +957,7 @@ def test_remove_invalid_ops_filters_aliased_list_returns(self) -> None: 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])