综合技术交流
直播中

ss

7年用户 8762经验值
擅长:电源/新能源 制造/封装 RF/无线
私信 关注
[经验]

TVM主要的编译过程解析

`  TVM主要的编译过程如下图:
  
  Import:将tensorflow,onnx,pytorch等构建的深度学习模型导入,转化成TVM的中间层表示IR。
  Lower:将高层IR表示转化成低阶tiR表示。
  Codegen:内存分配和硬件可执行程序生成。
  图导入
  通过一个tensorflow的reception网络来熟悉编译过程,其它深度学习框架也具有类似过程。从TVM官网可以下载tensorflow的编译程序
  https://tvm.apache.org/docs/tuto ... nsorflow.html#sphx.。主要代码如下:
   2.png
  模型的输入是一个后缀为pb的文件,它是神经网络模型图的protobuf格式存储文件。Pb是二进制形式,pbtxt是文本形式。Import_graph_def函数是导入pb,graph是tensorflow的图结构。
  From_tensorflow是将tensorflow的图结构转化成TVM的IR。这个函数在文件relay/frontend/tensorflow.py中。函数的调用关系为:
  From_tensorflow -》 GraphProto.from_tensorflow -》 self._get_relay_func。
  在get_relay_func中会遍历每个tensorflow的节点,转换成tvm的IR表示。重点关注_backtrack_construct函数。
   3.png
  继续深入和算子转化有关的函数调用为:_convert_operator -》 convert_map。Convert_map中对应了可支持tensorflow算子到tvm算子的转换关系。
   4.png
  完成了tensorflow到TVM算子转化后,我们就得到了一个IRModule。我们可以利用tvm的可视化来打印出转化后的图:
   5.png
  Main是主函数入口,在TVM中以函数形式反应了tensorflow的图结构。函数的调用关系反应了图的依赖关系。
  编译
  Python中主要代码位于relay/build_module.py文件中,调用关系为build -》 BuildModule -》 build。在build中通过字典获得了C++中的相应函数。
   6.png
  这里不明白如何通过self.mod[“build”]得到C++中函数的。_BuildModule()是C++中注册到环境中的一个函数。在src/relay/backend/build_module.cc中,
   7.png
  TVM_REGISTER_GLOBAL是将C++函数注册到一个全局map中。当python加载编译好的动态库时,会自动查询map中静态注册的函数,并添加到python模块当中。
  真正build操作位于RelayBuildModule类中,在其中有一个GetFunction函数,会通过名字查询要使用的函数,打包成PackedFunc返回,这个函数可能和self.mod[“build”]有关。PackedFunc是TVM中提供的python的一个接口,任何函数都可以封装成PackedFunc,并给python调用。更详细介绍可看:https://hjchen2.github.io/2020/0 ... %AE%9E%E7%8E%B0%E6.。。
  继续深入代码,Build -》 BuildRelay。这是编译的主要代码。其过程包括optimize,codgen。
   8.png
  Optimize就是执行一些优化passes,这些passes包括常数折叠,算符融合等。之后会调用graph_codegen-》Codegen。Codegen中实现了内存分配和硬件代码生成。

` 1.png

更多回帖

发帖
×
20
完善资料,
赚取积分