From 772746c0fc2cd42fe01ff81e09352f59cf20f2e5 Mon Sep 17 00:00:00 2001
From: lidanqing <danqing.li@intel.com>
Date: Sun, 28 Jun 2020 08:13:45 +0200
Subject: [PATCH] [oneDNN] Fix to inplace pass  (#24442) (#25182)

---
 .../ir/mkldnn/mkldnn_inplace_pass.cc          | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/paddle/fluid/framework/ir/mkldnn/mkldnn_inplace_pass.cc b/paddle/fluid/framework/ir/mkldnn/mkldnn_inplace_pass.cc
index 6590ef44f8..df2bc5af8b 100644
--- a/paddle/fluid/framework/ir/mkldnn/mkldnn_inplace_pass.cc
+++ b/paddle/fluid/framework/ir/mkldnn/mkldnn_inplace_pass.cc
@@ -32,6 +32,7 @@ void MKLDNNInPlacePass::ApplyImpl(ir::Graph* graph) const {
                           platform::errors::InvalidArgument(
                               "Pointer to graph argument should not be NULL."));
   std::unordered_map<std::string, std::string> original_output_names;
+  std::unordered_set<std::string> inplaced_vars;
   GraphPatternDetector gpd;
   patterns::MKLDNNInPlace mkldnn_inplace{gpd.mutable_pattern(),
                                          "mkldnn_inplace"};
@@ -94,6 +95,22 @@ void MKLDNNInPlacePass::ApplyImpl(ir::Graph* graph) const {
       VLOG(3) << "DNNL in-place pass FAIL: in-place var cannot "
                  "be an input to multiple operators";
       return;
+    } else {
+      // We will prevent in-place when
+      // input is used in other part of graph, unless it was a result of
+      // inplacing
+      // Allow to next op out reuse inpuit var, as this is the same chaing
+      if (inplaced_vars.find(current_op_in->Name()) == inplaced_vars.end()) {
+        for (const Node* n : graph->Nodes()) {
+          if ((n->id() != current_op_in->id()) &&
+              (n->id() != next_op_out->id()) &&
+              (n->Name() == current_op_in->Name())) {
+            VLOG(3) << "DNNL in-place pass FAIL var used in diffrent part of "
+                       "graph ";
+            return;
+          }
+        }
+      }
     }
 
     // If this op was alrady inplaced in previous pass placements
@@ -132,6 +149,8 @@ void MKLDNNInPlacePass::ApplyImpl(ir::Graph* graph) const {
     auto out_name = in_to_outs.begin()->second;
     current_op->Op()->SetOutput(
         out_name, std::vector<std::string>({current_op_out->Name()}));
+    // Record var name
+    inplaced_vars.insert(current_op_out->Name());
 
     // If next op in a line is doing inplace
     // then we need to update its output as well
-- 
GitLab