Arm backend: Fix meta propagation in some call passes

When a call pass creates a node, it won't get a meta["val"] field
until the graph is re-traced, unless explicitly set.

This is not usually a problem, but if a pass needs meta["val"]  from
an argument, and an  op is chained multiple times, the pass will crash
if the meta is not there/incorrectly set.

Fix a couple of passes where the field  was missing, or using the
incorrect value.

Signed-off-by: Erik Lundell <erik.lundell@arm.com>
Change-Id: I9b93e29a66bf9cce1fe0a371546020c2eb0828c2

Author: Erik-Lundell

backends/arm/_passes/arm_pass_utils.py

@@ -6,9 +6,10 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import operator
 import traceback
 from inspect import isclass
-from typing import List, Optional, Sequence, Tuple
+from typing import cast, List, Optional, Sequence, Tuple
 
 import torch
 import torch.fx
@@ -59,6 +60,17 @@ def is_get_attr_node(node: torch.fx.Node) -> bool:
     )
 
 
+def get_getitem_users(
+    source_node: torch.fx.Node, max_users: int
+) -> dict[int, torch.fx.Node | None]:
+    getitem_users: dict[int, torch.fx.Node | None] = {i: None for i in range(max_users)}
+    for user in source_node.users:
+        if user.target == operator.getitem:
+            getitem_users[cast(int, user.args[1])] = user
+
+    return getitem_users
+
+
 def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
     return (
         is_get_attr_node(node)

backends/arm/_passes/decompose_any_pass.py

@@ -7,7 +7,10 @@
 
 import torch
 from executorch.backends.arm._passes import ArmPass
-from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
+from executorch.backends.arm._passes.arm_pass_utils import (
+    create_node,
+    get_first_fake_tensor,
+)
 from executorch.exir.dialects._ops import (  # type: ignore[import-not-found]
     ops as exir_ops,
 )
@@ -91,16 +94,21 @@ def call(self, graph_module: torch.fx.GraphModule):
             with graph_module.graph.inserting_before(node):
                 for dim in dims_to_reduce:
                     args = (input_node, dim, True)
-                    input_node = graph_module.graph.create_node(
-                        "call_function", exir_ops.edge.aten.any.dim, args, node.kwargs
+                    input_node = create_node(
+                        graph_module.graph,
+                        exir_ops.edge.aten.any.dim,
+                        args=args,
+                        kwargs=node.kwargs,
+                        from_node=node,
                     )
 
                 if not keepdim:
                     output_shape = list(get_first_fake_tensor(node).shape)
-                    input_node = graph_module.graph.create_node(
-                        "call_function",
+                    input_node = create_node(
+                        graph_module.graph,
                         exir_ops.edge.aten.view_copy.default,
-                        (input_node, output_shape),
+                        args=(input_node, output_shape),
+                        from_node=node,
                     )
 
             node.replace_all_uses_with(input_node)

backends/arm/_passes/decompose_gru_pass.py

@@ -9,7 +9,10 @@
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmPass
-from executorch.backends.arm._passes.arm_pass_utils import create_node
+from executorch.backends.arm._passes.arm_pass_utils import (
+    create_node,
+    get_getitem_users,
+)
 from executorch.backends.arm._passes.insert_table_ops import InsertTableOpsPass
 from executorch.exir.pass_base import ExportPass, PassResult
 
@@ -149,6 +152,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                 or not self.allowed_to_transform(node.meta)
             ):
                 continue
+            getitem_users = get_getitem_users(node, 2)
 
             args = node.args
             input_node = args[0]
@@ -257,7 +261,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                             graph,
                             self._cat,
                             args=(merged, time_dim),
-                            from_node=node,
+                            from_node=getitem_users.get(0),
                         )
 
                         layer_final_hiddens.append(
@@ -281,15 +285,15 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                             graph,
                             self._cat,
                             args=(fw_outputs, time_dim),
-                            from_node=node,
+                            from_node=getitem_users.get(0),
                         )
 
                         layer_final_hiddens.append(
                             create_node(
                                 graph,
                                 self._unsqueeze,
                                 args=(fw_h_final, 0),
-                                from_node=node,
+                                from_node=getitem_users.get(1),
                             )
                         )
 
@@ -303,7 +307,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                         graph,
                         self._cat,
                         args=(layer_final_hiddens, 0),
-                        from_node=node,
+                        from_node=getitem_users.get(1),
                     )
 
                 output_node = current_input

backends/arm/_passes/decompose_lstm_pass.py

@@ -9,7 +9,10 @@
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmPass
-from executorch.backends.arm._passes.arm_pass_utils import create_node
+from executorch.backends.arm._passes.arm_pass_utils import (
+    create_node,
+    get_getitem_users,
+)
 from executorch.backends.arm._passes.insert_table_ops import InsertTableOpsPass
 from executorch.exir.pass_base import ExportPass, PassResult
 
@@ -142,6 +145,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                 or not self.allowed_to_transform(node.meta)
             ):
                 continue
+            getitem_users = get_getitem_users(node, 3)
 
             args = node.args
             input_node = args[0]
@@ -266,7 +270,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                             graph,
                             self._cat,
                             args=(merged, time_dim),
-                            from_node=node,
+                            from_node=(getitem_users.get(0)),
                         )
 
                         layer_final_hiddens.append(
@@ -306,23 +310,23 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                             graph,
                             self._cat,
                             args=(fw_outputs, time_dim),
-                            from_node=node,
+                            from_node=(getitem_users.get(0)),
                         )
 
                         layer_final_hiddens.append(
                             create_node(
                                 graph,
                                 self._unsqueeze,
                                 args=(fw_h_final, 0),
-                                from_node=node,
+                                from_node=(getitem_users.get(1)),
                             )
                         )
                         layer_final_cells.append(
                             create_node(
                                 graph,
                                 self._unsqueeze,
                                 args=(fw_c_final, 0),
-                                from_node=node,
+                                from_node=(getitem_users.get(2)),
                             )
                         )
 
@@ -336,7 +340,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                         graph,
                         self._cat,
                         args=(layer_final_hiddens, 0),
-                        from_node=node,
+                        from_node=getitem_users.get(1),
                     )
                 if len(layer_final_cells) == 1:
                     c_n = layer_final_cells[0]
@@ -345,7 +349,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                         graph,
                         self._cat,
                         args=(layer_final_cells, 0),
-                        from_node=node,
+                        from_node=getitem_users.get(2),
                     )
 
                 output_node = current_input

backends/arm/_passes/decompose_rnn_pass.py

@@ -9,7 +9,10 @@
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmPass
-from executorch.backends.arm._passes.arm_pass_utils import create_node
+from executorch.backends.arm._passes.arm_pass_utils import (
+    create_node,
+    get_getitem_users,
+)
 from executorch.backends.arm._passes.insert_table_ops import InsertTableOpsPass
 from executorch.exir.pass_base import ExportPass, PassResult
 
@@ -117,6 +120,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
 
             is_relu = node.target == torch.ops.aten.rnn_relu.input
             activation = self._relu if is_relu else self._tanh
+            getitem_users = get_getitem_users(node, 2)
 
             args = node.args
             input_node = args[0]
@@ -223,7 +227,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                             graph,
                             self._cat,
                             args=([fw_combined, bw_combined], -1),
-                            from_node=node,
+                            from_node=(getitem_users.get(0)),
                         )
 
                         layer_final_hiddens.append(
@@ -247,15 +251,15 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                             graph,
                             self._cat,
                             args=(fw_outputs, time_dim),
-                            from_node=node,
+                            from_node=(getitem_users.get(0)),
                         )
 
                         layer_final_hiddens.append(
                             create_node(
                                 graph,
                                 self._unsqueeze,
                                 args=(fw_h_final, 0),
-                                from_node=node,
+                                from_node=(getitem_users.get(1)),
                             )
                         )
 
@@ -269,7 +273,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa: C901
                         graph,
                         self._cat,
                         args=(layer_final_hiddens, 0),
-                        from_node=node,
+                        from_node=getitem_users.get(1),
                     )
 
                 output_node = current_input

backends/arm/_passes/rewrite_upsample.py

@@ -208,6 +208,7 @@ def call(self, graph_module):
                     rescale_node = create_node(
                         graph_module.graph,
                         exir_ops.backend.tosa.RESCALE.default,
+                        from_node=node,
                     )
                     tosa_resize_node.replace_all_uses_with(rescale_node)
                     if input_dtype == torch.int16:

backends/arm/test/ops/test_any.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -41,6 +41,14 @@ def forward(self, x: torch.Tensor):
         return torch.any(x)
 
 
+class ChainedAny(torch.nn.Module):
+    aten_op = "torch.ops.aten.any.default"
+    exir_op = "executorch_exir_dialects_edge__ops_aten_any_default"
+
+    def forward(self, x: torch.Tensor):
+        return torch.any(torch.any(x, dim=1, keepdim=True))
+
+
 input_t1 = Tuple[torch.Tensor]  # Input x
 
 
@@ -152,6 +160,19 @@ def test_any_tosa_INT(test_data: input_t1):
     pipeline.run()
 
 
+def test_any_tosa_FP_chained_reduce_all() -> None:
+    pipeline = TosaPipelineFP[Tuple[torch.Tensor]](
+        ChainedAny(),
+        (torch.randint(0, 2, (2, 3, 4), dtype=torch.bool),),
+        ChainedAny.aten_op,
+        ChainedAny.exir_op,
+        atol=0,
+        rtol=0,
+        qtol=0,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data)
 def test_any_u55_INT(test_data: input_t1):
     # Tests that we don't delegate these ops since they are not supported on U55.

backends/arm/test/ops/test_upsample_nearest2d.py

@@ -135,6 +135,16 @@ def forward(self, x):
         return self.upsample(x)
 
 
+class ChainedBilinearInterpolate(torch.nn.Module):
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(
+            x, scale_factor=2.0, mode="bilinear", align_corners=False
+        )
+        return torch.nn.functional.interpolate(
+            x, scale_factor=2.0, mode="bilinear", align_corners=False
+        )
+
+
 @common.parametrize(
     "test_data", test_data_suite | test_data_suite_bf16 | test_data_suite_fp16
 )
@@ -233,6 +243,16 @@ def test_upsample_nearest2d_vec_tosa_INT_interpolate(test_data: torch.Tensor):
     pipeline.run()
 
 
+def test_upsample_nearest2d_vec_tosa_INT_chained_bilinear_interpolate():
+    pipeline = TosaPipelineINT[input_t1](
+        ChainedBilinearInterpolate(),
+        (torch.rand(1, 3, 4, 4),),
+        "torch.ops.aten.upsample_bilinear2d.vec",
+        exir_op=[],
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_suite)
 def test_upsample_nearest2d_vec_tosa_INT_a16w8(test_data: torch.Tensor):
     """Test upsample_nearest2d vector op with int16 I/O quantization for TOSA

backends/arm/test/passes/test_decompose_gru_pass.py

@@ -8,7 +8,9 @@
 from typing import Tuple
 
 import torch
+from executorch.backends.arm._passes import DecomposeGruPass
 from executorch.backends.arm.test.tester.test_pipeline import (
+    PassPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
 )
@@ -62,6 +64,30 @@ def get_inputs(self) -> input_t:
         return (x, h)
 
 
+chain_input_t = Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+
+
+class SequentialGRU(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.gru1 = torch.nn.GRU(10, 12, batch_first=True)
+        self.gru2 = torch.nn.GRU(12, 8, batch_first=True)
+
+    def forward(
+        self, x: torch.Tensor, h1: torch.Tensor, h2: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        y, _ = self.gru1(x, h1)
+        z, h_n = self.gru2(y, h2)
+        return z, h_n
+
+    def get_inputs(self) -> chain_input_t:
+        return (
+            torch.randn(2, 5, 10),
+            torch.randn(1, 2, 12),
+            torch.randn(1, 2, 8),
+        )
+
+
 def _make_gru_fp_pipeline(module: GRU) -> TosaPipelineFP:
     pipeline = TosaPipelineFP[input_t](
         module,
@@ -139,3 +165,14 @@ def test_decompose_gru_tosa_FP_multilayer():
 def test_decompose_gru_tosa_INT_multilayer():
     """Test multi-layer GRU through quantized pipeline."""
     _make_gru_int_pipeline(GRU(num_layers=2)).run()
+
+
+def test_decompose_gru_pass_handles_chained_grus() -> None:
+    module = SequentialGRU()
+    pipeline = PassPipeline(
+        module,
+        module.get_inputs(),
+        quantize=True,
+        pass_list=[DecomposeGruPass],
+    )
+    pipeline.run()