Add DecomposeLstmPass for ARM backend (#17140)

[apullin]

Summary:

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

LSTM cell equations per timestep:
i_t = sigmoid(x_t @ W_ii.T + b_ii + h_{t-1} @ W_hi.T + b_hi)
f_t = sigmoid(x_t @ W_if.T + b_if + h_{t-1} @ W_hf.T + b_hf)
g_t = tanh(x_t @ W_ig.T + b_ig + h_{t-1} @ W_hg.T + b_hg)
o_t = sigmoid(x_t @ W_io.T + b_io + h_{t-1} @ W_ho.T + b_ho)
c_t = f_t * c_{t-1} + i_t * g_t
h_t = o_t * tanh(c_t)

Features:

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

Differential Revision: D92059277

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 Awaiting Approval, 2 Cancelled Jobs

As of commit 5c2da73 with merge base 3466332 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Presets (gh)
• Lint (gh)

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• pull / unittest / windows / windows-job (gh)
• pull / unittest-editable / windows / windows-job (gh)

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

[meta-codesync]

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

[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.

[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.

[gggekov]

What is the reason to decompose LSTM in the Arm backend rather than let torch.export.export decompose the LSTM ?

[gggekov]

Never mind- i see the torch.nn.LSTM is not decomposed in the torch.export.export as I thought initially.

[apullin]

@mansnils Trying to clear out the previous requested changes related to MLETORCH-1266, but that has long since been settled by others' PRs. Do I have to do anything from my end?

[gggekov]

Thanks for updating the PR and fixing the CI ! It's a great PR and would be really nice to add support for torch.nn.lstm/rnn/gru! I have added comments in the unit tests, I have two overall comments.

1. Could you please use TosaPipelineFP and TosaPipelineINT rather than the PassPipeline to test the decompositions for the three operators ?

2. We have very strict naming convention for how to name the unit tests. We are working on loosening the convention, but right now, if a unit test does not respect our naming convention, it is not picked up by the CI. I've added comments how to name the various tests if you keep them in backends/arm/test/passes.
   IMO, it would make more sense to move the tests under backends/arm/test/ops, that's where we test most operators that we decompose in transform_for_annotationn(e.g. test_adaptive_avg_pool2d.py, test_gelu.py, test_layer_norm.py, etc). In the case of LSTM, we also already have a unit test in executorch/backends/arm/test/models/test_lstm_arm.py for the rnn.LSTM, you can create another model with torch.nn.lstm and amend that file if you go down the route of moving the tests.

[apullin]

@gggekov The renaming is done.

And: I redid the tests to use TosaPipelineFP & TosaPipelineINT for all all three layer types (RNN, GRU, LSTM).

BUT: The LSTM TosaPipeline tests are currently skipped due to a TOSA backend shape mismatch on h_n/c_n (our decomposition pass outputs correct shapes, verified by the unit tests — the extra dimension is introduced during TOSA lowering). I'm pretty sure it's a TOSA backend issue. I can attempt to fix this directly as another commit, but that is some creep. Unless I am wrong about that as the failure origin.

I also tried adding EthosU55PipelineINT / EthosU85PipelineINT tests, and they ran fine locally, but GH CI doesn't seem to have the ethos-u-vela package installed?? So they fail there despite working locally.

[gggekov]

Hi @apullin ,
I am also getting a failure for the test_decompose_lstm_tosa_FP_e2e test case

FAILED backends/arm/test/passes/test_decompose_lstm_pass.py::test_decompose_lstm_tosa_FP_e2e - ValueError: Output needs to be of same shape: torch.Size([1, 1, 2, 20]) != torch.Size([1, 2, 20])

What is the reason for the shape mismatch for TOSA FP? Normally, when lowering an operator, the hard part is to get the numerics right for TOSA INT. In your case, the numerics for TOSA INT seem fine, but you get shape mismatch for TOSA FP. I wonder if something in another pass, run before the DecomposeLstmPass, doesn't cause the shape mismatch. That's why it's important to test via TosaPipelineINT & TosaPipelineFP because you run the full pipeline rather than a single pass in isolation.

Could you also add TosaPipelineINT & TosaPipelineFP tests for the other cases ( LSTM(bidirectional=True) , LSTM(bias=False), LSTM(num_layers=2)) ?

[gggekov]

Could you also add TosaPipelineINT & TosaPipelineFP tests for the other cases ( LSTM(bidirectional=True) , LSTM(bias=False), LSTM(num_layers=2)) ?

This applies also to the RNN & GRU test cases as well. It would be useful if you test the various attributes of RNN & GRU with TosaPipelineFP and TosaPipelineINT in order to make sure we run the full pipeline rather than just one pass. That's how we test all the operators we support via transform_for_annotation.

One idea about the shape mismatch for LSTM FP- it's possible that you may need to run the pass at a slightly different place in the arm_pass_manager.py - i did a few tests locally running the DecomposeLstmPass earlier in _tosa_pipeline, although i kept getting shape mismatch for TOSA FP :|

[apullin]

@gggekov ok, I THINK I have this down, now. Updated all three PRs with the following changes:

All three passes (GRU, RNN, LSTM):

• Added graph_module = super().call(graph_module).graph_module after recompile() to re-propagate FakeTensor shapes after decomposition. This was the root cause of the TOSA shape mismatch on LSTM h_n/c_n outputs.

Tests:

• All PassPipeline tests replaced with TosaPipelineFP + TosaPipelineINT across all three layer types.
• LSTM tests are no longer skipped — the shape fix plus a workaround (run_transform_for_annotation_pipeline before to_edge) resolves the extra-dimension
  issue. All 8 LSTM tests (4 FP + 4 INT) now run end-to-end through TosaPipeline.
• GRU: 8 tests (4 configs × FP + INT)
• RNN: 10 tests (5 configs × FP + INT)
• LSTM: 8 tests (4 configs × FP + INT)

All tests pass in a linux environment. But we will have to see if CI agrees with that.

[gggekov]

Hello @apullin ,
Thank you for updating the unit tests to use TosaPipelineFP/TosaPipelineINT! I think the PR is going in the right direction.

My main comment is that currently, your 3 unit tests test_decompose_<rnn/lstm/gru>_pass.py only work for TOSA INT(not sure if that is intentional?). However, for the TosaPipelineFP unit tests, your decomposition is never called. You are currently working around that in a slightly strange way by manually calling transform_for_annotation after the export from the ArmTester, but that is masking the underlying problem.

I think the proper fix is to make the pass decomposition include both EdgeIR ops(currently missing in your PR) and the ATen ops(that's what you currently do). See for example here (https://github.com/pytorch/executorch/blob/main/backends/arm/_passes/decompose_adaptive_avg_pool2d_pass.py#L23). In this way, for the TOSA FP case, you would use the EdgeIR decomposition and for the TOSA INT case, you'd use the ATen decomposition.

Why do you need the EdgeIR decomposition for the TOSA FP case ? Have a look at https://github.com/pytorch/executorch/blob/main/exir/program/_program.py#L1384, insert a breakpoint/print statement before and after the call to edge_manager = _gen_edge_manager_for_partitioners(....)
Do print(aten_programs['forward']) on line 1383 and print(edge_manager.exported_program("forward")) on line 1387. On line 1383, you still have the torch.ops.aten.lstm.input that you decompose. However, on line 1387, you no longer have the torch.ops.aten.lstm, you have EdgeIR ops instead. Note that this code is called before our backend has been called. Under the hood, the edge_manager = _gen_edge_manager_for_partitioners(....) API is calling the from torch._decomp import get_decompositions API to do the decomposition from ATen to EdgeIR(i mentioned that API at the very start of the code review a few weeks ago). I think for TosaPipelineFP, because the torch.ops.aten.lstm.input has been replaced with EdgeIR ops, your pipeline is never called and the unit test is passing relying on the decompositions taking place before our backend has been reached. We'd like to have consistency between TOSA FP & TOSA INT - have the same decomposition for both profiles. For other passes decomposed in transform_for_annotation, the TOSA FP case is managed via the EdgeIR operators.

Two more comments:

• I notice in the first commit for the GRU, you add test_decompose_recurrent_tosa_pipelines.py and then in the last commit(LSTM), you remove that file. Is it not possible to remove it altogether from the first commit ?
• Do you mind to do $ pip install -r requirements-lintrunner.txt & lintrunner init and then run the ./backends/arm/scripts/pre-commit script to run the lintrunner ? I had to manually fix a few minor errors to get your PR in the internal PR.

Well done, thanks for pushing with that!

[gggekov]

Hi @apullin ,
I had a second thought - i think for the TOSA FP case, it will be hard to do your own decomposition given that the op is already decomposed by ExecuTorch via the from torch._decomp import get_decompositions API before our backend is reached as per my comment above. I suggest to just remove the the _add_lstm_workaround function, and instead just mark the TOSA FP tests as expected failures with

@pytest.mark.skip(reason="Shape mismatch for TOSA FP")

and add appropriate comment that your pass is only handling the TOSA INT case. For TOSA INT, you shouldn't need the _add_lstm_workaround function.

It would be interesting to understand where is the shape mismatch stemming from for the TOSA FP case, maybe the from torch._decomp import get_decompositions API decomposes the op in a way that results in a slightly different shape compared to running the torch.nn.LSTM in eager mode, or maybe we are introducing an extra dimension somewhere in our FP passes by accident. I am not sure. I suggest to proceed the PR just managing the TOSA INT case, which should make it possible to run the LSTM on the Ethos-U85.

[apullin]

@gggekov done, updated across all three commits:

• Removed _add_lstm_workaround and test_decompose_recurrent_tosa_pipelines.py
• All tests use TosaPipelineFP / TosaPipelineINT (basic, bidirectional, no_bias, multilayer)
• LSTM FP tests skipped but present, INT is still faithfully tested
• No feature smear between commits