OneFlow 源码解析:算子签名的自动推断
撰文 | 郑建华
OneFlow 是一个原生支持分布式训练的、高性能的深度学习框架。最近读了一些 OneFlow 的源码、架构设计和代码实现的文章,简单梳理一下自己的理解。主要通过图形展示调用过程和类之间的关系,只对部分重要的代码作一下分析。
深度学习框架是一个复杂的系统,而用户使用最多的就是算子(op)。用户通过 op 构造模型,进行训练、预测。这个笔记就从 op 入手,看看从 Python 前端到 C++底层,OneFlow 如何执行算子的计算逻辑。
具体地说,以比较简单的 Relu 算子为例,分析如下代码怎么执行:
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 是源码目录。
编译比较耗时,可以把两个目录 mount 到容器,便于后续查看 build 目录中生成的文件。
在 cmake 配置、构建过程中,会下载很多第三方源码包,如果网络状况不好容易超时,直接重试 cmake/make 即可。
用 GDB 追踪 OneFlow 的执行过程
王益:Use GDB to Walkthrough OneFlow Source Code(https://quip.com/JuQ0AuodVJn4)
2、Python Binding
OneFlow 底层是 C++实现,通过 pybind11 实现 Python Binding。月踏在《从 Python 到 C++调用过程分析》对相关内容做了讲解。
2.1 Relu 的 Python 包路径
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):
其中 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):
从这段代码可以看出,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 为例,这个宏展开后的核心代码如下:
这段代码中的类似注册技巧,在 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:
下面就来看看 OneFlow 是怎么实现这个功能的。
Relu 算子的签名 Schema 如下所示:
先看一下从 PyFunction 开始的的调用顺序:
PyFunction 相关的代码如下(删掉了一些与核心逻辑无关的内容)。
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<typename tuple_element<I, typename ArgsType>>()...核心是拿到各个参数的实际类型并交给 As 处理,最终调用 ObjectAs 实现各种内部数据类型的转换。
unpack_call_dispatcher 返回的是 C++内部数据类型,最后要通过 CastToPyObject 转为 pybind11::object,主要是调用 pybind11::cast 函数。
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)会输出如下错误信息:
而 Relu 这种只支持一个签名的算子,如下面看到的,参数类型错误时的提示信息体现了单个签名的特点。如上所述,这是由 schema_size_ == 1 提示给 ParseArgs 的。
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)。宏展开后代码如下,通过定义一个静态变量来实现调用注册函数的目的。
稍微梳理一下就可以发现,FunctionLibrary 的 map 中的 value 是类似下面这样的 lambda:
注册的调用顺序如下:
那么,add_functor 的模版参数为何是变长的,内部又要展开呢?是因为 ScalarAdd 等名字对应多个 Functor。
4.2 算子 Functor 的执行
接下来看看 functional_api.yaml.cpp 中的 functional::Relu 函数。代码经过整理后如下所示。
核心逻辑就是 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 表达式:
find 返回的是 it->second(),也就是调用这个 lambda 表达式的返回值,即 PackedFunctorMaker::make 的返回值,类型是 PackedFunctor<F>,这就是 op__的类型。其中模版参数 F 的类型如 decltype(ReluFunctor::operator())。
PackedFunctor 构造时接受如下的 lambda 表达式,并保存到变量 impl_中:
所以__op->call(...)就是 PackedFunctor<Func>::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 执行函数签名特化,代价是要做一次调用转发(编译器有优化空间?)。
参考资料:
https://github.com/Oneflow-Inc/oneflow/tree/release/0.7.0
(本文经授权后发布,原文:https://segmentfault.com/a/1190000041843994)
其他人都在看
欢迎下载体验 OneFlow v0.7.0 最新版本:https://github.com/Oneflow-Inc/oneflow/
版权声明: 本文为 InfoQ 作者【OneFlow】的原创文章。
原文链接:【http://xie.infoq.cn/article/6ea618ed62ff09faefdd95e71】。文章转载请联系作者。
评论