Revert "Arm backend: Add real implementation for TOSA dialect ops (re-land)" (pytorch#20670)

Reverts pytorch#20537

This breaks some internal tests due to missing BUCK dependencies on the
tosa reference model. I tried to fix forward by just adding those deps
(pytorch#20642), but it looks like we'll need to do lazy imports here due to
some binary size expansion. I'll tackle the reland.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell

Author: rascani

backends/arm/_passes/fuse_constant_ops_pass.py

@@ -92,20 +92,6 @@ def _is_tosa_dialect_op(target) -> bool:
             or "<EdgeOpOverload: tosa." in target_str
         )
 
-    @staticmethod
-    def _has_real_tosa_dialect_impl(target) -> bool:
-        schema = getattr(target, "_schema", None)
-        op_name = getattr(schema, "name", None)
-        if op_name is None:
-            return False
-
-        try:
-            return torch._C._dispatch_has_kernel_for_dispatch_key(
-                op_name, "CompositeExplicitAutograd"
-            )
-        except RuntimeError:
-            return False
-
     @staticmethod
     def _arg_contains_symbolic_shape(arg) -> bool:
         if isinstance(arg, torch.fx.Node):
@@ -211,31 +197,19 @@ def resolve_arg(arg, arg_index=None):
 
         return True
 
-    def maybe_delete(self, input_nodes_to_maybe_delete):
-        for input_node in input_nodes_to_maybe_delete:
-            if input_node.meta.get("is_input", False):
-                # Never delete submodule inputs, they need to match the parameters from the outer module.
-                continue
-            if len(input_node.users) == 0:
-                self._delete_constant_placeholder(input_node)
-
     def call(self, graph_module):
         modified = False
         input_nodes_to_maybe_delete = set()
         for node in graph_module.graph.nodes:
             if node.op != "call_function":
                 continue
-            if node.target == exir_ops.backend.tosa.RESCALE.default:
-                # Leave fusing of RESCALES to the compiler.
-                continue
-            if self._is_tosa_dialect_op(
-                node.target
-            ) and not self._has_real_tosa_dialect_impl(node.target):
-                continue
-            # 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):
+            # 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)
+            ):
                 continue
 
             input_nodes = node.all_input_nodes
@@ -267,7 +241,10 @@ def call(self, graph_module):
 
         if modified:
             graph_module.graph.eliminate_dead_code()
-            self.maybe_delete(input_nodes_to_maybe_delete)
+            for input_node in input_nodes_to_maybe_delete:
+                if len(input_node.users) == 0:
+                    self._delete_constant_placeholder(input_node)
+
             graph_module = super().call(graph_module).graph_module
 
         return PassResult(graph_module, modified)

backends/arm/operators/node_visitor.py

@@ -31,7 +31,6 @@
     TosaSpecMapping,
 )
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -247,15 +246,3 @@ def get_node_visitors(*args) -> Dict[str, NodeVisitor]:
         node_visitors[target] = visitor(*args)
 
     return node_visitors
-
-
-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:
-            return node_visitor_cls(tosa_spec)
-
-    raise ValueError(f"No {target} NodeVisitor registered for {tosa_spec}")

backends/arm/test/misc/tosa_dialect/test_tosa_gather.py

@@ -13,7 +13,7 @@
 from torch._subclasses.fake_tensor import FakeTensorMode
 
 
-def test_identity_tosa_FP_fake() -> None:
+def test_identity_tosa_FP() -> None:
     sample_input = torch.randn((1, 2, 3, 4), dtype=torch.float32)
 
     with TosaLoweringContext(
@@ -23,12 +23,3 @@ def test_identity_tosa_FP_fake() -> None:
 
         assert output.dtype == sample_input.dtype
         assert tuple(output.shape) == tuple(sample_input.shape)
-
-
-def test_identity_tosa_FP_real() -> None:
-    sample_input = torch.randn((1, 2, 3, 4), dtype=torch.float32)
-
-    with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")):
-        output = exir_ops.backend.tosa.IDENTITY.default(sample_input)
-
-    torch.testing.assert_close(output, sample_input)

backends/arm/test/misc/tosa_dialect/test_tosa_identity.py

@@ -35,6 +35,7 @@ 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 = (
@@ -61,4 +62,5 @@ 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()