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)

编译环境

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

OneFlow提供了官网的编译镜像。用这个镜像能够十分不便地搭建编译环境。

我应用的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

Python Binding

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

relu的python包门路

# python/oneflow/__init__.pyfrom oneflow._C import relu# python/oneflow/_C/__init__.pyfrom oneflow._oneflow_internal._C import *

module解决逻辑的注册

Python代码次要在python/oneflow目录，C++实现的包次要在_oneflow_internal下，pybind11的绑定代码位于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是算子等模块的绑定；Pybind11ModuleRegistry应该是自定义的、相似protobuf的配置数据结构的绑定。

从OneflowModuleRegistry开始的具体调用流程如下：

把代码放到一起看看：

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& 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<void(pybind11::module&)>& BuildModule) {  // ...  BuildModule(m);  // ...}

从这段代码能够看出，python module的注册逻辑都保留在SubModuleMap中。它的key是module name；value是一组函数，BuildSubModule中调用这些函数、执行module注册逻辑。

GetSubModuleMap中保留map单例，Register函数设置map的值，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<functional::ReluSchema_TTB>);  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的定义。

多个接口签名的主动推断

从以上代码能够看到，relu算子被绑定到PyFunction这个函数执行计算逻辑，每次调用算子都会执行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<one::Tensor> (const std::shared_ptr<one::Tensor>& x, bool inplace);  using R = Maybe<one::Tensor>;  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_TTBtemplate<typename... SchemaT>class PyFunctionDispatcher { public:  // schema_t是第I个签名  template<size_t I>  using schema_t = typename std::tuple_element<I, std::tuple<SchemaT...>>::type;  // schema_size_是签名个数，比方relu是1，pow是4  PyFunctionDispatcher() : schema_size_(sizeof...(SchemaT)) {    signatures_.resize(schema_size_);    InitSignatures(std::make_index_sequence<sizeof...(SchemaT)>{});  }  template<size_t I0, size_t... I>  py::object call(const py::args& args, const py::kwargs& kwargs,                  std::index_sequence<I0, I...>) const {    // T是以后查看的签名，比方 ReluSchema_TTB    using T = schema_t<I0>;    std::vector<PythonArg> 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<I...>{});  }  py::object call(const py::args& args, const py::kwargs& kwargs, std::index_sequence<>) const {    // throw error ...    return py::none();  } private:  template<size_t... I>  void InitSignatures(std::index_sequence<I...>) {    __attribute__((__unused__)) int dummy[] = {        ((void)(signatures_[I] = schema_t<I>::signature), 0)...};  } private:  size_t schema_size_;  std::vector<const char*> signatures_;};// SchemaT如 ReluSchema_TTBtemplate<typename... SchemaT>inline py::object PyFunction(const py::args& args, const py::kwargs& kwargs) {  static PyFunctionDispatcher<SchemaT...> dispatcher;  return dispatcher.call(args, kwargs, std::make_index_sequence<sizeof...(SchemaT)>{});}// py module注册static void OneflowApiPythonModule9623(pybind11::module & m) {  m.def("relu", &functional::PyFunction<functional::ReluSchema_TTB>);  m.def("pow",  &functional::PyFunction<    functional::PowSchema_TTT,        functional::ScalarPowSchema_TTScB,    functional::ScalarPowSchema_TTSc, functional::ScalarReversePowSchema_TScT    >);}

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，校验失败时间接抛出错误信息。

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中未定义的参数

unpack_call: 开展算子函数的参数

在call函数中确定算子签名的Schema之后，间接调用unpack_call函数。这时曾经能够确定具体的算子执行函数了，对于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函数。

class PythonArg {  template<typename T>  T As() const {    return ObjectAsHelper<oneflow::detail::remove_cvref_t<T>>()(this).GetOrThrow();  }};template<typename F, typename R>struct unpack_call_dispatcher {  template<size_t... I>  static R apply(const F& f, const std::vector<PythonArg>& args, std::index_sequence<I...>) {    // 这里适当改写了一下，把ArgsType抽出来    using ArgsType = function_traits<F>::args_type;    return f(args[I]                 .As<oneflow::detail::remove_cvref_t<typename std::tuple_element<                     I, typename ArgsType>::type>>()...);  }};template<typename F>py::object unpack_call(const F& f, const std::vector<PythonArg>& args) {  constexpr size_t nargs = function_traits<F>::nargs;  using R = typename function_traits<F>::return_type;  return CastToPyObject(      unpack_call_dispatcher<F, R>::apply(f, args, std::make_index_sequence<nargs>{}));}

签名都有效时的错误处理

以上只是探讨了Python参数非法、能够找到匹配的函数签名的状况。如果传过来的参数是非法的，依据args/kwargs找不到匹配的签名怎么办？

如之前的探讨，PyFunctionDispatcher::call是递归模版参数，如果以后签名不匹配，就尝试下一个签名。如果所有签名都不匹配，就会进入call的模版参数列表为空的特化版本。这个函数会记录具体的错误信息。

例如，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

yaml cpp的生成

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

小结

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

  • 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。

外围是实现签名的主动校验推断，参数的对立解决，以及参数的合并、开展。整个过程环环相扣、天然晦涩。

算子Functor的注册与执行

算子Functor的注册

追踪一下functional::Relu的调用链路，容易发现最终会用到FunctionLibrary的动态map变量。先看看这个map是怎么初始化的。它在add_functor_creator中被增加元素，后者被add_functor间接调用。

搜寻一下add_functor和Relu，发现在activation_functor.cpp中调用宏ONEFLOW_FUNCTION_LIBRARY。宏开展后代码如下。通过定义一个动态变量来实现调用注册函数的目标。

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<impl::ReluFunctor>("Relu");};

略微梳理一下就能够发现，FunctionLibrary的map中的value是相似上面这样的lambda：

[=]() {  // Func如 impl::ReluFunctor  Func func;  // func_name来自lambda绑定，如Relu  return PackedFunctorMaker<func_type>::make(func_name, func);}

注册的调用程序如下：

多个Functor对应同一个名字

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

算子Functor的执行

接下来看看functional_api.yaml.cpp中的functional::Relu函数。代码通过整顿后如下所示。

Maybe<one::Tensor> Relu(const std::shared_ptr<one::Tensor>& x, bool inplace) {  static thread_local const auto& __op = CHECK_JUST(    FunctionLibrary::Global()->find      <        Maybe<one::Tensor>,        const std::shared_ptr<one::Tensor>&,        bool      > ("Relu"));  return __op->call(x, inplace);}

外围逻辑就是func_lib.find("Relu").call(x, inplace)。

获取__op并执行的调用程序如下（疏忽op的动态属性）：

依据下面的探讨以及调用链路容易发现，PackedFuncCreatorMap::Get内的动态map变量，其value理论是一个相似如下的lambda表达式：

[=]() {  // Func如 impl::ReluFunctor  Func func;  // func_name来自lambda绑定，如Relu  return PackedFunctorMaker<func_type>::make(func_name, func);}

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

PackedFunctor结构时承受如下的lambda表达式，并保留到变量impl_中：

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

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

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

须要留神的是，

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

小结

  • 同一个名字可能对应多个Functor。所以不能只用名字作为Functor的key，须要联合签名。
  • FunctionLibrary负责管理所有的Functor。然而单例不适宜作为模版类，所以通过内嵌的PackedFuncCreatorMap保留签名各异的Functor。
  • 每种签名都会特化一个PackedFuncCreatorMap模版类，再通过名字辨别不同的Functor。

PackedFunctor的作用

那么，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++调用过程剖析
  • OneFlow