Function SyncAccessSmallMem

oneflow/core/framework/instructions_builder.cpp:792–828

template<typename T, typename InstructionPolicyT>
Maybe<void> SyncAccessSmallMem(char* mem_ptr, size_t bytes, const T tensor) {
  static thread_local vm::InstructionList instruction_list;
  static thread_local InstructionsBuilder instructions_builder(&instruction_list);
  const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object = JUST(tensor->eager_blob_object());
  const Symbol<Stream> stream = JUST(GetAccessStream(tensor));
  if (eager_blob_object->last_used_stream().has_value()
      && stream != JUST(eager_blob_object->last_used_stream())) {
    // Synchronize stream.
    JUST(instructions_builder.SoftSyncStream({eager_blob_object}, stream));
  }
  InstructionPolicyT* instruction_policy = nullptr;
  {
    // Construct instruction.
    auto* instruction = JUST(MutThreadLocalInstruction<InstructionPolicyT>(stream));
    instruction_policy =
        static_cast<InstructionPolicyT*>(instruction->mut_instruction_policy());  // NOLINT
    instruction_policy->Reset(mem_ptr, bytes, eager_blob_object.get());
    instruction_list.PushBack(instruction);
  }
  // Dispatch instructions.
  JUST(vm::Run(&instruction_list));
  {
    // This thread should blocking wait if and only if there is a lot of workload on worker thread.
    // When workload is small, we want better performance by skipping cond_.notify_xxx which costs
    // about 2us to 3us.
    auto* virtual_machine = JUST(SingletonMaybe<VirtualMachine>());
    static constexpr int kSkipBlockingThreshold = 2;
    if (virtual_machine->flying_instruction_cnt() < kSkipBlockingThreshold) {
      // skip pthread_cond_broadcast on worker thread.
      instruction_policy->mut_btb()->mut_notifier()->Notify();
    }
  }
  // wait until done.
  JUST(instruction_policy->mut_btb()->WaitUntilCntEqualZero(
      VirtualMachine::GetPredicatorNoMoreInstructionsFinished()));
  return Maybe<void>::Ok();
}

template<typename T>
Maybe<void> SyncReadSmallMem(char* mem_ptr, size_t bytes, const T tensor) {