一个算子在深度学习框架中的旅程

撰文｜赵露阳

算子即Operator，这里简称op。op是深度学习的基础操作，任意深度学习框架中都包含了数百个op，这些op用于各种类型的数值、tensor运算。

在深度学习中，通过nn.Module这样搭积木的方式搭建网络，而op就是更基础的，用于制作积木的配方和原材料。

譬如如下的一个demo网络：

import oneflow as torch                  
class TinyModel(torch.nn.Module):

    def __init__(self):
        super(TinyModel, self).__init__()

        self.linear1 = torch.nn.Linear(100, 200)
        self.activation = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(200, 10)
        self.softmax = torch.nn.Softmax()

    def forward(self, x):
        x = self.linear1(x)
        x = self.activation(x)
        x = self.linear2(x)
        x = self.softmax(x)
        return xtinymodel = TinyModel()print('The model:')print(tinymodel)

从结构来看，这个网络是由各种nn.Module如Linear、ReLU、Softmax搭建而成，但从本质上，这些nn.Module则是由一个个基础op拼接，从而完成功能的。这其中就包含了Matmul、Relu、Softmax等op。 在OneFlow中，对于一个已有op，是如何完成从Python层->C++层的调用、流转和执行过程？本文将以

output = flow.relu(input)

为例，梳理一个op从Python -> C++执行的完整过程。

首先，这里给出一个流程示意图：

下面，将分别详细从源码角度跟踪其各个环节。

1

Binding

这里，binding是指Python和C++代码的绑定。通常，我们用Python搭建网络，训练模型，调用函数完成各种操作。实际上，这些函数通常在Python层只是一层wrapper，底层实现还是通过C++代码完成的，那么Python -> C++是如何调用的？这就需要用到Python和C++的绑定。

在深度学习框架的实现中，即可以用Python原生的C API，也可以通过pybind11来完成函数绑定，在OneFlow中，二者均有使用，譬如：

• oneflow/api/python/framework/tensor.cpp

• oneflow/api/python/framework/tensor_functions.cpp

中涉及到的 tensor.xxx 方法都是通过Python C API完成了函数绑定；

• oneflow/core/functional/functional_api.yaml

中定义的诸多 flow.xxx 方法则是通过pybind实现的绑定。这里关于Python C API和pybind不做过多介绍，具体用法可以参考相应文档：

• https://docs.python.org/zh-cn/3.8/c-api/index.html

• https://pybind11.readthedocs.io/en/stable/index.html

下面我们回到flow.relu方法，我们在Python层调用的flow.relu实际是调用了在

python/oneflow/__init__.py

中定义的oneflow._C.relu。 _C表示其实现位于底层C++。和PyTorch类似，我们也基于.yaml定义了一套接口导出及code gen的规则，譬如在 functional_api.yaml 中，我们可以看到Relu的导出接口的函数签名：

- name: "relu"
  signature: "Tensor (Tensor x, Bool inplace=False) => Relu"
  bind_python: True

从yaml定义可以看出，flow._C.relu 接收两个参数，tensor和一个bool值，其绑定了C++的Relu方法，函数返回值也是tensor。实际上，在OneFlow编译时，会通过执行

tools/functional/generate_functional_api.py

这个文件，对 functional_api.yaml 进行解析和代码生成，动态生成C++的.h和.cpp文件。

• build/oneflow/core/functional/functional_api.yaml.h

• build/oneflow/core/functional/functional_api.yaml.cpp

并在.cpp文件中调用相应的functor完成C++层面的函数调用。这里，还是以flow._C.relu为例，其对应的functor定义位于oneflow/core/functional/impl/activation_functor.cpp:

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 {
    ...
  }


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

‍

ReluFunctor通过

ONEFLOW_FUNCTION_LIBRARY(m) {
  m.add_functor<impl::ReluFunctor>("Relu");
  ...
}

‍

完成functor的注册，注册成functional接口后，在Python层flow._C.relu就完成了和“Relu”的绑定。同时，这个函数在C++中也可以通过functional::Relu直接调用。

2

Functor

Functor不仅是Python -> C++交互的核心，也是op调用、输入参数推导和检查的第一站。通常，各种op在functor层需要完成对输入tensor的shape、dtype、维度、元素个数等各种check，以及对op特有的逻辑进行解析和处理。Relu 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) {
      JUST(CheckInplaceValid(x));
      std::shared_ptr<TensorTuple> outputs = std::make_shared<TensorTuple>(1);
      outputs->at(0) = x;
      JUST(OpInterpUtil::Dispatch(*op_, {x}, outputs.get(), AttrMap{})