add warning and skip vars to mem opt passes (#16967)

test=release/1.4

Author: sneaxiy

paddle/fluid/framework/details/CMakeLists.txt

@@ -15,6 +15,8 @@ cc_library(alloc_continuous_space_for_grad_pass SRCS alloc_continuous_space_for_
 cc_library(fuse_adam_op_pass SRCS fuse_adam_op_pass.cc fuse_optimizer_op_pass.cc DEPS graph graph_helper)
 cc_library(fuse_sgd_op_pass SRCS fuse_sgd_op_pass.cc fuse_optimizer_op_pass.cc DEPS graph graph_helper)
 
+cc_library(record_skip_memory_opt_vars_pass SRCS record_skip_memory_opt_vars_pass.cc DEPS graph graph_helper)
+
 cc_library(variable_visitor SRCS variable_visitor.cc DEPS lod_tensor selected_rows)
 
 if(WITH_DISTRIBUTE)
@@ -114,4 +116,4 @@ cc_library(build_strategy SRCS build_strategy.cc DEPS
         fuse_relu_depthwise_conv_pass
         memory_optimize_pass lock_free_optimize_pass
         alloc_continuous_space_for_grad_pass fuse_all_reduce_op_pass
-        fuse_adam_op_pass fuse_sgd_op_pass)
+        fuse_adam_op_pass fuse_sgd_op_pass record_skip_memory_opt_vars_pass)

paddle/fluid/framework/details/build_strategy.cc

@@ -53,6 +53,9 @@ class ParallelExecutorPassBuilder : public ir::PassBuilder {
       viz_pass->Set<std::string>("graph_viz_path", new std::string(graph_path));
     }
 
+    // Note(zcd): record_skip_memory_opt_vars_pass should be the first pass.
+    AppendPass("record_skip_memory_opt_vars_pass");
+
     if (strategy_.enable_sequential_execution_) {
       VLOG(10) << "Add sequential_execution_pass";
       AppendPass("sequential_execution_pass");
@@ -320,3 +323,4 @@ USE_PASS(graph_to_program_pass);
 USE_PASS(fuse_adam_op_pass);
 USE_PASS(fuse_sgd_op_pass);
 USE_PASS(fuse_all_reduce_op_pass);
+USE_PASS(record_skip_memory_opt_vars_pass);

paddle/fluid/framework/details/inplace_op_pass.cc

@@ -303,7 +303,16 @@ void InplacePass::TryInplaceOpInputOutput(ir::Node* op,
     auto* in_node = view_.GetNodeByName(in_var_name, op->inputs);
     auto* out_node = view_.GetNodeByName(out_var_name, op->outputs);
 
-    VLOG(4) << "Try to inplace " << in_var_name << " with " << out_var_name;
+    VLOG(4) << "Try to replace: " << in_var_name << " => " << out_var_name;
+    if (view_.InSkipSet(in_var_name)) {
+      VLOG(4) << string::Sprintf("SKIP: %s is in skip set", in_var_name);
+      continue;
+    }
+
+    if (view_.InSkipSet(out_var_name)) {
+      VLOG(4) << string::Sprintf("SKIP: %s is in skip set", out_var_name);
+      continue;
+    }
 
     if (var_nodes_[in_var_name].back() != in_node) {
       VLOG(4) << "SKIP since " << in_var_name
@@ -318,11 +327,15 @@ void InplacePass::TryInplaceOpInputOutput(ir::Node* op,
               << out_var_name << " are the same";
     } else if (!NodeCanReused(in_node)) {
       can_replace = false;
-      VLOG(4) << "SKIP: Input varialbe " << in_var_name << "cannot be reused";
+      VLOG(4) << "SKIP: Input variable " << in_var_name << "cannot be reused";
     } else if (!NodeCanReused(out_node)) {
       can_replace = false;
       VLOG(4) << "SKIP: Output variable " << out_var_name
               << " cannot be reused";
+    } else if (in_node->Var()->GetType() != out_node->Var()->GetType()) {
+      can_replace = false;
+      VLOG(4) << "SKIP: Input type : " << in_node->Var()->GetType()
+              << " does not match Output type : " << out_node->Var()->GetType();
     } else if (details::NodeSize(*in_node->Var()) !=
                details::NodeSize(*out_node->Var())) {
       can_replace = false;
@@ -331,8 +344,8 @@ void InplacePass::TryInplaceOpInputOutput(ir::Node* op,
 
     if (!can_replace) continue;
 
-    // 2. there is no external pending op on the input node
-    // if (view_.PendingOpsOnVar(in_node).size() > 1) {
+    // 2. If the variable is the input of muliple ops, we need to make sure
+    // current op has dependecny on other ops use the same variable
     if (in_node->outputs.size() > 1 && !view_.CheckDeps(in_node, op)) {
       VLOG(4) << string::Sprintf(
           "Skiped pair %s => %s. %s input has external dependency."
@@ -341,17 +354,6 @@ void InplacePass::TryInplaceOpInputOutput(ir::Node* op,
       continue;
     }
 
-    // 3. if output has been memory optimize by python(fluid.memory_optmize()).
-    // this candidate  can not be inplaced. Will be deprecated in the future.
-    if (view_.InSkipSet(out_node->Name())) {
-      VLOG(4) << string::Sprintf(
-          "Skiped %s => %s reused previous memory block in python memory "
-          "optmize,"
-          "it inplace may generate a circle",
-          out_var_name, in_var_name, op->Name());
-      continue;
-    }
-
     // Debug Interface. Which would be skipped by the pass.
     if (out_node->Name() == FLAGS_memory_optimize_debug) {
       VLOG(3) << "Skiped var by force. FLAGS_memory_optimize_debug="
@@ -519,16 +521,22 @@ void GraphView::Build(ir::Graph* g) {
   // resolve data harzards depends on the var nodes in right order.
   TopoSort(g);
 
+  // fill the skip_set_
+  PADDLE_ENFORCE(g->Has(details::kMemOptSkipVars));
+  auto& mem_opt_whitelist = g->Get<MemOptSkipVars>(kMemOptSkipVars);
+  for (const auto& var : mem_opt_whitelist) skip_set_.emplace(var);
+
   // 2. track the nodes which used by parameter server.
   // these node can not be inplaced, otherwise trainer
   // pserver can not find each other name.
   auto update_skip_set = [&](ir::Node* node) {
     for (auto& in : node->inputs) {
-      if (in->IsVar() && in->Var() != nullptr) dup_nodes_.emplace(in->Name());
+      if (in->IsVar() && in->Var() != nullptr) {
+        skip_set_.emplace(in->Name());
+      }
     }
     for (auto& out : node->outputs) {
-      if (out->IsVar() && out->Var() != nullptr)
-        dup_nodes_.emplace(out->Name());
+      if (out->IsVar() && out->Var() != nullptr) skip_set_.emplace(out->Name());
     }
   };
   for (auto& node : g->Nodes()) {
@@ -545,7 +553,7 @@ void GraphView::Build(ir::Graph* g) {
 const std::vector<ir::Node*>& GraphView::AllOps() { return ops_; }
 
 bool GraphView::InSkipSet(const std::string& var) const {
-  return dup_nodes_.count(var);
+  return skip_set_.count(var);
 }
 
 }  // namespace details

paddle/fluid/framework/details/inplace_op_pass.h

@@ -57,7 +57,7 @@ class GraphView {
 
  private:
   std::vector<ir::Node*> ops_;
-  std::unordered_set<std::string> dup_nodes_;  // mem opt affect nodes
+  std::unordered_set<std::string> skip_set_;  // mem opt affect nodes
   std::map<ir::Node*, std::unordered_set<ir::Node*>> adj_list_;
   std::unordered_map<ir::Node*, uint32_t> op_level_;
 };

paddle/fluid/framework/details/memory_optimize_helper.h

@@ -21,6 +21,7 @@
 #include <set>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 #include "paddle/fluid/framework/data_type.h"
@@ -30,6 +31,11 @@ namespace paddle {
 namespace framework {
 namespace details {
 
+/// this attribute is used to avoid some core variables removed/reused
+/// in memory optimize related passes
+constexpr char kMemOptSkipVars[] = "@MEM_OPT_SKIP_VARS@";
+typedef std::unordered_set<std::string> MemOptSkipVars;
+
 std::vector<ir::Node*> SortOpLikeDescOrder(const ir::Graph& graph);
 
 // NOTE(dzh): A ordered set for node reuse in memory optimize.

paddle/fluid/framework/details/memory_optimize_pass.cc

@@ -45,8 +45,7 @@ namespace framework {
 namespace details {
 
 void MemoryOptimizePass::ApplyImpl(ir::Graph* graph) const {
-  auto nodes = graph->Nodes();
-  CollectSkipVarsSet(nodes);
+  CollectSkipVarsSet(graph);
 
   cfg_.reset(new details::ControlFlowGraph(*graph));
   cfg_->LiveVariableAnalysis();
@@ -204,14 +203,20 @@ void MemoryOptimizePass::SubGraphOptimize(OpDesc* op_desc) const {
   }
 }
 
-void MemoryOptimizePass::CollectSkipVarsSet(
-    const std::unordered_set<ir::Node*>& nodes) const {
+void MemoryOptimizePass::CollectSkipVarsSet(ir::Graph* graph) const {
+  // fill skip_set_
+  PADDLE_ENFORCE(graph->Has(details::kMemOptSkipVars));
+  auto& mem_opt_whitelist = graph->Get<MemOptSkipVars>(kMemOptSkipVars);
+  for (const auto& var : mem_opt_whitelist) skip_set_.emplace(var);
+
   auto update_skip_set = [&](OpDesc* op_desc) {
     auto inputs = op_desc->InputArgumentNames();
     auto outputs = op_desc->OutputArgumentNames();
     skip_set_.insert(inputs.begin(), inputs.end());
     skip_set_.insert(outputs.begin(), outputs.end());
   };
+
+  auto nodes = graph->Nodes();
   for (auto& op : nodes) {
     if (!op->IsOp() || op->Op() == nullptr) continue;
     auto* op_desc = op->Op();

paddle/fluid/framework/details/memory_optimize_pass.h

@@ -53,7 +53,8 @@ class MemoryOptimizePass : public ir::Pass {
   // 1. scan op with subblock and collect the output/input vars.
   // while, while_grad, conditional_block
   // 2. scan distributed ops and collect the output/input vars
-  void CollectSkipVarsSet(const std::unordered_set<ir::Node*>&) const;
+  // 3. op_role_vars
+  void CollectSkipVarsSet(ir::Graph* graph) const;
 
  private:
   // Reuse Node Pool, Owned.

paddle/fluid/framework/details/record_skip_memory_opt_vars_pass.cc

@@ -0,0 +1,64 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed 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.
+
+#include <string>
+#include "paddle/fluid/framework/details/memory_optimize_helper.h"
+#include "paddle/fluid/framework/ir/graph.h"
+#include "paddle/fluid/framework/ir/graph_helper.h"
+#include "paddle/fluid/framework/ir/pass.h"
+#include "paddle/fluid/framework/op_proto_maker.h"
+
+namespace paddle {
+namespace framework {
+namespace details {
+
+class RecordSkipMemoryOptVarsPass : public ir::Pass {
+ protected:
+  void ApplyImpl(ir::Graph* graph) const override {
+    PADDLE_ENFORCE(!graph->Has(kMemOptSkipVars));
+    graph->Set(kMemOptSkipVars, new MemOptSkipVars);
+    auto& skip_vars = graph->Get<MemOptSkipVars>(kMemOptSkipVars);
+
+    // NOTE(zcd): Insert OpRoleVars to SkipVarSet to prevent the vars are rename
+    // in memory optimize pass.
+    InsertOpRoleVarsToSkipVarSet(graph, &skip_vars);
+  }
+
+  void InsertOpRoleVarsToSkipVarSet(const ir::Graph* graph,
+                                    MemOptSkipVars* skip_vars) const {
+    for (auto& node : graph->Nodes()) {
+      PADDLE_ENFORCE_NOT_NULL(node, "The node should not be nullptr.");
+      if (node->IsOp() && node->Op()) {
+        try {
+          auto op_role_vars =
+              boost::get<std::vector<std::string>>(node->Op()->GetNullableAttr(
+                  OpProtoAndCheckerMaker::OpRoleVarAttrName()));
+          PADDLE_ENFORCE_EQ(op_role_vars.size() % 2, 0);
+          for (size_t i = 0; i < op_role_vars.size(); i += 2) {
+            auto& g_name = op_role_vars[i + 1];
+            skip_vars->insert(g_name);
+          }
+        } catch (boost::bad_get e) {
+        }
+      }
+    }
+  }
+};
+
+}  // namespace details
+}  // namespace framework
+}  // namespace paddle
+
+REGISTER_PASS(record_skip_memory_opt_vars_pass,
+              paddle::framework::details::RecordSkipMemoryOptVarsPass);

paddle/fluid/framework/inplace_op_inference_test.cc

@@ -19,6 +19,7 @@
 #include <vector>
 #include "gtest/gtest.h"
 #include "paddle/fluid/framework/details/inplace_op_pass.h"
+#include "paddle/fluid/framework/details/memory_optimize_helper.h"
 #include "paddle/fluid/framework/ir/pass_builder.h"
 #include "paddle/fluid/framework/op_info.h"
 #include "paddle/fluid/framework/op_registry.h"
@@ -217,6 +218,7 @@ TEST(InferInplace, SingleOpInplaceInToOut) {
 
   FakeSuccData(&prog);
   std::unique_ptr<ir::Graph> g(new ir::Graph(prog));
+  g->Set(details::kMemOptSkipVars, new std::unordered_set<std::string>());
   g = test_SingleOpInplaceInToOut(std::move(g));
   auto op_node = GetNodeFromGraph(g.get(), "single_op");
 
@@ -232,6 +234,7 @@ TEST(InferInplace, SingleOpInplaceInToOutNoInplace) {
 
   FakeNoInplaceData(&prog);
   std::unique_ptr<ir::Graph> g(new ir::Graph(prog));
+  g->Set(details::kMemOptSkipVars, new std::unordered_set<std::string>());
   g = test_SingleOpInplaceInToOut(std::move(g));
   auto op_node = GetNodeFromGraph(g.get(), "single_op");
 
@@ -264,6 +267,7 @@ TEST(InferInplace, MultiOutInplaceInToOut) {
   prog.MutableBlock(0)->Var("z0")->SetShape({32, 16, 1024, 1024});
 
   std::unique_ptr<ir::Graph> g(new ir::Graph(prog));
+  g->Set(details::kMemOptSkipVars, new std::unordered_set<std::string>());
   std::unique_ptr<details::InplacePass> pass(new details::InplacePass());
   pass->Apply(g.get());
   auto op_node = GetNodeFromGraph(g.get(), "multi_out_op");
@@ -299,6 +303,7 @@ TEST(InferInplace, MultiGradInplaceInToOut) {
   prog.MutableBlock(0)->Var("z0")->SetShape({32, 15, 1024, 1024});
 
   std::unique_ptr<ir::Graph> g(new ir::Graph(prog));
+  g->Set(details::kMemOptSkipVars, new std::unordered_set<std::string>());
   std::unique_ptr<details::InplacePass> pass(new details::InplacePass());
   pass->Apply(g.get());
   auto op_node = GetNodeFromGraph(g.get(), "multi_out_grad");

paddle/fluid/pybind/const_value.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/pybind/const_value.h"
+#include "paddle/fluid/framework/details/memory_optimize_pass.h"
 #include "paddle/fluid/framework/ir/node.h"
 #include "paddle/fluid/framework/op_proto_maker.h"
 #include "paddle/fluid/framework/operator.h"
@@ -28,6 +29,7 @@ void BindConstValue(pybind11::module* m) {
   m->def("kControlDepVarName",
          [] { return framework::ir::Node::kControlDepVarName; });
   m->def("kNewGradSuffix", [] { return framework::kNewGradSuffix; });
+  m->def("kMemOptSkipVars", [] { return framework::details::kMemOptSkipVars; });
 
   auto op_proto_and_checker_maker =
       m->def_submodule("op_proto_and_checker_maker");