OneFlow 源码解析：Eager 模式下 Tensor 的存储管理
作者｜郑建华
1、不同 Tensor 类型的存储管理方式
Lazy Tensor 的存储是由 Runtime 和 Actor 等对象管理的。静态图完成编译后，需要多少个对象、多少存储空间都是确定的，Runtime 等在初始化时会分配存储，在退出时回收资源。
Eager 模式下，Global Tensor 可以视为对 Local Tensor 的分布式封装，EagerGlobalTensorImpl 在本地的数据是一个 EagerLocalTensorImpl 对象。可以通过考察 EagerLocalTensorImpl 来理解 eager 模式下 tensor 的存储管理。
参考的示例代码如下：
2、Tensor 存储相关类的关系
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
UpdateTensorStorage
创建 one::TensorStorage 对象
设置 tensor 存储释放的回调函数
上述对象的创建，都只是记录相关信息，还不涉及 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 析构时会执行如下操作：
在 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()
5、reshape 等场景的存储管理
在 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 断点示例
参考资料
OneFlow（https://github.com/Oneflow-Inc/oneflow/tree/b51cb72430619f6088e47bbb8b8226f37299573a）
其他人都在看
欢迎 Star、试用 OneFlow 最新版本：https://github.com/Oneflow-Inc/oneflow/
版权声明: 本文为 InfoQ 作者【OneFlow】的原创文章。
原文链接:【http://xie.infoq.cn/article/18d471e5d185f996dac5b11db】。文章转载请联系作者。
不至于成为世界上最快的深度学习框架。 2022-03-23 加入
★ OneFlow深度学习框架：github.com/Oneflow-Inc/oneflow ★ OF云平台：oneflow.cloud
评论