5

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

 2 years ago
source link: https://blog.csdn.net/OneFlow_Official/article/details/124013895
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

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

9cae8c5d7f9eab06d4e0e58390735065.png

撰文|月踏

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

虽然调用在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:

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

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

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

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

可以看到里面最和调用流程相关的是一个指向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:

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

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

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

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

2

C++ glue functor

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

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

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

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

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

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

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

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

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

其他人都在看

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


About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK