Add DecomposeLstmPass for ARM backend (#17140)

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)
---
> Generated by [Confucius Code Assist (CCA)](https://www.internalfb.com/wiki/Confucius/Analect/Shared_Analects/Confucius_Code_Assist_(CCA)/)
[Confucius Session](https://www.internalfb.com/confucius?host=62602.od.fbinfra.net&port=8086&tab=Chat&session_id=e1d1ac52-0014-11f1-9d55-75b7d4e71d8a&entry_name=Code+Assist), [Trace](https://www.internalfb.com/confucius?session_id=e1d1ac52-0014-11f1-9d55-75b7d4e71d8a&tab=Trace)

Differential Revision: D92059277

Author: Andrew Pullin

backends/arm/_passes/__init__.py

@@ -64,6 +64,7 @@
 from .decompose_linear_pass import DecomposeLinearPass  # noqa
 from .decompose_log1p_pass import DecomposeLog1pPass  # noqa
 from .decompose_logit_pass import DecomposeLogitPass  # noqa
+from .decompose_lstm_pass import DecomposeLstmPass  # noqa
 from .decompose_masked_fill_pass import DecomposeMaskedFillPass  # noqa
 from .decompose_matmul import DecomposeMatmulPass  # noqa
 from .decompose_maxpool2d_with_dilation_pass import DecomposeMaxPool2dPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -65,6 +65,7 @@
     DecomposeLinearPass,
     DecomposeLog1pPass,
     DecomposeLogitPass,
+    DecomposeLstmPass,
     DecomposeMaskedFillPass,
     DecomposeMatmulPass,
     DecomposeMaxPool2dPass,
@@ -239,6 +240,7 @@ def _tosa_pipeline(
                 DecomposeTOSAUnsupportedClampPass(),
                 DecomposeGroupNormPass(),
                 DecomposeGruPass(),
+                DecomposeLstmPass(),
                 DecomposeRnnPass(),
                 DecomposeLayerNormPass(),
                 DecomposeVarPass(),
@@ -429,6 +431,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 DecomposeAddSubAlphaPass(tfa_pass=True),
                 DecomposeGroupNormPass(tfa_pass=True),
                 DecomposeGruPass(tfa_pass=True),
+                DecomposeLstmPass(tfa_pass=True),
                 DecomposeRnnPass(tfa_pass=True),
                 DecomposeLayerNormPass(tfa_pass=True),
                 DecomposeVarPass(tfa_pass=True),