Arm backend: Support for identity for TOSA 1.0

Signed-off-by: Per Åstrand <[email protected]>
Change-Id: I3f133d55feac211be6196f4d6811feff97ade98c

Author: per

backends/arm/operators/ops_identity.py

@@ -5,20 +5,50 @@
 
 # pyre-unsafe
 
-from typing import List
+from typing import Any, List
 
 import torch
 import torch.fx
 
-import tosa_tools.v0_80.serializer.tosa_serializer as ts
-
 from executorch.backends.arm.operators.node_visitor import (
     NodeVisitor,
     register_node_visitor,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 
 
+def identity_operator_factory_v0_80(identity_target: str):
+    """
+    Creates and registers NodeVisitors for operators that map directly
+    to a TOSA IDENTITY op.
+    """
+
+    class IdentityOperatorVisitor(NodeVisitor):
+        target = identity_target
+
+        tosa_specs = NodeVisitor.tosa_specs_0_80
+
+        def define_node(
+            self,
+            node: torch.fx.Node,
+            tosa_graph: Any,
+            inputs: List[TosaArg],
+            output: TosaArg,
+        ) -> None:
+            import tosa_tools.v0_80.serializer.tosa_serializer as ts
+
+            # Simply add an identityOp
+            tosa_graph.addOperator(
+                ts.TosaOp.Op().IDENTITY, [inputs[0].name], [output.name]
+            )
+
+    register_node_visitor(IdentityOperatorVisitor)
+
+
+identity_operator_factory_v0_80("getitem")
+identity_operator_factory_v0_80("aten.alias_copy.default")
+
+
 def identity_operator_factory(identity_target: str):
     """
     Creates and registers NodeVisitors for operators that map directly
@@ -28,13 +58,17 @@ def identity_operator_factory(identity_target: str):
     class IdentityOperatorVisitor(NodeVisitor):
         target = identity_target
 
+        tosa_specs = NodeVisitor.tosa_specs_1_00
+
         def define_node(
             self,
             node: torch.fx.Node,
-            tosa_graph: ts.TosaSerializer,
+            tosa_graph: Any,
             inputs: List[TosaArg],
             output: TosaArg,
         ) -> None:
+            import serializer.tosa_serializer as ts
+
             # Simply add an identityOp
             tosa_graph.addOperator(
                 ts.TosaOp.Op().IDENTITY, [inputs[0].name], [output.name]