Arm backend: support list-of-tensors args in FuseConstantArgsPass (pytorch#13037)

FuseConstantArgsPass is changed so that it no longer unpacks lists or
tuples in node.args and node.kwargs. Those sequences are now preserved
when the operator is invoked. That means ops like aten::cat, which
expect their first argument to be a `List[Tensor]`, still get exactly
that. Previously, the pass would flatten the list into separate tensor
arguments, leading to this runtime error:

[WARNING 2025-07-24 11:37:13,749 fuse_constant_ops_pass.py:133] Failed
to fuse constant op aten_cat_default due to exception: aten::cat()
Expected a value of type 'List[Tensor]' for argument 'tensors' but
instead found type 'Tensor'.

Signed-off-by: Sebastian Larsson <[email protected]>

Author: Sebastian-Larsson

backends/arm/_passes/fuse_constant_ops_pass.py

@@ -6,6 +6,7 @@
 import logging
 
 import torch._export.utils
+import torch.fx
 from executorch.backends.arm._passes.arm_pass_utils import (
     get_constant_placeholder_kind,
     get_first_fake_tensor,
@@ -50,22 +51,26 @@ def _fuse_nodes(self, node) -> bool:
         the operations already carried out on the data.
         """
 
-        # Extract tensors and args from the node
-        data_list = [
-            get_param_tensor(self.exported_program, input_node)
-            for input_node in node.all_input_nodes
-        ]
-
-        args = node.args[len(node.all_input_nodes) :]
-        kwargs = node.kwargs
-
-        if "input_qparams" in node.meta and len(node.meta["input_qparams"]) > 0:
-            for i in range(len(node.all_input_nodes)):
-                q_params = node.meta["input_qparams"][i]
-                data_list[i] = q_params.dequantize_value(data_list[i])
-
-        # Run the op on the extracted tensor
-        data = node.target(*data_list, *args, **kwargs)
+        input_nodes = list(node.all_input_nodes)
+        qparams = node.meta.get("input_qparams", None)
+
+        def resolve_arg(arg):
+            if isinstance(arg, torch.fx.Node) and arg in input_nodes:
+                idx = input_nodes.index(arg)
+                t = get_param_tensor(self.exported_program, arg)
+                if qparams:
+                    t = qparams[idx].dequantize_value(t)
+                return t
+            if isinstance(arg, tuple):
+                return tuple(resolve_arg(x) for x in arg)
+            if isinstance(arg, list):
+                return [resolve_arg(x) for x in arg]
+            return arg
+
+        new_args = tuple(resolve_arg(a) for a in node.args)
+        new_kwargs = {k: resolve_arg(v) for k, v in node.kwargs.items()}
+
+        data = node.target(*new_args, **new_kwargs)
 
         # Only fuse if the tensor does not get bigger.
         if data.numel() > get_first_fake_tensor(node).numel():

backends/arm/test/passes/test_fuse_constant_ops_pass.py

@@ -15,6 +15,7 @@
 from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
 
 input_t = Tuple[torch.Tensor]  # Input x
+input_t2 = Tuple[torch.Tensor, torch.Tensor]
 
 
 class FuseParameter(torch.nn.Module):
@@ -86,12 +87,32 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return operator.add(sliced, x)
 
 
+class CatConst(torch.nn.Module):
+    ops_before_pass = {
+        "executorch_exir_dialects_edge__ops_aten_cat_default": 1,
+    }
+    ops_after_pass = {
+        "executorch_exir_dialects_edge__ops_aten_cat_default": 1,
+    }
+    ops_not_after_pass = []
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, a, b):
+        return torch.cat((a, b), dim=0)
+
+
 modules = {
     "fuse_parameter": FuseParameter(),
     "fuse_buffer": FuseBuffer(),
     "fuse_const_tensor": FuseLiftedTensor(),
 }
 
+cat_module = {
+    "fuse_cat": CatConst(),
+}
+
 
 @common.parametrize("module", modules)
 def test_fuse_const_ops_tosa_MI(module: torch.nn.Module):
@@ -118,3 +139,16 @@ def test_fuse_const_ops_tosa_BI(module: torch.nn.Module):
         passes_with_exported_program=[ComputeConstantOpsAOT, FuseConstantArgsPass],
     )
     pipeline.run()
+
+
+@common.parametrize("module", cat_module)
+def test_fuse_const_ops_tosa_BI_cat(module: torch.nn.Module):
+    pipeline = PassPipeline[input_t2](
+        module,
+        (torch.rand(3), torch.rand(2)),
+        quantize=True,
+        ops_before_pass=module.ops_before_pass,
+        ops_after_pass=module.ops_after_pass,
+        passes_with_exported_program=[ComputeConstantOpsAOT, FuseConstantArgsPass],
+    )
+    pipeline.run()