乐趣区

关于深度学习:OneFlow源码阅读4tensor类型体系与local-tensor

jiezi

tensor 和 op 是神经网络模型最根本的组件：op 是模型的节点，tensor 是连贯节点的边。

然而，构建一个 tensor 并不仅仅是结构一个对象那么简略，至多要思考上面这些问题：

  • 要反对节点本地的 local tensor，以及分布式的 global tensor。
  • 要反对 eager 和 lazy 执行模式。
  • 要反对不同的数据类型，包含 float、double、int 等。
  • 要反对不同设施。

1 创立 tensor 的办法

从 init.py 看，有两个办法能够创立 tensor 对象，一个是 Tensor，另一个是 tensor。这两种形式最终都会通过 PyFunction 转发到特定的 Functor。

1.1 Tensor 类型

Tensor是在 tensor.py 中引入的，构造函数被绑定为 C ++ 的 ApiNewTensor，通过 RegisterMethods 为 Tensor 注册了一些 Python 实现的办法（如将 get_item/set_item 等转发给对应的 C ++ 函数），在包初始化时会通过 RegisterMethod4Class 实现这些办法的注册。

RegisterMethod4Class 的调用流程如下：

1.2 tensor 函数

tensor 是一个函数，其绑定定义在 tensor_api.yaml.pybind.cpp 中，这是构建阶段主动生成的文件。tensor函数间接绑定到PyFunction。

1.3 手动构建 tensor 的两种形式

剖析 Tensor 和 tensor 的 PyFunction 签名，能够通过如下形式结构 local tensor，也就是只能在节点外部应用的 tensor。其中只有 tensor 能够指定 dtype 参数。

import oneflow as flow
flow.tensor([[1,2,3],[4,5,6]])
flow.tensor([1, 2, 3], dtype=flow.int64)
flow.Tensor([[1,2,3],[4,5,6]])
# error
# flow.Tensor([1, 2, 3], dtype=flow.int64)

2 oneflow 的 tensor 类型体系

ApiNewTensor 函数返回 Tensor 类型。这是一个抽象类接口。通过其继承和子类的字段蕴含关系，能够失去如下的类图：

以上次要是 Tensor 相干的接口定义。MirroredTensor即节点内的 local tensor，ConsistentTensor即一致性视角的、分布式的 global tensor。

Tensor 应用了 Bridge 模式，每个 Tensor 子类外部有一个 TensorImpl 字段。TensorImpl 相干的类图如下：

3 local tensor 的结构

咱们以 flow.Tensor([[1,2,3],[4,5,6]]) 为例，看一下 Tensor 对象结构的过程。次要的流程如下：

在这个例子中，TensorWithDataCtorFunctor 最终会调用 MakeLocalTensorFromData，次要的逻辑都在这个函数中。其中大量调用 Python 和 numpy 的接口，查看 PyObject 的数据类型，获取 Shape 和 DataType，如果用户没有制订 device，默认会设置为 CPU 设施。

前面次要是调用 EmptyFunctor 和 SwitchCopyMirroredTensorFromUntypedArray。前者为 tensor 分配内存，后者进行数据拷贝，两个步骤都会通过虚拟机指令实现。

为什么要通过虚拟机指令实现呢？无论是内存资源的调配，还是数据拷贝，CPU 和 CUDA 等不同设施上的操作都不一样。之前探讨 Op/Kernel 时曾经看到，虚拟机和 InstructionType 反对不同的设施，所以内存调配和数据拷贝也通过虚拟机执行。

3.1 分配内存：EmptyFunctor

matmul和 relu（inplace=false 时）等操作在执行过程中也会创立 output tensor。之前探讨 relu 时重点关注了 op 和 kernel 的计算逻辑，而疏忽了 tensor 相干的内容。

而这里只须要结构一个 tensor 对象，不须要其它计算，所以是一个 Empty 操作，EmptyKernel 没有实质性的计算逻辑。

因为是 eager 模式下的 local tensor，EmptyFunctor 会进入 NaiveInterpret 执行。在这里会先结构 EagerMirroredTensorImpl 和 MirroredTensor 对象，用于寄存 tensor 后果。但这只是一个壳子，还没有为 tensor 的数据调配存储空间。

之后会初始化 EagerBlobObject、创立 TensorStorage，这样 tensor 次要的字段根本构建结束。

而后结构指令、提交虚拟机执行。EmptyFunctor 是 UserOp，最终会进入 LocalCallOpKernelUtil: Compute，其中 AllocateOutputBlobsMemory 实现内存分配任务。

EmptyFunctor 的调用流程如下：

AllocateOutputBlobsMemory 的调用流程如下。BlobDesc::ByteSizeOfBlobBody 提供内存 size，即 elem_cnt * SizeOf(data_type。CPU 环境下，CpuAllocator 通过aligned_alloc 申请内存资源。

3.2 拷贝数据：SwitchCopyMirroredTensorFromUntypedArray

SwitchCopyMirroredTensorFromUntypedArray其实是 MAKE_SWITCH_ENTRY 宏开展后的函数名。宏开展后的代码如下。理论会调用 CopyMirroredTensorFromUntypedArray。

template<typename... Args>
static Maybe<void> SwitchCopyMirroredTensorFromUntypedArray(const std::tuple<DataType>& switch_tuple, Args&& ... args) {static const std::map<std::tuple<DataType>, std::function<Maybe<void>(Args && ...)>>
      case_handlers {{SwitchCase(DataType::kFloat),
           [](Args&&... args) {return CopyMirroredTensorFromUntypedArray<float>(std::forward<Args>(args)...);
           }},
           // ...
      };
  return case_handlers.at(switch_tuple)(std::forward<Args>(args)...);
};

数据拷贝的调用流程如下：

根据上述宏开展后的代码，CopyMirroredTensorFromUntypedArray 的模版参数是 tensor 的 dtype，如DataType::kFloat。在 tensor 结构的场景下，函数 CopyBetweenMirroredTensorAndNumpy 的模版参数如 BlobNumpyCopyUtil<DataType::kFloat>::From。

CopyBetweenMirroredTensorAndNumpy中会结构指令提交虚拟机执行。PhysicalRun的逻辑相似如下代码：

    vm::InstructionMsgList instruction_list;
    InstructionsBuilder instructions_builder(std::make_shared<vm::PhysicalIdGenerator>(),
                                            &instruction_list);
    // JUST(Build(&instructions_builder));
    builder->AccessBlobByCallback(
        tensor,
        [array_ptr, Copy](uint64_t ofblob_ptr) {CHECK_JUST(Copy(ofblob_ptr, array_ptr)); },
        modifier);
    JUST(vm::Run(instructions_builder.mut_instruction_list()));

lambda 表达式中的 Copy 就是 BlobNumpyCopyUtil<DataType::kFloat>::From；array_ptr示意 Python 端传过来的数组数据指针；前面咱们会看到，ofblob_ptr就是 tensor 的 Blob 中的指针。

InstructionsBuilder::AccessBlobByCallback中创立 AccessBlobArgCbPhyInstrOperand 对象，对应的指令类型是 AccessBlobByCallbackInstructionType。所以虚拟机执行指令时，会进入 AccessBlobByCallbackInstructionType::Compute 执行。理论的执行逻辑相似如下代码：

const auto* ptr =
  dynamic_cast<const vm::AccessBlobArgCbPhyInstrOperand*>(phy_instr_operand.get());
OfBlob ofblob(device_ctx->stream(), ptr->eager_blob_object()->mut_blob());
// ptr->callback()(reinterpret_cast<uint64_t>(&ofblob));
BlobNumpyCopyUtil<DataType::kFloat>::From(&ofblob, array_ptr);

ptr->callback()就是上述 lambda 表达式。OfBlob 是对 tensor 的 Blob 的封装。一路追踪上来，CPU 环境下最终会调用 std::memcpy 拷贝数据。

参考资料

  • oneflow v0.7.0
  • OneFlow 源码浏览 1：算子签名的主动推断
  • OneFlow 源码浏览 2：Op、Kernel 与解释器
  • OneFlow 源码浏览 3：Op 指令在虚拟机中的执行
  • 一个 Tensor 在深度学习框架中的执行过程简略梳理
退出移动版