Function GetStreamId

oneflow/core/graph/task_graph.cpp:297–318

namespace {

StreamId GetStreamId(const OpNode* op_node, int64_t parallel_id, TaskType task_type) {
  const ParallelDesc& parallel_desc = op_node->parallel_desc();
  int64_t machine_id = CHECK_JUST(parallel_desc.MachineId4ParallelId(parallel_id));
  int64_t dev_phy_id = CHECK_JUST(parallel_desc.DeviceId4ParallelId(parallel_id));

  DeviceId::device_index_t device_index = parallel_desc.device_type() == DeviceType::kCPU
                                              ? 0
                                              : static_cast<DeviceId::device_index_t>(dev_phy_id);
  DeviceId device_id{static_cast<DeviceId::rank_t>(machine_id), parallel_desc.device_type(),
                     device_index};
  StreamId::stream_index_t stream_index = 0;
  if (op_node->op().op_conf().has_stream_name_hint()) {
    const std::string& stream_name_hint = op_node->op().op_conf().stream_name_hint();
    VLOG(3) << "set op: " << op_node->op().op_name() << " to stream: " << stream_name_hint;
    stream_index = Singleton<TaskStreamIndexManager>::Get()->GetNamedTaskStreamIndex(
        device_id, stream_name_hint);
  } else {
    stream_index =
        Singleton<TaskStreamIndexManager>::Get()->GetTaskStreamIndex(task_type, device_id);
  }
  return StreamId{device_id, stream_index};
}

TaskType TaskType4OpNode(const OpNode* op_node) {
  std::unique_ptr<CompTaskNode> comp_task_node(NewCompTaskNode4OpNode(op_node));