Plan rank compiler

oneflow/core/graph/graph.h

@@ -45,6 +45,7 @@ class Graph {
       std::function<Maybe<void>(NodeType*)> NodeHandler) const;
   void ReverseTopoForEachNode(std::function<void(NodeType*)> NodeHandler) const;
   void ForEachEdge(std::function<void(EdgeType*)> EdgeHandler) const;
+  Maybe<void> MaybeForEachEdge(std::function<Maybe<void>(EdgeType*)> EdgeHandler) const;
 
   void SortedTopoForEachNode(std::function<bool(const EdgeType* lhs, const EdgeType* rhs)> LessThan,
                              std::function<void(NodeType*)> NodeHandler) const;
@@ -292,6 +293,16 @@ void Graph<NodeType, EdgeType>::ForEachEdge(std::function<void(EdgeType*)> EdgeH
   }
 }
 
+template<typename NodeType, typename EdgeType>
+Maybe<void> Graph<NodeType, EdgeType>::MaybeForEachEdge(
+    std::function<Maybe<void>(EdgeType*)> EdgeHandler) const {
+  for (auto& x : edges_) {
+    if (x->src_node() == nullptr && x->dst_node() == nullptr) { continue; }
+    JUST(EdgeHandler(x.get()));
+  }
+  return Maybe<void>::Ok();
+}
+
 template<typename NodeType, typename EdgeType>
 NodeType* Graph<NodeType, EdgeType>::SoleNode() const {
   CHECK_EQ(nodes_.size(), 1);

oneflow/core/graph/op_graph.cpp

@@ -17,11 +17,37 @@ limitations under the License.
 #include "oneflow/core/job/job_builder.h"
 #include "oneflow/core/job/local_sig_infer_hint.h"
 #include "oneflow/core/job/lazy_mode.h"
+#include "oneflow/core/common/container_util.h"
+#include "oneflow/core/persistence/tee_persistent_log_stream.h"
 #include "oneflow/core/auto_parallel/algorithm_util.h"
 #include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/framework/sbp_infer_util.h"
 
 namespace oneflow {
 
+bool OpEdge::NeedBoxing() const {
+  if (src_node()->parallel_desc_sym() != dst_node()->parallel_desc_sym()) { return true; }
+  if (src_node()->parallel_desc().parallel_num() == 1) { return false; }
+  for (const auto& lbi : *lbis_) {
+    Shape src_reduced_hierarchy;
+    Shape dst_reduced_hierarchy;
+    NdSbp src_reduced_nd_sbp;
+    NdSbp dst_reduced_nd_sbp;
+
+    InOutParallelDimReduce(*src_node()->parallel_desc().hierarchy(),
+                           *dst_node()->parallel_desc().hierarchy(), src_node()->NdSbp4Lbi(lbi),
+                           dst_node()->NdSbp4Lbi(lbi), &src_reduced_hierarchy,
+                           &dst_reduced_hierarchy, &src_reduced_nd_sbp, &dst_reduced_nd_sbp,
+                           src_node()->LogicalBlobDesc4Lbi(lbi).shape());
+    if (src_reduced_hierarchy != dst_reduced_hierarchy
+        || src_reduced_nd_sbp != dst_reduced_nd_sbp) {
+      // Not one to one
+      return true;
+    }
+  }
+  return false;
+}
+
 std::string OpEdge::VisualStr() const {
   std::string str;
   int32_t idx = 0;
@@ -54,12 +80,11 @@ const NdSbp& OpNode::NdSbp4Lbi(const LogicalBlobId& lbi) const {
   return it->second;
 }
 
-OpNode::OpNode(const std::shared_ptr<const ParallelDesc>& parallel_desc,
-               const OperatorConf& op_conf)
+OpNode::OpNode(Symbol<ParallelDesc> parallel_desc, const OperatorConf& op_conf)
     : parallel_desc_(parallel_desc),
       op_(CHECK_JUST(ConstructOp(op_conf, parallel_desc->device_type()))),
       ibns_(op_->input_bns().begin(), op_->input_bns().end()) {
-  CHECK_JUST(op_->FillOpParallelDesc(parallel_desc));
+  CHECK_JUST(op_->FillOpParallelDesc(parallel_desc.shared_from_symbol()));
 }
 
 std::string OpNode::VisualStr() const {
@@ -194,16 +219,14 @@ void OpGraph::CheckIsDAG() const {
 
 namespace {
 
-std::function<std::shared_ptr<const ParallelDesc>(const std::string&)>
-MakeGetterParallelDesc4OpName(const Job& job) {
+std::function<Symbol<ParallelDesc>(const std::string&)> MakeGetterParallelDesc4OpName(
+    const Job& job) {
   const Placement& placement = job.placement();
-  auto op_name2parallel_desc =
-      std::make_shared<HashMap<std::string, std::shared_ptr<const ParallelDesc>>>();
+  auto op_name2parallel_desc = std::make_shared<HashMap<std::string, Symbol<ParallelDesc>>>();
   op_name2parallel_desc->reserve(job.net().op_size());
   for (const auto& placement_group : placement.placement_group()) {
     const ParallelConf& parallel_conf = placement_group.parallel_conf();
-    std::shared_ptr<const ParallelDesc> parallel_desc =
-        std::make_shared<const ParallelDesc>(parallel_conf);
+    Symbol<ParallelDesc> parallel_desc = SymbolOf(ParallelDesc(parallel_conf));
     for (const std::string& op_name : placement_group.op_set().op_name()) {
       CHECK(op_name2parallel_desc->emplace(op_name, parallel_desc).second)
           << "op_name: " << op_name;
@@ -566,6 +589,11 @@ Maybe<void> OpGraph::ForEachOpNode(const std::function<Maybe<void>(const OpNode&
   return Maybe<void>::Ok();
 }
 
+std::function<bool(const OpNode* src, const OpNode* dst)> OpGraph::CreatePredicatorIsReachable()
+    const {
+  return MakePredicatorIsReachable();
+}
+
 // Print the graph with SBP in order
 void OpGraph::PrintSBPGraphDebugInfo() const {
   // test debug
@@ -622,4 +650,17 @@ void OpGraph::PrintSBPGraphDebugInfo() const {
   }
 }
 
+OpGraphSingletonGuard::OpGraphSingletonGuard(const Job& job) {
+  // new Singleton<OpGraph> and set log configs.
+  Singleton<OpGraph>::New(job);
+  const JobDesc& job_desc = GlobalJobDesc();
+  if (Singleton<ResourceDesc, ForSession>::Get()->enable_debug_mode()) {
+    TeePersistentLogStream::Create(StrCat("optimized_job", job_desc.job_id()))->Write(job);
+    Singleton<OpGraph>::Get()->ToDotWithFilePath(
+        "optimized_dlnet_" + std::to_string(job_desc.job_id()) + "_op_graph.dot");
+  }
+}
+
+OpGraphSingletonGuard::~OpGraphSingletonGuard() { Singleton<OpGraph>::Delete(); }
+
 }  // namespace oneflow

oneflow/core/graph/op_graph.h

@@ -34,15 +34,15 @@ class OpGraph;
 class OpNode final : public Node<OpNode, OpEdge> {
  public:
   OF_DISALLOW_COPY_AND_MOVE(OpNode);
-  explicit OpNode(const std::shared_ptr<const ParallelDesc>& parallel_desc,
-                  const OperatorConf& op_conf);
+  explicit OpNode(Symbol<ParallelDesc> parallel_desc, const OperatorConf& op_conf);
   ~OpNode() = default;
 
   // Getters
   bool IsTimeShapeIdentity() const;
   const Operator& op() const { return *op_; }
   std::shared_ptr<const Operator> shared_op() const { return op_; }
   const ParallelDesc& parallel_desc() const { return *parallel_desc_; }
+  Symbol<ParallelDesc> parallel_desc_sym() const { return parallel_desc_; }
   const SbpSignature& sbp_signature() const { return *CHECK_JUST(op().sbp_signature()); }
   const NdSbpSignature& nd_sbp_signature() const { return *CHECK_JUST(op().nd_sbp_signature()); }
   const SbpParallel& SbpParallel4Lbi(const LogicalBlobId& lbi) const;
@@ -67,7 +67,7 @@ class OpNode final : public Node<OpNode, OpEdge> {
   void InitLbi2SourceNode();
   void InitLbi2NdSbp();
 
-  std::shared_ptr<const ParallelDesc> parallel_desc_;
+  Symbol<ParallelDesc> parallel_desc_;
   std::shared_ptr<Operator> op_;
   HashSet<std::string> ibns_;
   HashMap<LogicalBlobId, OpNode*> lbi2source_node_;
@@ -88,6 +88,8 @@ class OpEdge final : public Edge<OpNode, OpEdge> {
   const std::vector<LogicalBlobId>& lbis() const { return *lbis_; }
   const HashMap<LogicalBlobId, std::string>& lbi2obn() const { return *lbi2obn_; }
   const HashMap<LogicalBlobId, std::vector<std::string>>& lbi2ibns() const { return *lbi2ibns_; }
+
+  bool NeedBoxing() const;
   std::string VisualStr() const override;
 
  private:
@@ -130,6 +132,7 @@ class OpGraph final : public Graph<OpNode, OpEdge> {
 
   Maybe<void> Init(const Job& job);
 
+  std::function<bool(const OpNode* src, const OpNode* dst)> CreatePredicatorIsReachable() const;
   // Print the graph with SBP in order
   void PrintSBPGraphDebugInfo() const;
 
@@ -155,6 +158,13 @@ class OpGraph final : public Graph<OpNode, OpEdge> {
   HashMap<std::string, HashSet<std::string>> producer_op_name2ctrl_consumer_op_names_;
 };
 
+class OpGraphSingletonGuard {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(OpGraphSingletonGuard);
+  explicit OpGraphSingletonGuard(const Job& job);
+  ~OpGraphSingletonGuard();
+};
+
 }  // namespace oneflow
 
 #endif  // ONEFLOW_CORE_GRAPH_OP_GRAPH_H_