OneFlow学习笔记:从Python到C++调用过程分析

Posted OneFlow深度学习框架

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了OneFlow学习笔记:从Python到C++调用过程分析相关的知识,希望对你有一定的参考价值。

撰文|月踏

在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>
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)>);

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

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...>);
  
  ...
;

这里把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/

以上是关于OneFlow学习笔记:从Python到C++调用过程分析的主要内容,如果未能解决你的问题,请参考以下文章

OneFlow学习笔记:从OpExprInterpreter到OpKernel

Global View的概念和实现|OneFlow学习笔记

BBuf的CUDA笔记八,对比学习OneFlow 和 FasterTransformer 的 Softmax Cuda实现

C++异常处理的学习笔记

BBuf的CUDA笔记一,解析OneFlow Element-Wise 算子实现

BBuf的CUDA笔记一,解析OneFlow Element-Wise 算子实现