Arm backend: Move TOSA operators to dialect (#13408)

### Summary
Move rescale, table and transpose TOSA operators to new implementation.

### Test plan
Tested through existing CI unit tests.

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

Author: per

backends/arm/_passes/annotate_channels_last_dim_order_pass.py

@@ -14,36 +14,12 @@
 from executorch.backends.arm.tosa_utils import is_consumer_node_depthwise_conv2d
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
-from torch.library import impl, Library
-
-# Define lib with passthrough operators. The operators have no real meaning in edge IR
-# except for argument validaiton and a passthrough output. The operators will be used
-# when lowering to TOSA, e.g. a passthrough_to_tosa._transpose will not affect
-# the edge IR graph but will be lowered to a TOSA-TRANSPOSE.
-lib = Library("passthrough_to_tosa", "DEF")
-# For certain operators we need the data in a specific data format. Changing tosa_dim_order
-# is not sufficient as we also need transpose the data.
-# By utilizing an edge IR passthrough operator we can keep the edge program in
-# channels-first/contiguous and get the desired behavior in the TOSA lowering.
-lib.define("_transpose(Tensor self, int[] dim_order) -> Tensor")
-
-
-@impl(lib, "_transpose")
-def _transpose_impl(*args, **kwargs):
-    # Validate length of dim_order array
-    dim = args[1]
-    if len(dim) != 4 and len(dim) != 5:
-        raise ValueError(
-            f"Dim order length must be either 4 or 5, got {len(dim)}: {dim}"
-        )
-    # Pass-through in edge-IR
-    return args[0]
 
 
 class AnnotateChannelsLastDimOrder(ExportPass):
     """
     Annotates each node with a tosa_dim_order. tosa_dim_order can be seen as a channels-last dim-order
-    that in most cases will be (0, 2, 3, 1) for nodes with 4D-shapes. The pass also inserts passthrough_to_tosa._transpose
+    that in most cases will be (0, 2, 3, 1) for nodes with 4D-shapes. The pass also inserts backend.tosa.TRANSPOSE
     when a transition between 3D and 4D/5D tensors happen.
     The annotated tosa_dim_order is used to permute the node's shape such that it gives a TOSA-compliant shape.
     """
@@ -119,7 +95,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.default,
+                exir_ops.backend.tosa.TRANSPOSE.default,
                 args=(
                     input_node,
                     list(
@@ -141,7 +117,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.default,
+                exir_ops.backend.tosa.TRANSPOSE.default,
                 args=(
                     node,
                     list(

backends/arm/_passes/fuse_constant_ops_pass.py

@@ -107,7 +107,7 @@ def call(self, graph_module):
         for node in graph_module.graph.nodes:
             if node.op != "call_function":
                 continue
-            if node.target == torch.ops.tosa._table.default:
+            if node.target == exir_ops.backend.tosa.TABLE.default:
                 continue
 
             input_nodes = node.all_input_nodes

backends/arm/_passes/insert_rescales_pass.py

@@ -3,70 +3,25 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import logging
 from copy import copy
 from typing import cast
 
-import torch
 from executorch.backends.arm._passes.arm_pass_utils import create_node
 from executorch.backends.arm._passes.quant_args import QuantArgs
 from executorch.backends.arm.constants import DQ_OPS, Q_OPS
+from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
-from torch import Tensor
 from torch.fx import GraphModule, Node
-from torch.library import custom_op, register_fake
-
-logger = logging.getLogger(__name__)
-
-
-@custom_op("tosa::_rescale", mutates_args=())  # type: ignore[misc]
-def rescale(
-    x: Tensor, dtype: torch.dtype, scale: float, in_zp: int, out_zp: int
-) -> Tensor:
-    logger.warning(
-        "Ran default implementation of tosa::_rescale."
-        "This op is meant to always be inserted inside a partition and a correct default implementation is not implemented."
-    )
-    # Clone is needed to not return reference when rescaling to same dtype.
-    # This is a neccessary requirement for non-mutating custom ops.
-    return x.to(dtype=dtype).clone()
-
-
-@register_fake("tosa::_rescale")  # type: ignore[misc]
-def rescale_fake(
-    x: Tensor, dtype: torch.dtype, scale: float, in_zp: int, out_zp: int
-) -> Tensor:
-    """Casts the input tensor to dtype `dtype` to produce the correct tensor meta for a _rescale op.
-    Additionally validates TOSA constraints of a RESCALE op.
-    """
-    if dtype not in (torch.int32, torch.int8, torch.int16):
-        raise NotImplementedError(
-            f"tosa::rescale currently only supports int32, int16 and int8, not {dtype}"
-        )
-    if dtype in (torch.int32, torch.int16) and out_zp != 0:
-        raise ValueError(
-            f"TOSA requires output_zp to be zero when the output dtype is {dtype}."
-        )
-    if x.dtype in (torch.int32, torch.int16) and in_zp != 0:
-        raise ValueError(
-            f"TOSA requires input_zp to be zero when the input dtype is {dtype}"
-        )
-    if x.dtype == torch.int8 and not -128 <= in_zp <= 127:
-        raise ValueError(f"{in_zp=} outside valid range (-128,127) for int8.")
-    if dtype == torch.int8 and not -128 <= out_zp <= 127:
-        raise ValueError(f"{out_zp=} outside valid range (-128,127) for int8.")
-
-    return x.to(dtype=dtype).clone()
 
 
 class InsertRescalePass(ExportPass):
     """Finds patterns of dq -> q, and replaces them
-    with passthrough_to_tosa::rescales.
+    with backend dialect tosa::RESCALE op.
 
-    Does not garantuee that the dtypes and zero points are valid
+    Does not guarantee that the dtypes and zero points are valid
     in TOSA, that is the job of the quantization annotator that
     produced the dq and q nodes. The TOSA constraints are validated
-    in the fake implementation of passthrough_to_tosa:rescale.
+    in the fake implementation of.
     """
 
     def fold_dq_q_to_rescale(self, node: Node, user: Node, graph_module: GraphModule):
@@ -77,7 +32,7 @@ def fold_dq_q_to_rescale(self, node: Node, user: Node, graph_module: GraphModule
         with graph_module.graph.inserting_before(node):
             rescale_node = create_node(
                 graph_module.graph,
-                torch.ops.tosa._rescale.default,
+                exir_ops.backend.tosa.RESCALE.default,
                 (
                     node.all_input_nodes[0],
                     q_args.dtype,

backends/arm/_passes/insert_table_ops.py

@@ -11,26 +11,17 @@
 import torch
 from executorch.backends.arm._passes.arm_pass_utils import create_node
 from executorch.backends.arm._passes.quant_args import QuantArgs
+from executorch.backends.transforms.utils import create_constant_placeholder
+
 from executorch.exir import ExportedProgram
 
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 
 from executorch.exir.pass_base import ExportPass, PassResult
+from torch.export.graph_signature import InputKind
 from torch.fx import GraphModule
 from torch.fx.node import Node
-from torch.library import impl, Library
-
-lib = Library("tosa", "DEF")
-lib.define("_table(Tensor self) -> Tensor")
-
-
-@impl(lib, "_table")
-def _table_impl(*args, **kwargs):  # pyre-ignore
-    in_dtype = args[0].dtype
-    if in_dtype == torch.int8:
-        return args[0]
-    return args[0].to(dtype=torch.int32)
 
 
 class TableOps:
@@ -242,13 +233,8 @@ def call(self, graph_module: GraphModule) -> PassResult:
                 # We only want to replace the node if it's quantized
                 continue
             # Create table node
-            with graph_module.graph.inserting_before(node):
-                table_node = create_node(
-                    graph=graph_module.graph,
-                    op_target=torch.ops.tosa._table.default,
-                    args=(node.args[0],),
-                )
-                output_node = table_node
+            insert_pos = list(node.graph.nodes)[0]
+            with graph_module.graph.inserting_before(insert_pos):
                 # Expect exactly one quantization parameter for input and output
                 if len(input_qparams) != 1:
                     raise ValueError(
@@ -268,27 +254,37 @@ def call(self, graph_module: GraphModule) -> PassResult:
                     out_quantargs=output_qparams[0],
                 )
                 # Register buffer in self.exported_program.state_dict
-                # 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
+                const_table_node = create_constant_placeholder(
+                    exp_program=self.exported_program,
+                    graph=node.graph,
+                    kind=InputKind.BUFFER,
+                    name=node.name + "_table_constant",
+                    data=buffer,
+                    persistent_buffer=True,
                 )
 
+            # Create table node
+            with graph_module.graph.inserting_before(node):
+                table_op_node = create_node(
+                    graph=graph_module.graph,
+                    op_target=exir_ops.backend.tosa.TABLE.default,
+                    args=(node.args[0], const_table_node),
+                )
+                output_node = table_op_node
+
                 if lshift != 0:
                     scale = 2.0**lshift
                     rescale_node = create_node(
                         graph=graph_module.graph,
-                        op_target=torch.ops.tosa._rescale.default,
-                        args=(table_node, output_qparams[0].dtype, scale, 0, 0),
+                        op_target=exir_ops.backend.tosa.RESCALE.default,
+                        args=(table_op_node, output_qparams[0].dtype, scale, 0, 0),
                     )
                     output_node = rescale_node
 
                 node.replace_all_uses_with(output_node)
-
             graph_module.graph.erase_node(node)
-
-            output_node.meta["input_qparams"] = input_qparams
-            output_node.meta["output_qparams"] = output_qparams
+            table_op_node.meta["input_qparams"] = input_qparams
+            table_op_node.meta["output_qparams"] = output_qparams
             modified = True
 
         if modified:

backends/arm/operators/op_rescale.py

@@ -23,8 +23,8 @@
 
 
 @register_node_visitor
-class RescaleVisitor_INT(NodeVisitor):
-    target = "_rescale.default"
+class RescaleVisitor(NodeVisitor):
+    target = "tosa.RESCALE.default"
 
     tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+INT")]
 

backends/arm/operators/op_table.py

@@ -23,7 +23,7 @@
 
 @register_node_visitor
 class TableVisitor(NodeVisitor):
-    target = "_table.default"
+    target = "tosa.TABLE.default"
 
     tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+INT")]
 
@@ -36,7 +36,7 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts  # type: ignore
 
-        validate_num_inputs(self.target, inputs, 1)
+        validate_num_inputs(self.target, inputs, 2)
         validate_valid_dtype(
             self.target, inputs, [ts.DType.INT8, ts.DType.INT16], output.tosa_spec
         )
@@ -45,12 +45,12 @@ def define_node(
         if inputs[0].dtype == ts.DType.INT16:
             validate_valid_dtype(self.target, output, ts.DType.INT32, output.tosa_spec)
 
-        if node.name not in self._exported_program.state_dict.keys():  # type: ignore[union-attr]
+        if inputs[1].name not in self._exported_program.state_dict.keys():  # type: ignore[union-attr]
             raise RuntimeError(
                 f"Did not find key {node.name} in state_dict {self._exported_program.state_dict.keys()}."
             )
 
-        table = self._exported_program.state_dict[node.name]
+        table = self._exported_program.state_dict[inputs[1].name]  # type: ignore[union-attr]
 
         table_tensor_name = node.name + "_table"
         tosa_graph.addConst(

backends/arm/operators/op_transpose.py

@@ -24,12 +24,12 @@
 @register_node_visitor
 class TransposeVisitor(NodeVisitor):
     """
-    This node visitor targets the _transpose op defined in the
-    passthrough_to_tosa library. Used when switching between tosa_dim_orders.
+    This node visitor targets the tosa::TRANSPOSE op defined in the
+    TOSA backend dialect. Used when switching between tosa_dim_orders.
     Inserts a TOSA TRANSPOSE.
     """
 
-    target = "_transpose.default"
+    target = "tosa.TRANSPOSE.default"
 
     tosa_specs = NodeVisitor.tosa_specs
 

backends/arm/test/passes/test_insert_table_ops_pass.py

@@ -33,13 +33,13 @@ def test_insert_table_tosa_INT(test_data: input_t):
         module,
         test_data,
         quantize=True,
-        ops_before_pass={},
+        ops_before_pass={"executorch_exir_dialects_edge__ops_aten_sigmoid_default": 1},
         ops_after_pass={
             "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 1,
             "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 1,
-            "tosa._table": 1,
+            "backend__ops_tosa_TABLE_default": 1,
         },
-        ops_not_after_pass=["aten_sigmoid_default"],
+        ops_not_after_pass=["executorch_exir_dialects_edge__ops_aten_sigmoid_default"],
         pass_list=[FoldAndAnnotateQParamsPass],
         passes_with_exported_program=[InsertTableOpsPass],
     )