Method CreateThreadCtx

oneflow/core/vm/virtual_machine.cpp:425–462

}

Maybe<vm::ThreadCtx*> VirtualMachine::CreateThreadCtx(Symbol<Device> device, StreamType stream_type,
                                                      size_t thread_uid) {
  std::unique_lock<std::recursive_mutex> lock(stream_and_thread_ctx_mutex_);
  // thread_ctx_ptr may be used after timout.
  auto thread_ctx_ptr = std::make_shared<vm::ThreadCtx*>(nullptr);
  {
    auto bc = std::make_shared<BlockingCounter>(1);
    engine_->InsertProbe([thread_ctx_ptr, bc](vm::VirtualMachineEngine* engine) {
      auto thread_ctx = intrusive::make_shared<vm::ThreadCtx>();
      engine->mut_thread_ctx_list()->PushBack(thread_ctx.Mutable());
      *thread_ctx_ptr = thread_ctx.Mutable();
      bc->Decrease();
      return true;
    });
    JUST(NotifyOrRunScheduler());
    JUST(bc->WaitUntilCntEqualZero(VirtualMachine::GetPredicatorNoMoreInstructionsFinished()));
  }
  auto* thread_ctx = *thread_ctx_ptr;
  {
    const std::string thread_tag = [&] {
      std::string device_tag = *CHECK_JUST(DeviceTag4DeviceType(device->enum_type()));
      if (StreamOnIndependentThread::Visit(stream_type)) {
        return device_tag + GetStreamTypeName::Visit(stream_type);
      } else {
        return std::to_string(thread_uid);
      }
    }();
    const auto& WorkerInitializer = [thread_tag](vm::ThreadCtx* thread_ctx) {
      OF_PROFILER_NAME_THIS_HOST_THREAD("_VM::Worker_" + thread_tag);
    };
    auto thread = std::make_unique<std::thread>(&WorkerLoop, thread_ctx, WorkerInitializer);
    {
      std::unique_lock<std::mutex> lock(worker_threads_mutex_);
      worker_threads_.push_back(std::move(thread));
    }
  }
  return thread_ctx;
}

Maybe<vm::Stream*> VirtualMachine::CreateStream(vm::ThreadCtx* thread_ctx, Symbol<Stream> stream) {
  std::unique_lock<std::recursive_mutex> lock(stream_and_thread_ctx_mutex_);