撰文｜月踏

在OneFlow中，从Python端我们可以使用各种Op进行相关操作，下面是一个最最简单的relu op的使用示例：

>>> import oneflow as of>>> x=of.tensor([-3,-2,-1,0,1,2,3], dtype=of.float)>>> of.relu(x)tensor([0., 0., 0., 0., 1., 2., 3.], dtype=oneflow.float32)

虽然调用在Python端，但具体的实现是在C++端，那么OneFlow是怎么样一步步从Python端调到C++中的呢，本文以最最简单的Relu这个Op作为例子，来追溯一下在OneFlow中从Python端到C++中的大致调用过程，具体过程大概总结为Python wrapper和C++ glue functor两部分，下面是两部分的具体细节。

1

Python wrapper

Python的代码都在python/oneflow文件夹中，在分析Python wrapper的过程中，也会涉及很多C++代码，主要是和pybind11绑定相关的，也一并归类到Python wrapper这部分了。

先看本文开头示例中的relu接口的直接来源，在python/oneflow/__init__.py中可以找到下面这一行：

from oneflow._C import relu

可以看到relu是从_C这个module中导出来的，所以继续看oneflow/_C/__init__.py这个文件：

from oneflow._oneflow_internal._C import *

可见relu接口来自_oneflow_internal这个module，_oneflow_internal是pybind11定义的一个module，位于oneflow/api/python/init.cpp：

PYBIND11_MODULE(_oneflow_internal, m) { ... ::oneflow::cfg::Pybind11ModuleRegistry().ImportAll(m); ::oneflow::OneflowModuleRegistry().ImportAll(m);}

继续看上面代码中的OneflowModuleRegistry，它是注册过的Op暴露到Python层的关键，它位于oneflow/api/python/of_api_registry.h：

class OneflowModuleRegistry { ...void Register(std::string module_path, std::function BuildModule);void ImportAll(pybind11::module& m);};

这个类提供了一个Register接口，被封装进了下面这个注册宏里，代码位于oneflow/api/python/of_api_registry.h：

#define ONEFLOW_API_PYBIND11_MODULE(module_path, m)   struct OfApiRegistryInit {                                                     OfApiRegistryInit() {       ::oneflow::OneflowModuleRegistry()                                         .Register(module_path, &OF_PP_CAT(OneflowApiPythonModule, __LINE__));     }                                                                            }; OfApiRegistryInit of_api_registry_init; static void OF_PP_CAT(OneflowApiPythonModule, __LINE__)(pybind11::module & m)

知道了ONEFLOW_API_PYBIND11_MODULE这个宏，继续搜哪里会用到它，在build/oneflow/api/python/functional/functional_api.yaml.pybind.cpp这个自动生成的文件中，可以搜到它被用到：

ONEFLOW_API_PYBIND11_MODULE("_C", m) { py::options options;  options.disable_function_signatures(); ... m.def("relu", &functional::PyFunction<:reluschema_ttb>); ... options.enable_function_signatures();}

由此可知本节刚开头的from oneflow._C import relu这句代码中的_C这个module和Relu这个算子是从哪来的了，在这里Relu被映射到了functional::PyFunction<:reluschema_ttb>这个函数，这是一个模板函数，先看其中的模板参数ReluSchema_TTB的定义：

struct ReluSchema_TTB {using FType = Maybe<:tensor>(const std::shared_ptr<:tensor>& x, bool inplace);using R = Maybe<: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;};

可以看到里面最和调用流程相关的是一个指向functional::Relu的函数指针成员，functional::Relu这个系列的函数非常重要，它是一个自动生成的全局C++ 接口，可以认为是Python和C++之间的分水岭，细节在第二节会详细讲，下面继续来看functional::PyFunction<:reluschema_ttb>这个模板函数，是它决定了怎么样去调用functional::ReluSchema_TTB中的func这个指向functional::Relu的函数指针，functional::PyFunction模板函数定义位于oneflow/api/python/functional/py_function.h：

templateinline py::object PyFunction(const py::args& args, const py::kwargs& kwargs) {static PyFunctionDispatcher dispatcher;return dispatcher.call(args, kwargs, std::make_index_sequence{});}

这里又继续调用了PyFunctionDispatcher中的call函数：

templateclass PyFunctionDispatcher { ...template py::object call(const py::args& args, const py::kwargs& kwargs,std::index_sequence) const {std::cout using T = schema_t;std::vector parsed_args(T::max_args);    if (ParseArgs(args, kwargs, &parsed_args, T::function_def, T::max_pos_args, schema_size_ == 1)) {return detail::unpack_call(*T::func, parsed_args); }return call(args, kwargs, std::index_sequence{}); } ...};

这里把functional::ReluSchema_TTB中的func这个指向functional::Relu的函数指针作为参数，继续调用了oneflow/api/python/functional/unpack_call.h中的unpack_call：

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

这里又把functional::ReluSchema_TTB中的func这个指向functional::Relu的函数指针作为参数，继续调用了同一个文件中的unpack_call_dispatcher::apply：

templatestruct unpack_call_dispatcher {templatestatic R apply(const F& f, const std::vector& args, std::index_sequence) {    return f(args[I].As<:detail::remove_cvref_t>::args_type>::type>>()...); }};

至此完成了对全局C++接口functional::Relu的调用，下一节具体讲functional::Relu这个全局C++接口怎么生成的。

2

C++ glue functor

先看oneflow/core/functional/impl/activation_functor.cpp中的一个类，它对下是通往Relu底层实现的大门，通往底层的实现是OneFlow框架的精髓，我还没有往里看，以后时机到了会继续总结出来，对上则提供了上层调用的接口，本文只关注接口部分：

class ReluFunctor { ... Maybe operator()(const std::shared_ptr& x, bool inplace) const {    ...    return OpInterpUtil::Dispatch(*op_, {x});  }};

ReluFunctor提供了一个函数调用符的重载函数，所以它对应的对象是可调用对象，它会被下面的代码进行注册：

ONEFLOW_FUNCTION_LIBRARY(m) { m.add_functor<:relufunctor>("Relu"); ...};

继续看ONEFLOW_FUNCTION_LIBRARY的定义，它通过定义一个静态变量的办法来在OneFlow的初始化阶段把上面的类似ReluFunctor的这些funtor通过add_functor接口全部注册到FunctionLibrary这个单例类中：

#define ONEFLOW_FUNCTION_LIBRARY(m) ONEFLOW_FUNCTION_LIBRARY_IMPL(m, __COUNTER__)#define ONEFLOW_FUNCTION_LIBRARY_IMPL(m, uuid)                                  static int OF_PP_CAT(_oneflow_function_library_dummy_, uuid) = []() { FunctionLibrary* library = FunctionLibrary::Global(); OF_PP_CAT(_oneflow_function_library_, uuid)(*library); return 0; }(); void OF_PP_CAT(_oneflow_function_library_, uuid)(FunctionLibrary & m)

FunctionLibrary的主要数据结构和接口如下，其中PackedFuncMap是一个用于存放注册对象的数据结构，add_functor用于注册，find用于查找已经注册过的对象， Global是单例接口：

class FunctionLibrary { templatestruct PackedFuncMap {static HashMap<:string functorcreator>* Get() { using FunctorCreator = typename std::function()>;static HashMap<:string functorcreator> functors;return &functors; } };
template  void add_functor(const std::string& func_name) { ... }
template auto find(const std::string& func_name) -> Maybe::FType>> { ... }
static FunctionLibrary* Global() {static FunctionLibrary global_function_library;return &global_function_library; }};

再继续看上面代码中的数据结构部分中用到的PackedFunctor，位于oneflow/core/functional/packed_functor.h，它通过call接口封装了functor的调用：

templateclass PackedFunctor {public:  PackedFunctor(const std::string& func_name, const std::function& impl) : func_name_(func_name), impl_(impl) {}  R call(Args&&... args) const {    return impl_(std::forward(args)...); }
private:std::string func_name_;std::function impl_;};

前面这部分都是functor的定义和注册部分，它们是提供全局C++接口的基石，下面继续看全局的C++接口functional::Relu是怎么来的，在code base中，有一个oneflow/core/functional/functional_api.yaml的配置文件，与Relu相关的内容如下：

- name: "relu" signature: "Tensor (Tensor x, Bool inplace=False) => Relu"  bind_python: True

这是一个yaml配置脚本，最终的functional::Relu这个全局C++接口就是通过前面的functor的定义、注册、yaml配置，最后再通过tools/functional/generate_functional_api.py这个python脚本自动生成出来，精简代码如下：

if __name__ == "__main__": g = Generator("oneflow/core/functional/functional_api.yaml") g.generate_cpp_header_file(header_fmt, "oneflow/core/functional/functional_api.yaml.h") g.generate_cpp_source_file(source_fmt, "oneflow/core/functional/functional_api.yaml.h") ...

可见具体的接口被生成到了上面指定的文件中，具体的生成过程在generator.py中，内容比较trivial，主要是通过hard code的方式来自动生成全局C++接口，下面是functional::Relu这个全局C++接口的示例：

namespace oneflow {namespace one {namespace functional {...Maybe<:tensor> Relu(const std::shared_ptr<:tensor>& x, bool inplace) {static thread_local const auto& op = CHECK_JUST(FunctionLibrary::Global()->find, const std::shared_ptr<:tensor>&, bool>("Relu"));return op->call(x, inplace);}...} // namespace functional} // namespace one} // namespace oneflow

可以看到上面的Relu接口通过注册类的find接口找到了注册过的ReluFunctor，然后用PackedFunctor中的call接口进行了调用，至此，我们终于知道了functional::Relu这个全局C++接口的前因后果。

其他人都在看

• 颠覆式编程：软件2.0

• “远见者”特斯拉AI主管Karpathy

• TVM：成为深度学习领域的“Linux”

• 深度学习崛起十年：“开挂”的OpenAI革新者

• 计算机架构新黄金时代，GPU能否继续保持辉煌

• 芯片设计“花招”已耍完？无指令集架构颠覆旧套路

欢迎下载体验OneFlow新一代开源深度学习框架：https://github.com/Oneflow-Inc/oneflow/

 

---

本文分享自微信公众号 - OneFlow（OneFlowTechnology）。
如有侵权，请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。