文章详情

短信预约-IT技能 免费直播动态提醒

请输入下面的图形验证码

提交验证

短信预约提醒成功

OneFlow源码解析之Eager模式下Tensor存储管理

2023-05-16 14:35

关注

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执行。对于示例代码,这个函数的输入参数如下:

因为 outputs 中的 tensor 指针都是空的,所以需要创建一个 EagerLocalTensorImpl 对象,其 one::TensorStorage 成员变量是空指针。

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

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

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

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

分配 tensor 存储的过程如下:

在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 方法执行具体的存储释放逻辑,过程如下:

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(github.com/Oneflow-Inc… )

OneFlow源码解析:Tensor类型体系与Local Tensor

欢迎 Star、试用 OneFlow 最新版本:

github.com/Oneflow-Inc…

以上就是OneFlow源码解析:Eager模式下Tensor的存储管理的详细内容,更多关于OneFlow源码解析:Eager模式下Tensor的存储管理的资料请关注编程网其它相关文章!

阅读原文内容投诉

免责声明:

① 本站未注明“稿件来源”的信息均来自网络整理。其文字、图片和音视频稿件的所属权归原作者所有。本站收集整理出于非商业性的教育和科研之目的,并不意味着本站赞同其观点或证实其内容的真实性。仅作为临时的测试数据,供内部测试之用。本站并未授权任何人以任何方式主动获取本站任何信息。

② 本站未注明“稿件来源”的临时测试数据将在测试完成后最终做删除处理。有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341

软考中级精品资料免费领

  • 历年真题答案解析
  • 备考技巧名师总结
  • 高频考点精准押题
  • 2024年上半年信息系统项目管理师第二批次真题及答案解析(完整版)

    难度     807人已做
    查看
  • 【考后总结】2024年5月26日信息系统项目管理师第2批次考情分析

    难度     351人已做
    查看
  • 【考后总结】2024年5月25日信息系统项目管理师第1批次考情分析

    难度     314人已做
    查看
  • 2024年上半年软考高项第一、二批次真题考点汇总(完整版)

    难度     433人已做
    查看
  • 2024年上半年系统架构设计师考试综合知识真题

    难度     221人已做
    查看

相关文章

发现更多好内容

猜你喜欢

AI推送时光机
位置:首页-资讯-后端开发
咦!没有更多了?去看看其它编程学习网 内容吧
首页课程
资料下载
问答资讯