TaskNode::order_in_chain

oneflow/core/graph/plan_task_graph.cpp

@@ -17,11 +17,6 @@ limitations under the License.
 
 namespace oneflow {
 
-int64_t PlanTaskNode::chain_id() const {
-  int64_t chain_id = task_proto_->task_set_info().chain_id();
-  return chain_id;
-}
-
 PlanTaskGraph::PlanTaskGraph(const Plan& plan) : plan_(&plan) {
   InitNodes();
   InitEdges();

oneflow/core/graph/plan_task_graph.h

@@ -38,8 +38,8 @@ class PlanTaskNode final : public Node<PlanTaskNode, PlanTaskEdge> {
 
   const TaskProto* task_proto() const { return task_proto_; }
   int64_t task_id() const { return task_proto_->task_id(); }
-  int64_t chain_id() const;
-  int64_t order_in_graph() const { return task_proto_->task_set_info().order_in_graph(); }
+  int64_t chain_id() const { return task_proto_->chain_id(); }
+  int64_t order_in_chain() const { return task_proto_->order_in_chain(); }
 
  private:
   const TaskProto* task_proto_;

oneflow/core/graph/straighten_nodes.cpp

@@ -522,9 +522,6 @@ int32_t MaximumOverlapNum(StraightenAlgorithmTag sat, bool nccl_use_compute_stre
 
 void StraightenNodes(TaskGraph* task_graph, std::vector<TaskNode*>* ordered_task_nodes,
                      bool nccl_use_compute_stream) {
-  // The function for settle the order in the graph
-  int64_t order_in_graph = 0;
-
   // Generate topological data structure for each task node
   HashMap<TaskNode*, TopoStruct> task_node2topo_struct;
   // Determine the same nodes which should run simultaneously by the keys
@@ -606,11 +603,7 @@ void StraightenNodes(TaskGraph* task_graph, std::vector<TaskNode*>* ordered_task
 
   std::vector<int32_t> remain_task_nums(num_classifier, 0);
 
-  auto SetOrderInGraph = [&](TaskNode* task_node) {
-    task_node->set_order_in_graph(order_in_graph);
-    ordered_task_nodes->emplace_back(task_node);
-    ++order_in_graph;
-  };
+  auto AddOrderedNodes = [&](TaskNode* task_node) { ordered_task_nodes->emplace_back(task_node); };
 
   // wait in the list
   auto wait = [&](TaskNode* node) {
@@ -709,7 +702,7 @@ void StraightenNodes(TaskGraph* task_graph, std::vector<TaskNode*>* ordered_task
     remain_task_nums[list_classifier] -= execution_list.size();
     // Set the order and then remove from the execution list
     for (auto* node : execution_list) {
-      SetOrderInGraph(node);
+      AddOrderedNodes(node);
       finish_execution(node);
     }
   };
@@ -743,7 +736,7 @@ void StraightenNodes(TaskGraph* task_graph, std::vector<TaskNode*>* ordered_task
         move2execution_list(waiting_lists[TaskClassifier::kWaitingOverlapNode],
                             overlap_execution_list);
         remain_task_nums[TaskClassifier::kWaitingOverlapNode] -= overlap_execution_list.size();
-        for (auto* overlap_node : overlap_execution_list) { SetOrderInGraph(overlap_node); }
+        for (auto* overlap_node : overlap_execution_list) { AddOrderedNodes(overlap_node); }
         // Overlap the node with computation from the trunk
         execute(TaskClassifier::kWaitingMainComputation, computation_num);
 

oneflow/core/graph/task_graph.cpp

@@ -207,7 +207,7 @@ MakePredicatorIsLbiAllConsumersReachable(
                                                              const TaskNode* dst_node) -> bool {
     if (IsValidChainId(src_node->chain_id()) && IsValidChainId(dst_node->chain_id())
         && src_node->chain_id() == dst_node->chain_id()
-        && src_node->order_in_graph() <= dst_node->order_in_graph()) {
+        && src_node->order_in_chain() <= dst_node->order_in_chain()) {
       return true;
     }
     const CompTaskNode* comp_src_node = dynamic_cast<const CompTaskNode*>(src_node);
@@ -614,14 +614,10 @@ void TaskGraph::MergeChainAndAddOrderingCtrlEdgeInSameChain() {
   BuildCtrlRegstDescInSameChain();
 }
 
-void TaskGraph::SetOrderInGraphForEachNode() {
-  int64_t order_in_graph = 0;
-  auto SetOrderInGraph = [&](TaskNode* task_node) {
-    task_node->set_order_in_graph(order_in_graph);
-    ordered_task_nodes_.emplace_back(task_node);
-    ++order_in_graph;
-  };
-  TopoForEachNode(SetOrderInGraph);
+void TaskGraph::InitOrderedTaskNodes() {
+  // NOTE(chengcheng): Warning, ordered_task_nodes_ by topo is NOT valid in process
+  //  parallel compile, because the current rank task graph is Incomplete.
+  TopoForEachNode([&](TaskNode* task_node) { ordered_task_nodes_.emplace_back(task_node); });
 }
 
 void TaskGraph::MergeChainByPhysicalTaskGraph() {
@@ -638,15 +634,31 @@ void TaskGraph::MergeChainByPhysicalTaskGraph() {
 
     ++chain_id;
   }
-  for (auto* node : ordered_task_nodes_) { CHECK(IsValidChainId(node->chain_id())); }
+
+  // set order_in_chain by ordered_task_nodes_
+  HashMap<int64_t, int64_t> chain_id2order;
+  for (auto* node : ordered_task_nodes_) {
+    CHECK(IsValidChainId(node->chain_id()));
+    int64_t this_chain_id = node->chain_id();
+    if (chain_id2order.find(this_chain_id) == chain_id2order.end()) {
+      chain_id2order.emplace(this_chain_id, 0);
+    }
+    node->set_order_in_chain(chain_id2order.at(this_chain_id)++);
+  }
 }
 
 void TaskGraph::MergeChainByLogicalChainId() {
   for (TaskNode* this_node : ordered_task_nodes_) {
     CompTaskNode* comp_node = dynamic_cast<CompTaskNode*>(this_node);
     if (!comp_node) { continue; }
-    const int64_t logical_chain_id = comp_node->op()->op_conf().logical_chain_id();
-    if (IsValidChainId(logical_chain_id)) { this_node->set_chain_id(logical_chain_id); }
+    const OperatorConf& conf = comp_node->op()->op_conf();
+    if (conf.has_logical_chain_id()) {
+      const int64_t logical_chain_id = conf.logical_chain_id();
+      CHECK(IsValidChainId(logical_chain_id));
+      this_node->set_chain_id(logical_chain_id);
+      CHECK(conf.has_order_in_logical_chain());
+      this_node->set_order_in_chain(conf.order_in_logical_chain());
+    }
   }
 }
 
@@ -980,7 +992,7 @@ void TaskGraph::DecideExecutionOrder() {
   if (straighten_algorithm_tag == StraightenAlgorithmTag::kDisableStraighten
       || (straighten_algorithm_tag == StraightenAlgorithmTag::kOverlap4Transfer
           && GlobalProcessCtx::WorldSize() == 1)) {
-    SetOrderInGraphForEachNode();
+    InitOrderedTaskNodes();
   } else {
     StraightenNodes(this, &ordered_task_nodes_,
                     Singleton<ResourceDesc, ForSession>::Get()->nccl_use_compute_stream());

oneflow/core/graph/task_graph.h

@@ -82,7 +82,7 @@ class TaskGraph final : public Graph<TaskNode, TaskEdge> {
   void ConnectCtrlEdges(const std::vector<CompTaskNode*>& src_task_nodes,
                         const std::vector<CompTaskNode*>& dst_task_nodes);
 
-  void SetOrderInGraphForEachNode();
+  void InitOrderedTaskNodes();
   void MergeChainByPhysicalTaskGraph();
   void MergeChainByLogicalChainId();
   void BuildCtrlRegstDescInSameChain();

oneflow/core/graph/task_node.cpp

@@ -40,7 +40,7 @@ void ForEachDataEdge(const std::unordered_set<TaskEdge*>& edges,
 }  // namespace
 
 TaskNode::TaskNode()
-    : machine_id_(-1), thrd_id_(-1), task_id_(-1), chain_id_(-1), order_in_graph_(-1) {}
+    : machine_id_(-1), thrd_id_(-1), task_id_(-1), chain_id_(-1), order_in_chain_(-1) {}
 
 std::shared_ptr<RegstDesc> TaskNode::GetProducedRegst(const std::string& name) {
   auto produced_regsts_it = produced_regsts_.find(name);
@@ -88,9 +88,9 @@ void TaskNode::set_chain_id(int64_t val) {
   chain_id_ = val;
 }
 
-void TaskNode::set_order_in_graph(int64_t val) {
-  CHECK_EQ(order_in_graph_, -1);
-  order_in_graph_ = val;
+void TaskNode::set_order_in_chain(int64_t val) {
+  CHECK_EQ(order_in_chain_, -1);
+  order_in_chain_ = val;
 }
 
 void TaskNode::PinConsumedRegst() {
@@ -209,8 +209,8 @@ void TaskNode::ToProto(TaskProto* task_proto) const {
   task_proto->set_thrd_id(thrd_id_);
   task_proto->set_task_id(task_id_);
   task_proto->set_job_id(GlobalJobDesc().job_id());
-  task_proto->mutable_task_set_info()->set_chain_id(chain_id_);
-  task_proto->mutable_task_set_info()->set_order_in_graph(order_in_graph_);
+  task_proto->set_chain_id(chain_id_);
+  task_proto->set_order_in_chain(order_in_chain_);
 
   // Step2: process exec_gph.
   exec_gph_.ToExecSequence(parallel_ctx(), task_proto->mutable_exec_sequence());

oneflow/core/graph/task_node.h

@@ -56,7 +56,7 @@ class TaskNode : public Node<TaskNode, TaskEdge> {
   int64_t task_id() const { return task_id_; }
   const StreamId& stream_id() const;
   int64_t chain_id() const { return chain_id_; }
-  int64_t order_in_graph() const { return order_in_graph_; }
+  int64_t order_in_chain() const { return order_in_chain_; }
   const ExecGraph& exec_gph() const { return exec_gph_; }
   std::shared_ptr<RegstDesc> GetProducedRegst(const std::string& name);
   const std::list<std::shared_ptr<RegstDesc>>& GetConsumedRegst(const std::string& name);
@@ -74,7 +74,7 @@ class TaskNode : public Node<TaskNode, TaskEdge> {
   void set_machine_id(int64_t val);
   void set_thrd_id(int64_t val);
   void set_chain_id(int64_t val);
-  void set_order_in_graph(int64_t val);
+  void set_order_in_chain(int64_t val);
 
   // Build
   virtual void ProduceAllRegstsAndBindEdges() = 0;
@@ -149,7 +149,7 @@ class TaskNode : public Node<TaskNode, TaskEdge> {
   int64_t thrd_id_;
   int64_t task_id_;
   int64_t chain_id_;
-  int64_t order_in_graph_;
+  int64_t order_in_chain_;
   std::unique_ptr<TaskId> new_task_id_;
 
   ExecGraph exec_gph_;

oneflow/core/job/intra_job_mem_sharing_util.cpp

@@ -91,11 +91,10 @@ void InitMemoryChains(Plan* plan,
     DeviceType device_type = stream_id.device_id().device_type();
     // TODO(zwx): eliminate this special 'is cpu' determine
     if (device_type == DeviceType::kCPU) { continue; }
-    if (!IsValidChainId(task->task_set_info().chain_id())) { continue; }
+    if (!IsValidChainId(task->chain_id())) { continue; }
     int64_t device_id = stream_id.device_id().device_index();
     int64_t device_unique_id = GenDeviceUniqueId(machine_id, device_id);
-    MemoryChain* mem_chain =
-        &((*device2chain2mem_chain)[device_unique_id][task->task_set_info().chain_id()]);
+    MemoryChain* mem_chain = &((*device2chain2mem_chain)[device_unique_id][task->chain_id()]);
     mem_chain->sorted_tasks.emplace_back(task);
     for (auto& pair : *(task->mutable_produced_regst_desc())) {
       RegstDescProto* regst_desc = &pair.second;
@@ -131,10 +130,10 @@ void InitMemoryChains(Plan* plan,
       MemoryChain* mem_chain = &pair.second;
       std::sort(mem_chain->sorted_tasks.begin(), mem_chain->sorted_tasks.end(),
                 [&](const TaskProto* lhs, const TaskProto* rhs) {
-                  int64_t lhs_order_in_graph = lhs->task_set_info().order_in_graph();
-                  int64_t rhs_order_in_graph = rhs->task_set_info().order_in_graph();
-                  CHECK_NE(lhs_order_in_graph, rhs_order_in_graph);
-                  return lhs_order_in_graph < rhs_order_in_graph;
+                  int64_t lhs_order_in_chain = lhs->order_in_chain();
+                  int64_t rhs_order_in_chain = rhs->order_in_chain();
+                  CHECK_NE(lhs_order_in_chain, rhs_order_in_chain);
+                  return lhs_order_in_chain < rhs_order_in_chain;
                 });
     }
   }

oneflow/core/job/plan_util.cpp

@@ -248,8 +248,8 @@ void PlanUtil::MergeMemBlockIdByLogicalChainId(Plan* plan, const Job& job) {
     DeviceType device_type = stream_id.device_id().device_type();
     // TODO(zwx): eliminate this special 'is cpu' determine
     if (device_type == DeviceType::kCPU) { continue; }
-    if (!IsValidChainId(task->task_set_info().chain_id())) { continue; }
-    int64_t logical_chain_id = task->task_set_info().chain_id();
+    if (!IsValidChainId(task->chain_id())) { continue; }
+    int64_t logical_chain_id = task->chain_id();
 
     for (auto& pair : *(task->mutable_produced_regst_desc())) {
       RegstDescProto* regst_desc = &pair.second;
@@ -300,7 +300,7 @@ void PlanUtil::MergeMemBlockIdByLogicalChainId(Plan* plan, const Job& job) {
     DeviceType device_type = stream_id.device_id().device_type();
     // TODO(zwx): eliminate this special 'is cpu' determine
     if (device_type == DeviceType::kCPU) { continue; }
-    if (!IsValidChainId(task->task_set_info().chain_id())) { continue; }
+    if (!IsValidChainId(task->chain_id())) { continue; }
 
     for (auto& pair : *(task->mutable_produced_regst_desc())) {
       RegstDescProto* regst_desc = &pair.second;
@@ -1066,13 +1066,6 @@ struct RankDeviceMemoryInfo {
 }  // namespace
 
 void PlanUtil::PlanMemoryLog(Plan* plan, const std::string& plan_name) {
-  std::vector<const TaskProto*> ordered_tasks;
-  for (const TaskProto& task : plan->task()) { ordered_tasks.push_back(&task); }
-  auto CompTask = [](const TaskProto* a, const TaskProto* b) {
-    return a->task_set_info().order_in_graph() < b->task_set_info().order_in_graph();
-  };
-  std::sort(ordered_tasks.begin(), ordered_tasks.end(), CompTask);
-
   std::vector<RankDeviceMemoryInfo> rank_device_memory_infos(GlobalProcessCtx::WorldSize(),
                                                              RankDeviceMemoryInfo());
   HashMap<int64_t, MemBlockMemoryInfo> mem_block_id2info;
@@ -1111,8 +1104,8 @@ void PlanUtil::PlanMemoryLog(Plan* plan, const std::string& plan_name) {
     }
   }
 
-  for (const auto* task : ordered_tasks) {
-    for (const auto& pair : task->produced_regst_desc()) {
+  for (const auto& task : plan->task()) {
+    for (const auto& pair : task.produced_regst_desc()) {
       const auto& regst = pair.second;
       if (regst.regst_desc_type().has_data_regst_desc()
           && mem_block_id2info.find(regst.mem_block_id()) != mem_block_id2info.end()) {
@@ -1214,10 +1207,11 @@ void PlanUtil::PlanMemoryLog(Plan* plan, const std::string& plan_name) {
 }
 
 void PlanUtil::GenLightPlan(Plan* plan, const std::string& plan_name) {
+  // NOTE(chengcheng): ordered_tasks is NOT exec order, just task id order.
   std::vector<const TaskProto*> ordered_tasks;
   for (const TaskProto& task : plan->task()) { ordered_tasks.push_back(&task); }
   auto CompTask = [](const TaskProto* a, const TaskProto* b) {
-    return a->task_set_info().order_in_graph() < b->task_set_info().order_in_graph();
+    return a->task_id() < b->task_id();
   };
   std::sort(ordered_tasks.begin(), ordered_tasks.end(), CompTask);
 
@@ -1288,8 +1282,10 @@ void PlanUtil::GenLightPlan(Plan* plan, const std::string& plan_name) {
           << " task_id2name cannot find" << task_id;
       int64_t thrd_id = task->thrd_id();
       StreamId stream_id = DecodeStreamIdFromInt64(thrd_id);
-      file_stream << "order : " << std::to_string(i) << " , actor id : " << std::to_string(task_id)
-                  << " name : " << task_id2name.at(task_id) << " thrd : " << std::to_string(thrd_id)
+      file_stream << "i : " << std::to_string(i) << " , actor id : " << std::to_string(task_id)
+                  << " thrd : " << std::to_string(thrd_id) << " name : " << task_id2name.at(task_id)
+                  << "\n  chain_id : " << std::to_string(task->chain_id())
+                  << " order_in_chain : " << std::to_string(task->order_in_chain())
                   << " device_type : " << DeviceType_Name(stream_id.device_type())
                   << " stream_index : " << std::to_string(stream_id.stream_index()) << " {\n";
       for (const auto& key2consume_regst : task->consumed_regst_desc_id()) {

oneflow/core/job/task.proto

@@ -44,23 +44,19 @@ message RegstDescIdSet {
   repeated int64 regst_desc_id = 1;
 }
 
-message TaskSetInfo {
-  required int64 chain_id = 4;
-  required int64 order_in_graph = 5;
-}
-
 message TaskProto {
   // common
   required TaskType task_type = 1;
   required int64 machine_id = 2;
   required int64 thrd_id = 3;
   required int64 task_id = 4;
   required int64 job_id = 5;
-  required TaskSetInfo task_set_info = 6;
   required ExecSequence exec_sequence = 7;
   map<string, RegstDescProto> produced_regst_desc = 8;
   map<string, RegstDescIdSet> consumed_regst_desc_id = 9;
   optional bool all_register_num_eq_one_hint = 10 [default = false];
+  required int64 chain_id = 20;
+  required int64 order_in_chain = 21;
   // compute task
   optional ParallelContext parallel_ctx = 1000; // CompTask
 };

oneflow/core/job_rewriter/logical_chain_pass.cpp

@@ -625,9 +625,11 @@ Maybe<void> LogicalChainPass::Apply(const OpGraph& op_graph, JobBuilder* job_bui
 
   auto InsertLogicalChainId = [&](const std::vector<const OpNode*>& ordered_op_nodes,
                                   const int64_t logical_chain_id) {
+    int64_t order = 0;
     for (const OpNode* op_node : ordered_op_nodes) {
-      CHECK_JUST(MapAt(mut_op_name2conf, op_node->op().op_name()))
-          .set_logical_chain_id(logical_chain_id);
+      auto& conf = CHECK_JUST(MapAt(mut_op_name2conf, op_node->op().op_name()));
+      conf.set_logical_chain_id(logical_chain_id);
+      conf.set_order_in_logical_chain(order++);
     }
   };
 

oneflow/core/operator/op_conf.proto

@@ -398,6 +398,7 @@ message OperatorConf {
   optional string pass_tag = 10;
   optional string loc = 11 [default = ""];
   optional int64 logical_chain_id = 12 [default = -1];
+  optional int64 order_in_logical_chain = 13 [default = -1];
   oneof op_type {
     // system op
     CopyCommNetOpConf copy_comm_net_conf = 106;