Arm backend: Add real impls to all TOSA dialect ops

Additionally,
- Remove special case in ComputeOpsAOT pass for such ops,
since they can now be executed.
- Start running the model in tests were this was previously
impossible due to ops not having a real impl.

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

Author: Erik-Lundell

backends/arm/_passes/fuse_constant_ops_pass.py

@@ -174,13 +174,10 @@ def call(self, graph_module):
         for node in graph_module.graph.nodes:
             if node.op != "call_function":
                 continue
-            # Don't fuse TOSA dialect ops as they do not have eager forward functions.
-            # Also don't fuse ops whose explicit args/kwargs include symbolic shape values.
-            if (
-                self._is_tosa_dialect_op(node.target)
-                or self._arg_contains_symbolic_shape(node.args)
-                or self._arg_contains_symbolic_shape(node.kwargs)
-            ):
+            # Don't fuse ops whose explicit args/kwargs include symbolic shape values.
+            if self._arg_contains_symbolic_shape(
+                node.args
+            ) or self._arg_contains_symbolic_shape(node.kwargs):
                 continue
 
             input_nodes = node.all_input_nodes

backends/arm/operators/node_visitor.py

@@ -250,6 +250,9 @@ def get_node_visitors(*args) -> Dict[str, NodeVisitor]:
 
 
 def get_node_visitor(target: str, tosa_spec: TosaSpecification):
+    # Ensure all operator modules are imported so visitors are registered.
+    import executorch.backends.arm.operators  # noqa: F401
+
     node_visitor_tuples = _node_visitor_tuples.get(tosa_spec)
     for target_name, node_visitor_cls in node_visitor_tuples:
         if target_name == target:

backends/arm/test/passes/test_ensure_unique_output_nodes_pass.py

@@ -35,7 +35,6 @@ def test_ensure_unique_output_nodes_no_target_inserts_identity_per_repeated_outp
             "executorch_exir_dialects_backend__ops_tosa_IDENTITY_default": 2,
         },
     )
-    pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.run()
 
     graph_module = (
@@ -62,5 +61,4 @@ def test_ensure_unique_output_nodes_no_target_keeps_unique_outputs_unchanged() -
             "executorch_exir_dialects_backend__ops_tosa_IDENTITY_default",
         ],
     )
-    pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.run()

backends/arm/test/passes/test_rewrite_conv_pass.py

@@ -213,8 +213,6 @@ def test_rewrite_conv_tosa_FP():
     pipeline = PassPipeline(
         module, module.get_inputs(), passes_with_exported_program=[RewriteConvPass]
     )
-    # We cannot run TOSA backend dialect operators in eager mode.
-    pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.run()
 
 

backends/arm/test/passes/test_rewrite_max_pool2d_pass.py

@@ -37,15 +37,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class MaxPool2dWithoutStride(torch.nn.Module):
     def get_inputs(self) -> input_t:
-        return (torch.rand(1, 3, 8, 8),)
+        return (torch.rand(1, 3, 9, 9),)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return torch.nn.functional.max_pool2d(x, kernel_size=3)
 
 
 class MaxPool2dListKernel(torch.nn.Module):
     def get_inputs(self) -> input_t:
-        return (torch.rand(1, 3, 8, 8),)
+        return (torch.rand(1, 3, 8, 9),)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return torch.nn.functional.max_pool2d(x, kernel_size=[2, 3])
@@ -56,7 +56,7 @@ def get_inputs(self) -> input_t:
         return (torch.rand(1, 3, 8, 8),)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return torch.nn.functional.max_pool2d(x, kernel_size=[2, 3], stride=[])
+        return torch.nn.functional.max_pool2d(x, kernel_size=[2, 2], stride=[])
 
 
 class MaxPool2dDynamic(torch.nn.Module):
@@ -94,9 +94,6 @@ def test_rewrite_max_pool2d_tosa(module: ModuleWithInputs) -> None:
         },
         pass_list=[RemoveGetItemPass, RewriteMaxPool2dPass],
     )
-    pipeline.pop_stage(
-        "run_method_and_compare_outputs"
-    )  # Cannnot run aten graph with tosa dialect ops
     pipeline.run()
 
 
@@ -131,11 +128,10 @@ def test_rewrite_max_pool2d_tosa_empty_stride_uses_kernel_size() -> None:
         },
         pass_list=[RemoveGetItemPass, RewriteMaxPool2dPass],
     )
-    pipeline.pop_stage("run_method_and_compare_outputs")
     pipeline.run()
 
     tosa_node = _get_tosa_max_pool2d_node(pipeline)
-    assert tosa_node.args[2] == [2, 3]
+    assert tosa_node.args[2] == [2, 2]
 
 
 def test_rewrite_max_pool2d_tosa_dynamic_shape() -> None:

backends/arm/tosa/dialect/ops/avg_pool2d_adaptive.py

@@ -11,7 +11,7 @@
     compute_avg_pool2d_output_shape,
     validate_avg_pool2d_dtype,
 )
-from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op
+from executorch.backends.arm.tosa.dialect.ops_registration import register_tosa_op
 from executorch.backends.arm.tosa.specification import (
     get_context_shape_env,
     get_context_spec,
@@ -36,7 +36,7 @@ def _is_directly_representable(
     return remainder in (0, 1)
 
 
-@register_fake_tosa_op(
+@register_tosa_op(
     "AVG_POOL2D_ADAPTIVE(Tensor input, Tensor input_zp, Tensor output_zp, SymInt[2] kernel, SymInt[2] stride, SymInt[4] pad, ScalarType acc_type) -> Tensor",
     TosaSpecification.all_profiles_for_version("1.1"),
 )

backends/arm/tosa/dialect/ops/conv2d.py

@@ -8,7 +8,7 @@
 
 import torch
 from executorch.backends.arm.tosa.dialect.lib import TosaValueError
-from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op
+from executorch.backends.arm.tosa.dialect.ops_registration import register_tosa_op
 from executorch.backends.arm.tosa.specification import (
     get_context_spec,
     TosaSpecification,
@@ -77,7 +77,7 @@ def validate_conv2d_args_dtypes(
     return output_dtype
 
 
-@register_fake_tosa_op(
+@register_tosa_op(
     "CONV2D(Tensor input, "
     "Tensor weight, "
     "Tensor bias, "

backends/arm/tosa/dialect/ops/conv3d.py

@@ -9,7 +9,7 @@
 import torch
 from executorch.backends.arm.tosa.dialect.lib import TosaValueError
 from executorch.backends.arm.tosa.dialect.ops.conv2d import validate_conv2d_args_dtypes
-from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op
+from executorch.backends.arm.tosa.dialect.ops_registration import register_tosa_op
 from executorch.backends.arm.tosa.specification import (
     get_context_spec,
     TosaSpecification,
@@ -30,7 +30,7 @@ def validate_conv3d_args_dtypes(
     return validate_conv2d_args_dtypes(tosa_spec, x, weight, bias, op="CONV3D")
 
 
-@register_fake_tosa_op(
+@register_tosa_op(
     "CONV3D(Tensor input, "
     "Tensor weight, "
     "Tensor bias, "

backends/arm/tosa/dialect/ops/custom.py

@@ -31,7 +31,7 @@
 from collections.abc import Callable
 
 import torch
-from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op
+from executorch.backends.arm.tosa.dialect.ops_registration import register_tosa_op
 
 from executorch.backends.arm.tosa.specification import (
     get_context_spec,
@@ -132,7 +132,7 @@ def run_registered_fake_tosa_impl(
     return outputs
 
 
-@register_fake_tosa_op(
+@register_tosa_op(
     "CUSTOM(Tensor[] inputs, str operator_name, str domain_name, int[] implementation_attrs) -> Tensor[]",
     TosaSpecification.all_versions_and_profiles(),
 )

backends/arm/tosa/dialect/ops/depthwise_conv2d.py

@@ -7,15 +7,15 @@
 
 import torch
 from executorch.backends.arm.tosa.dialect.ops.conv2d import validate_conv2d_args_dtypes
-from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op
+from executorch.backends.arm.tosa.dialect.ops_registration import register_tosa_op
 
 from executorch.backends.arm.tosa.specification import (
     get_context_spec,
     TosaSpecification,
 )
 
 
-@register_fake_tosa_op(
+@register_tosa_op(
     "DEPTHWISE_CONV2D(Tensor input, "
     "Tensor weight, "
     "Tensor bias, "