一个Tensor在深度学习框架中的执行过程简单梳理

撰文：BBuf。审稿：王迎港。

0x0. 前言

相信看到这篇文章的人都对深度学习框架是有所了解和熟悉的，也多多少少会使用Python写一些神经网络相关的代码。例如我们可以在PyTorch写出下面的代码：

import torch
x = torch.tensor([-1.0, 2.0], device="cuda")
y = torch.relu(x)
print(y)

使用PyTorch运行之后我们会获得如下结果：

tensor([0., 2.], device='cuda:0')

对于x这个输入Tensor来说，它被喂给了relu这个Op，然后输出结果，一切看起来都很简单和正常。但如果有人问你是否清楚这背后到底发生了什么，relu这个Op对应的Cuda Kernel是在什么时候被GPU调用的，相信一部分人是不会很清楚的。因为包括我的大多数人习惯在舒适区使用深度学习框架，对背后的原理可能没有深入了解，所以回答不了也很正常。

这篇文章我就将尝试解开这个问题，但我并不是以PyTorch为例来讲解，而是以OneFlow为例子。为什么以OneFlow为例子呢？首先我在OneFlow工作，对这背后的执行机制比PyTorch要清楚一些，在调用链跟踪的时候会更流畅。其次，OneFlow背后这套运行机制含有挺多PyTorch不存在的设计思想，相信读者看完之后对深度学习框架系统设计方面有更多的思考和启发。

所以，接下来就一起看看一个Tensor在OneFlow深度学习框架中的执行过程吧。为了简单起见，本文只考虑单机单卡模式下的Op执行过程，不涉及OneFlow特有的consistent模式（和分布式相关），如果你对这部分感兴趣可以自行查看。

0x1. Python和C++的桥梁

当我们敲下如下代码并将其移交给OneFlow执行时：

import oneflow as flow
x = flow.tensor([-1.0, 2.0], device="cuda")
y = flow.relu(x)
print(y)

系统首先创建了一个在GPU上的输入Tensor，然后调用了导出到python端的c++ functional接口relu。这里涉及到pybind11绑定相关的Python wrapper和C++ relu functor。这个交互的上层，同事在OneFlow学习笔记：python到C++调用过程分析 这篇文章有解析过了，感兴趣可以看看。我们上面Python代码中的flow.relu这个Op最终调用的是ReLU C++ Functor的实现，我们看一下代码。

class ReluFunctor {
   
   
 public:
  ReluFunctor() {
   
    op_ = CHECK_JUST(one::OpBuilder("relu").Input("x", 1).Output("y", 1).Build()); }
  Maybe<Tensor> operator()(const std::shared_ptr<Tensor>& x, bool inplace) const {
   
   
    if (inplace) {
   
   
      ...
    } else {
   
   
      return OpInterpUtil::Dispatch<Tensor>(*op_, {
   
   x});
    }
  }

 private:
  std::shared_ptr<OpExpr> op_;
};

这段代码里面的op_是一个OpExpr的指针，然后在构造函数里面调用了OpBuilder函数来创建了一个新的OpExpr。从后面的实际调用代码OpInterpUtil::Dispatch<Tensor>(*op_, {x});可以发现这里的算子构建和执行是分开的（因为Dispatch函数是同时将OpExpr和输入Tensor等分发出去，没有直接分发执行的结果Tensor出去，所以这里还没有真正的执行Op），这里的OpInterpUtil::Dispatch是负责将OpExpr，输入Tensor和其它参数（ReLU这个算子没有除输入外的参数）分发出去，还没有真正的执行。

OpExpr可以简单理解为是OneFlow算子的统一抽象。OpExpr大体可以分为BuiltinOpExpr、FunctionOpExpr和其他类别的OpExpr，其中BuiltinOpExpr又可以细分为UserOpExpr和其他非UserOpExpr，用户可以通过OpBuilder构建出UserOpExpr。

不需要完全理解OpExpr的定义，我们只需要知道这里是通过OpBuilder类构造了一个新的OpExpr，这个OpExpr有Op name，UserOpConf proto_这个序列化Op信息的ProtoBuf对象，以及输入输出Tensor的名字等关键信息。然后顺着这个Dispatch函数可以发现最后在oneflow/core/framework/op_interpreter/op_interpreter_util.cpp中调用到了GetInterpreter函数的Apply方法：

/* static */ Maybe<void> OpInterpUtil::Dispatch(const OpExpr& op_expr, const TensorTuple& inputs,
                                                TensorTuple* outputs,
                                                const OpExprInterpContext& ctx) {
   
   
  return JUST(GetInterpreter(inputs, ctx, op_expr))->Apply(op_expr, inputs, outputs, ctx);
}

这里的OpExprInterpContext对象会存储Op的动态属性，设备信息，分布式信息等，对于Relu Functor来说，这里为空，所以我们这里不关注这个对象。再往下跟就属于InterPreter的内容了，新开一节来讲。

0x2. Interpreter

从上面的Op调用流程可以看出，我们在Python层的Op实际上是调用的导出到Python的Functor接口，而Functor接口会将OpExpr，输入Tensor和动态属性attr递交给Interpreter来处理，因为上面的GetInterpreter函数获取的就是一个Interpreter对象。Interpreter这个类就是专门用来解释Op执行过程的，上一节在Relu Functor里面的Dispatch就是把任务分发到Interpreter来执行。OneFlow的Interpreter又分为几种类型，如Eager Mirrored Interpreter，Eager Consistent Interpreter和LazyInterpreter，我们这篇文章的例子没有考虑分布式信息，所以输入Tensor都是Eager Mirroed Tensor，所以走的是Eager Mirrored Interpreter这个调用链。Mirrored Tensor和PyTorch的Tensor类似，在各个Rank上是独立的。

再往下跟一下我们发现上面的Apply实际上调用的是oneflow/core/framework/op_interpreter/eager_mirrored_op_interpreter.cpp文件中的NaiveInterpret函数，这个函数接收OpExpr对象，输入输出Tensor和一个OpExprInterpContext对象来对Op的device，输出dtype，输出shape等进行推导，然后根据推导的元信息（元信息对应TensorMeta类对象，把 Tensor 的基本信息：shape, dtype, stride 等抽出来一个类型，放一起方便管理）构造分别对应输入输出的BlobObject对象input_eager_blob_objects和output_eager_blob_objects（可理解为输入输出Tensor的数据指针），另外还会根据OpExpr和推导后的device构造一个特定执行kernel。最后将执行kernel，输入输出Tensor的数据指针以及OpExprInterpContext对象以指令的方式发给OneFlow的虚拟机（VM，可以理解为OneFlow的Eager运行时，后面会细讲）执行并获得结果。

这里我们分段看一下NaiveInterpret的实现。第一段：

Maybe<void> NaiveInterpret(const UserOpExpr& user_op_expr, const TensorTuple& inputs,
                           const Symbol<Device>& default_device, TensorTuple* outputs,
                           const OpExprInterpContext& ctx) {
   
   
  const auto& attrs = ctx.attrs;
  std::shared_ptr<EagerBlobObjectList> input_eager_blob_objects =
      std::make_shared<EagerBlobObjectList>(inputs.size());
  for (int i = 0; i < inputs.size(); i++) {
   
   
    const auto& input_device = JUST(inputs.at(i)->device());
    if (i > 0) {
   
   
      CHECK_OR_RETURN(*default_device == *input_device) << Error::InputDeviceNotMatchError();
    }
    input_eager_blob_objects->at(i) = JUST(inputs.at(i)->eager_blob_object());
  }

上面这段代码遍历输入Tensor的列表，将每一个输入Tensor的device和函数传入的默认device进行比较，如果发现输入Tensor的device和默认device不一致就抛出异常。可以对类似输入Tensor在CPU上，但nn.Module在GPU上的例子进行错误检查，输出设备不匹配的错误信息。如果设备都匹配上了，这个时候会将输入Tensor的eager_blob_object添加到input_eager_blob_objects这个列表中。输入Tensor的eager_blob_object是一个EagerBlobObject类型的指针，是输入Tensor的数据指针，后续通过它和OneFlow的虚拟机（VM）进行交互。

这里要补充说明一下OneFlow中Tensor，TensorImpl，TensorMeta和BlobObject的关系。 Tensor 和 TensorImpl 用了桥接设计模式，Tensor 负责向上和 python 接口、autograd 的对接；TensorImpl 是向下负责真实数据这部分。TensorMeta 就是把 Tensor 的基本信息：shape, dtype, stride 等抽出来一个类型，放一起方便管理。BlobObject是真正的数据对象，数据指针在这个对象中，这个类被虚拟机使用来完成指令的计算任务。

第二段：

std::shared_ptr<EagerBlobObjectList> output_eager_blob_objects =
      std::make_shared<EagerBlobObjectList>(outputs->size());
  auto* output_tensor_metas = ThreadLocalDefaultOutputMutTensorMetas(outputs->size());
  for (int i = 0; i < outputs->size(); i++) {
   
   
    if (!outputs->at(i)) {
   
   
      const auto& tensor_impl = std::make_shared<EagerMirroredTensorImpl>();
      outputs->at(i) = std::make_shared<MirroredTensor>(tensor_impl);
      output_tensor_metas->at(i) = tensor_impl->mut_tensor_meta();
    } else {
   
   
      bool has_eager_blob_object = JUST(outputs->at(i)->has_eager_blob_object());
      CHECK_OR_RETURN(has_eager_blob_object);
      output_eager_blob_objects->at(i) = JUST