Remove no-op clones in xnnpack (pytorch#15884)

Summary: Pull Request resolved: pytorch#15884

Differential Revision: D87405074

Author: GregoryComer

backends/transforms/remove_clone_ops.py

@@ -25,8 +25,9 @@ class RemoveCloneOpsTransform(ExportPass):
         exir_ops.edge.dim_order_ops._clone_dim_order.default,
     }
 
-    def __init__(self) -> None:
+    def __init__(self, preserve_input_output_copies: bool = True) -> None:
         super().__init__()
+        self._preserve_input_output_copies = preserve_input_output_copies
 
     def _remove(self, graph_module: torch.fx.GraphModule) -> None:
         dequant_nodes = []
@@ -38,6 +39,11 @@ def _remove(self, graph_module: torch.fx.GraphModule) -> None:
             if self._is_non_identity_clone(n):
                 continue
 
+            # If preserve_input_output_copies is set, don't remove clones that directly
+            # copy from input to output.
+            if self._is_input_output_copy(n) and self._preserve_input_output_copies:
+                continue
+
             to_be_removed = n
             for user_n in list(n.users.keys()):
                 user_n.replace_input_with(n, n.args[0])
@@ -76,3 +82,16 @@ def _is_non_identity_clone(self, node: torch.fx.Node) -> bool:
             )
 
         return False
+
+    def _is_input_output_copy(self, node: torch.fx.Node) -> bool:
+        """Return True if the node input is a graph input and output goes into an output node."""
+
+        input_node = node.args[0]
+        if input_node.op != "placeholder":
+            return False
+
+        for users in node.users:
+            if users.op == "output":
+                return True
+
+        return False

backends/xnnpack/_passes/TARGETS

@@ -8,6 +8,7 @@ runtime.python_library(
     deps = [
         "//caffe2:torch",
         "//executorch/backends/transforms:addmm_mm_to_linear",
+        "//executorch/backends/transforms:remove_clone_ops",
         "//executorch/backends/transforms:lib",
         "//executorch/backends/xnnpack/partition:partitioner_graphs",
         "//executorch/backends/xnnpack/serialization:xnnpack_schema",

backends/xnnpack/_passes/__init__.py

@@ -6,6 +6,8 @@
 
 from typing import List, Optional, Type
 
+from executorch.backends.transforms.remove_clone_ops import RemoveCloneOpsTransform
+
 from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass
 
 from executorch.backends.xnnpack._passes.channels_last_tagged_reshape_pass import (
@@ -58,6 +60,7 @@ def __init__(
         if not passes:
             # All the XNNPACK passes
             self.passes = [
+                RemoveCloneOpsTransform,
                 # TODO - remove this pass once we have a better support for dim_order ops lowering
                 DimOrderOpsRevertPass,
                 ConvertToUpsampleBilinear2d,

backends/xnnpack/operators/__init__.py

@@ -14,6 +14,7 @@
     op_cat,
     op_ceiling,
     op_clamp,
+    op_clone,
     op_conv2d,
     op_div,
     op_dynamic_dequantize_ops,

backends/xnnpack/operators/op_clone.py

@@ -0,0 +1,58 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict
+
+import torch
+from executorch.backends.xnnpack.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    XNNCopy,
+    XNNGraph,
+    XNode,
+)
+from executorch.backends.xnnpack.utils.utils import get_input_node
+
+
+@register_node_visitor
+class CloneVisitor(NodeVisitor):
+    target = "aten.clone.default"
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        xnn_graph: XNNGraph,
+        vals_to_ids: Dict[torch.fx.Node, int],
+        debug_handle: int,
+    ) -> None:
+        self.define_nodes_tensor_inputs_outputs(node, xnn_graph, vals_to_ids)
+
+        # Sanity check that the input and output dim order are the same. We don't
+        # handle dim order conversions yet.
+        dim_order = node.kwargs.get("dim_order", None)
+        input_meta = node.args[0].meta["val"]
+        assert(dim_order is None or list(input_meta.dim_order() == dim_order))
+
+        # input
+        input_id = vals_to_ids[get_input_node(node, 0)]
+
+        # output
+        output_id = vals_to_ids[node]
+
+        ser_node = XNode(
+            xnode_union=XNNCopy(
+                input_id=input_id,
+                output_id=output_id,
+                flags=0,
+            ),
+            debug_handle=debug_handle,
+        )
+        xnn_graph.xnodes.append(ser_node)

backends/xnnpack/partition/config/__init__.py

@@ -22,6 +22,7 @@
     CatConfig,
     CeilConfig,
     ClampConfig,
+    CloneDimOrderConfig,
     ConstantPadConfig,
     DeQuantizedPerTensorConfig,
     DivConfig,
@@ -77,6 +78,7 @@
     BMMConfig,
     CatConfig,
     CeilConfig,
+    CloneDimOrderConfig,
     ConstantPadConfig,
     ConvolutionConfig,
     ClampConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -643,3 +643,24 @@ class SinConfig(GenericNodePartitionerConfig):
 
     def supported_precision_types(self) -> List[ConfigPrecisionType]:
         return [ConfigPrecisionType.FP32]
+
+class CloneDimOrderConfig(GenericNodePartitionerConfig):
+    target_name = "_clone_dim_order.default"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+    def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
+        if not self.check_common_constraints(node, ep):
+            return False
+
+        # Only partition no-op _clone_dim_order nodes (output dim order = input).
+        # We can relax this in the future.
+        # This is also a conservative check and doesn't consider ambiguity.
+        dim_order = node.kwargs.get("dim_order", None)
+        input_meta = node.args[0].meta["val"]
+        if dim_order is not None and list(input_meta.dim_order()) != dim_order:
+            why(node, reason="Only dim-order preserving clones are supported.")
+            return False
+
+        return True

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -1459,6 +1459,34 @@ Error defineBatchMatrixMultiplyNode(
   return Error::Ok;
 }
 
+/*
+ * Defines a code node in the XNN subgraph.
+ */
+Error defineCopyNode(
+    xnn_subgraph_t subgraph_ptr,
+    const std::unordered_map<uint32_t, uint32_t>& remapped_ids,
+    const NodePtr node,
+    const fb_xnnpack::XNNGraph* graph) noexcept {
+  MAYBE_UNUSED(graph);
+
+  auto graph_node = node->xnode_union_as_XNNCopy();
+
+  xnn_status status = xnn_define_copy(
+      subgraph_ptr,
+      remapped_ids.at(graph_node->input_id()),
+      remapped_ids.at(graph_node->output_id()),
+      graph_node->flags());
+
+  ET_CHECK_OR_RETURN_ERROR(
+      status == xnn_status_success,
+      Internal,
+      "Failed to create copy node %i with code: %s",
+      node->debug_handle(),
+      xnn_status_to_string(status));
+
+  return Error::Ok;
+}
+
 /*
 Returns not Implemented Error code. This function is meant to be
 called when the compiler encountes a XNodeType from the flatbuffer
@@ -1763,6 +1791,7 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) {
     _DEFINE(Concatenate5)
     _DEFINE(StaticSlice)
     _DEFINE(BatchMatrixMultiply)
+    _DEFINE(Copy)
     case fb_xnnpack::XNodeUnion::NONE:
     default: // Adding here as a catch all, just in case
       return &defineNotImplementedNode;

backends/xnnpack/serialization/runtime_schema.fbs

@@ -157,6 +157,7 @@ union XNodeUnion {
   XNNTanh: _XNNNode1x1,
   XNNExp: _XNNNode1x1,
   XNNSin: _XNNNode1x1,
+  XNNCopy: _XNNNode1x1,
 }
 
 union XValueUnion {

backends/xnnpack/serialization/schema.fbs

@@ -153,6 +153,7 @@ union XNodeUnion {
   XNNTanh: _XNNNode1x1,
   XNNExp: _XNNNode1x1,
   XNNSin: _XNNNode1x1,
+  XNNCopy: _XNNNode1x1,
 }
 
 union XValueUnion {