Arm backend: Improve support for multiple output ops (#15695)

The pattern of using .name from the TosaArg output is used in
essentially all node visitors. Refactoring output to possibly be a list
of TosaArgs would be very cumbersome. Instead, use a new
multiple_output_names field to provide the names of all getitem tensors.
This also means that we don't have to insert identity ops for getitem
operators, they just need to insert a tensor.


Signed-off-by: Erik Lundell <[email protected]>

Author: Erik-Lundell

backends/arm/operators/op_cond_if.py

@@ -56,6 +56,6 @@ def define_node(
                 inputs[0].name,
                 *(subgraph_input.name for subgraph_input in inputs[-1].special),
             ],
-            [output.name],
+            output.multiple_output_names,
             attr,
         )

backends/arm/operators/ops_identity.py

@@ -40,7 +40,7 @@ def define_node(
             inputs: List[TosaArg],
             output: TosaArg,
         ) -> None:
-            validate_num_inputs(self.target, inputs, [1, 2])
+            validate_num_inputs(self.target, inputs, 1)
             validate_same_dtype(self.target, [inputs[0], output], ts)
 
             # Simply add an identityOp
@@ -58,5 +58,4 @@ def define_node(
     register_node_visitor(IdentityOperatorVisitor)
 
 
-identity_operator_factory("getitem")
 identity_operator_factory("aten.alias_copy.default")

backends/arm/process_node.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 #
 
+import operator
 from typing import Any, cast, Dict
 
 import numpy as np
@@ -45,9 +46,15 @@ def process_call_function(
             f"Failed processing call_function: {node.name}. "
             "Is the original torch function supported?"
         ) from e
-    tosa_graph.currRegion.currBasicBlock.addTensor(
-        output.name, tosa_shape(output.shape, output.dim_order), output.dtype
-    )
+
+    if not output.multiple_output_names:
+        tosa_graph.currRegion.currBasicBlock.addTensor(
+            output.name, tosa_shape(output.shape, output.dim_order), output.dtype
+        )
+
+    # Get item nodes just add tensors, no node visitor is needed.
+    if node.target == operator.getitem:
+        return
 
     # Visiting each Node
     # pyre-ignore[16]: Undefined attribute.

backends/arm/test/ops/test_cond.py

@@ -174,9 +174,7 @@ def _create() -> tuple[torch.nn.Module, input_t2]:
     "case",
     test_cases,
     xfails={
-        "one_arg_two_outputs": "Multiple outputs is not supported.",
         "one_arg_and_scalar_one_output": "Scalars become get_attr nodes that are not supported.",
-        "two_args_two_outputs": "Nodes with multiple outputs are not properly supported.",
         "multiple_one_arg_one_output": "Scalars become get_attr nodes that are not supported.",
     },
 )

backends/arm/tosa/mapping.py

@@ -10,6 +10,7 @@
 
 """
 
+import operator
 from enum import Enum
 from typing import Any, Optional, Sequence
 
@@ -136,7 +137,7 @@ class TosaArg:
         special (list | None): Captured list when the argument is a sequence.
         number (float | int | None): Captured numeric value when given.
         tosa_spec (TosaSpecification): Active specification used for mapping.
-
+        multiple_output_name (list[str]): Output node names when node has multiple outputs; empty otherwise.
     """
 
     def __process_node(self, argument: torch.fx.Node):
@@ -146,7 +147,8 @@ def __process_node(self, argument: torch.fx.Node):
             argument (torch.fx.Node): FX node to inspect.
 
         """
-        self.name = argument.name + argument.meta.get(TOSA_TENSOR_NAME_META, "")
+        suffix = argument.meta.get(TOSA_TENSOR_NAME_META, "")
+        self.name = argument.name + suffix
 
         if "val" in argument.meta:
             output_dtype, self.shape, self.dim_order = extract_tensor_meta(
@@ -158,6 +160,16 @@ def __process_node(self, argument: torch.fx.Node):
 
             self.dtype = output_dtype
 
+        # If all users of the node are getitems, node visitors should connect the output of this node directly to the getitem tensors.
+        # Add a new attribute 'multiple_output_names' instead of making 'name' a list to avoid ambiguity regarding the type of 'name'.
+        # Make name of the output is the first getitem since we in most cases only handle that output.
+        users = list(argument.users)
+        if len(users) > 0 and all(user.target == operator.getitem for user in users):
+            self.multiple_output_names: list = [user.name + suffix for user in users]
+            self.name = self.multiple_output_names[0]
+        else:
+            self.multiple_output_names = []
+
     def __process_list(self, argument):
         """Capture a sequence argument as ``special``.
 
@@ -244,4 +256,6 @@ def __repr__(self):
             attrs.append(f"number={self.number!r}")
         if hasattr(self, "tosa_spec") and self.tosa_spec is not None:
             attrs.append(f"tosa_spec={self.tosa_spec!r}")
+        if hasattr(self, "names"):
+            attrs.append(f"names={self.multiple_output_names!r}")
         return f"{self.__class__.__name__}({', '.join(attrs)})"