1.为什么需要深度学习编译器
深度学习编译器主要为解决不同框架下训练的模型部署到指定的某些设备上时所遇到的一系列复杂的问题,即将各种深度学习训练框架的模型部署到各种硬件所面临的问题;
首先深度学习领域,从训练框架看,当前可选的框架有pytorch、TensorFlow、Mxnet、paddle,oneflow、caffe/caffe2、mindspored等,具体选择哪个,不尽相同,但如果项目要部署落地,则面临很多问题,即从推理框架角度来看,无论我们选择何种训练框架训练模型,我们最终都是要将训练好的模型部署到实际场景的,在模型部署的时候我们会发现我们要部署的设备可能是五花八门的,例如Intel CPU/Nvidia GPU/Intel GPU/Arm CPU/Arm GPU/FPGA/NPU(华为海思)/BPU(地平线)/MLU(寒武纪),如果我们要手写一个用于推理的框架在所有可能部署的设备上都达到良好的性能并且易于使用是一件非常困难的事。
为了解决上面的问题,科学家为编译器抽象出了编译器前端,编译器中端,编译器后端等概念,并引入IR (Intermediate Representation)的概率。解释如下:
编译器前端:接收C/C++/Java等不同语言,进行代码生成,吐出IR
编译器中端:接收IR,进行不同编译器后端可以共享的优化,如常量替换,死代码消除,循环优化等,吐出优化后的IR
编译器后端:接收优化后的IR,进行不同硬件的平台相关优化与硬件指令生成,吐出目标文件
因此我们可以将各个深度学习框架训练出来的模型看做各种编程语言,传入深度学习编译器,之后吐出IR,由于深度学习的IR其实就是计算图,所以可以叫做Graph IR,针对这些Graph IR可以做一些计算图优化在吐出IR分发给各种硬件使用,这样就解决了上述很多繁琐的问题,如下图所示:
2. TVM
基于上面深度学习编译器的思想,陈天奇领衔的TVM横空出世。TVM就是一个基于编译优化的深度学习推理框架(暂且说是推理吧,训练功能似乎也开始探索和接入了),我们来看一下TVM的架构图,图片来自:https://tvm.apache.org/2017/10/06/nnvm-compiler-announcement
从这个图中我们可以看到,TVM架构的核心部分就是NNVM编译器(注意一下最新的TVM已经将NNVM升级为了Realy,所以后面提到的Relay也可以看作是NNVM)。NNVM编译器支持直接接收深度学习框架的模型,如TensorFlow/Pytorch/Caffe/MxNet等,同时也支持一些模型的中间格式如ONNX、CoreML。这些模型被NNVM直接编译成Graph IR,然后这些Graph IR被再次优化,吐出优化后的Graph IR,最后对于不同的后端这些Graph IR都会被编译为特定后端可以识别的机器码完成模型推理。比如对于CPU,NNVM就吐出LLVM可以识别的IR,再通过LLVM编译器编译为机器码到CPU上执行