Fix identity fuse for oneDNN (#20767)

* Fix identity fuse

* add new line

* apply review comments

* fix

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

Author: bgawrych

src/operator/subgraph/dnnl/dnnl_identity_property.h

@@ -40,6 +40,7 @@ namespace op {
 class SgDNNLIdentitySelector : public SubgraphSelectorV2 {
  private:
   std::vector<const BiDirectedNode*> matched_list_;
+  bool pattern_found = false;
 
  public:
   bool Select(const BiDirectedNode& seed_node,
@@ -65,10 +66,11 @@ class SgDNNLIdentitySelector : public SubgraphSelectorV2 {
   }
 
   bool SelectInput(const BiDirectedNode& n, const BiDirectedNode& input_node) override {
-    if (input_node.node->is_variable()) {
+    if (pattern_found || input_node.node->is_variable()) {
       return false;
     } else if (input_node.node->op()) {
       matched_list_.emplace_back(&input_node);
+      pattern_found = true;
       return true;
     }
     return false;
@@ -80,7 +82,8 @@ class SgDNNLIdentitySelector : public SubgraphSelectorV2 {
 
   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()) {
+    if (pattern_found && candidates.size() == matched_list_.size()) {
+      CHECK_EQ(candidates.size(), 2);
       return candidates;
     } else {
       return std::vector<BiDirectedNode*>(0);
@@ -134,8 +137,10 @@ class SgDNNLIdentityProperty : public SubgraphProperty {
     // 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));
+    CHECK(n->op()) << "WRTF";
+    if (n->op()->attr_parser) {
+      n->op()->attr_parser(&(n->attrs));
+    }
     return n;
   }
 

tests/python/dnnl/subgraphs/test_fc_subgraph.py

@@ -208,18 +208,27 @@ 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.fc1 = nn.Dense(units=64, use_bias=False, weight_initializer=None, flatten=True)
+      self.fc2 = 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)
+      out = self.fc1(x)
       if self.identity_node == 'copy':
-        fc_out = mx.np.copy(fc_out)
+        out = mx.np.copy(out)
       else:
-        fc_out = mx.npx.dropout(fc_out)
-      out = mx.npx.activation(fc_out, act_type='relu')
+        out = mx.npx.dropout(out)
+      out = mx.npx.activation(out, act_type='relu')
+      out = self.fc2(out)
+      if self.identity_node == 'copy':
+        out = mx.np.copy(out)
+      else:
+        out = mx.npx.dropout(out)
+      out = mx.npx.activation(out, act_type='relu')
       return out
 
   data_shape = (64, 4, 10, 10)
-  attrs = {'fc': {'with_eltwise': 'true'}}
+  attrs = {'sg_onednn_fully_connected_eltwise_0' : {'with_eltwise': 'true'},
+           'sg_onednn_fully_connected_eltwise_1' : {'with_eltwise': 'true'}}
   net = FCIdentityEltwise(identity_node)
   check_fusion(net, data_shape, attrs, check_quantization=False)