作者｜鄭建華
更新｜趙露陽

通過這篇筆記，希望能初步瞭解 OneFlow 在 Eager 模式下對裝置的管理方式、裝置執行計算的過程以及如何充分利用裝置計算能力。這裡的裝置主要指類似 CUDA 這樣的平行計算加速裝置。

1

裝置、流相關型別及關係

框架通過流（Stream）向裝置（Device）提交計算任務。一個 Stream 是一個命令序列，可以類比 CUDA Stream（https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#streams），或者 CPU Thread 的指令序列。同一個 Stream 中的命令按順序執行；不同 Stream 之間的命令有依賴關係時，需要同步。不同的任務，比如 kernel 計算、host2device、device2host 等都有自己獨立的 Stream，可以併發執行，從而在 Eager 模式下儘可能充分利用裝置的非同步併發執行能力。

OneFlow 中Device和Stream相關的部分類結構如下所示：

Device相關型別

oneflow::Device

oneflow::Device是用於表示裝置的基礎型別，例如：構建tensor時 flow.tensor(shape, device="cuda:1")就會在內部構造出這個基礎的Device型別，其中裝置編號為1、裝置型別為CUDA。

oneflow/core/framework/device.h：

  
  
  

   
   
   
class Device final { public:  ...
 private:  Device(const std::string& type, int64_t device_id);  Maybe<void> Init();
  const std::string type_;  DeviceType enum_type_;  const int64_t device_id_;  const size_t hash_value_;  std::shared_ptr<MemoryCase> mem_case_;};
  
  
  

oneflow::Device中最重要的兩個成員變數分別是用於表示裝置型別的DeviceType；用於表示裝置編號的device_id_。

DeviceType

DeviceType是一個列舉類，不同的值代表不同的計算裝置型別，其定義位於 oneflow/core/common/device_type.proto：

  
  
  

   
   
   
enum DeviceType {  kInvalidDevice = 0;  // 無效裝置  kCPU = 1;            // cpu裝置  kCUDA = 2;           // cuda裝置  kMockDevice = 3;     // pseudo device for test.}
  
  
  

目前在oneflow master分支中，主要有kCPU表示cpu裝置；kCUDA表示nvidia cuda裝置；在其他多裝置支援的分支中，這裡還增加了更多的裝置型別。

oneflow::ep::Device

oneflow::Device是oneflow中對裝置的基礎封裝型別，而oneflow::ep::Device則是一個抽象類，屬於oneflow的ep模組（execution provider），是對裝置行為的封裝，ep模組為多硬體裝置提供了更高層次的抽象，方便oneflow支援和相容多硬體裝置提供了更高的靈活性和可拓展性。

oneflow::ep::Device不僅提供了表示裝置型別的device_type()方法、表示裝置編號的device_index()方法，還提供了建立/銷燬ep::Stream、建立/銷燬Event、在裝置上申請/釋放記憶體的各種方法。

oneflow/core/ep/include/device.h：

  
  
  

   
   
   
class Device { public:  OF_DISALLOW_COPY_AND_MOVE(Device);  Device() = default;  virtual ~Device() = default;
  virtual void SetAsActiveDevice() = 0;
  virtual DeviceType device_type() const = 0;  virtual size_t device_index() const = 0;  virtual DeviceManager* device_manager() const = 0;
  virtual Stream* CreateStream() = 0;  virtual void DestroyStream(Stream* stream) = 0;
  virtual Event* CreateEvent();  virtual void DestroyEvent(Event* event);  virtual void CreateEvents(Event** events, size_t count) = 0;  virtual void DestroyEvents(Event** events, size_t count) = 0;
  virtual Maybe<void> Alloc(const AllocationOptions& options, void** ptr, size_t size) = 0;  virtual void Free(const AllocationOptions& options, void* ptr) = 0;  virtual Maybe<void> AllocPinned(const AllocationOptions& options, void** ptr, size_t size) = 0;  virtual void FreePinned(const AllocationOptions& options, void* ptr) = 0;  virtual bool IsStreamOrderedMemoryAllocationSupported() const;};
  
  
  

oneflow::ep::Device有如下子類實現：

Stream相關型別

oneflow::Stream和cuda device以及stream的關係類似，oneflow中也存在類似的基礎Stream型別。

oneflow/core/framework/stream.h：

  
  
  

   
   
   
class Stream final { .... private:  Stream(Symbol<Device> device, StreamType stream_type, size_t thread_uid);
  static Maybe<Symbol<Stream>> RawNew(Symbol<Device> device, StreamType stream_type,                                      size_t thread_uid);
  Maybe<void> Init(size_t unique_stream_id);
  Symbol<Device> device_;  StreamType stream_type_;  size_t thread_uid_;  size_t unique_stream_id_;};
  
  
  

可以看見Stream類中的成員變數：

• device_  表示該Stream物件將在何種裝置上執行
• streamtype_  表示該Stream的型別，是用於計算的compute stream還是用於資料搬運的host2device、device2host stream等
• threaduid_ 表示負責啟動該Stream的執行緒id
• unique_streamid_ 表示這個stream自身的unique id

StreamType

和 DeviceType 分為kCpu和kCuda類似， Stream 也有各種型別之分，具體如下：

oneflow/core/common/stream_type.h

  
  
  

   
   
   
enum class StreamType {    kInvalid = 0,    // 無效    kCompute,        // kernel計算流    kHost2Device,    // 資料搬運（host -> device）流    kDevice2Host,    // 資料搬運（device -> host）流    kCcl,            // 集合通訊流    kBarrier,        // 執行緒屏障流    kCriticalSection,// 臨界區流    kLazyJobLauncher,// job啟動流（lazy mode）    kPinnedCompute   // pinned memory kernel計算流};
  
  
  
  
  
  

   
   
   

  
  
  

oneflow::ep::Stream

oneflow中的ep模組提供了一個更高層次的對Stream的抽象類，除了可以獲取裝置的 device() 、獲取裝置型別的 device_type() 方法外，還提供了一系列虛方法如：

• 同步Sync()
• 執行Event事件RecordEvent()

oneflow/core/ep/include/stream.h ：

  
  
  

   
   
   
class Stream { public:  OF_DISALLOW_COPY_AND_MOVE(Stream);  Stream() = default;  virtual ~Stream() = default;
  virtual DeviceType device_type() const = 0;  virtual Device* device() const = 0;  virtual Maybe<void> Sync() = 0;  virtual void RecordEvent(Event* event) = 0;  virtual Maybe<void> GetAsyncError() { return Maybe<void>::Ok(); }
  virtual Maybe<void> AllocAsync(void** ptr, size_t size) { UNIMPLEMENTED_THEN_RETURN(); }  virtual Maybe<void> FreeAsync(void* ptr) { UNIMPLEMENTED_THEN_RETURN(); }  template<typename T>  Maybe<void> AllocAsync(T** ptr, size_t size) {    return AllocAsync(reinterpret_cast<void**>(ptr), size);  }
  virtual Maybe<void> OnExecutionContextSetup() { return Maybe<void>::Ok(); }  virtual Maybe<void> OnExecutionContextTeardown() { return Maybe<void>::Ok(); }
  template<typename T>  T* As() {    return static_cast<T*>(this);  }};
  
  
  
  
  
  

   
   
   

  
  
  

oneflow::ep::Stream 有如下子類實現：

oneflow::vm::Stream

oneflow vm（virtual machine）中的 oneflow::vm::Stream 型別，用於vm內部維護stream極其依賴關係、StreamPolicy、排程執行緒等。

oneflow/core/vm/stream.h ：

  
  
  

   
   
   
class Stream final : public intrusive::Base { public:  ... private:  ...  // fields  ThreadCtx* thread_ctx_;  Symbol<Device> device_;  StreamType stream_type_;  std::shared_ptr<StreamPolicy> stream_policy_;  bool on_scheduler_thread_;  std::unique_ptr<char, std::function<void(char*)>> small_pinned_mem_ptr_;  ...};
  
  
  
  
  
  

   
   
   

  
  
  

StreamPolicy

StreamPolicy 是oneflow vm中獨有的概念，提供了一系列虛方法如：

• stream() 獲取oneflow::ep::Stream指標

• mut_allocator() 獲取vm::Allocator指標（用於tensor記憶體管理）

• device_type() 獲取device裝置型別

除此之外，提供了一系列vm相關的指令狀態初始化、查詢、刪除等方法。

oneflow/core/vm/stream_policy.h ：

  
  
  

   
   
   
class StreamPolicy { public:  virtual ~StreamPolicy() = default;
  virtual ep::Stream* stream() = 0;  virtual vm::Allocator* mut_allocator() = 0;  virtual DeviceType device_type() const = 0;
  virtual void InitInstructionStatus(const Stream& stream,                                     InstructionStatusBuffer* status_buffer) const = 0;  virtual void DeleteInstructionStatus(const Stream& stream,                                       InstructionStatusBuffer* status_buffer) const = 0;  virtual bool QueryInstructionStatusLaunched(      const Stream& stream, const InstructionStatusBuffer& status_buffer) const = 0;  virtual bool QueryInstructionStatusDone(const Stream& stream,                                          const InstructionStatusBuffer& status_buffer) const = 0;  virtual bool OnSchedulerThread(StreamType stream_type) const;  virtual bool SupportingTransportInstructions() const = 0;
  void RunIf(Instruction* instruction) const;
 protected:  StreamPolicy() = default;
 private:  virtual void Run(Instruction* instruction) const = 0;};
  
  
  
  
  
  

   
   
   

  
  
  

StreamPolicy有如下子類實現：

2

Eager Local模式下的Device和Stream推導

下面，梳理一下普通的eager模式（eager local mode）下，運算元執行全過程中device和stream相關的推導流程。

2.1 推導Device

首先，對於一個運算元（op）來說，要為其設定一個預設的device用於實際計算，這一步在：

Symbol<Device> default_device = JUST(GetDefaultDevice(inputs, ctx))

這裡 GetDefaultDevice 的邏輯是：

• 1.如果inputs tensor非空，則根據第一個input tensor的device來設定default的device
• 2.如果inputs tensor為空，則優先從OpExprInterpContext中獲取device，若OpExprInterpContext中未設定，則會通過
  Device::New("cpu")
  ;預設給一個cpu device

值得說明的是，在1.種情況時，如果input tensor建立時指定了device為cuda裝置，則這裡推匯出的default device同樣為相同的cuda device；如果未顯示指定，則預設還是cpu device。

2.2 推導Stream

oneflow::Stream的推導主要在：

• JUST(user_op_expr.mut_local_tensor_infer_cache()->GetOrInfer(infer_args))
  ）；
  Symbol<Stream> stream = JUST(InferDeviceAndStream(... ));

InferDeviceAndStream 中， Stream 推導的邏輯是會根據user_op_expr是否定義了device_and_stream_infer_fn而有所區別

• （少數情況）如果該op定義了推導函式，則呼叫此推導函式來推導Stream，例如 tensor.cuda() 方 法，inputs 在 CPU 上， output s 在 CUDA，二者的裝置型別不同。 這時就不會預設推導而是利用 op 註冊的推導函式獲取 oneflow::Stream（
  。 例如 CopyOp::InferDeviceA ndStream 。
• （多數情況）否則會通過
  stream = JUST(GetDefaultStreamByDevice(default_device))
  ; 來推導。

GetDefaultStreamByDevice 的具體實現：

  
  
  

   
   
   
Maybe<Symbol<Stream>> RawGetDefaultStreamByDevice(Symbol<Device> device) {  return Stream::New(device, StreamType::kCompute);}
  
  
  
  
  
  

   
   
   
  
  
  
  
  
  

   
   
   

  
  
  

可以看見，根據傳入的 device 、 StreamType::kCompute ，new了一個 oneflow::Stream 。

2.3 InstructionsBuilder::Call和vm::Stream推導

在上述device和stream推導完成後，會通過InstructionsBuilder呼叫Call方法 ：

  
  
  

   
   
   
JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {    return builder->Call(kernel, std::move(input_eager_blob_objects),                         std::move(output_eager_blob_objects), ctx, result->stream());}));
  
  
  
  
  
  

   
   
   
  
  
  
  
  
  

   
   
   

  
  
  

Call方法中會通過

• JUST(SoftSyncStream(output_eager_blob_objects, stream)) ;
• JUST(SoftSyncStream(input_eager_blob_objects, stream)) ;
• auto* vm_stream = JUST(Singleton<VirtualMachine>::Get()->GetVmStream(stream)) ;

完成outputs inputs tensor的流同步（SoftSyncStream）過程以及 vm::Stream 的推導，然後通過構造 OpCallInstructionPolicy 指令派發至vm執行。

SoftSyncStream的同步這裡省略，具體過程見第4節。

2.3.1 構造 ThreadCtx 物件，啟動執行指令的執行緒

ThreadCtx 物件指標儲存在 VirtualMachine 的 HashMap  中。每個 DeviceType （CPU或CUDA）對應一個 ThreadCtx 物件；臨界區和 LazyJob  有自己的 ThreadCtx 物件 。

首次訪問 HashMap  時得到的是零值（空指標），需要呼叫 CreateThreadCtx  建立物件。實際通過虛擬機器指令建立物件 ，ThreadCtx 物件儲存在 VirtualMachineEngine::thread_ctx_list_ 中 。

ThreadCtx 物件構造後，會建立一個 worker 執行緒、執行 WorkerLoop 方法 ，並新增到 worker_threads_ 。所以 worker_threads_ 是與 ThreadCtx 物件一一對應的。

這個執行緒負責其所歸屬的指令 的執行：

• WorkerLoop 在收到通知 後，會呼叫 ThreadCtx::TryReceiveAndRun   處理指令。
• 在這個函式中，將 ThreadCtx 的指令挪到臨時列表 、通過 StreamPolicy 執行每個指令 。
• ThreadCtx 的指令，是 VirtualMachineEngine 在 DispatchInstruction 時新增進去 的。

ThreadCtx建立完成後，將持有vm::Stream物件 。 oneflow::vm::Stream 和 oneflow::Stream 的數量是一一對應的 ，vm::Stream 按照<DeviceType, StreamRole>分組儲存在對應的 ThreadCtx 中 。

vm::Stream 的推導流程細節如下：

• auto* vm_stream = JUST(Singleton<VirtualMachine>::Get()->GetVmStream(stream)) ;

• VirtualMachine::GetVmStream()
• Maybe<vm::Stream*>      VirtualMachine::CreateStream(Symbol<Stream> stream)
• Stream::__Init__(ThreadCtx* thread_ctx, Symbol<Device> device, StreamType stream_type...)
• stream_policy_ = CHECK_JUST(CreateStreamPolicy::Visit(stream_type, device)) ;

2.4 執行OpCall指令和ep::Stream推導

有幾個場景會建立（獲取） ep::Stream 物件。比如 kernel 執行時。

OpCall指令在構造時，指令策略型別是 OpCallInstructionPolicy 。虛擬機器在 DispatchInstruction 時，無論哪個分支，後續都會呼叫 EpStreamType::Run ，最終通過

• EpStreamPolicyBase::Run()

• instruction->Compute()
• OpCallInstructionPolicy::Compute()
• OpCallInstructionUtil::Compute()
• OpCallInstructionUtil::OpKernelCompute()
• op_call_instruction_policy->mut_opkernel()->Compute()執行 kernel 的 Compute 方法

例如 GpuL2NormalizeKernel::Compute ，最終在其kernel的Compute方法中，會通過ctx->stream()建立 （獲取）ep::Stream 物件，launch kernel 執行計算。

2.4.1 獲取/建立ep::Stream

下面，我們重點看一下OpCall指令實際執行時，呼叫的OpCallInstructionUtil::Compute() 方法：

  
  
  

   
   
   
static inline Maybe<void> Compute(OpCallInstructionPolicy* op_call_instruction_policy,                                    Instruction* instruction) {    Allocator* allocator = instruction->mut_stream()->mut_stream_policy()->mut_allocator();    JUST(AllocateOutputBlobsMemory(op_call_instruction_policy, allocator, instruction));    if (unlikely(op_call_instruction_policy->need_temp_storage())) {      JUST(TryAllocateTempStorage(op_call_instruction_policy, allocator));    }    ep::Stream* stream = instruction->mut_stream()->mut_stream_policy()->stream();    user_op::OpKernelState* state = nullptr;    user_op::OpKernelCache* cache = nullptr;    if (op_call_instruction_policy->user_opkernel()->has_state_or_cache()) {      TryInitOpKernelStateAndCache(op_call_instruction_policy, stream, &state, &cache);    }    OpKernelCompute(op_call_instruction_policy, stream, state, cache);    if (unlikely(op_call_instruction_policy->need_temp_storage())) {      DeallocateTempStorage(op_call_instruction_policy, allocator);    }    return Maybe<void>::Ok();  }
  
  
  
  
  
  

   
   
   

  
  
  

其中會通過 ep::Stream* stream = instruction->mut_stream()->mut_stream_policy()->stream() ;完成 ep::Stream 的推導，之後在 OpKernelCompute() 方法中實際完成op/kernel的執行。

ep::Stream* stream = instruction->mut_stream()->mut_stream_policy()->stream() ;

• ep::Stream* stream() override { return GetOrCreateEpStream(); }

• GetOrCreateEpStream()
• ep_stream_ = GetOrCreateEpDevice()->CreateStream() ;

這裡->stream()會呼叫ep_stream_policy_base.h中的：

ep::Stream* stream() override { return GetOrCreateEpStream(); }

這是一個private方法：

  
  
  

   
   
   
private:  ep::Stream* GetOrCreateEpStream() const {    if (unlikely(ep_stream_ == nullptr)) {      ep_stream_ = GetOrCreateEpDevice()->CreateStream();      CHECK(ep_stream_ != nullptr);    }    return ep_stream_;  }
  
  
  
  
  
  

   
   
   
  
  
  
  
  
  

   
   
   

  
  
  

可以看到，如果成員變數 ep_stream_ 非空，則直接返回；否則，通過 ep_stream_ = GetOrCreateEpDevice()->CreateStream() ; 來建立建立 ep::Stream 。

2.4.2 獲取/建立ep::Device

而這裡的 GetOrCreateEpDevice 方法如下：

  
  
  

   
   
   
ep::Device* GetOrCreateEpDevice() const {    if (unlikely(ep_device_ == nullptr)) {      ep_device_ = Singleton<ep::DeviceManagerRegistry>::Get()->GetDevice(device_->enum_type(),                                                                          device_->device_id());      CHECK(ep_device_);    }    return ep_device_.get();  }
  
  
  
  
  
  

   
   
   

  
  
  

根據 oneflow::Device 中拿到的device id和device type，去全域性單例的 ep::DeviceManagerRegistry 中取出對應的 oneflow::ep::Device

oneflow::vm::StreamPolicy和oneflow::vm::EpStreamPolicy推導

• stream_policy_ = CHECK_JUST(CreateStreamPolicy::Visit(stream_type, device)) ;

• std::shared_ptr<vm::StreamPolicy>(new vm::EpStreamPolicy(device)) ;

3

Eager Global模式下的Device和Stream推導

eager global模式下，device資訊隱藏在placement中，placement不僅包括了device type資訊還包括其tensor具體分佈在哪些ranks上的資訊， placement 在 C++ 中的對應型別是 ParallelDesc 。

所以device以及stream的部分推導過程和eager local模式下有所區別，但OpCall指令執行；device、vm::Stream和ep::Stream的推導過程都和eager local模式下是類似的。

3.1 推導Device

3.1.1 placement 的 parallel_id

oneflow中的placement表示tensor存放的裝置叢集（device group），如：

p = flow.placement(type="cuda", ranks=[0, 1, 2, 3])

表示tensor分佈於1臺機器上，cuda device 0、1、2、3四個裝置上；

p = flow.placement(type="cuda", ranks=[[0, 1], [2, 3]])

則表示tensor分佈於2臺機器上，host1的device0、1以及host2的device2、3。

在 oneflow 的分散式環境下，各個 host 上需要有相同數量的device，每個程序使用一個device。這樣根據環境變數 RANK 可以得出 machine_id， LOCAL_RANK 就是程序在 制定host 上的 rank序號。

如果 input tensor 的 placement 與當前程序無關，可以省掉很多不必要的計算。通過 placement 的 parallel_id 可以判斷計算任務是否與當前程序相關。

placement 在 C++ 中的對應型別是 ParallelDesc ，其中並沒有 parallel_id 欄位，這個資訊隱含在其它欄位中。

ParallelDesc 在構造時會呼叫 ClearUp 函式，從中可以看到

• ParallelDesc::parallel_id2machine_id_ 是 placement 分佈的 machine。
• ParallelDesc::parallel_id2device_id_ 是 placement 分佈的 device_id。
• parallel_id 是上述 2 個數組的索引，一個 parallel_id 對應一個 machine_id:device_id 組合。這樣，根據parallel_id可以查到對應的 machine_id 和 device_id。
• 反過來，根據 machine_id:device_id 也可以從 machine_id2device_id2parallel_id_ 查到 parallel_id。

  
  

3.1.2 eager 模式下根據 parallel_id 忽略無關計算任務

在 eager 分佈時場景處理計算任務時，會呼叫 GetTensorDevice4CurrentProcessCtx ，推導得到輸出tensor的device，以及獲取當前程序的 machine_id、device_id 在 placement 中的 parallel_id 值。

如果當前程序與該 placement 無關，parallel_id 就是空，後續處理時就可以忽略一些計算：

• EagerGlobalTensorImpl::New  中只需要用 functional::Empty 構造一個 shape 為 0 的空的 tensor。
• GetBoxingOutput 計算時，如果parallel_id為空則表示當前rank程序無效，無需計算直接返回。
• Interpret 可以不給 vm 提交指令、提前返回 。

  
  

3.2 推導Stream

在 ConsistentTensorInferCache 中推導 SBP Signature 時，也會同時推匯出當前的 tensor 計算任務、在當前程序所用的device。推導時，會先確認所有 inputs 的 placement 是一致的 ，都分佈在相同的device上。如前所述，如果計算任務與當前程序無關，會提前返回；而一個程序只使用一個device。

這裡和eager local模式下stream的推導類似，通過 JUST(InferDeviceAndStream(user_op_expr, infer_args)) 推匯出 oneflow::Stream 物件，StreamRole 是 kCompute 。區別在於eager global模式下

3.2.1 unique_stream_id

unique_stream_id 表示 oneflow::Stream 物件的建立次序。

所有的 oneflow::Stream 物件都儲存在全域性的 StreamMgr::stream2unique_stream_id_ 中 。 unique_stream_id2stream_symbol_ 可看作是引用型別的副本， unique_stream_id 就是 Stream 物件在這個陣列中的索引。與 parallel_id 不同，unique_stream_id 是 Stream 物件在程序內的唯一標識。

並不是每次都需要加鎖訪問 StreamMgr 。 oneflow::Stream 包含的都是描述性資訊，其引用是以 ThreadLocal  的方式儲存的，可以提升後續讀取的效率。虛擬機器在執行指令時，也會用 unique_stream_id 進行邏輯判斷。

4

Eager模式下的Stream同步——SoftSyncStream

設想以下場景：將 CPU 下的 tensor 拷貝到 CUDA 裝置，然後在 CUDA 上再進行 tensor add 的計算。這涉及到兩個流，一個是 Host2Device，一個是 CUDA Compute。這兩個流的計算任務是併發執行的。需要有同步措施，才能保證拷貝完再執行 add 計算。

Eager 模式下，在 InstructionsBuilder::Call 中構造指令時，對 SoftSyncStream 的呼叫會在必要時向指令列表插入同步指令。

SoftSyncStream 中，幾個重要概念：

• tensor在oneflow記憶體中的實際承載者是 eager_blob_object
• last_used_stream 表示一個tensor(blob)上一次使用到的stream，可能是compute stream、h2d stream、d2h stream、集合通訊ccl stream等
• 如果 last_used_stream 與當前計算執行的流 stream 相同，則可以忽略，因為相同stream間天然順序執行所以無需同步，否則就需要進行後續的同步處理

SoftSyncStream 程式碼如下：

  
  
  

   
   
   
Maybe<void> InstructionsBuilder::SoftSyncStream(const vm::EagerBlobObjectList& eager_blob_objects,                                                Symbol<Stream> stream) {  JUST(ForEachEagerBlobObjectsNeedingSoftSync(      eager_blob_objects, stream,      [&](Symbol<Stream> last_used_stream, auto&& dep_objects) -> Maybe<void> {        return SoftSyncStreamBetween(std::move(dep_objects), last_used_stream, stream);      }));  for (const auto& eager_blob_object : eager_blob_objects) {    eager_blob_object->set_last_used_stream(stream);  }  return Maybe<void>::Ok();}
  
  
  
  
  
  

   
   
   

  
  
  

主體邏輯是，會在 ForEachEagerBlobObjectsNeedingSoftSync 方法中遍歷每一個tensor物件（eager blob object），對於每一個需要同步的blob運用lambda方法並最終呼叫 SoftSyncStreamBetween 完成stream間的同步。

這裡，我們看一下ForEachEagerBlobObjectsNeedingSoftSync 的邏輯：

首先if/else的主體業務邏輯是類似的，主要區別在於，當blob的size <= kOpArgsReservedSize時（預設為4）會使用small vector來存放LocalDepObject變數，效率會更快一些（否則會走到else分支，主體邏輯類似，這裡就不看了）。

• const auto& opt_last_used_stream = eager_blob_object->last_used_stream()
  ;
• if (unlikely(!opt_last_used_stream.has_value())) { continue; }

這兩句是查詢該tensor(blob)上一次被使用時用到的stream——last_used_stream，如果為空，則直接continue跳過，因為如果此tensor之前並未被任何stream使用，則無需進行stream間的同步操作，因為在當前stream上不會有關於該tensor的其他依賴關係；

  
  
  

   
   
   
if (last_used_stream != stream) {    small_vector<intrusive::shared_ptr<LocalDepObject>, kOpArgsReservedSize> dep_objects{        intrusive::shared_ptr<LocalDepObject>(            JUST(eager_blob_object->compute_local_dep_object()))};    JUST(DoEach(last_used_stream, std::move(dep_objects)));  }
  
  
  
  
  
  

   
   
   

  
  
  

如果 last_used_stream!=stream 則表示需要在兩個stream間進行同步，則會應用傳入的lambda函式DoEach進行處理，在這裡lambda函式即：

  
  
  

   
   
   
[&](Symbol<Stream> last_used_stream, auto&& dep_objects) -> Maybe<void> {    return SoftSyncStreamBetween(std::move(dep_objects), last_used_stream, stream);  }));
  
  
  
  
  
  

   
   
   

  
  
  

既實際呼叫的是SoftSyncStreamBetween 來完成實際的stream間同步，這裡主要有3個變數：

• dep_objects
  儲存了tensor間的依賴關係
• last_used_stream
  則是該tensor上一次使用的stream
• stream
  該tensor當前使用的stream

  
  

SoftSyncStreamBetween 的程式碼如下：

  
  
  

   
   
   
Maybe<void> InstructionsBuilder::SoftSyncStreamBetween(    small_vector<intrusive::shared_ptr<LocalDepObject>, kOpArgsReservedSize>&& dependences,    Symbol<Stream> from_stream, Symbol<Stream> to_stream) {  CHECK(from_stream != to_stream) << "synchronization is unnecessary";  if (SupportingStreamWait(from_stream, to_stream)) {    JUST(StreamWait(std::move(dependences), from_stream, to_stream));  } else {    JUST(RecordEvent(std::move(dependences), from_stream));  }  return Maybe<void>::Ok();}
  
  
  
  
  
  

   
   
   

  
  
  

SoftSyncStreamBetween 的主要邏輯如下：

• 先額外做了一次check，檢測如果待同步的兩個stream相同，則check會報錯並提示"synchronization is unnecessary"
• 通過SupportingStreamWait判斷from 和 to stream間是否支援stream wait，是則呼叫StreamWait方法；否則，直接呼叫RecordEvent 方法
  
  
• SupportingStreamWait 的主要邏輯是，通過stream的device、以及StreamType的Visit方法來判斷。簡單來說，如果from 和 to stream之間是不同的device（譬如cpu stream <-> cuda stream之間的同步），或者from stream的device為cpu，則SupportingStreamWait一定是false；如果是相同的，則繼續通過其他判斷條件進行判斷。
  
  

SupportingStreamWait為True

SupportingStreamWait 為True時，即from to stream同為Cuda Stream間的同步情況，在這種情況下會走到StreamWait 的函式，該函式最終會派發一個 StreamWaitEventInstructionPolicy 的指令給vm執行，StreamWaitEventInstructionPolicy 的執行邏輯主要是兩個cuda event：

• cudaEventRecord
• cudaStreamWaitEvent

  
  

• 對於from_stream來說，插入一個
  cudaEventRecord
  ，用於標誌from stream是否完成該stream上的event事件；

  
  

• 對於to_stream來說，插入一個
  cudaStreamWaitEvent
  等待from stream上的事件完成後，再繼續執行to_stream。

SupportingStreamWait為False

SupportingStreamWait 為False時，會直接呼叫 JUST(RecordEvent(std::move(dependences) , from_stream)); 其內部實現會從物件池中獲取可複用的cuda event物件並執行event。

這裡有個細節，由於cuda event的建立和銷燬都會引發cuda kernel的launch由非同步轉同步，所以基於物件池的cuda event可以避免這個開銷。

實際上最終呼叫的還是 cudaEventRecord ，而 cudaEventRecord 本身只是起到一個“佔位符”的作用，並不能起到（保證該stream上其他kernel全部執行完）的作用，真正能保證stream同步作用的是oneflow vm（vitual machine）控制下的指令間依賴關係/執行順序。

5

CPU 下的平行計算

CpuStream 只有一個執行緒。CPU kernel 應該是通過 OpenMP  或者 Intel OneApi 等實現平行計算加速。

參考資料

1.https://github.com/Oneflow-Inc/oneflow/tree/845595e2c0abc3d384ff047e188295afdc41faaa

其他人都在看

• GPT-4，大增長時代的序幕
  

• GPT-4創造者：第二次改變AI浪潮的方向
  

• OpenAI創始人：GPT-4的研究起源和構建心法

• ChatGPT作者Schulman：我們成功的祕密武器

• 比快更快，開源Stable Diffusion重新整理作圖速度
• OneEmbedding:單卡 訓練TB級推薦模型不是夢
• GLM訓練加速：效能最高提升3倍，視訊記憶體節省1/3

試用OneFlow: github.com/Oneflow-Inc/oneflow/

---

​