Skip to content

Conversation

MartinPavella
Copy link
Collaborator

@MartinPavella MartinPavella commented Aug 28, 2025

Summary

This PR introduces a pre-processing pass on the aten dialect level, which splits gru nodes with num_layers > 1 into an equivalent sequence of single layer gru nodes.

Test plan

Unit tests provided in backends/nxp/tests/test_gru_splitting.py.

cc @robert-kalmar @roman-janik-nxp @StrycekSimon @jirioc

Copy link

pytorch-bot bot commented Aug 28, 2025

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 1 New Failure

As of commit b0f6ecb with merge base f8a422c (image):

NEW FAILURE - The following job has failed:

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

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 28, 2025
@MartinPavella
Copy link
Collaborator Author

@pytorchbot label "module: nxp" "release notes: nxp"

@pytorch-bot pytorch-bot bot added module: nxp Issues related to NXP Neutron NPU delegation and code under backends/nxp/ release notes: nxp Changes to the NXP Neutron backend delegate labels Aug 28, 2025
@MartinPavella MartinPavella force-pushed the upstream/main-nxp/EIEX-501-upstream-pass-to-decompose-multilayer-gru branch from ffc75b4 to 325ee60 Compare August 28, 2025 12:51
passes: list[PassType] = passes or [
FuseBatchNormWithConvPass(),
FuseBatchNormWithLinearPass(),
SplitGRUBasedOnNumLayers(),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we do it here in aten is because we want to support quantization?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We do it in the aten dialect because it allows the pass to focus only of the transformation, and not worry about the quantization. This way, the implementation is simpler, and if the quantization requirements of GRU (or the other ops) ever change, we only need to update the quantizer (otherwise we would also have to update this pass).

Copy link
Collaborator

@robert-kalmar robert-kalmar Sep 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We do it here for 2 reasons:

  • Neutron NPU supports GRU as primitive operation Neutron IR ==> for Neutron NPU we do not need and want to decompose it to primitive ops by the ExecuTorch.
  • Neutron NPU supports GRU, but only single layer. Therefore we transform the multilayer GRU to a sequence of single layer GRU, perform the quantization and obtain a graph what can be represented in Neutron IR, including the proper quantization parameters for inputs and outputs of individual GRU layers.
    Note: we will preserve the GRU op in to_edge tranformation - coming change

@MartinPavella MartinPavella force-pushed the upstream/main-nxp/EIEX-501-upstream-pass-to-decompose-multilayer-gru branch 2 times, most recently from 5d6e466 to fbb061c Compare September 1, 2025 08:22
…ltiple operators.

The `aten.gru.input` has a parameter `num_layers`. For values > 1, it represents multiple `aten.gru.input` operators chained together. The introduced pass can split the original GRU into a chain of simpler GRU nodes.
@MartinPavella MartinPavella force-pushed the upstream/main-nxp/EIEX-501-upstream-pass-to-decompose-multilayer-gru branch from fbb061c to b0f6ecb Compare September 3, 2025 09:40
@robert-kalmar robert-kalmar merged commit 5abad6c into pytorch:main Sep 3, 2025
112 of 113 checks passed
@robert-kalmar robert-kalmar deleted the upstream/main-nxp/EIEX-501-upstream-pass-to-decompose-multilayer-gru branch September 3, 2025 15:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: nxp Issues related to NXP Neutron NPU delegation and code under backends/nxp/ release notes: nxp Changes to the NXP Neutron backend delegate
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants