撰文｜月踏
在 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<void(pybind11::module&)> 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<functional::ReluSchema_TTB>); ... options.enable_function_signatures();}
复制代码
由此可知本节刚开头的 from oneflow._C import relu 这句代码中的_C 这个 module 和 Relu 这个算子是从哪来的了，在这里 Relu 被映射到了 functional::PyFunction<functional::ReluSchema_TTB>这个函数，这是一个模板函数，先看其中的模板参数 ReluSchema_TTB 的定义：
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;};
复制代码
可以看到里面最和调用流程相关的是一个指向 functional::Relu 的函数指针成员，functional::Relu 这个系列的函数非常重要，它是一个自动生成的全局 C++ 接口，可以认为是 Python 和 C++之间的分水岭，细节在第二节会详细讲，下面继续来看 functional::PyFunction<functional::ReluSchema_TTB>这个模板函数，是它决定了怎么样去调用 functional::ReluSchema_TTB 中的 func 这个指向 functional::Relu 的函数指针，functional::PyFunction 模板函数定义位于 oneflow/api/python/functional/py_function.h：
template<typename... SchemaT>class PyFunctionDispatcher { ...template<size_t I0, size_t... I> py::object call(const py::args& args, const py::kwargs& kwargs,std::index_sequence<I0, I...>) const {std::cout << I0 << std::endl;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, schema_size_ == 1)) {return detail::unpack_call(*T::func, parsed_args); }return call(args, kwargs, std::index_sequence<I...>{}); } ...};
复制代码
这里又继续调用了 PyFunctionDispatcher 中的 call 函数：
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>{}));}
复制代码
这里把 functional::ReluSchema_TTB 中的 func 这个指向 functional::Relu 的函数指针作为参数，继续调用了 oneflow/api/python/functional/unpack_call.h 中的 unpack_call：
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>{}));}
复制代码
这里又把 functional::ReluSchema_TTB 中的 func 这个指向 functional::Relu 的函数指针作为参数，继续调用了同一个文件中的 unpack_call_dispatcher<F, R>::apply：
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...>) { return f(args[I].As<oneflow::detail::remove_cvref_t<typename std::tuple_element<I, typename function_traits<F>::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<Tensor> operator()(const std::shared_ptr<Tensor>& x, bool inplace) const { ... return OpInterpUtil::Dispatch<Tensor>(*op_, {x}); }};
复制代码
ReluFunctor 提供了一个函数调用符的重载函数，所以它对应的对象是可调用对象，它会被下面的代码进行注册：
ONEFLOW_FUNCTION_LIBRARY(m) { m.add_functor<impl::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 { template<typename R, typename... Args>struct PackedFuncMap<R(Args...)> {static HashMap<std::string, FunctorCreator>* Get() { using FunctorCreator = typename std::function<PackedFunctor<R(Args...)>()>;static HashMap<std::string, FunctorCreator> functors;return &functors; } };
template<typename... Fs> void add_functor(const std::string& func_name) { ... }
template<typename R, typename... Args> auto find(const std::string& func_name) -> Maybe<PackedFunctor<typename PackedFunctorMaker<R(Args...)>::FType>> { ... }
static FunctionLibrary* Global() {static FunctionLibrary global_function_library;return &global_function_library; }};
复制代码
再继续看上面代码中的数据结构部分中用到的 PackedFunctor，位于 oneflow/core/functional/packed_functor.h，它通过 call 接口封装了 functor 的调用：
template<typename R, typename... Args>class PackedFunctor<R(Args...)> {public: PackedFunctor(const std::string& func_name, const std::function<R(Args...)>& impl) : func_name_(func_name), impl_(impl) {} R call(Args&&... args) const { return impl_(std::forward<Args>(args)...); }
private:std::string func_name_;std::function<R(Args...)> 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<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);}...} // namespace functional} // namespace one} // namespace oneflow
复制代码
可以看到上面的 Relu 接口通过注册类的 find 接口找到了注册过的 ReluFunctor，然后用 PackedFunctor 中的 call 接口进行了调用，至此，我们终于知道了 functional::Relu 这个全局 C++接口的前因后果。
欢迎下载体验 OneFlow 新一代开源深度学习框架：https://github.com/Oneflow-Inc/oneflow/
评论