热文：oneflow源码解析：算子签名的自动推断-ag真人官方网

撰文 | 郑建华

(资料图片仅供参考)

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

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

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

# import会触发一系列初始化工作，暂时忽略import oneflow as flow# tensor的实现其实很复杂，因为要融合local和分布式的global tensort = 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 bashcd /mnt/buildcmake -s /mnt/oneflowcmake --build . # --parallel 8cd ../oneflow/pythonpython3 setup.py bdist_wheelpip 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/oneflowcmake --build . --parallel 8source /mnt/build/source.shgdb python3b oneflow::one::makelocaltensorfromdatarunimport oneflow as flowflow.tensor([[1,2,3],[4,5,6]])

2

python binding

oneflow底层是c 实现，通过pybind11实现python binding。月踏在《从python到c 调用过程分析》对相关内容做了讲解。

2.1 relu的python包路径

# python/oneflow/__init__.pyfrom oneflow._c import relu# python/oneflow/_c/__init__.pyfrom 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>>;submodulemap* getsubmodulemap() { static submodulemap sub_module_map; return &sub_module_map;}// 修改map，执行注册void oneflowmoduleregistry::register(std::string module_path, std::function buildmodule) { (*getsubmodulemap())[module_path].emplace_back(buildmodule);}void oneflowmoduleregistry::importall(pybind11::module& m) { for (const auto& pair : (*getsubmodulemap())) { for (const auto& buildmodule : pair.second) { buildsubmodule(pair.first, m, buildmodule); } }}void oneflowmoduleregistry::buildsubmodule( const std::string& module_path, pybind11::module& m, const std::function& buildmodule) { // ... buildmodule(m); // ...}

从这段代码可以看出，python module的注册逻辑都保存在submodulemap中。它的key是module name；value是一组函数，buildsubmodule中调用这些函数、执行module注册逻辑。

getsubmodulemap中保存map单例，register函数设置map的值，of_api_registry.h（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/api/python/of_api_registry.h）中的宏oneflow_api_pybind11_module调用register函数处理module注册逻辑。搜索一下可以知道relu的注册逻辑在build/oneflow/api/python/functional/functional_api.yaml.pybind.cpp中，这个文件中注册了很多算子（user_op）。以relu和pow为例，这个宏展开后的核心代码如下：

static void oneflowapipythonmodule9623(pybind11::module&);namespace { struct ofapiregistryinit { ofapiregistryinit() { ::oneflow::oneflowmoduleregistry().register("_c", &oneflowapipythonmodule9623); } }; ofapiregistryinit of_api_registry_init;}static void oneflowapipythonmodule9623(pybind11::module & m) { m.def("relu", &functional::pyfunction); m.def("pow", &functional::pyfunction< functional::powschema_ttt, functional::scalarpowschema_ttscb, functional::scalarpowschema_ttsc, functional::scalarreversepowschema_tsct >);}

这段代码中的类似注册技巧，在oneflow中的很多地方都被用到。

module注册逻辑在函数oneflowapipythonmodule9623中（9623来自宏定义中的line以避免名字冲突），ofapiregistryinit在构造对象时将这个函数注册到submodulemap，匿名空间中的变量of_api_registry_init就是为了通过构造对象、在构造函数中调用注册逻辑（而这个对象不占用任何空间）。这样在系统加载时就通过静态对象的初始化实现了module处理逻辑的注册，再通过pybind11的调用完成对python binding的定义。

3

多个接口签名的自动推断

从以上代码可以看到，relu算子被绑定到pyfunction（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#l120）这个函数执行计算逻辑，每次调用算子都会执行pyfunction这个函数。

从签名看，pyfunction是一个模版函数，给python前端返回py::object作为算子执行结果。

relu只有一个模版参数，pow有4个模版参数。每个模版参数表示算子支持的一种调用接口签名。oneflow可以根据python传过来的arguments类型，自动推断合适的签名，调用相关函数。

例如下面的代码，算子pow的指数参数既支持标量，也支持tensor：

import oneflow as flowr = flow.randn(1, 10)flow.pow(r, 2)flow.pow(r, flow.ones(1, 10))

下面就来看看oneflow是怎么实现这个功能的。

relu算子的签名schema如下所示：

struct reluschema_ttb { using ftype = maybe (const std::shared_ptr& x, bool inplace); using r = maybe; static constexpr ftype* func = &functional::relu; static constexpr size_t max_args = 2; static constexpr size_t max_pos_args = 2; static constexpr char const* signature = "tensor (tensor x, bool inplace=false)"; static functiondef function_def;};

先看一下从pyfunction开始的的调用顺序：

pyfunction相关的代码如下（删掉了一些与核心逻辑无关的内容）。

// schemat如 reluschema_ttbtemplateclass pyfunctiondispatcher { public: // schema_t是第i个签名 template using schema_t = typename std::tuple_element>::type; // schema_size_是签名个数，比如relu是1，pow是4 pyfunctiondispatcher() : schema_size_(sizeof...(schemat)) { signatures_.resize(schema_size_); initsignatures(std::make_index_sequence{}); } template py::object call(const py::args& args, const py::kwargs& kwargs, std::index_sequence) const { // t是当前检查的签名，比如 reluschema_ttb using t = schema_t; std::vector parsed_args(t::max_args); if (parseargs(args, kwargs, &parsed_args, t::function_def, t::max_pos_args, /*raise_exception*/ schema_size_ == 1)) { return detail::unpack_call(*t::func, parsed_args); } return call(args, kwargs, std::index_sequence{}); } py::object call(const py::args& args, const py::kwargs& kwargs, std::index_sequence<>) const { // throw error ... return py::none(); } private: template void initsignatures(std::index_sequence) { __attribute__((__unused__)) int dummy[] = { ((void)(signatures_[i] = schema_t::signature), 0)...}; } private: size_t schema_size_; std::vector signatures_;};// schemat如 reluschema_ttbtemplateinline py::object pyfunction(const py::args& args, const py::kwargs& kwargs) { static pyfunctiondispatcher dispatcher; return dispatcher.call(args, kwargs, std::make_index_sequence{});}// py module注册static void oneflowapipythonmodule9623(pybind11::module & m) { m.def("relu", &functional::pyfunction); m.def("pow", &functional::pyfunction< functional::powschema_ttt, functional::scalarpowschema_ttscb, functional::scalarpowschema_ttsc, functional::scalarreversepowschema_tsct >);}

3.1 dispatcher: 算子接口签名的自动推断

pyfunction是一个模版函数，每个模版参数表示算子的一个接口签名。

pyfunction及其后续执行链路的最重要的功能，就是实现这些签名的自动筛选。自动筛选的实质，就是通过index_sequence逐个检查签名与pyfunction的参数args/kwargs是否匹配。函数内的静态变量dispatcher实现了这个自动筛选功能。

每个算子都会特化一个pyfunction和pyfunctiondispatcher实例，也有一个算子自己的dispatcher变量。pyfunction直接将请求转发给dispatcher.call，顺带加上一个index_sequence模版参数，正是依靠这个模版参数实现了签名的自动筛选。

在call函数中，先确定当前检查的签名类型t（例如reluschema_ttb），然后通过parseargs检查python传过来的参数args/kwargs与签名t是否匹配。如果不匹配，就去掉当前签名t，将剩余的签名类型作为模版参数、继续递归调用call函数。

如果算子只有一个签名，就通过schema_size_ == 1通知parseargs（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.cpp#l48），校验失败时直接抛出错误信息。

3.2 parseargs: 签名与参数的匹配

python的keyword arguments是类似map的结构，在c 中不方便直接用，需要转为positional arguments，同时按顺序保存到parsed_args中供后续执行使用。而这个顺序只能是签名指定的顺序，所以parseargs中只能按function_def的顺序循环校验。

函数的参数可能是各种类型，parseargs统一转为pythonarg类型，并通过pyobject*类型的成员读取python的变量值。

参数校验不一致的情况主要包括：

positional与keyword参数类型冲突

签名中的keyword参数名在kwargs中不存在且不接受默认值

参数类型不符合pythonargcheck规定的内部类型检查要求

kwargs包含function_def中未定义的参数

3.3 unpack_call: 展开算子函数的参数

在call函数中确定算子签名的schema之后，直接调用unpack_call（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/unpack_call.h#l69）函数。这时已经可以确定具体的算子执行函数了，对于relu来说就是functional::relu，同时将python传过来的参数都整理到args中。

unpack_call的模版参数是函数类型，例如functional::relu，在函数体内利用function_traits推导出函数的参数个数和返回值类型。

unpack_call_dispatcher内主要是调用f，也就是functional::relu。但还不能直接调用这个函数。因为每个算子对应函数的签名都不一样，又不能把vector args直接传给这些函数。

oneflow通过如下步骤完成模版的特化适配：

将args展开为各个pythonarg元素，通过index_sequence和变长模版参数包的展开实现；

利用function_traits推导得到函数参数类型列表argstype；

as函数调用可简化为as>()...核心是拿到各个参数的实际类型并交给as处理，最终调用objectas实现各种内部数据类型的转换。

unpack_call_dispatcher返回的是c 内部数据类型，最后要通过casttopyobject转为pybind11::object，主要是调用pybind11::cast函数。

class pythonarg { template t as() const { return objectashelper>()(this).getorthrow(); }};templatestruct unpack_call_dispatcher { template static r apply(const f& f, const std::vector& args, std::index_sequence) { // 这里适当改写了一下，把argstype抽出来 using argstype = function_traits::args_type; return f(args[i] .as::type>>()...); }};templatepy::object unpack_call(const f& f, const std::vector& args) { constexpr size_t nargs = function_traits::nargs; using r = typename function_traits::return_type; return casttopyobject( unpack_call_dispatcher::apply(f, args, std::make_index_sequence{}));}

3.4 签名都无效时的错误处理

以上只是讨论了python参数合法，可以找到匹配的函数签名的情况。如果传过来的参数是非法的，根据args/kwargs找不到匹配的签名怎么办？

如之前的讨论，pyfunctiondispatcher::call（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#l58c）是递归模版参数，如果当前签名不匹配，就尝试下一个签名。如果所有签名都不匹配，就会进入call的模版参数列表为空的特化版本（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/api/python/functional/py_function.h#l69）。这个函数会记录详细的错误信息。

例如，flow.pow("abc", 123)会输出如下错误信息：

file ".../oneflow/api/python/functional/py_function.h", line 76, in call typeerror: pow(): received an invalid combination of arguments. the valid signatures are: *0: tensor (tensor input, tensor exponent) *1: tensor (tensor input, scalar exponent, *, bool inplace=false) *2: tensor (tensor input, scalar exponent) *3: tensor (scalar exponent, tensor input)

而relu这种只支持一个签名的算子，如下面看到的，参数类型错误时的提示信息体现了单个签名的特点。如上所述，这是由schema_size_ == 1提示给parseargs的。

flow.relu(1)typeexception: file ".../oneflow/api/python/functional/py_function.cpp", line 98, in parseargs typeerror: relu(): argument "x" must be tensor, not int

3.5 yaml cpp的生成

functional_api.yaml的相关代码是在cmake构建过程中生成的，对应的cmake脚本是cmake/functional.cmake。

3.6 小结

总结一下上述几个主要组件的作用：

pyfunction是pybind11的def定义的入口函数，并为算子保存一个dispatcher对象用于推断合适的签名；

pyfunctiondispatcher通过模版函数的递归调用实现了签名的自动筛选，通过成员变量为参数校验和异常提示保存必要的信息；

unpack_call在编译期就确定了具体执行的算子函数类型，这一点在pyfunctiondispatcher中是无法做到的；

unpack_call_dispatcher的作用是将vector展开为多个元素、作为调用算子函数的参数，这在unpack_call中也是无法做到的；

pythonarg是python与c 类型转换的桥梁，同时承担类型检查的职能；

基于yaml生成的2组文件，yaml.pybind.cpp中调用pybind11的m.def指定模块调用的函数，并定义了函数签名的schema结构作为pyfunction的模版参数。yaml.cpp中则定义了具体的执行函数，如relu。将二者衔接起来的就是schema的字段func，对于relu算子来说，签名schema的func字段就是函数functional:relu。

核心是实现签名的自动校验推断，参数的统一处理以及参数的合并、展开。整个过程环环相扣、自然流畅。

4

算子functor的注册与执行

4.1 算子functor的注册

追踪一下functional::relu（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#l40）的调用链路，容易发现最终会用到functionlibrary的静态map变量。先看看这个map是怎么初始化的。它在add_functor_creator（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#l93）中被添加元素，后者被add_functor（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#l63）间接调用。

搜索一下add_functor和relu，发现在activation_functor.cpp中调用宏oneflow_function_library（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/core/functional/impl/activation_functor.cpp#l444）。宏展开后代码如下，通过定义一个静态变量来实现调用注册函数的目的。

static void _oneflow_function_library_0(functionlibrary & m);// 以定义一个静态变量的方式调用注册函数static int _oneflow_function_library_dummy_0 = []() { functionlibrary* library = functionlibrary::global(); _oneflow_function_library_0(*library); return 0; }();void _oneflow_function_library_0(functionlibrary & m) { m.add_functor("relu");};

稍微梳理一下就可以发现，functionlibrary的map中的value是类似下面这样的lambda：

[=]() { // func如 impl::relufunctor func func; // func_name来自lambda绑定，如relu return packedfunctormaker::make(func_name, func);}

注册的调用顺序如下：

那么，add_functor的模版参数为何是变长的，内部又要展开呢？是因为scalaradd等名字对应多个functor。

4.2 算子functor的执行

接下来看看functional_api.yaml.cpp中的functional::relu函数。代码经过整理后如下所示。

maybe relu(const std::shared_ptr& x, bool inplace) { static thread_local const auto& __op = check_just( functionlibrary::global()->find < maybe, const std::shared_ptr&, bool > ("relu")); return __op->call(x, inplace);}

核心逻辑就是func_lib.find("relu").call(x, inplace)。

获取__op并执行的调用顺序如下（忽略op的静态属性）：

根据上面的讨论以及调用链路容易发现，packedfunccreatormap::get内的静态map变量（https://github.com/oneflow-inc/oneflow/blob/release/0.7.0/oneflow/core/functional/function_library.h#l40），其value实际是一个类似如下的lambda表达式：

[=]() { // func如 impl::relufunctor func func; // func_name来自lambda绑定，如relu return packedfunctormaker::make(func_name, func);}

find返回的是it->second()，也就是调用这个lambda表达式的返回值，即packedfunctormaker::make的返回值，类型是packedfunctor，这就是op__的类型。其中模版参数f的类型如decltype(relufunctor::operator())。

packedfunctor构造时接受如下的lambda表达式，并保存到变量impl_中：

// func是一个函数变量，类型如 impl::relufunctor[func](const remove_cvref_t&... args) -> r { return func(std::forward&>(args)...);}

所以__op->call(...)就是packedfunctor::call(...)，最终相当于调用impl::relufunctor::operator()(args)。

也就是说，relu的操作就由impl::relufunctor执行。

需要注意的是，这里整个链路的分析，最关键的是模版参数的梳理和推导。模版参数确定后，整个逻辑还是比较清楚的。

4.3 小结

同一个名字可能对应多个functor。所以不能只用名字作为functor的key，需要结合签名。

functionlibrary负责管理所有的functor。但是单例不适合作为模版类，所以通过内嵌的packedfunccreatormap保存签名各异的functor。

每种签名都会特化一个packedfunccreatormap模版类，再通过名字区分不同的functor。

那么，packedfunctor类的作用是什么？或者换个角度，如果没有这个类，能否实现需求？答案是不能。

首先，yaml生成的2个cpp文件，都没有functor信息，只有relu这个名字、以及functor的签名信息。functor是在各个模块根据名字注册的。yaml与functionlibrary通过名字和签名进行交互。

其次，functionlibrary::find返回的packedfunctor是带模版参数的（参数就是functor签名）。find能否直接返回functor对象呢？主要是map不便存储不同类型的functor。即使functor都有共同的虚基类、map的value存储指针，但不能要求所有functor的执行接口是一致的，虚函数不满足这个场景的需求。所以find不能直接返回functor对象。

packedfunctor的作用就在于，它把真正的functor包在自己的结构里面；它的模版参数与functor的调用接口一致；它的call方法将op的所有入参通过lambda转发给functor。

functor能直接作为packedfunctor的成员变量吗？应该是可以的。packedfunctormaker::make的模版参数也包含functor。但是这样每个functor都要特化一个packedfunctor，编译后的可执行程序容易膨胀。而现在的实现，packedfunctor只根据functor执行函数签名特化，代价是要做一次调用转发（编译器有优化空间？）。

参考资料

从python到c 调用过程分析

https://github.com/oneflow-inc/oneflow/tree/release/0.7.0

（本文经授权后发布，原文：https://segmentfault.com/a/1190000041843994）

其他人都在看

深度学习概述

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

手把手推导分布式矩阵乘的最优并行策略

训练千亿参数大模型，离不开四种并行策略

解读pathways（二）：向前一步是oneflow

关于并发和并行，go和erlang之父都弄错了？

oneflow v0.7.0发布：全新分布式接口，libai、serving等一应俱全

欢迎体验oneflow v0.7.0：github - oneflow-inc/oneflow: oneflow is a performance-centered and open-source deep learning framework.oneflow is a performance-centered and open-source deep learning framework. - github - oneflow-inc/oneflow: oneflow is a performance-centered and open-source deep learning framework.https://github.com/oneflow-inc/oneflow/