赞
踩
最近对深度学习编译器TVM学习了一下,了解各部分的功能,根据自己的理解进行总结,如有不正确请指正,话不多说上脑图:
TVM为解决各种深度学习训练框架的模型部署到各种硬件而诞生
TVM的前端,完成各种深度学习框架的计算图到Relay IR的转化
Import(fronted): 各种深度学习框架的模型导入
Transform: IRModule等价转换达到优化的效果
上层relay::Function(一个支持控制流,递归等计算图),在编译阶段,一个relay::Function可能会被lower成多个tir::PrimFunc
下层tir::PrimFunc(底层threading,vector/tensor的指令等OP执行单元)
前端的Relay IR:定义:high-level 计算图的描述,用IR表示来描述神经网络结构
LLVM的底层IR(Tir):借助Object元类实现统一的AST Node表示
Topi:Tensor计算库,包含了很多神经网络通用的算子(矩阵乘法、卷积)
Te:Te表示Tensor Expression,用户可以通过调用te中的函数来构建Tir(可直接通过Te来写神经网络)
Node:可以允许用户对一些函数进行访问
Arith:这个和Tir有关,可以在Tir优化时进行一些分析
Relay Passes到TIR Passes的过程
1. 调度原语(scheduling primitives),比如循环变化(split, unrool..),内联,向量化,target搜索
2. 可选的自动调优模块。不断的跑程序,记录性能,调整调度选择,从log文件中选择最优的性能对应的调度方案来进行执行
定义:pass是对计算图的一些优化和转换,比如常量折叠,算符融合,死代码消除等。
常规:如constant folding、dead-code elimination
张量计算:如transformation,scaling factor folding
主要功能是lower和一些optimization
如多维数据扁平化、后端intrinsics扩展
Module Level Pass 基于全局信息进行优化,可以删减Function
Function Level Pass 对Module中的每个Function进行优化,只有局部信息
Sequential Level Pass 顺序执行一系列的Pass
通过遍历AST,修改node来实现的,通过TVM_REGISTER_GLOBAL宏来注册和暴露支持的Pass
1. 新增一个AST Traversers,用来确定哪些node是需要修改的
2. 新增一个Expression Mutators,用于修改和替换满足条件的node
编译器将TIR变换为目标硬件上可执行的格式(即代码生成)
对于X86和Arm CPU,TVM使用LLVM IR Builder在内存中构建llvm ir
生成源代码级别的语言,比如生成CUDA C或者OpenCL的源码
支持直接从Relay Function到特定后端的Codegen
封装编译DSO的核心单元,它包含了很多PackedFunc,可以根据name来获取
模式
模型是静态shape,没有控制流,则lower到graph runtime
模型是动态shape,有控制流,可以使用「virtual machine backend」
直接将子图级别的程序转换为executable and generated primitive functions的级别
后端生成的函数,类似于深度学习框架中的核函数
封装了执行期Tensor的结构
仅支持静态Tensor Shape以及没有控制流的程序
解决动态Tensor shape和控制流的问题(动态Tensor shape只支持典型的推理场景,不支持训练)
只用来debug,因为解释器是通过遍历AST来执行程序,效率较低
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。