NXP backend: added aten.split support #16276
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16276
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 Unrelated FailureAs of commit 033bf11 with merge base df626bd ( NEW FAILURE - The following job has failed:
UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
|
@pytorchbot label "release notes: nxp" |
pytorch-bot
bot
added
the
release notes: nxp
|
@pytorchbot label "module: nxp" |
pytorch-bot
bot
added
the
module: nxp
There was a problem hiding this comment.
Pull request overview
This PR adds support for the aten.split operator in the NXP backend by introducing a new pass that decomposes split operations into slice operations. This enables the backend to handle tensor splitting operations that were previously unsupported.
- Implements
DecomposeSplitToSlicesPassto convert split operations into equivalent slice operations - Adds comprehensive test coverage for different split scenarios including split with size, split with sections, GRU-based splits, and single-chunk edge cases
- Integrates the new pass into the default NXP backend pass pipeline
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| backends/nxp/aten_passes/decompose_split_to_slices_pass.py | Implements the core pass logic to decompose split operations into slice operations |
| backends/nxp/aten_passes/neutron_aten_pass_manager.py | Integrates the new decompose split pass into the default pass pipeline |
| backends/nxp/tests/models.py | Adds test model classes (GRUModel, SplitWithSize, SplitWithSections) to support split operation testing |
| backends/nxp/tests/test_decompose_split_to_slices.py | Provides comprehensive test coverage for the split decomposition functionality |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| raise RuntimeError("Faulty split chunks") | ||
|
|
||
| # Get split dim | ||
| dim = -1 |
There was a problem hiding this comment.
Why are you guessing the dim? Is it not possible to read it from the arguments?
| split_node.replace_all_uses_with(input_node) | ||
| self.graph_module.graph.erase_node(split_node) | ||
|
|
||
| return True |
There was a problem hiding this comment.
Why return True? I don't see any other possible return value.
|
|
||
| # Check if split is even necessary - if not, remove it | ||
| if len(split_nodes_chunks) == 1: | ||
| getitem_node = split_node.next |
There was a problem hiding this comment.
The split_node.next is merely the node that is stored right after the split_node in the graph.nodes. The nodes don't have to be connected at all. Please use something like list(split_node.users)[0] instead.
| ) | ||
| def test_decompose_split_with_size(mocker, input_shape, split_size, dim): | ||
| model = SplitWithSize(split_size, dim) | ||
| example_input = torch.rand(input_shape) |
There was a problem hiding this comment.
Please set a seed to make the tests deterministic. Add something like
@pytest.fixture(autouse=True)
def reseed_model_per_test_run():
torch.manual_seed(23)
np.random.seed(23)to the top of the file.
| @pytest.mark.parametrize( | ||
| "input_shape, size_or_sections, dim", | ||
| [ | ||
| # pytest.param((8, 4), 4, 1, id="2D, one chunk using split size."), |
There was a problem hiding this comment.
Why is this commented out?
There was a problem hiding this comment.
Very nice tests!
Would it be possible to add some end to end test? I would like to see a case where you use our regular pipeline (to_quantized_edge_program()), where the split gets correctly decomposed (without calling the NeutronAtenPassManager manually like you did here), and the slice nodes then get delegated.
2 participants
Summary
adds support for aten.split operator
Test plan
tests can be manually run using
pytest -c /dev/null backends/nxp/tests/cc @robert-kalmar @MartinPavella