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 flowflow.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在深度学习框架中的执行过程简略梳理