graph_to_program topology sort (#33949)

See #33949 for details

Author: thisjiang

paddle/fluid/framework/ir/graph.cc

@@ -56,10 +56,12 @@ Graph::Graph(const ProgramDesc &program, const int64_t start_op_index,
     // sub_graph.
     std::unique_ptr<Graph> first_sub_graph = std::make_unique<Graph>(
         program_.Block(0), this, start_op_index, end_op_index);
+    first_sub_graph->block_id_ = 0;
     sub_graphs_.push_back(std::move(first_sub_graph));
     for (size_t idx = 1; idx < program_.Size(); ++idx) {
       std::unique_ptr<Graph> sub_graph =
           std::make_unique<Graph>(program_.Block(idx), this);
+      sub_graph->block_id_ = idx;
       sub_graphs_.push_back(std::move(sub_graph));
     }
   } else {
@@ -90,14 +92,32 @@ std::map<std::string, std::vector<ir::Node *>> Graph::InitFromProgram(
 std::map<std::string, std::vector<ir::Node *>> Graph::InitFromBlock(
     const BlockDesc &block, const int64_t start_op_index,
     const int64_t end_op_index) {
-  std::unordered_map<std::string, VarDesc *> all_vars;
+  std::unordered_map<std::string, std::pair<VarDesc *, int>>
+      name_to_desc_block_id;
+
+  const BlockDesc *block_var_visible = &block;
+  while (block_var_visible != nullptr) {
+    for (auto *var : block_var_visible->AllVars()) {
+      name_to_desc_block_id.emplace(
+          var->Name(), std::make_pair(var, block_var_visible->ID()));
+    }
+    const BlockDesc *forward_block = block_var_visible->ForwardBlock();
+    if (forward_block != nullptr) {
+      for (auto *var : forward_block->AllVars()) {
+        name_to_desc_block_id.emplace(var->Name(),
+                                      std::make_pair(var, forward_block->ID()));
+      }
+    }
+    block_var_visible = block_var_visible->ParentBlock();
+  }
   // var nodes for each var name, will have multiple versions in SSA
   std::map<std::string, std::vector<ir::Node *>> var_nodes;
+  std::unordered_map<std::string, VarDesc *> not_visited_vars;
   for (auto *var : block.AllVars()) {
-    all_vars.emplace(var->Name(), var);
+    not_visited_vars.emplace(var->Name(), var);
   }
 
-  auto not_visited_vars = all_vars;
+  int desc_order = 0;
   auto all_ops = block.AllOps();
   PADDLE_ENFORCE_LE(
       end_op_index, all_ops.size(),
@@ -109,15 +129,18 @@ std::map<std::string, std::vector<ir::Node *>> Graph::InitFromBlock(
     auto *op = all_ops[i];
     VLOG(3) << "create OpNode by " << op->Type();
     ir::Node *node = CreateOpNode(op);
+    node->SetDescOrder(desc_order);
+    ++desc_order;
     // 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));
+      } else if (name_to_desc_block_id.count(each_var_name) != 0) {
+        auto desc_and_block_id = name_to_desc_block_id.at(each_var_name);
+        var = CreateVarNode(desc_and_block_id.first, desc_and_block_id.second);
         var_nodes[each_var_name].push_back(var);
       } else {
         // Operation input var can be optional (dispensable). Which means
@@ -143,8 +166,9 @@ std::map<std::string, std::vector<ir::Node *>> Graph::InitFromBlock(
       }
 
       ir::Node *var = nullptr;
-      if (all_vars.count(each_var_name) != 0) {
-        var = CreateVarNode(all_vars.at(each_var_name));
+      if (name_to_desc_block_id.count(each_var_name) != 0) {
+        auto desc_and_block_id = name_to_desc_block_id.at(each_var_name);
+        var = CreateVarNode(desc_and_block_id.first, desc_and_block_id.second);
       } else {
         // Operation output vars can be @EMPTY@. For example, while_grad
         // can have multi @EMPTY@ outputs with no VarDesc.
@@ -270,6 +294,7 @@ std::shared_ptr<Graph> Graph::Clone() {
     auto cloned_graph = std::make_shared<Graph>(this->program_);
     cloned_graph->ReleaseNodes();
     cloned_graph->num_node_created_ = 0;
+    cloned_graph->block_id_ = this->block_id_;
     std::unordered_map<ir::Node *, ir::Node *> origin_to_cloned;
     for (auto *n : this->node_set_) {
       PADDLE_ENFORCE_NOT_NULL(n, platform::errors::InvalidArgument(
@@ -313,6 +338,7 @@ std::unique_ptr<Graph> Graph::CloneSubGraph(const size_t idx) {
       std::make_unique<Graph>(this->program_.Block(idx), this);
   cloned_sub_graph->ReleaseNodes();
   cloned_sub_graph->num_node_created_ = 0;
+  cloned_sub_graph->block_id_ = idx;
   std::unordered_map<ir::Node *, ir::Node *> origin_to_cloned;
   for (auto *n : this->sub_graphs_.at(idx)->Nodes()) {
     PADDLE_ENFORCE_NOT_NULL(n, platform::errors::InvalidArgument(

paddle/fluid/framework/ir/graph.h

@@ -104,7 +104,14 @@ class Graph {
     attr_dels_.clear();
   }
 
-  bool IsConstructedByPartialProgram() const { return is_partial_; }
+  bool IsConstructedByPartialProgram() const {
+    if (FLAGS_convert_all_blocks) {
+      if (IsMainGraph()) {
+        return GetSubGraph(0)->IsConstructedByPartialProgram();
+      }
+    }
+    return is_partial_;
+  }
 
   bool Has(const std::string &attr_name) const {
     if (FLAGS_convert_all_blocks) {
@@ -210,7 +217,7 @@ class Graph {
   }
 
   // Create a normal variable with non-null VarDesc.
-  ir::Node *CreateVarNode(VarDesc *var_desc) {
+  ir::Node *CreateVarNode(VarDesc *var_desc, int block_id = -1) {
     if (FLAGS_convert_all_blocks) {
       if (IsMainGraph()) {
         return GetSubGraph(0)->CreateVarNode(var_desc);
@@ -219,7 +226,8 @@ class Graph {
     PADDLE_ENFORCE_NOT_NULL(
         var_desc, platform::errors::InvalidArgument(
                       "The VarDesc used to create variable node is null."));
-    auto *x = AddNode(new ir::Node(var_desc));
+    auto *x =
+        AddNode(new ir::Node(var_desc, block_id == -1 ? block_id_ : block_id));
     x->SetId(num_node_created_++);
     return x;
   }
@@ -252,7 +260,7 @@ class Graph {
     const std::string name = string::Sprintf(
         "%s@%llu", static_cast<const char *>(ir::Node::kControlDepVarName),
         num_node_created_);
-    auto *x = AddNode(new ir::Node(name, ir::Node::Type::kVariable));
+    auto *x = AddNode(new ir::Node(name, ir::Node::Type::kVariable, block_id_));
     x->SetId(num_node_created_++);
     return x;
   }
@@ -265,7 +273,7 @@ class Graph {
         return GetSubGraph(0)->CreateEmptyNode(name, type);
       }
     }
-    auto *x = AddNode(new ir::Node(name, type));
+    auto *x = AddNode(new ir::Node(name, type, block_id_));
     x->SetId(num_node_created_++);
     return x;
   }
@@ -365,6 +373,15 @@ class Graph {
     return sub_graphs_.at(idx).get();
   }
 
+  int GetBlockId() const {
+    if (FLAGS_convert_all_blocks) {
+      if (IsMainGraph()) {
+        return GetSubGraph(0)->block_id_;
+      }
+    }
+    return block_id_;
+  }
+
   size_t SubGraphsSize() const {
     PADDLE_ENFORCE_EQ(
         this->IsMainGraph(), true,
@@ -394,6 +411,9 @@ class Graph {
     PADDLE_ENFORCE_EQ(
         this->IsMainGraph(), true,
         platform::errors::InvalidArgument("This graph is not main_graph"));
+    PADDLE_ENFORCE_EQ(sub_graphs_.size(), sub_graph->block_id_,
+                      platform::errors::InvalidArgument(
+                          "sub_graph idx is not equal to block_id_"));
     sub_graphs_.push_back(std::move(sub_graph));
   }
 
@@ -416,6 +436,8 @@ class Graph {
   // parts: forward graph and backward graph, which can be executed
   // independently.
   bool is_partial_{false};
+  // The block this SubGraph belongs to.
+  int block_id_{0};
 };
 
 bool IsControlDepVar(const ir::Node &var);

paddle/fluid/framework/ir/graph_helper.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/framework/ir/graph_helper.h"
+#include <queue>
 #include <stack>
 
 DEFINE_string(print_sub_graph_dir, "",
@@ -395,6 +396,85 @@ std::vector<Node *> TopologyVarientSort(const Graph &graph,
   }
 }
 
+class DescOrderComparator {
+ public:
+  bool operator()(const Node *n1, const Node *n2) {
+    return (n1->DescOrder() > n2->DescOrder()) ||
+           ((n1->DescOrder() == n2->DescOrder()) &&
+            (n1->ToString() > n2->ToString()));
+  }
+};
+
+std::vector<ir::Node *> TopologySortGraphByDescOrder(const Graph &graph) {
+  std::vector<ir::Node *> sorted_ops;
+  std::priority_queue<Node *, std::vector<Node *>, DescOrderComparator> q;
+  std::unordered_map<Node *, std::unordered_set<Node *>> in_ops;
+  std::unordered_map<Node *, std::unordered_set<Node *>> out_ops;
+
+  // ensure all op node in 'in_ops' and 'out_ops'
+  for (const auto &n : graph.Nodes()) {
+    if (!n->IsOp()) continue;
+
+    in_ops.emplace(n, std::unordered_set<Node *>());
+    out_ops.emplace(n, std::unordered_set<Node *>());
+  }
+
+  // record all op's input op and output op
+  for (const auto &n : graph.Nodes()) {
+    if (!n->IsOp()) continue;
+
+    // traverse all input op
+    for (const auto &var : n->inputs) {
+      for (const auto &in : var->inputs) {
+        // use at instead of [] to prevent no unrecorded op node
+        in_ops.at(n).insert(in);
+        out_ops.at(in).insert(n);
+      }
+    }
+  }
+
+  // find topology entrance
+  for (const auto &n : graph.Nodes()) {
+    if (!n->IsOp()) continue;
+
+    if (in_ops.at(n).empty()) {
+      q.push(n);
+    }
+  }
+
+  // topological sorting
+  while (!q.empty()) {
+    // Do not get by reference!!! The element will pop later.
+    const auto cur_op = q.top();
+    q.pop();
+
+    sorted_ops.push_back(cur_op);
+    for (const auto &out : out_ops.at(cur_op)) {
+      PADDLE_ENFORCE_GT(in_ops.at(out).count(cur_op), 0,
+                        platform::errors::InvalidArgument(
+                            "We find %s in %s's output list, "
+                            "but cannot find %s in %s's input list. "
+                            "Please ensure graph completely.",
+                            out->Name().c_str(), cur_op->Name().c_str(),
+                            cur_op->Name().c_str(), out->Name().c_str()));
+      in_ops.at(out).erase(cur_op);
+
+      // push if in-degree is 0
+      if (in_ops.at(out).empty()) {
+        q.push(out);
+      }
+    }
+  }
+
+  PADDLE_ENFORCE_EQ(
+      sorted_ops.size(), in_ops.size(),
+      platform::errors::InvalidArgument("Topological sorting incompletely, "
+                                        "only sorted %zd op but total %zd.",
+                                        sorted_ops.size(), in_ops.size()));
+
+  return sorted_ops;
+}
+
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle

paddle/fluid/framework/ir/graph_helper.h

@@ -87,6 +87,8 @@ std::vector<T *> FilterByNodeWrapper(const Graph &graph) {
   return ret;
 }
 
+std::vector<ir::Node *> TopologySortGraphByDescOrder(const Graph &graph);
+
 }  // namespace ir
 }  // namespace framework
 }  // namespace paddle