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)

Differential Revision: D92059277

Author: Andrew Pullin

backends/arm/_passes/__init__.py

@@ -71,6 +71,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

@@ -72,6 +72,7 @@
     DecomposeLinearPass,
     DecomposeLog1pPass,
     DecomposeLogitPass,
+    DecomposeLstmPass,
     DecomposeMaskedFillPass,
     DecomposeMatmulPass,
     DecomposeMaxPool2dPass,
@@ -363,6 +364,7 @@ def _tosa_pipeline(
                 DecomposeTOSAUnsupportedClampPass(),
                 DecomposeGroupNormPass(),
                 DecomposeGruPass(),
+                DecomposeLstmPass(),
                 DecomposeRnnPass(),
                 DecomposeLayerNormPass(),
                 DecomposeVarPass(),
@@ -583,6 +585,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
             [
                 NormalizeWhileInitialArgsPass(use_exir_clone=False, tfa_pass=True),
                 DecomposeGruPass(tfa_pass=True),
+                DecomposeLstmPass(tfa_pass=True),
                 DecomposeRnnPass(tfa_pass=True),
                 DecomposeNotEqualPass(tfa_pass=True),
                 DecomposeCosineSimilarityPass(tfa_pass=True),