OneFlow源码解析：算子签名的自动推断

撰文 | 郑建华

OneFlow是一个原生支持分布式训练的、高性能的深度学习框架。最近读了一些OneFlow的源码、架构设计和代码实现的文章，简单梳理一下自己的理解。主要通过图形展示调用过程和类之间的关系，只对部分重要的代码作一下分析。

深度学习框架是一个复杂的系统，而用户使用最多的就是算子（op）。用户通过op构造模型，进行训练、预测。这个笔记就从op入手，看看从Python前端到C++底层，OneFlow如何执行算子的计算逻辑。

具体地说，以比较简单的Relu算子为例，分析如下代码怎么执行：

# import会触发一系列初始化工作，暂时忽略
import oneflow as flow
# tensor的实现其实很复杂，因为要融合local和分布式的global tensor
t = flow.tensor([-1, 0, 1])
r = flow.relu(t)

1

编译环境

在开始分析之前，需要搭建环境编译OneFlow的源码，因为有些代码是在编译构建过程中自动生成的。在分析的过程中，这些自动生成的代码也是必要的环节。

OneFlow提供了官方的编译镜像（https://hub.docker.com/r/oneflowinc/manylinux2014_x86_64_cuda11.2）。用这个镜像可以非常方便地搭建编译环境（https://github.com/Oneflow-Inc/oneflow#option-2-build-in-docker-container-recommended）。

我使用的OneFlow版本是v0.7.0。本地编译环境目录结构如下，build是

cmake的构建目录，oneflow是源码目录。

.
├── build
└── oneflow

编译比较耗时，可以把两个目录mount到容器，便于后续查看build目录中生成的文件。

在cmake配置、构建过程中，会下载很多第三方源码包，如果网络状况不好容易超时，直接重试cmake/make即可。

# docker run -itd -v $PWD/oneflow:/mnt/oneflow -v $PWD/build:/mnt/build \
#   manylinux2014_x86_64_cuda11.2 bash
cd /mnt/build
cmake -S /mnt/oneflow
cmake --build . # --parallel 8
cd ../oneflow/python
python3 setup.py bdist_wheel
pip install ./dist/oneflow-0.7.0+cpu-cp38-cp38-linux_x86_64.whl

用GDB追踪OneFlow的执行过程

王益：Use GDB to Walkthrough OneFlow Source Code（https://quip.com/JuQ0AuodVJn4）

CMAKE_BUILD_TYPE=Debug cmake -S /mnt/oneflow
cmake --build . --parallel 8


source /mnt/build/source.sh
gdb python3


b oneflow::one::MakeLocalTensorFromData


run


import oneflow as flow


flow.Tensor([[1,2,3],[4,5,6]])

2

Python Binding

OneFlow底层是C++实现，通过pybind11实现Python Binding。月踏在《从Python到C++调用过程分析》对相关内容做了讲解。

2.1 Relu的Python包路径

# python/oneflow/__init__.py
from oneflow._C import relu


# python/oneflow/_C/__init__.py
from oneflow._oneflow_internal._C import *

2.2 module处理逻辑的注册

Python代码主要在python/oneflow目录，C++实现的包主要在_oneflow_internal下，pybind11的绑定代码位于init.cpp（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/init.cpp）：

PYBIND11_MODULE(_oneflow_internal, m) {
  // ...
  py::class_<::oneflow::cfg::Message, std::shared_ptr<::oneflow::cfg::Message>>(m, "CfgMessage");
  ::oneflow::cfg::Pybind11ModuleRegistry().ImportAll(m);
  ::oneflow::OneflowModuleRegistry().ImportAll(m);
}

其中OneflowModuleRegistry（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/init.cpp#L106）是算子等模块的绑定；Pybind11ModuleRegistry（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/init.cpp#L105）应该是自定义的、类似protobuf的配置数据结构的绑定。

从OneflowModuleRegistry开始的详细调用流程如下：

把代码放到一起看看（https://github.com/Oneflow-Inc/oneflow/blob/release/0.7.0/oneflow/api/python/of_api_registry.cpp）：

using SubModuleMap = std::map<std::string, std::vector<std::function<void(pybind11::module&)>>>;


SubModuleMap* GetSubModuleMap() {
  static SubModuleMap sub_module_map;
  return &sub_module_map;
}


// 修改map，执行注册
void OneflowModuleRegistry::Register(std::string module_path,
                                     std::function<void(pybind11::module&)> BuildModule) {
  (*GetSubModuleMap())[module_path].emplace_back(BuildModule);
}


void OneflowModuleRegistry::ImportAll(pybind11::module&am