关于人工智能:OneFlow源码解析Eager模式下Tensor的存储管理

3次阅读

共计 4088 个字符,预计需要花费 11 分钟才能阅读完成。

作者|郑建华

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)

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 析构时会执行如下操作:


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 办法

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 断点示例


 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(https://github.com/Oneflow-Inc/oneflow/tree/b51cb72430619f6088e47bbb8b8226f37299573a)
  • OneFlow 源码解析:Tensor 类型体系与 Local Tensor

欢送 Star、试用 OneFlow 最新版本:https://github.com/Oneflow-Inc/oneflow/

正文完
 0