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

56次阅读

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

撰文|月踏

在 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::PyFunctionfunctional::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::PyFunctionfunctional::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-In…

正文完
 0