Add DecomposeGruPass for ARM backend (#17137)

[apullin]

Summary:

Adds a decomposition pass that transforms aten.gru.input into elementary
ops supported by TOSA (matmul, sigmoid, tanh, mul, add, slice, cat).

GRU cell equations per timestep:
r_t = sigmoid(x_t @ W_ir.T + b_ir + h_{t-1} @ W_hr.T + b_hr)
z_t = sigmoid(x_t @ W_iz.T + b_iz + h_{t-1} @ W_hz.T + b_hz)
n_t = tanh(x_t @ W_in.T + b_in + r_t * (h_{t-1} @ W_hn.T + b_hn))
h_t = n_t + z_t * (h_{t-1} - n_t)

Features:

• Multi-layer GRU support
• Bidirectional GRU support
• With/without bias
• batch_first support
• Batched gate computation (2 mm ops per timestep instead of 6)

Differential Revision: D92058313

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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:

• viable/strict has been blocked for 5+ days

✅ You can merge normally! (2 Unrelated Failures)

As of commit e4f7631 with merge base 2a68e74 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / unittest-editable / windows / windows-job (gh) (matched win rule in flaky-rules.json)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but was 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.

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[meta-codesync]

@apullin has exported this pull request. If you are a Meta employee, you can view the originating Diff in D92058313.

[pytorch-bot]

To add the ciflow label ciflow/trunk please first approve the workflows that are awaiting approval (scroll to the bottom of this page).

This helps ensure we don't trigger CI on this PR until it is actually authorized to do so. Please ping one of the reviewers if you do not have access to approve and run workflows.

[apullin]

@AdrianLundell @martinlsm I think we re seeking an approval from someone inside ARM for a PR like this. Please have a look.

[AdrianLundell]

I'm a bit confused it seems like this is added also in #17140? And did you look into what @gggekov said about using torch._decomp.get_decompositions?

[apullin]

added also in #17140
@AdrianLundell
This is my mistake - I have several changes in a stack, and PR's got exported for each one, although they are 3 instances of the same technique, adding explicte decomp passes for common RNN layers.

See comment here:
#17140 (comment)

If we want to hit all 3, yes, #17140 would have that. instead of #17137 (this) and #17139 separately. Please advise.

[JakeStevens]

per @gggekov comment on a different pull request:

apologies for the radio silence over the last weeks. Happy to go with your approach of manually decomposing rather than the get_decompositions route. I see you have a few failing tests and a merge conflict, could you resolve these ? Your patch looks good to me, I haven't verified the numerics of the decomposition, I assume they must be alright since you pass the unit tests. It would be great to have the LSTM, GRU & RNN available in the Arm backend.

based on this and internal testing LGTM

[JakeStevens]

you need to rebase and resolve merge conflicts

[apullin]

you need to rebase and resolve merge conflicts

Done now, and the other failed test is expected to be green now. However - do we want to approve this one, or #17140, which includes the RNN & LSTM support, following the same pattern?

[AdrianLundell]

I would say go with #17140 and get everything at the same time