引言
有将近两个月没有学习一些新东西,更新一下博客了。一直在忙公司的一个项目,是做一款支持LSTM和RNN的通用架构加速IP。自己恰好负责指令编译工作,虽然开始的指令比较粗糙,没有一套完整的编译器架构。但是这其中也去思考了一下基于FPGA加速器的编译器架构。在FPGA深度学习加速器中,编译器除了需要自动化生成指令外,还要优化指令的结构,来最大化加速器性能。TVM是一个支持GPU、CPU、FPGA指令生成的开源编译器框架,虽然在我们自己的加速IP上无法直接拿过来用,但是其中的很多方法和思想还是很值得借鉴的。TVM最大的特点是基于图和算符结构来优化指令生成,最大化硬件执行效率。其中使用了很多方法来改善硬件执行速度,包括算符融合、数据规划、基于机器学习的优化器等。它向上对接Tensorflow、Pytorch等深度学习框架,向下兼容GPU、CPU、ARM、TPU等硬件设备。
1. 整体结构
TVM是一个端到端的指令生成器。它从深度学习框架中接收模型输入,然后进行图的转化和基本的优化,最后生成指令完成到硬件的部署。整个架构是基于图描述结构,不论是对指令的优化还是指令生成,一个图结构清晰的描述了数据流方向,操作之间的依赖关系等。基于机器学习的优化器是优化过程中的重点,指令空间很大,通过优化函数来寻找最优值是一个很合理的想法。它的主要特点如下:
1) 基于GPU、TPU等硬件结构,将张量运算作为一个基本的算符,通过把一个深度学习网络描述成图结构来抽象出数据计算流程。在这样的图结构基础上,更方便记性优化。同时能够有更好的向上向下兼容性,同时支持多种深度学习框架和硬件架构。
2) 巨大的优化搜索空间。在优化图结构方面,其不再局限于通过某一种方式,而是通过机器学习方法来搜索可能的空间来最大化部署效率。这种方式虽然会导致编译器较大的计算量,但是更加通用。
TVM提供了一个非常简单的端到端用户接口,通过调用TVM的API可以很方便的进行硬件部署。比如:
以上就是将Keras的模型输入到TVM,指定部署的硬件GPU,然后进行优化和代码生成。
TVM也提供了Java、C++和Python界面供用户调用。
2. 图结构基本优化
图结构是大多数深度学习框架中普遍采用的描述方式。这种图是一种高层次的描述,将一个张量运算用一个算符描述,而不是拆分的更细。这样更有利于优化,而且也符合GPU、TPU的硬件架构,在这些芯片中计算核算力很大,通常可以完成一个较大的计算,比如卷积、矩阵运算等。以下是一个卷积计算的例子,整个计算图包含了2D卷积、ReLu,dense、softmas等。这样的图结构也正好符合FPGA加速器的结构,在FPGA中也是用一个计算核来专门计算某个大的计算。TVM图中节点描述一个张量数据或者算符,而边表示了不同计算的依赖关系。
基于图结构,TVM采用了很多图优化策略。包括算符融合,将可以在硬件上用一个算符完成的多个连续运算合并;常量折叠,将可以预先计算的数据放在编译器中完成,减少硬件计算;存储规划,预先为中间数据分配存储空间来储存中间值,避免中间数据无法存储在片上而增加片外存储开销;数据规划,重新排列数据有利于硬件计算。
1) 算符融合
TVM中将运算划分为4种:1对1运算,比如加法、点乘;降运算,比如累加;复杂运算,比如2D卷积,融合了乘法和累加;不透明的,比如分类、数据排列等,这些不能被融合。算符融合可以减少存储开销,实现pipeline,特别是在FPGA中更有利。比如我们目前的项目中是开发一款通用RNN架构IP,其中涉及到矩阵乘法,加法。其中加法就可以融合到矩阵乘中,这样就减少了单独加法模块计算开销以及读写cache开销。
2) 数据规划
以我们XRNN来说,片上有一个矩阵运算阵列,由于阵列大小固定,一次计算矩阵大小也是固定的。比如计算一个32x32对应32x1的矩阵向量乘法,那么就要求权重和向量必须要按照32倍数进行对齐,这就需要对权重数据等进行规划。
3. 张量计算
TVM中使用的张量描述语言是透明的,可以根据硬件需要进行修改。这样更加灵活和有利于进行优化。但是这样可能增加了编译器优化的复杂性。TVM描述例子如下:
描述算符中包含结果大小,计算方式。但是这其中没有涉及到循环结构和更多数据操作细节。TVM采用了Halide思想,通过使用schedule来对张量计算进行等价变换,从中计算出执行效率最高的schedule结构。整个schedule流程如下图:
从中可以看出,TVM除了采用了Halide的schedule方式外,还增加了三种针对GPU和TPU的schedule方式:specile memory scope,tensorization,latency hiding。这些schedule方式可以对一个张量运算进行等价变换,产生多种代码结构,从中选择出最有利于硬件执行的代码结构。
3.1 并行优化
并行计算是提高硬件执行效率的重要一步,因为诸如卷积、矩阵计算等都是大量的可以并发进行的计算,如何优化并行结构对改善硬件性能很关键。在这里需要考虑两点问题:一个是并行度,另外一个是数据共享。如果数据不共享那么就会增加数据读写消耗。而尽量利用可共享数据,则需要尽心设计一个计算结构。
TVM提出了memory scope的概念,其将数据计算进行可并行和不可并行分类,对于可以并行计算的,就可以使用多线程来并行计算,而不可并行,则需要等待被依赖数据计算完成。比如一个矩阵乘法例子:
在XRNN中也会遇到类似的问题,对于不同的计算比如矩阵乘法和激活函数,如果没有依赖的话,就是可以并行执行的。
3.2 存储读写优化
在FPGA中读写cache或者external ddr也是一笔开销。如何将存储读写开销降低也有利于提高硬件执行效率。比如在我们的XRNN中,save数据到ddr会消耗很多时间,而这个数据下一次又会被利用,同时又增加了load的时间。如果将数据缓存到片上,那么就减少了load和save开销。还有比如片上cache也会来回读写数据,如果可以将一个计算完成的数据直接送给下一个计算核,实现流水,那么读写cache开销也节省了。
对于隐藏存储读写开销还有一种方法,在下一次计算开始前而这一次计算正在进行的时候,就进行片外数据加载,就将加载和计算重合了,减少额外加载消耗的时间。
4. 自动优化器
TVM在丰富的schedule方式基础上,提出了一个机器学习模型来寻找最优化的schedule结构。其包含两部分:一部分是基于schedule方式产生所有可能的计算结构;另外一个是机器学习代驾模型来预测性能。
Schedule空间是巨大的,它可能产生很多种计算流结构,而对其中进行探索,找到最合适的结构。这会产生大量的计算。不过在我们的XRNN结构中,因为受到硬件内核的限制,还是可以对空间进行缩减的。只采用最可能影响性能schedule方式。比如矩阵乘、加法、点乘、激活这些不同的计算可以进行并发和非并发安排,这样的空间相对小,有利于加速编译器生成工作。
机器学习代价模型主要考虑了不同操作的延时来预测性能:存储访问方式,数据重利用,pipeline等。TVM中对代价函数的求解不是通过随机统计方式,而是使用实时的配置数据进行训练,周期性更替代码结构。人为的探寻更合理的优化结构,然后提供给模型,让其不断更替。这种方式避免了探索大量schedule空间的时间消耗,同时也能更接近于实际情况。
总结
以上简要介绍了TVM的整体架构和基本方法,其实还挺符合FPGA加速硬件结构的。很多方法还是可以借用的。而且TVM是一个兼容更广的编译器架构,针对我们自身的FPGA特性,也会有很多不一样的设计。
审核编辑:汤梓红
全部0条评论
快来发表一下你的评论吧 !