Convert all blocks in program into SSAgraphs.

paddle/fluid/framework/ir/graph.cc

@@ -17,15 +17,114 @@ limitations under the License. */
 #include "paddle/fluid/framework/ir/graph.h"
 #include "paddle/fluid/framework/operator.h"
 
+DEFINE_bool(convert_all_blocks, false,
+            "Convert all blocks in program into SSAgraphs");
+
 namespace paddle {
 namespace framework {
 namespace ir {
 
-Graph::Graph(const ProgramDesc &program) : program_(program) {
-  auto var_nodes = InitFromProgram(program_);
+Graph::Graph(const ProgramDesc &program)
+    : program_(program), main_graph_(nullptr) {
+  if (FLAGS_convert_all_blocks) {
+    PADDLE_ENFORCE_GE(
+        program_.Size(), 1,
+        platform::errors::InvalidArgument("Can't construct a graph from this "
+                                          "program, it doesn't have a block"));
+    for (size_t idx = 0; idx < program_.Size(); ++idx) {
+      std::unique_ptr<Graph> sub_graph =
+          std::make_unique<Graph>(program_.Block(idx), this);
+      sub_graphs_.push_back(std::move(sub_graph));
+    }
+  } else {
+    auto var_nodes = InitFromProgram(program_);
+    ResolveHazard(var_nodes);
+  }
+}
+
+Graph::Graph(const BlockDesc &block, const Graph *main_graph)
+    : main_graph_(main_graph) {
+  auto var_nodes = InitFromBlock(block);
   ResolveHazard(var_nodes);
 }
 
+std::map<std::string, std::vector<ir::Node *>> Graph::InitFromBlock(
+    const BlockDesc &block) {
+  std::unordered_map<std::string, VarDesc *> all_vars;
+  // var nodes for each var name, will have multiple versions in SSA
+  std::map<std::string, std::vector<ir::Node *>> var_nodes;
+  for (auto *var : block.AllVars()) {
+    all_vars.emplace(var->Name(), var);
+  }
+
+  auto not_visited_vars = all_vars;
+
+  for (auto *op : block.AllOps()) {
+    ir::Node *node = CreateOpNode(op);
+    // For input args, reuse the same var name if it was created before.
+    // Otherwise, create a new one.
+    for (auto &each_var_name : op->InputArgumentNames()) {
+      not_visited_vars.erase(each_var_name);
+      ir::Node *var = nullptr;
+      if (var_nodes.find(each_var_name) != var_nodes.end()) {
+        var = var_nodes.at(each_var_name).back();
+      } else if (all_vars.count(each_var_name) != 0) {
+        var = CreateVarNode(all_vars.at(each_var_name));
+        var_nodes[each_var_name].push_back(var);
+      } else {
+        // Operation input var can be optional (dispensable). Which means
+        // the operation doesn't really need the var at runtime. In this
+        // case, the no-existed var is ready at the beginning.
+        var = CreateEmptyNode(each_var_name, ir::Node::Type::kVariable);
+        var_nodes[each_var_name].push_back(var);
+      }
+      node->inputs.push_back(var);
+      var->outputs.push_back(node);
+    }
+    // For output args, always create a new var.
+    std::unordered_set<std::string> out_arg_set;
+    for (auto &each_var_name : op->OutputArgumentNames()) {
+      not_visited_vars.erase(each_var_name);
+      if (each_var_name != kEmptyVarName) {
+        PADDLE_ENFORCE_EQ(out_arg_set.count(each_var_name), 0,
+                          platform::errors::InvalidArgument(
+                              "The input Program is invalid. Variable %s occurs"
+                              " in output of %s multiple times.",
+                              each_var_name, op->Type()));
+        out_arg_set.insert(each_var_name);
+      }
+
+      ir::Node *var = nullptr;
+      if (all_vars.count(each_var_name) != 0) {
+        var = CreateVarNode(all_vars.at(each_var_name));
+      } else {
+        // Operation output vars can be @EMPTY@. For example, while_grad
+        // can have multi @EMPTY@ outputs with no VarDesc.
+        // TODO(panyx0718): Add a test.
+        var = CreateEmptyNode(each_var_name, ir::Node::Type::kVariable);
+      }
+      var_nodes[each_var_name].push_back(var);
+      node->outputs.push_back(var);
+      var->inputs.push_back(node);
+    }
+  }
+
+  for (auto &pair : not_visited_vars) {
+    const auto &var_name = pair.first;
+    auto *var_desc = pair.second;
+    if (var_name != kEmptyVarName) {
+      VLOG(10) << "Create isolated var node " << var_name;
+      var_nodes[var_name].push_back(CreateVarNode(var_desc));
+    }
+  }
+
+  Set<const std::vector<OpDesc *>>(details::kStaleProgramOpDescs,
+                                   new std::vector<OpDesc *>(block.AllOps()));
+  return var_nodes;
+}
+
+// TODO(levi): delete this interface after when we can convert all
+// blocks into sub_graphs.
 std::map<std::string, std::vector<ir::Node *>> Graph::InitFromProgram(
     const ProgramDesc &program) {
   VLOG(3) << "block in program:" << program_.Size();
@@ -176,38 +275,93 @@ void Graph::ResolveHazard(
 }
 
 std::shared_ptr<Graph> Graph::Clone() {
-  auto cloned_graph = std::make_shared<Graph>(this->program_);
-  cloned_graph->ReleaseNodes();
-  cloned_graph->num_node_created_ = 0;
+  PADDLE_ENFORCE_EQ(
+      this->IsMainGraph(), true,
+      platform::errors::InvalidArgument(
+          "This graph is a sub_graph, and can't be cloned individually"));
+  if (FLAGS_convert_all_blocks) {
+    auto cloned_graph = std::make_shared<Graph>(this->program_);
+    cloned_graph->ReleaseSubGraphs();
+    for (size_t idx = 0; idx < this->program_.Size(); ++idx) {
+      cloned_graph->AddSubGraph(this->CloneSubGraph(idx));
+    }
+    return cloned_graph;
+  } else {
+    auto cloned_graph = std::make_shared<Graph>(this->program_);
+    cloned_graph->ReleaseNodes();
+    cloned_graph->num_node_created_ = 0;
+    std::unordered_map<ir::Node *, ir::Node *> origin_to_cloned;
+    for (auto *n : this->node_set_) {
+      PADDLE_ENFORCE_NOT_NULL(n, platform::errors::InvalidArgument(
+                                     "The node to be cloned is nullptr."));
+      ir::Node *cloned_node = nullptr;
+      if (n->IsCtrlVar()) {
+        cloned_node = cloned_graph->CreateControlDepVar();
+      } else if (!n->var_desc_ && !n->op_desc_) {  // empty node
+        cloned_node = cloned_graph->CreateEmptyNode(n->Name(), n->NodeType());
+      } else if (n->IsVar()) {
+        cloned_node = cloned_graph->CreateVarNode(n->Var());
+      } else if (n->IsOp()) {
+        cloned_node = cloned_graph->CreateOpNode(n->Op());
+      }
+      PADDLE_ENFORCE_NOT_NULL(
+          cloned_node,
+          platform::errors::InvalidArgument(
+              "Failed to clone new node from original node in graph."));
+      origin_to_cloned[n] = cloned_node;
+    }
+    for (auto *n : this->node_set_) {
+      for (auto it = n->inputs.begin(); it != n->inputs.end(); it++) {
+        origin_to_cloned[n]->inputs.push_back(origin_to_cloned[*it]);
+      }
+      for (auto it = n->outputs.begin(); it != n->outputs.end(); it++) {
+        origin_to_cloned[n]->outputs.push_back(origin_to_cloned[*it]);
+      }
+    }
+    return cloned_graph;
+  }
+}
+
+std::unique_ptr<Graph> Graph::CloneSubGraph(const size_t idx) {
+  PADDLE_ENFORCE_EQ(
+      this->IsMainGraph(), true,
+      platform::errors::InvalidArgument("This graph is not main_graph"));
+  PADDLE_ENFORCE_LT(
+      idx, this->sub_graphs_.size(),
+      platform::errors::InvalidArgument("Invalid sub_graph index"));
+  std::unique_ptr<Graph> cloned_sub_graph =
+      std::make_unique<Graph>(this->program_.Block(idx), this);
+  cloned_sub_graph->ReleaseNodes();
+  cloned_sub_graph->num_node_created_ = 0;
   std::unordered_map<ir::Node *, ir::Node *> origin_to_cloned;
-  for (auto *n : this->node_set_) {
+  for (auto *n : this->sub_graphs_.at(idx)->Nodes()) {
     PADDLE_ENFORCE_NOT_NULL(n, platform::errors::InvalidArgument(
                                    "The node to be cloned is nullptr."));
     ir::Node *cloned_node = nullptr;
     if (n->IsCtrlVar()) {
-      cloned_node = cloned_graph->CreateControlDepVar();
+      cloned_node = cloned_sub_graph->CreateControlDepVar();
     } else if (!n->var_desc_ && !n->op_desc_) {  // empty node
-      cloned_node = cloned_graph->CreateEmptyNode(n->Name(), n->NodeType());
+      cloned_node = cloned_sub_graph->CreateEmptyNode(n->Name(), n->NodeType());
     } else if (n->IsVar()) {
-      cloned_node = cloned_graph->CreateVarNode(n->Var());
+      cloned_node = cloned_sub_graph->CreateVarNode(n->Var());
     } else if (n->IsOp()) {
-      cloned_node = cloned_graph->CreateOpNode(n->Op());
+      cloned_node = cloned_sub_graph->CreateOpNode(n->Op());
     }
     PADDLE_ENFORCE_NOT_NULL(
         cloned_node,
         platform::errors::InvalidArgument(
             "Failed to clone new node from original node in graph."));
     origin_to_cloned[n] = cloned_node;
   }
-  for (auto *n : this->node_set_) {
+  for (auto *n : this->sub_graphs_.at(idx)->Nodes()) {
     for (auto it = n->inputs.begin(); it != n->inputs.end(); it++) {
       origin_to_cloned[n]->inputs.push_back(origin_to_cloned[*it]);
     }
     for (auto it = n->outputs.begin(); it != n->outputs.end(); it++) {
       origin_to_cloned[n]->outputs.push_back(origin_to_cloned[*it]);
     }
   }
-  return cloned_graph;
+  return cloned_sub_graph;
 }
 
 bool IsControlDepVar(const ir::Node &var) {