Remove identity operators from oneDNN optimized graph (#20712)

* Remove identity operators from inference graph

* Add new line at EOF

* review fixes

* Small refactor & review

* remove commented fragment

Co-authored-by: Bartlomiej Gawrych <barlomiej.gawrych@intel.com>

Author: bgawrych

src/operator/subgraph/dnnl/dnnl_identity_property.h

@@ -0,0 +1,168 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file dnnl_identity_property.cc
+ * \brief Graph property for removing identity operators
+ */
+
+#ifndef MXNET_OPERATOR_SUBGRAPH_DNNL_DNNL_IDENTITY_PROPERTY_H_
+#define MXNET_OPERATOR_SUBGRAPH_DNNL_DNNL_IDENTITY_PROPERTY_H_
+#if MXNET_USE_ONEDNN == 1
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include "../common.h"
+#include "../../nn/dropout-inl.h"
+#include "dnnl_subgraph_base-inl.h"
+
+namespace mxnet {
+namespace op {
+
+class SgDNNLIdentitySelector : public SubgraphSelectorV2 {
+ private:
+  std::vector<const BiDirectedNode*> matched_list_;
+
+ public:
+  bool Select(const BiDirectedNode& seed_node,
+              const std::shared_ptr<NodeAttr>& node_attr) override {
+    bool status = false;
+    if (seed_node.node->op() == Op::Get("_npi_copy")) {
+      status = true;
+    }
+
+    if (seed_node.node->op() == Op::Get("Dropout")) {
+      auto const& dropout_param = nnvm::get<DropoutParam>(seed_node.node->attrs.parsed);
+      if (dropout_param.mode == dropout::kTraining) {
+        status = true;
+      }
+    }
+
+    if (status) {
+      matched_list_.clear();
+      matched_list_.emplace_back(&seed_node);
+      return true;
+    }
+    return false;
+  }
+
+  bool SelectInput(const BiDirectedNode& n, const BiDirectedNode& input_node) override {
+    if (input_node.node->is_variable()) {
+      return false;
+    } else if (input_node.node->op()) {
+      matched_list_.emplace_back(&input_node);
+      return true;
+    }
+    return false;
+  }
+
+  bool SelectOutput(const BiDirectedNode& n, const BiDirectedNode& output_node) override {
+    return false;
+  }
+
+  std::vector<BiDirectedNode*> Filter(const std::vector<BiDirectedNode*>& candidates) override {
+    // candidates should contain only two nodes - custom node and identity node
+    if (candidates.size() == 2 && candidates.size() == matched_list_.size()) {
+      return candidates;
+    } else {
+      return std::vector<BiDirectedNode*>(0);
+    }
+  }
+
+  void Reset() override {
+    CHECK_GE(matched_list_.size(), 1);
+    auto new_selector = SgDNNLIdentitySelector();
+    new_selector.Select(*matched_list_[0], nullptr);
+    *this = new_selector;
+  }
+};
+
+inline bool IsIdentityNode(const nnvm::ObjectPtr node) {
+  return node->op() && (node->op() == Op::Get("_npi_copy") || node->op() == Op::Get("Dropout"));
+}
+
+class SgDNNLIdentityProperty : public SubgraphProperty {
+ public:
+  SgDNNLIdentityProperty() {}
+
+  static SubgraphPropertyPtr Create() {
+    static const std::string& name = "DNNL Identity optimization pass";
+    auto property                  = std::make_shared<SgDNNLIdentityProperty>();
+    property->SetAttr<std::string>("property_name", name);
+    property->SetAttr<bool>("inference_only", true);
+    return property;
+  }
+
+  nnvm::ObjectPtr CreateSubgraphNode(const nnvm::Symbol& sym,
+                                     const int subgraph_id = 0) const override {
+    nnvm::NodeEntry identity_node_entry;
+    for (auto entry : sym.outputs) {
+      if (IsIdentityNode(entry.node)) {
+        identity_node_entry = entry;
+      }
+    }
+
+    auto last_node = identity_node_entry.node;
+    nnvm::Symbol new_sym;
+    new_sym.outputs.emplace_back(last_node);
+
+    nnvm::ObjectPtr org_node;
+    DFSVisit(new_sym.outputs, [&](const nnvm::ObjectPtr& node) {
+      if (!IsIdentityNode(node)) {
+        org_node = node;
+      }
+    });
+
+    // Create copy of original node
+    nnvm::ObjectPtr n = nnvm::Node::Create();
+    n->attrs          = org_node->attrs;
+    CHECK(n->op());
+    n->op()->attr_parser(&(n->attrs));
+    return n;
+  }
+
+  void ConnectSubgraphOutputs(const nnvm::ObjectPtr n,
+                              std::vector<nnvm::NodeEntry*>* output_entries) const override {
+    // output of identity must be connected as output of operator before identity
+    // e.g. for:        /--index 0--> custom_op
+    //         (n) slice
+    //                  \--index 1--> Dropout --index 0--> OUT_NODE
+    //  for OUT_NODE index 0 must be changed to index 1
+    for (int i = 0; i < output_entries->size(); ++i) {
+      auto out_node = output_entries->at(i)->node;
+      if (IsIdentityNode(out_node)) {
+        output_entries->at(i)->index = out_node->inputs[0].index;
+      }
+      output_entries->at(i)->node = n;
+    }
+  }
+
+  SubgraphSelectorV2Ptr CreateSubgraphSelectorV2() const override {
+    auto selector = std::make_shared<SgDNNLIdentitySelector>();
+    return selector;
+  }
+};
+
+}  // namespace op
+}  // namespace mxnet
+
+#endif  // if MXNET_USE_ONEDNN == 1
+#endif  // MXNET_OPERATOR_SUBGRAPH_DNNL_DNNL_IDENTITY_PROPERTY_H_

src/operator/subgraph/dnnl/dnnl_subgraph_base-inl.h

@@ -31,7 +31,7 @@ static inline bool SupportDNNLAttr(const std::shared_ptr<NodeAttr>& node_attr) {
     return (node_attr->dispatch_mode == DispatchMode::kFComputeEx) &&
            (node_attr->itype[0] == mshadow::kFloat32 ||
             node_attr->itype[0] == mshadow::kBfloat16) &&
-           (ndim == 1 || ndim == 2 || ndim == 4 || ndim == 5);
+           (ndim >= 1 && ndim <= 5);
   } else {
     return true;
   }

src/operator/subgraph/dnnl/dnnl_subgraph_property.cc

@@ -23,6 +23,7 @@
 #include "dnnl_bn_relu_property.h"
 #include "dnnl_conv_property.h"
 #include "dnnl_fc_property.h"
+#include "dnnl_identity_property.h"
 #include "dnnl_post_quantize_align_scale_property.h"
 #include "dnnl_post_quantize_property.h"
 #include "dnnl_transformer_qk_property.h"
@@ -35,6 +36,7 @@ MXNET_REGISTER_SUBGRAPH_BACKEND(ONEDNN)
     .set_attr("enable", DNNLEnvSet())
     .set_attr("context", Context::CPU());
 
+MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN, SgDNNLIdentityProperty);
 MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN, SgDNNLConvProperty);
 MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN, SgDNNLFCProperty);
 MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN, SgDNNLBNReLUProperty);
@@ -44,6 +46,7 @@ MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN, SgDNNLBatchDotProperty);
 
 MXNET_REGISTER_SUBGRAPH_BACKEND(ONEDNN_QUANTIZE).set_attr("context", Context::CPU());
 
+MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN_QUANTIZE, SgDNNLIdentityProperty);
 MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN_QUANTIZE, SgDNNLConvProperty).set_attr("quantize", true);
 MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN_QUANTIZE, SgDNNLFCProperty).set_attr("quantize", true);
 MXNET_REGISTER_SUBGRAPH_PROPERTY(ONEDNN_QUANTIZE, SgDNNLTransformerQKProperty);

tests/python/dnnl/subgraphs/test_fc_subgraph.py

@@ -200,3 +200,26 @@ def forward(self, x):
   attrs = {'fc': {}}
   net = MultiOutputFC()
   check_fusion(net, data_shape, attrs, check_quantization=flatten)
+
+
+@mx.util.use_np
+@pytest.mark.parametrize('identity_node', ['dropout', 'copy'])
+def test_fc_identity_eltwise(identity_node):
+  class FCIdentityEltwise(nn.HybridBlock):
+    def __init__(self, identity_node, **kwargs):
+      super(FCIdentityEltwise, self).__init__(**kwargs)
+      self.fc = nn.Dense(units=64, use_bias=False, weight_initializer=None, flatten=True)
+      self.identity_node = identity_node
+    def forward(self, x):
+      fc_out = self.fc(x)
+      if self.identity_node == 'copy':
+        fc_out = mx.np.copy(fc_out)
+      else:
+        fc_out = mx.npx.dropout(fc_out)
+      out = mx.npx.activation(fc_out, act_type='relu')
+      return out
+
+  data_shape = (64, 4, 10, 10)
+  attrs = {'fc': {'with_eltwise': 'true'}}
+  net = FCIdentityEltwise(identity_node)
+  check_fusion(net, data_shape, attrs, check_quantization=False)