OneFlow源码阅读1:算子签名的自动推断

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)

编译环境

在开始分析之前,需要搭建环境编译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 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

Python Binding

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

relu的python包路径

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

# python/oneflow/_C/__init__.py
from 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开始的详细调用流程如下:
OneFlow源码阅读1:算子签名的自动推断_第1张图片

代码放到一起看看:

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中的宏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的定义。

多个接口签名的自动推断

从以上代码可以看到,relu算子被绑定到PyFunction这个函数执行计算逻辑,每次调用算子都会执行PyFunction这个函数。

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

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

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

import oneflow as flow
r = 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开始的的调用顺序:
OneFlow源码阅读1:算子签名的自动推断_第2张图片

PyFunction相关的代码如下(删掉了一些与核心逻辑无关的内容)。

// SchemaT如 ReluSchema_TTB
template
class 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_TTB
template
inline 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
    >);
}

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

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

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

每个算子都会特化一个PyFunctionPyFunctionDispatcher实例,也有一个算子自己的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>()...,核心是拿到各个参数的实际类型并交给As处理,最终调用ObjectAs实现各种内部数据类型的转换。

unpack_call_dispatcher返回的是C++内部数据类型,最后要通过CastToPyObject转为pybind11::object,主要是调用pybind11::cast函数。

class PythonArg {
  template
  T As() const {
    return ObjectAsHelper>()(this).GetOrThrow();
  }
};

template
struct 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>>()...);
  }
};

template
py::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{}));
}

签名都无效时的错误处理

以上只是讨论了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_functorRelu,发现在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("Relu");
};

稍微梳理一下就可以发现,FunctionLibrary的map中的value是类似下面这样的lambda:

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

注册的调用顺序如下:
OneFlow源码阅读1:算子签名的自动推断_第3张图片

多个Functor对应同一个名字

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

算子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的静态属性):
OneFlow源码阅读1:算子签名的自动推断_第4张图片

根据上面的讨论以及调用链路容易发现,PackedFuncCreatorMap::Get内的静态map变量,其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执行。

需要注意的是,

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

小结

  • 同一个名字可能对应多个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执行函数签名特化,代价是要做一次调用转发(编译器有优化空间?)。

参考资料

你可能感兴趣的