OneFlow源碼解析之Eager模式下Tensor存儲管理方法是什么

今天小編給大家分享一下OneFlow源碼解析之Eager模式下Tensor存儲管理方法是什么的相關知識點，內容詳細，邏輯清晰，相信大部分人都還太了解這方面的知識，所以分享這篇文章給大家參考一下，希望大家閱讀完這篇文章后有所收獲，下面我們一起來了解一下吧。

1不同Tensor類型的存儲管理方式

Lazy Tensor 的存儲是由 Runtime 和 Actor 等對象管理的。靜態圖完成編譯后，需要多少個對象、多少存儲空間都是確定的，Runtime 等在初始化時會分配存儲，在退出時回收資源。

Eager 模式下，Global Tensor 可以視為對 Local Tensor 的分布式封裝，EagerGlobalTensorImpl 在本地的數據是一個EagerLocalTensorImpl 對象。可以通過考察 EagerLocalTensorImpl 來理解 eager 模式下 tensor 的存儲管理。

參考的示例代碼如下：

import numpy as np
 import oneflow as flow
 a = np.random.randn(1, 4)
 flow.tensor(a, device=flow.device("cpu"), dtype=flow.float)

2Tensor 存儲相關類的關系

EagerLocalTensorImpl 的存儲相關的類關系如下。

后續會順著示例代碼的執行過程，看看圖中的對象都是在何時、如何構造的，存儲被誰持有、如何分配并釋放。

3通過虛擬機指令為 Tensor 分配存儲

tensor 的構造函數通過 Python C API 注冊為 PyTensorObject_init，由
functional::_legacy_tensor_ctor根據簽名進行轉發。

示例代碼對應的是 TensorWithDataFunctor，調用 MakeLocalTensorFromData 構造 tensor，在這個函數中通過調用 functional::Empty以及 EmptyFunctor分配存儲。在 EmptyFunctor 中把相關屬性都存到 attrs，然后調用 OpInterpUtil::Dispatch在 vm 指令的執行準備過程中分配存儲。

EmptyFunctor 返回的 tensor 是一個只有存儲空間、不含數據的對象。數據拷貝在后面由
CopyLocalTensorFromUntypedArray完成。

3.1 存儲相關對象的構造

因為是 eager 模式下的 local tensor，OpInterpUtil::Dispatch 會被轉發到 NaiveInterpret執行。對于示例代碼，這個函數的輸入參數如下：

• inputs 是一個空數組

• outputs 只有一個元素、且是空指針

因為 outputs 中的 tensor 指針都是空的，所以需要創建一個 EagerLocalTensorImpl 對象，其 one::TensorStorage 成員變量是空指針。

因為 output_eager_blob_objects 中的元素尚未初始化，會調用 tensor_impl->InitEagerBlobObject進行初始化。因為 tensor_storage_ 還是空的，這個過程會執行如下操作：

• 創建 vm::TensorStorage 對象

• 創建 EagerBlobObject 對象

• set_eager_blob_object

• 
  

• 創建 one::TensorStorage 對象

• 設置 tensor 存儲釋放的回調函數

• UpdateTensorStorage

• 
  

上述對象的創建，都只是記錄相關信息，還不涉及 tensor 的存儲分配。

需要注意的是，注冊到 one::TensorStorage 的回調函數被賦值給了成員變量 releaser_hook_，這個函數會通過虛擬機指令釋放 tensor。

3.2 在指令執行過程中分配 tensor 存儲

分配 tensor 存儲的過程如下：

• vm::Instruction::Compute

• vm::InstructionPolicy::ComputeIf

• vm::OpCallInstructionPolicy::Compute

• OpCallInstructionUtil::Compute

• 獲取內存分配器

• OpCallInstructionUtil::AllocateOutputBlobsMemory

• blob_object->TryAllocateBlobBodyMemory

• allocator->Allocate

在EagerBlobObject::TryAllocateBlobBodyMemory 中，allocator 分配的存儲地址會賦值給 dptr，存儲地址 dptr 和 Free 函數一起構造一個智能指針，并賦值給 vm::TensorStorage 的 blob_dptr_ 變量。

4通過虛擬機指令釋放 Tensor 存儲

在前面的 3.1 節提到，EagerLocalTensorImpl 在初始化 EagerBlobObject、創建 one::TensorStorage 的同時，會設置一個釋放 tensor 的回調函數，回調函數保存在變量 releaser_hook_ 中，one::TensorStorage 析構時調用這個回調函數。把這些信息綜合整理一下，one::TensorStorage 析構時會執行如下操作：

vm::InstructionList instruction_list;
 InstructionsBuilder instructions_builder(&instruction_list);
 // JUST(Build(&instructions_builder));
 if (eager_blob_object->producer_stream().has_value()) {
   JUST(instructions_builder->ReleaseTensor(eager_blob_object));
 }
 JUST(vm::Run(instructions_builder.mut_instruction_list()));

在InstructionsBuilder::ReleaseTensor 中，如果有其它 stream 最近使用了 eager_blob_object，會通過 SoftSyncStreamBetween 進行同步。通過這種方式解決存儲的依賴問題。

一般情況下，通過 tensor 的 producer_stream 釋放存儲，根據這個對象獲取對應的 vm::Stream 對象，并據此構造指令 instruction（包含 eager_blob_object 和 vm_stream），示例代碼對應的指令類型是 FastReleaseTensorInstructionPolicy，其 Compute 方法執行具體的存儲釋放邏輯，過程如下：

• ReleaseTensorInstructionPolicy::Release()

• eager_blob_object->DeallocateBlobDataPtr()

• tensor_storage_->Release()

• tensor_storage_->_Release()

• blob_dptr_.reset()

• 
  

• 智能指針重置，調用分配存儲時指定的 Free 方法

5reshape 等場景的存儲管理

在 reshape、slice、transpose 等場景中，調用的 EagerLocalTensorImpl 構造函數的參數包括 input 的 tensor_storage，所以這個 tensor 的 tensor_storage_ 變量不是空的，在執行 InitEagerBlobObject 時，只創建 EagerBlobObject以提供 shape、stride等信息；但不會再創建 one::TensorStorage，而是復用 input 的存儲。

6兩個 TensorStorage 類型可以合并嗎？

為什么在 one::TensorStorage 析構時、由它保存的回調函數來觸發釋放 vm::TensorStorage 中的存儲呢？

one::TensorStorage 只多了一個 releaser，這兩個 Storage 類型是否可以合并呢？

在當前的設計下，這兩個類型不能合并。因為
one::TensorStorage::releaser_hook_ 中持有 EagerBlobObject 的智能指針，EagerBlobObject 中也持有 vm::TensorStorage 的智能指針。如果兩個 Storage 類型合并為一個，就會出現循環引用、對象無法析構而導致內存泄漏。

所以，vm::TensorStorage 只是單純的存儲，可以在多個 tensor 之間共享。EagerBlobObject 既包括存儲、也包括 shape、stride、data_type 等獨特的對象信息。而 one::TensorStorage 是為了避免循環引用而引入的、專門負責釋放存儲的角色。

7附錄

GDB 斷點示例

break oneflow::one::MakeLocalTensorFromData
 break oneflow::one::NaiveInterpret
 break oneflow::vm::VirtualMachineEngine::DispatchInstruction
 break oneflow::vm::OpCallInstructionUtil::Compute
 break oneflow::vm::OpCallInstructionUtil::AllocateOutputBlobsMemory
 break oneflow::vm::EagerBlobObject::TryAllocateBlobBodyMemory
 break oneflow::vm::ReleaseTensorInstructionPolicy::Release
 break oneflow/core/eager/eager_blob_object.cpp:107

以上就是“OneFlow源碼解析之Eager模式下Tensor存儲管理方法是什么”這篇文章的所有內容，感謝各位的閱讀！相信大家閱讀完這篇文章都有很大的收獲，小編每天都會為大家更新不同的知識，如果還想學習更多的知識，請關注億速云行業資訊頻道。