OneFlow 源码阅读 10：eager 模式下 tensor 的存储管理

tensor 对象涵盖的内容非常多，这篇笔记着重梳理一下存储的管理。

1 不同 tensor 类型的存储管理方式

lazy tensor 的存储是由 Runtime 和 Actor 等对象管理的。静态图完成编译后，需要多少个对象、多少存储都是确定的，Runtime 等在初始化时会分配存储，在退出时回收资源。

eager 模式下，global tensor 可以视为对 local tensor 的分布式封装，EagerConsistentTensorImpl 在本地的数据是一个 EagerMirroredTensorImpl 对象。可以通过考察 EagerMirroredTensorImpl 来理解 eager 模式下 tensor 的存储管理。

参考的示例代码如下：

import oneflow as flow
flow.Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

2 tensor 存储相关类的关系

EagerMirroredTensorImpl 的存储相关的类关系如下。v0.8.0 的 tensor 存储实现比之前更简洁。
后续会顺着示例代码的执行过程，看看图中的对象都是在何时、如何构造的，存储被谁持有、如何分配并释放。

3 通过虚拟机指令为 tensor 分配存储

tensor 的构造函数通过 Python C API 注册为 PyTensorObject_init，由 functional::_legacy_tensor_ctor 根据签名进行转发。
示例代码对应的是 TensorWithDataCtorFunctor，调用 MakeLocalTensorFromData 构造 tensor，在这个函数中通过调用 functional::Empty 以及 EmptyFunctor 分配存储。在 EmptyFunctor 中把相关属性都存到 attrs，然后调用 OpInterpUtil::Dispatch 在 vm 指令的执行准备过程中分配存储。

EmptyFunctor 返回的 tensor 是一个只有存储空间、不含数据的对象。数据拷贝在后面由 SwitchCopyMirroredTensorFromUntypedArray 完成。

3.1 存储相关对象的构造

因为是 eager 模式下的 local tensor，OpInterpUtil::Dispatch 会被转发到 NaiveInterpret 执行。对于示例代码，这个函数的输入参数如下：

inputs 是一个空数组
outputs 只有一个元素、且是空指针

因为 outputs 中的 tensor 指针都是空的，所以需要创建一个 EagerMirroredTensorImpl 对象，调用的是无参的构造函数，其 one::TensorStorage 成员变量是空指针。后续推导设备时，因为EmptyOp 定义了自己的设备推导函数，会通过 user_op_expr 调用这个推导函数、从 attrs 中获取设备信息。

因为 output_eager_blob_objects 中的元素尚未初始化，会调用 tensor_impl->InitEagerBlobObject 进行初始化。因为 tensor_storage_ 还是空的，这个过程会执行如下操作：

上述对象的创建，都只是记录相关信息，还不涉及 tensor 的存储分配。

需要注意的是，one::TensorStorage 注册的回调函数被赋值给了成员变量 releaser_hook_，这个函数会通过虚拟机指令释放 tensor。

3.2 在指令执行过程中分配 tensor 存储

分配 tensor 存储的过程如下：

在 EagerBlobObject::TryAllocateBlobBodyMemory 中，通过 DeviceCtx（这里是 EpDeviceCtx）获取 allocator。allocator 分配的存储地址会赋值给 dptr，存储地址 dptr 和 Free 函数一起构造一个智能指针，并赋值给 vm::TensorStorage 的 blob_dptr_ 变量。

3.3 Allocator 相关类的关系

EpStreamType 在初始化 EpDeviceCtx 前，先创建一个 EpBackendAllocator 对象。EpDeviceCtx 在构造时，根据这个基础 allocator 构造一个线程安全、具有内存池功能的复合对象，对象关系图示如下：

实际的内存分配是由 EpBackendAllocator 完成的，它通过调用设备的 Alloc 方法完成存储的分配。CUDA 设备会调用 cudaMalloc，CPU 会调用 aligned_alloc。

4 通过虚拟机指令释放 tensor 存储

在前面的 3.1 节提到，EagerMirroredTensorImpl 在初始化 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，会通过 SoftSyncStream 进行同步。通过这种方式解决存储的依赖问题。

一般情况下，通过 tensor 的 producer_stream 释放存储，根据这个对象获取对应的 vm::Stream 对象，并据此构造指令操作数 ReleaseTensorArgPhyInstrOperand（包含 eager_blob_object 和 vm_stream），对应的指令类型是 ReleaseTensorInstructionType。

ReleaseTensorInstructionType::Compute 或 Prepare 的执行流程如下：

GetEagerBlobObject
- 获取上述 ReleaseTensorArgPhyInstrOperand 中保存的 eager_blob_object
ReleaseTensorInstructionType::Release
- eager_blob_object->DeallocateBlobDataPtr
  - vm::TensorStorage::Release
    - blob_dptr_.reset()
      - 智能指针重置，调用分配存储时指定的 Free 方法
  - tensor_storage_.reset(new TensorStorage): 让 vm::TensorStorage 的 stream 都为空，避免不必要的释放。

这样，在 one::TensorStorage 析构时会通过虚拟机指令释放 vm::TensorStorage 中的存储。

5 reshape 等场景的存储管理

在 reshape、slice、transpose 等场景中，调用的 EagerMirroredTensorImpl 构造函数的参数包括 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 应该是为了避免循环引用而引入的、专门负责释放存储的角色。