关于机器学习:OneFlow源码阅读1算子签名的自动推断

12次阅读

共计 12482 个字符,预计需要花费 32 分钟才能阅读完成。

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 开始的具体调用流程如下:

把代码放到一起看看:

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 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<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_TTB
template<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_TTB
template<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 实现了这个主动筛选性能。

每个算子都会特化一个 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<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_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<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

正文完
 0