Function GenCompTaskNode

oneflow/core/graph/task_graph.cpp:327–342

}  // namespace

CompTaskNode* GenCompTaskNode(
    const OpNode* op_node, int64_t parallel_id,
    const std::function<StreamId(const OpNode* op_node, int64_t parallel_id, TaskType task_type)>&
        GetOrCreateStreamId) {
  const ParallelDesc& parallel_desc = op_node->parallel_desc();
  int64_t parallel_num = parallel_desc.parallel_num();
  CompTaskNode* comp_task_node = NewCompTaskNode4OpNode(op_node);
  int64_t machine_id = CHECK_JUST(parallel_desc.MachineId4ParallelId(parallel_id));
  comp_task_node->set_machine_id(machine_id);
  comp_task_node->mut_parallel_ctx()->set_parallel_id(parallel_id);
  comp_task_node->mut_parallel_ctx()->set_parallel_num(parallel_num);
  StreamId stream_id = GetOrCreateStreamId(op_node, parallel_id, comp_task_node->GetTaskType());
  comp_task_node->set_thrd_id(EncodeStreamIdToInt64(stream_id));
  comp_task_node->set_op_node(op_node);
  return comp_task_node;
}

void GenSortedCompTaskNodes(const OpNode* op_node, std::vector<CompTaskNode*>* sorted_comp_tasks) {
  int64_t parallel_idx = 0;