如何学习TVM的代码?

问题:
对陈天奇团队的开源深度学习编译器TVM很感兴趣,特别是看到18年发的论文中提到的在FPGA上的部署。对于基础知识薄弱(如体系架构编译等方面)的学生,应该如何学习TVM的代码呢?

如提问有误请指正,求轻喷。

相关链接:
TVM: End-to-End Optimization Stack for Deep Learning
自动生成硬件优化内核:陈天奇等人发布深度学习编译器TVM
dmlc/tvm

作者:蓝色
链接:https://www.zhihu.com/question/268423574/answer/506008668
来源:知乎
著作权归作者所有。商业转载请联系作者获得授权,非商业转载请注明出处。

或许和很多人不同,以我的经验来看,觉得理解TVM,或者推理框架一定要从前端开始。即你从一个Tensorflow模型 / MXNet模型等,是如何转为NNVM的,然后再应该是后续的图优化,以及后续的TVM Tensor,LLVM代码生成等东西。为什么我会这么强调从前端开始呢?因为不理解前端模型,就很难理解后续TVM为什么是这样,而且出了错以后很难知道到底是什么原因,比如很多时候找了半天,其实只是你忘记了模型输入图片的预处理,却误认为是后续卷积的调度优化做的有问题,**所以我强烈建议先从一个模型前端开始,在tvm/nnvm/frontend里面选取一个前端。**而选取前端开始不应该仅仅是看,Bug / 需求驱动永远是最好学习源代码的方式,建议从一个固化好的模型开始,然后补足NNVM算子,比如Mobilenet / Resnet50等,这里也是让你熟悉工具,熟悉NNVM的开始,可能会遇到很多问题,但是一个一个克服会收获很多,这里面推荐一个看模型的好工具: https://github.com/lutzroeder/Netron 我也是看苹果公司一个人用了以后发现的,确实是好东西。

接下来你应该首先理解TOPI,这是架设在NNVM与TVM之间的东西(首先忽略图优化,你后面再去看),因为你需要理解NNVM Symbol (其它模型在转为NNVM前端表示时会以Symbol形式的Api表示) 如何与TVM之间是如何连接起来的,在这里面你会有点迷糊,因为TVM是C++和Python混合的工程,这里面你需要在这两者跳来跳去,但是你这一步你最重要的是抓住两个核心: FTVMCompute (@reg.register_compute) / @reg.register_schedule,这一个你需要分别在nnvm/top里面的C++ / Python去找,top里面会告诉你是如何从NNVM进入topi的。

这一步完成以后,你则需要进入topi里面的任意一个后端Target去看,我暂时推荐x86后端,因为这一个后端还没有被AutoTVM改造。对于你来说,理解起来更容易。在这里你会遇到topi/nn里面的@tvm.target.generic_func到类似具体@generic.schedule_conv2d_nchw.register([“cpu”])的改变,这是TVM的核心所在,对于卷积这样的数据负载处理,为了优化而沿用Halide的思想: 计算与调度分离。为了理解这个,你最好参考一下这个文档: https://docs.tvm.ai/tutorials/optimize/opt_gemm.html#sphx-glr-tutorials-optimize-opt-gemm-py

到这一步理解好以后,后续的TVM底层API大部分情况下你都不需要去动,包括后续的LLVM自动生成,优化等你也大部分不需要去动,因为类似CNN这样的网络,大部分你要做的工作就是在调度上,如何减少Cache Miss ,如何更好的让数据Locality是更关键的地方。

到这一步以后,你可以再回过头去理解图优化的部分,如Operator Fusion / FoldScaleAxis等,以及包括TVM目前最核心最与众不同的地方: AutoTVM(https://docs.tvm.ai/tutorials/autotvm/tune_nnvm_arm.html#sphx-glr-tutorials-autotvm-tune-nnvm-arm-py),这是TVM去击败NCNN等用手写汇编的推理框架的关键一环,用机器学习去解决机器学习的问题,让你从调度参数的设置中解放出来,而专心写调度算法。这里面目前ARM CPU的调度算法并非是最优的,但是从测试来看,至少在测试中使用硬件和环境来看,已经超过能找到的推理框架。后续我将撰写一篇文章到TVM社区,将我在ARM CPU的工作写出来,这将改善目前ARM CPU的官方调度版本,这将在Mobilenet等模型中有很好的提升,敬请关注!

TVM是很好的一个项目,这种基于编译优化思想的深度学习推理框架正是我赞同的,虽然还有很多工作需要做,但是我认为它已经走在一个很好的方向上了。

手把手带你遨游TVM
在这里插入图片描述

人工智能无疑是目前风口浪尖的话题,大部分提到人工智能都关注在了算法,然而若要让人工智能落脚在实际应用中,工程化是至关重要的。

让我们再细化一点这个问题,我们目前要让一个算法落地下来,或者说我们要产生一个模型,然后部署到目标设备上进行运行。我们需要分两步:Training + Inference,目前对于Training来说,比较流行的是TensorFlow,PyTorch等,市场主流框架基本上已经确定下来了。而对于Inference来说,以我之所见,其实是“群雄逐鹿”,因为模型Training一次,但是会跑到的设备可能会是多种多样的,如Intel CPU / Intel GPU / ARM CPU / ARM GPU / NV GPU / FPGA / AI芯片等,而要在这多种多样的设备中都保持一个高效的Inference性能,其实是一件很有挑战的事情,也是作为程序员很喜欢的一件事情,知难而上,挑战难题达到颅内高潮的那种感觉也不是任何一件事情能达到的。

让我们图例化这一点:
在这里插入图片描述

其实要达到这一个目标并不容易,各种硬件设备的特性千差万别,要如何保持一个统一的高效执行,是一个非常难做到的事情。而各大硬件厂商针对这样的情况都推出了自己的Inference 框架(相比TensorFlow等这样的框架孱弱的Inference性能,各大设备厂商的Inference框架性能都比较不错),比如Intel的OpenVINO,ARM的ARM NN,NV的TensorRT等,但是这里面有一个问题,各大设备厂商的框架并不具备通用性,比如对训练框架模型产生的算子支持不全(尤其是像TensorFlow这种算子很多的),通常在一个设备厂商的Inference框架能跑,但是不一定在另外一个设备厂商的Inference框架上能跑。同时,对于业务开发来说也是非常痛苦的事情,我在这个硬件上要用这个,我在另外一个硬件上要用另外一个,而且两者还没有统一的使用体验,算子支持也不一样,性能还不一定是最好的…其实对于业务方来说,也是想要一个统一的Inference框架,然后我业务场景的各种硬件设备都能高效的跑,使用体验都是一致的,我只要换一个device_target参数就好了。

以我之所见,其实我们并不是第一次遇到这样的问题,我们曾经出现了很多种编程语言,有很多种硬件,历史上最开始也是一种语言对应一种硬件,从而造成编译器的维护困难与爆炸。
在这里插入图片描述

而编译器后面解决了这个问题,其具体解决办法是这样的:抽象出编译器前端,编译器中端,编译器后端等概念,引入IR (Intermediate Representation)

  • 编译器前端:接收C / C++ / Fortran等不同语言,进行代码生成,吐出IR
  • 编译器中端:接收IR,进行不同编译器后端可以共享的优化,如常量替换,死代码消除,循环优化等,吐出优化后的IR
  • 编译器后端:接收优化后的IR,进行不同硬件的平台相关优化与硬件指令生成,吐出目标文件

类似这样的架构:
在这里插入图片描述

其实对于Inference框架也是类似的道理,我们也想要得到类似的架构:
在这里插入图片描述

我们把各种模型抽象看成各种编程语言,于是我们引入一个新的编译器,负责把这些编程语言识别,吐出IR
在这里插入图片描述

接下来我们就可以对这种中间的IR进行优化,而深度学习中是计算图,所以我们可以称为Graph IR
在这里插入图片描述

于是,我们得到这样的架构:
在这里插入图片描述

这就是我们想要解决的问题了。所以,就来到我们今天的重点了,我们能不能做一个基于编译优化思想的推理框架呢?答案就是:TVM (https://tvm.ai/).
在这里插入图片描述

让我们给出一副更详尽,更漂亮的一张图:
在这里插入图片描述

这就是TVM官方给的架构图。在我们图中的CPU目标是通过什么来做到的呢?LLVM。当然,也包括了AMD GPU,也是LLVM来做到的,但是这张图里面没有列举出来。在这张图里面,对于NV GPU的支持是产生CUDA,但是从我的角度来看,其实也是可以产生NVVM的,然后走LLVM的路线。

TVM基本上就是基于编译优化思想的深度学习推理框架的完美体现,自从加入了AutoTVM(https://arxiv.org/pdf/1805.08166.pdf)机制以后,可以用一个词来形容,就是如虎添翼,击败NCNN等框架是没有任何问题的,可参考这篇文章的对比:Automatic Kernel Optimization for Deep Learning on All Hardware Platforms,并且这里的性能还没有应用上我这个PR:

Improve ARM CPU depthwise convolution performance by FrozenGene · Pull Request #2345 · dmlc/tvm

我的这个PR可以将深度卷积的性能提高近2倍。

而AutoTVM的存在最厉害的是具有非常强的适应性,很多Inference框架在一些标准模型上表现的很好,但是一遇到业务模型就性能急剧下降,比如来个257 * 513这种输入,但是AutoTVM不会存在这个问题。

深度学习推理引擎的一些思考
在这里插入图片描述

作者:蓝色
链接:https://zhuanlan.zhihu.com/p/87392811
来源:知乎
著作权归作者所有。商业转载请联系作者获得授权,非商业转载请注明出处。

我想熟悉我背景的人应该知道我以前是做传统编译器的,而加入阿里巴巴以后我开始做深度学习推理引擎,而在我们部门走编译优化这条路来做深度学习推理引擎也是我主推起来的。在很早以前的一篇文章:蓝色:手把手带你遨游TVM 我曾提到了我为什么认为这是一个编译器的问题,也为什么认为编译优化才是解决这个问题的最佳方案与未来,这也是我说服老板的技术缘由。

在18年初我刚加入阿里巴巴的时候,基于编译优化的深度学习推理引擎其实并不是主流,那时候更流行的解决方案是TFLite这样的框架,即调用高效的GEMM加速库,抑或着NCNN这样的框架,针对ARM CPU手写汇编。但是,我很坚定的对我的老板表示,这都不是未来,用编译器来解决这个问题才是未来。当时我有多坚定呢,基本上可以理解为如果不这么搞,我就离职那种。很庆幸的是,我的老板非常力挺我的想法,于是从我最开始调研到最后实际推动起来,并没有遇到很多阻力,而最后也回报了老板的信任,我们取得了很好的落地效果。而在这两年时间,无论是TVM,GLOW,MLIR等,走编译优化来做深度学习推理引擎越来越受到认可,以今年 TVM Conf Talk来看,可以看到很多公司的身影,Amazon,Facebook,Microsoft,ARM,Intel,Qualcomm,Xilinx等,当然也包括了没有列举在这上面的华为,阿里巴巴,以及很多基于TVM来做NPU或者应用的公司(题外话,据我所知,华为或许是中国投入做TVM力量最大的团队了),这里面其实很多人都是像我一样是传统编译器出身,比如参加过TVM Conf 2018的Qualcomm的Krzysztof Parzyszek,也是做LLVM编译器的,我和他就TVM如何支持Hexagon DSP也做过很多交流。说到这里,或许会问为什么我们会选取TVM作为Base。在18年初的时候,TVM其实很不成熟(而即使放到现在,我认为TVM也还有很多改善的地方,但是相比18年好很多,后文我会说到我的一些看法),但是相比那时候调研的Tensor Comprehensions, XLA, ngraph等等,TVM是最接近我想法的框架。而为什么没有完全从0开始?我当时的看法是除了开发者的技术情怀以外,我找不到任何理由从0开始,因为你最后要做出来的也是类似TVM这种,没有任何的必要,不如在这个基础上来做,很多脏活累活也都已经帮你做完了。如做编译器,现在从0开始做编译器的已经基本上少很多了,基于LLVM的基本上算是通用做法了。在底层软件这一块儿,我个人的看法是拥抱开源生态远比一个人玩强很多,一起把这个蛋糕做大。

上面说了这么多,其实还没有说到有关技术一些东西。接下来,我会就我观测的,以及实际开发的经验来说一下我的一些感想,或许会有不对的地方,权当抛砖引玉。

做深度学习推理引擎难不难?我个人认为没有想象的那么难,也没有想象的那么简单。看起来好像很废话,但是是针对不同的情景。如我刚接触的时候,我其实还是把它想象的还是很神秘的,模型跑起来后到底是怎么就把猫识别出来的呢?我现在看法就不一样了,你其实可以不管它怎么识别出来的,对于你引擎开发人员来说,这里面就是矩阵的计算,你如何把它算的又快又准就可以了,至于它为什么可以识别出来猫,可以去补补算法的知识。所以,我认为做深度学习推理引擎很适合做编译器、体系结构、HPC的人。而没有想象的那么简单是指什么呢?深度学习推理引擎要真正做好,还是需要下很多功夫的,也需要很多方面知识的。以TVM支持Hexagon DSP为例,这里面如果你不了解Hexagon的体系结构知识与编程模型,你基本上无从下手。而为了生成Hexagon DSP的指令,你还需要LLVM的知识来进行LLVM CodeGen,这又回到了传统编译器领域。而深度学习推理引擎也不仅仅如此,还包括量化压缩,图优化,子图分离,异构执行等,后续我也会谈到。

那我们从头开始说起。对于深度学习推理引擎,无论是你基于编译优化这条路,还是传统的调用加速库这条路,第一步不可避免的就是接收模型,如Tensorflow, TFLite, MXNet等。无论是TVM,还是其它框架,这方面做的其实都不够好。很多时候,框架使用者第一步尝试的就是拿你这个框架跑一下我的模型,然后发现各种算子不支持,更别提后面的性能了。如果你是基于TVM这种开源的引擎来做还好,遇到不支持的可以添加,但是如果你是使用类似高通的SNPE,Intel的OpenVINO这样的闭源方案,那你基本上除了改模型没有其他的办法。那么,对于没有开发引擎的人来说,或许很难理解这件事情,为什么各大推理引擎不能都支持完了。这里我们先限制在CNN领域的算子,其实并不是说这个多难,我个人认为这其实是一个体力活。以我做的TVM TFLite前端为例,我把这个前端做完了,并加上了convolution, relu等算子,我不能添加类似add算子的支持吗?显然不是我不能,而是我个人精力有限,我以Mobilenet V1模型为驱动,验证了我这条路和框架是通的,然后我开源出去。如果你遇到了不支持的算子,你可以在我的框架上添加,我可以帮忙Review,然后我们一起把这个壮大,这也是开源的伟大之处。然而,抛出开发者这个角度来说,这其实是最影响用户体验的一环,目前大家的做法基本上都是以用户报哪个算子不支持或者跑哪个模型发现不支持然后来增加,并不是像传统编译器一样,拿着一份表挨着挨着来做,比如Clang这样http://clang.llvm.org/cxx_status.html

在这里就不得不提ONNX,这是很伟大的构想,如果真的成功了,我们深度学习推理引擎只需要支持ONNX即可,其余的任何框架模型都先变为ONNX即可。但是ONNX属于一把好牌打得稀烂,目前在这方面投入还算比较积极的就属于微软了,至于原因是为什么,也是很显然能想到的,这里就不再赘述了。

对于深度学习推理引擎来说,解析完模型以后都会变为自己的计算图表示,在这里面大家都会做事情,但是做的都不尽相同。不过一些通用的图优化大家都会做,如算子融合。在图这方面,TVM从NNVM变为了Relay,并引入了类似LLVM的Pass机制,并支持了异构。这方面的优势暂且不谈,因为其实你能做的,我也能做,比如异构,算子融合等。这里我谈谈我个人觉得TVM这里还可以改进的一块,那就是自动化的异构子图分离。对于设备上有多个硬件的话(以高通的硬件为例,即有CPU,GPU,DSP),如何高效的自动把计算图分离成多个计算子图,并且异构的执行,从而更高效的执行模型,这一块儿其实TVM基本上没有探索。在这方面,Training当然是有更强烈的需求,但是不是说Inference就没有。那么如果进行子图分离,哪些算子放在DSP,哪些放在GPU,哪些放在CPU,实际运行的时候,如CPU有巨大的波动如何处理等。业界有一些在基于传统机器学习的Cost Model来做这个事情,也有一些在基于强化学习来做,我觉得都是可以学习的一点。

在计算图这一块,其实还有一块儿我觉得是可以考虑的,那就是如何在计算图对接外部推理引擎。这一块儿或许听起来很矛盾,因为类似TVM这样的本就是推理引擎,你为什么要考虑对接外部推理引擎。然而我认为只有真正的去接触了业务,你才能知道大家想要的到底是什么。对于业务部门来说,想要的就两点:我的模型可以很轻易地跑起来;我的模型可以跑得很快。以GPU为例,TensorRT是一个很强的高性能推理框架,现在我业务部门有一个模型我需要你在GPU跑起来,并且我需要在XX ms 跑完。而很遗憾的是TensorRT不支持这个模型的一些算子,并且在一些层比TVM慢,一些层比TVM快。而要在规定时间完成这个业务目标,最直接的办法就是TVM + TensorRT,前端统一走TVM,然后在Relay这里进行子图分离,把一些层交给TensorRT,一些层走TVM。这样,算子都由你来控制,同时TensorRT执行它高效的一部分,TVM执行不支持的算子以及TVM本身高效的一部分从而达到目的。这一块儿其实可以推广到NPU的支持,如果NPU不暴露指令集,但是暴露了计算图的API,你不能走编译生成这条路,你也可以这样做,这样你不再受限于NPU不支持算子的问题,你可以异构的跑在TVM CPU端。

经历了计算图,我们来到算子层面。这一块儿也是大家费工夫优化的点,以CNN为例,基本上就是卷积算子。我为什么说做编译器、体系结构、HPC的同学很适合呢,就是因为我们很大一部分优化都在这里,说白了,就是如何在目标平台更高效的执行多重循环。你以前学到的Cache、Loop优化完全能无缝的派上用场。在这一块儿,我觉得各大框架都差不多,Winograd, Spatial Pack, im2col,layout变换等,大家都会做,区别点就在于谁做的更好一些,但是都半斤八两,没道理你会的我不会。针对这个,我认为TVM跟很多传统框架有几个很不同的点。第一个就是Auto TVM,这一点在业务模型中非常厉害。因为无论业务模型的卷积形状如何变,Auto TVM总能给你找到很好的参数让你的计算有很好的局部性,这一点其实很影响性能。这一点其实也被很多框架看到了,也有很多框架开始学习Auto TVM来做Auto Tuning。第二个是TVM借助Halide思想,进行了计算与调度分离,并提供了DSL。你在ARM CPU做的优化,很容易移植到x86 CPU,你不需要生写NEON与SSE,你所关注的点在于是不是应该.vectorize,是不是应该.unroll,是不是应该.parallel。这个带来的工作效率的提升,我觉得只有真正做过优化的才能体会到,人就应该干人擅长的事情,关注用哪些优化技术,而不是痛苦的把NEON移植成SSE,或者CUDA,OpenCL,Metal等。

到后面代码指令生成部分,TVM走的就是编译生成,如CPU就是走的LLVM。其实相比传统编译器,这部分的编译生成简单很多,你要去做生成的TVM IR也很少很少,如Load, Store, Add等,就几十个。如你有一个新的硬件,你从TVM IR到硬件的CodeGen工作量不会很大,如果这个硬件是基于LLVM Toolchain的,这个工作量更少,你可以复用TVM CodeGenLLVM的很大一部分。很多人会说这样的编译效率不高,我个人是持反对意见的,在这部分只要能生成你想要的指令,如mla,那么无论是走LLVM生成,还是你手写汇编生成,抑或着你去调用高效的加速库,其实没有差别。而如果你觉得做得不好,其实在TVM中提供了一个机制,叫做Tensorize,该机制允许你替代自动生成的部分,而是自己手写的代码。这是我觉得TVM在这里做得很好的地方,主推编译优化,却允许容纳手工微内核。

接下来谈谈量化压缩,对于量化有两种,一种是以TensorFlow为代表的,quantization-aware training,其完整的路线是Tensorflow -> TFLite Convert -> TFLite,其在工业界应用很广泛,还有一种是类似TensonRT这种,深度学习推理引擎做量化。我认为这两种都应该支持,如MNN目前就是支持这两种,而TVM也是。第一种抛开不谈,这种支持没有什么好说的,就是解析TFLite量化模型,然后底层需要支持INT8的运算,以及需要更高效的支持INT8的卷积运算,包括类似ARM v8.2的dot指令产生等,大家都会做。而这里想谈谈第二点,这里面其实也是很自由发挥的点,第一是如何支持类似INT N 的量化,而不仅仅是INT 8,如FPGA,其实可以跑类似INT4这样的类型。第二点是如何更好地量化,得到更好地准确度。这里面其实各家做的不同,目前TVM社区就有很多用户发帖说TVM的这一点其实做的不好,经常说准确度不好。而目前业界一些框架的做法是会加入快速的重训练,不仅仅是单纯的Post Training Quantization,这一套会抽象出来一个工具,供算法业务团队使用。而这套工具往往还会包括模型压缩,甚至告诉你怎么压缩可以在这个硬件跑的更快,而不仅仅是降低模型GFLOPS这一简单的指标,我觉得这或许都是类似TVM等引擎学习的一点。

好了,其实想说的还有很多,但是时间比较晚了,希望下次有机会再来谈谈。但是,无论如何,基于编译优化的深度学习推理引擎我认为才是最佳解决方案以及未来,而基于这一思想的解决方案终将会出现一个类似LLVM的东西出来,目前现状依然是群雄逐鹿,但是我个人觉得TVM目前是看起来最有希望的。:-)延伸文章讨论:
谈谈对深度学习编译技术的一些思考
也谈TVM和深度学习编译器

猜你喜欢

转载自blog.csdn.net/qq_33287871/article/details/113805752