Use default overloads when calling custom ops

If a node is created without specifying an overload,
A OpOverloadPacket is created, rather than an OpOverload.
This works in a GraphModule, but the OpOverloadPacket is not a valid
operator type in the _EXIREdgeDialectVerifier, which means
that Edge ExportedPrograms can't contain a GraphModule with
such ops.

In short, specifying using the default overload seems
to be the more correct way of creating a custom
operator.

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

Author: Erik-Lundell

backends/arm/_passes/annotate_channels_last_dim_order_pass.py

@@ -118,7 +118,7 @@ def insert_input_transpose(node, input_node, graph_module):
         with graph_module.graph.inserting_before(node):
             permute_node = create_node(
                 graph_module.graph,
-                torch.ops.passthrough_to_tosa._transpose,
+                torch.ops.passthrough_to_tosa._transpose.default,
                 args=(
                     input_node,
                     list(AnnotateChannelsLastDimOrder.NHWC_inverse_order),
@@ -137,7 +137,7 @@ def insert_output_transpose(node, graph_module):
         with graph_module.graph.inserting_after(node):
             permute_node = create_node(
                 graph_module.graph,
-                torch.ops.passthrough_to_tosa._transpose,
+                torch.ops.passthrough_to_tosa._transpose.default,
                 args=(node, list(AnnotateChannelsLastDimOrder.NHWC_order)),
             )
             permute_node.meta["tosa_dim_order"] = (

backends/arm/_passes/insert_table_ops.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
+# Copyright 2024-2025 Arm Limited and/or its affiliates.
 # All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -92,7 +92,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
             with graph_module.graph.inserting_before(node):
                 table_node = create_node(
                     graph=graph_module.graph,
-                    op_target=torch.ops.tosa._table,
+                    op_target=torch.ops.tosa._table.default,
                     args=(node.args[0],),
                 )
                 assert len(input_qparams) == 1
@@ -104,7 +104,11 @@ def call(self, graph_module: GraphModule) -> PassResult:
                     out_quantargs=output_qparams[0],
                 )
                 # Register buffer in self.exported_program.state_dict
-                self.register_buffer(buffer_name=table_node.name, buffer=buffer)
+                # When the graph is retraced, the implementation _table is used and the suffix _default disappears from the node name
+                # Remove it here to make it possible to find in the node_visitor
+                self.register_buffer(
+                    buffer_name=table_node.name.replace("_default", ""), buffer=buffer
+                )
                 node.replace_all_uses_with(table_node)
             graph_module.graph.erase_node(node)
             table_node.meta["input_qparams"] = input_qparams

backends/arm/operators/op_table.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
+# Copyright 2024-2025 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -21,7 +21,7 @@
 
 @register_node_visitor
 class TableVisitor(NodeVisitor):
-    target = "_table"
+    target = "_table.default"
 
     def define_node(
         self,

backends/arm/operators/op_transpose.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Arm Limited and/or its affiliates.
+# Copyright 2024-2025 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -25,7 +25,7 @@ class TransposeVisitor(NodeVisitor):
     Inserts a TOSA TRANSPOSE.
     """
 
-    target = "_transpose"
+    target = "_transpose.default"
 
     def define_node(
         self,