Arm Backend: Add Silero VAD Ethos-U example and recurrent-state lowering (pytorch#22172)

### Summary

Silero VAD is a compact streaming model that demonstrates a realistic
recurrent-state workload on Ethos-U85 rather than another stateless
classification example.

This PR adds an in-tree Silero VAD Ethos-U example covering model
preparation, PT2E quantization, Arm lowering, bare-metal execution on
Corstone-320, and host-side output validation. It also adds the reusable
Arm backend support needed to lower the LSTMCell gate/update pattern
while preserving model-owned mutable hidden and cell state between audio
frames.

The runtime accepts one 576-sample audio input per frame, keeps
recurrent state inside the model, prints speech probabilities and
detected segments over serial, and validates the FVP results against
export-time reference probabilities.

Tracking: MLETORCH-2112

This PR was developed with OpenAI Codex assistance.

### Test plan

- `python -m pytest -q backends/arm/test/passes/test_mutable_buffers.py
backends/arm/test/misc/test_silero_vad_example.py` — 14 passed.
- `backends/arm/test/test_arm_ootb.sh run_silero_vad_e2e_ethos_u` —
exports the PTE, builds the runtime, runs Corstone-320 / Ethos-U85-256
FVP, and compares 79 output frames.
- FVP validation: maximum absolute error `0.156800`, mean absolute error
`0.018459`, and zero threshold-decision mismatches at `0.55`.
- Repository pre-push lint, license, commit-message, and Arm
documentation checks passed.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

---------

Signed-off-by: Usamah Zaheer <usamah.zaheer@arm.com>

Author: usamahz

backends/arm/_passes/arm_pass.py

@@ -11,6 +11,10 @@
 from typing import Any, List, Optional, Set, Type
 
 import torch
+from executorch.backends.arm._passes.mutable_buffer_utils import (
+    collect_mutable_buffer_infos,
+    restore_mutable_buffer_targets,
+)
 from executorch.backends.arm.constants import DISALLOW_TFA_META_KEY
 from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -217,8 +221,14 @@ def __call__(self, graph_module: GraphModule) -> PassResult | None:
         if not self.should_run_pass(graph_module):
             self.ensures(graph_module)
             return PassResult(graph_module, False)
+        mutable_buffers = collect_mutable_buffer_infos(graph_module)
         res = self.call(graph_module)
         self.ensures(graph_module)
+        if res is not None:
+            res = PassResult(
+                restore_mutable_buffer_targets(res.graph_module, mutable_buffers),
+                res.modified,
+            )
         return res
 
 

backends/arm/_passes/arm_pass_manager.py

@@ -182,7 +182,6 @@
     TosaLoweringContext,
     TosaSpecification,
 )
-
 from executorch.exir import ExportedProgram
 from executorch.exir._program_utils import _get_updated_graph_signature
 from executorch.exir.pass_base import (
@@ -822,4 +821,4 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 ]
             )
 
-            return self._transform_graph_module(graph_module)
+            return GraphModulePassManager(self.passes)(graph_module).graph_module

backends/arm/_passes/constant_folding_pass.py

@@ -8,7 +8,7 @@
 import torch
 from executorch.backends.arm._passes import ArmPass
 from executorch.exir.pass_base import ExportPass, PassResult
-from torch._export.passes.constant_folding import constant_fold
+from torchao.quantization.pt2e.constant_fold import constant_fold
 
 
 class ConstantFoldingPass(ArmPass):
@@ -52,19 +52,8 @@ def _ensure_param_attr(
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         before_graph = str(graph_module.graph)
-
-        # The ConstantFolder checks for this attribute which does not always exist
-        faked_dequantize_affine = False
-        if not hasattr(torch.ops.pt2e_quant, "dequantize_affine"):
-            torch.ops.pt2e_quant.dequantize_affine = None  # type: ignore[attr-defined]
-            faked_dequantize_affine = True
-
         constant_fold(graph_module)
 
-        # Remove attribute again if added to not affect global scope
-        if faked_dequantize_affine:
-            del torch.ops.pt2e_quant.dequantize_affine  # type: ignore[attr-defined]
-
         modified = before_graph != str(graph_module.graph)
 
         # Constant folding may have replaced some parameters with tensors, so we need to ensure they are still parameters for the backend to work correctly.

backends/arm/_passes/decompose_lstm_pass.py

@@ -6,7 +6,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import operator
-from typing import List, Set, Tuple, Type
+from typing import cast, List, Optional, Set, Tuple, Type
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmPass
@@ -18,7 +18,8 @@
 
 
 class DecomposeLstmPass(ArmPass):
-    """Decomposes aten.lstm.input into elementary ops supported by TOSA.
+    """Decomposes aten.lstm.input and aten.lstm_cell.default into elementary ops
+    supported by TOSA.
 
     LSTM cell equations per timestep:
         i_t = sigmoid(x_t @ W_ii.T + b_ii + h_{t-1} @ W_hi.T + b_hi)
@@ -32,13 +33,17 @@ class DecomposeLstmPass(ArmPass):
     the result is sliced into i/f/g/o components. This yields 2 mm ops per
     timestep instead of 8.
 
-    Supports multi-layer, bidirectional, with/without bias, and batch_first.
+    ``lstm_cell`` applies the equations once. ``lstm.input`` unrolls the
+    sequence and reuses the same cell-step decomposition at each timestep.
+    Full LSTM supports multi-layer, bidirectional, with/without bias, and
+    batch_first.
 
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
     _TARGET = torch.ops.aten.lstm.input
+    _TARGET_CELL = torch.ops.aten.lstm_cell.default
 
     _linear = torch.ops.aten.linear.default
     _add = torch.ops.aten.add.Tensor
@@ -59,8 +64,8 @@ def _build_direction(
         c_prev: torch.fx.Node,
         weight_ih: torch.fx.Node,
         weight_hh: torch.fx.Node,
-        bias_ih,
-        bias_hh,
+        bias_ih: Optional[torch.fx.Node],
+        bias_hh: Optional[torch.fx.Node],
         hidden_size: int,
         seq_len: int,
         time_dim: int,
@@ -82,50 +87,21 @@ def _build_direction(
                 from_node=node,
             )
 
-            gates_x = create_node(
-                graph, self._linear, args=(x_t, weight_ih, bias_ih), from_node=node
-            )
-            gates_h = create_node(
-                graph, self._linear, args=(h_prev, weight_hh, bias_hh), from_node=node
-            )
-
-            gates = create_node(
-                graph, self._add, args=(gates_x, gates_h), from_node=node
-            )
-
-            H = hidden_size
-            i_pre = create_node(
-                graph, self._slice, args=(gates, 1, 0, H), from_node=node
-            )
-            f_pre = create_node(
-                graph, self._slice, args=(gates, 1, H, 2 * H), from_node=node
-            )
-            g_pre = create_node(
-                graph, self._slice, args=(gates, 1, 2 * H, 3 * H), from_node=node
-            )
-            o_pre = create_node(
-                graph, self._slice, args=(gates, 1, 3 * H, 4 * H), from_node=node
+            h_prev, c_prev = self._build_cell_step(
+                graph,
+                node,
+                x_t,
+                h_prev,
+                c_prev,
+                weight_ih,
+                weight_hh,
+                bias_ih,
+                bias_hh,
+                hidden_size,
             )
 
-            i_t = create_node(graph, self._sigmoid, args=(i_pre,), from_node=node)
-            f_t = create_node(graph, self._sigmoid, args=(f_pre,), from_node=node)
-            g_t = create_node(graph, self._tanh, args=(g_pre,), from_node=node)
-            o_t = create_node(graph, self._sigmoid, args=(o_pre,), from_node=node)
-
-            # c_t = f_t * c_{t-1} + i_t * g_t
-            f_c = create_node(graph, self._mul, args=(f_t, c_prev), from_node=node)
-            i_g = create_node(graph, self._mul, args=(i_t, g_t), from_node=node)
-            c_t = create_node(graph, self._add, args=(f_c, i_g), from_node=node)
-
-            # h_t = o_t * tanh(c_t)
-            tanh_c = create_node(graph, self._tanh, args=(c_t,), from_node=node)
-            h_t = create_node(graph, self._mul, args=(o_t, tanh_c), from_node=node)
-
-            h_prev = h_t
-            c_prev = c_t
-
             h_t_expanded = create_node(
-                graph, self._unsqueeze, args=(h_t, time_dim), from_node=node
+                graph, self._unsqueeze, args=(h_prev, time_dim), from_node=node
             )
             timestep_outputs.append(h_t_expanded)
 
@@ -134,17 +110,129 @@ def _build_direction(
 
         return timestep_outputs, h_prev, c_prev
 
+    def _build_cell_step(
+        self,
+        graph: torch.fx.Graph,
+        node: torch.fx.Node,
+        input_node: torch.fx.Node,
+        h_prev: torch.fx.Node,
+        c_prev: torch.fx.Node,
+        weight_ih: torch.fx.Node,
+        weight_hh: torch.fx.Node,
+        bias_ih: Optional[torch.fx.Node],
+        bias_hh: Optional[torch.fx.Node],
+        hidden_size: int,
+    ) -> Tuple[torch.fx.Node, torch.fx.Node]:
+        gates_x = create_node(
+            graph, self._linear, args=(input_node, weight_ih, bias_ih), from_node=node
+        )
+        gates_h = create_node(
+            graph, self._linear, args=(h_prev, weight_hh, bias_hh), from_node=node
+        )
+        gates = create_node(graph, self._add, args=(gates_x, gates_h), from_node=node)
+        gate_dim = -1
+
+        i_pre = create_node(
+            graph,
+            self._slice,
+            args=(gates, gate_dim, 0, hidden_size),
+            from_node=node,
+        )
+        f_pre = create_node(
+            graph,
+            self._slice,
+            args=(gates, gate_dim, hidden_size, 2 * hidden_size),
+            from_node=node,
+        )
+        g_pre = create_node(
+            graph,
+            self._slice,
+            args=(gates, gate_dim, 2 * hidden_size, 3 * hidden_size),
+            from_node=node,
+        )
+        o_pre = create_node(
+            graph,
+            self._slice,
+            args=(gates, gate_dim, 3 * hidden_size, 4 * hidden_size),
+            from_node=node,
+        )
+
+        i_t = create_node(graph, self._sigmoid, args=(i_pre,), from_node=node)
+        f_t = create_node(graph, self._sigmoid, args=(f_pre,), from_node=node)
+        g_t = create_node(graph, self._tanh, args=(g_pre,), from_node=node)
+        o_t = create_node(graph, self._sigmoid, args=(o_pre,), from_node=node)
+
+        f_c = create_node(graph, self._mul, args=(f_t, c_prev), from_node=node)
+        i_g = create_node(graph, self._mul, args=(i_t, g_t), from_node=node)
+        c_t = create_node(graph, self._add, args=(f_c, i_g), from_node=node)
+        tanh_c = create_node(graph, self._tanh, args=(c_t,), from_node=node)
+        h_t = create_node(graph, self._mul, args=(o_t, tanh_c), from_node=node)
+        return h_t, c_t
+
+    def _decompose_cell(
+        self,
+        graph: torch.fx.Graph,
+        node: torch.fx.Node,
+    ) -> Tuple[torch.fx.Node, torch.fx.Node]:
+        args = node.args
+        input_node = cast(torch.fx.Node, args[0])
+        h_prev, c_prev = cast(Tuple[torch.fx.Node, torch.fx.Node], args[1])
+        weight_ih = cast(torch.fx.Node, args[2])
+        weight_hh = cast(torch.fx.Node, args[3])
+        bias_ih = cast(Optional[torch.fx.Node], args[4]) if len(args) > 4 else None
+        bias_hh = cast(Optional[torch.fx.Node], args[5]) if len(args) > 5 else None
+
+        hidden_size = h_prev.meta["val"].shape[-1]
+
+        return self._build_cell_step(
+            graph,
+            node,
+            input_node,
+            h_prev,
+            c_prev,
+            weight_ih,
+            weight_hh,
+            bias_ih,
+            bias_hh,
+            hidden_size,
+        )
+
     def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
         graph = graph_module.graph
         modified = False
 
         for node in list(graph.nodes):
             if (
                 node.op != "call_function"
-                or node.target != self._TARGET
+                or node.target not in (self._TARGET, self._TARGET_CELL)
                 or not self.allowed_to_transform(node.meta)
             ):
                 continue
+
+            if node.target == self._TARGET_CELL:
+                users = list(node.users)
+                if any(
+                    user.target != operator.getitem
+                    or user.args[1] not in (-2, -1, 0, 1)
+                    for user in users
+                ):
+                    continue
+
+                with graph.inserting_before(node):
+                    h_t, c_t = self._decompose_cell(graph, node)
+
+                getitem_nodes = []
+                for user in users:
+                    idx = cast(int, user.args[1])
+                    user.replace_all_uses_with(h_t if idx in (-2, 0) else c_t)
+                    getitem_nodes.append(user)
+
+                for gi in getitem_nodes:
+                    graph.erase_node(gi)
+                graph.erase_node(node)
+                modified = True
+                continue
+
             getitem_users = get_getitem_users(node, 3)
 
             args = node.args