【导读】近年来深度学习芯片的研究发展是业界关注的焦点。专知之前报道过Google AI Jeff Dean独自撰文的关于AI时代芯片的历史发展趋势-【Google Jeff Dean独自署名论文】深度学习革命及其对计算机架构和芯片设计的影响,讲述AI芯片发展历程与未来,但学术业界一直缺乏对当前深度学习编译器框架的综述。最近北航的学者撰写了第一篇关于当下深度学习编译器的综述论文《The Deep Learning Compiler: A Comprehensive Survey》36页pdf1,从各个方面对现有的DL编译器进行了全面的比较。此外,还详细分析了多级红外设计和编译优化技术。最后,指出了DL编译器的研究方向,很好的论文。
摘要
在不同的DL硬件上部署各种深度学习(DL)模型的困难,促进了社区中DL编译器的研究和开发。工业界和学术界已经提出了几个DL编译器,如Tensorflow XLA和TVM。类似地,DL编译器将不同DL框架中描述的DL模型作为输入,然后为不同的DL硬件生成优化代码作为输出。然而,现有的综述论文没有一个全面地分析DL编译器的独特设计。在本文中,我们通过剖析通常采用的设计细节,对现有的DL编译器进行了全面的调研,重点是面向DL的多级IRs和前端/后端优化。具体来说,我们从各个方面对现有的DL编译器进行了全面的比较。此外,还详细分析了多级红外设计和编译优化技术。最后,指出了DL编译器的研究方向。这是第一篇关于DL编译器的独特设计的综述论文,我们希望能够为将来DL编译器的研究铺平道路。
概述
深度学习(DL)的发展对各个科学领域产生了深远的影响。它不仅在自然语言处理(NLP) (Manning et al., 1999)和计算机视觉(CV) (Forsyth and Ponce, 2002)等人工智能领域显示出显著的价值,而且在电子商务(Ha et al., 2016)、智慧城市(Mohammadi et al., 2017)和药物发现(Chen et al., 2018a)等更广泛的应用领域也取得了巨大的成功。随着卷积神经网络(convolutional neural network, CNN) (LeCun et al., 1998)、递归神经网络(neural network, RNN) (Rumelhart et al., 1986)、长短时记忆(long short-term memory, LSTM) (Hochreiter and Schmidhuber, 1997)、生成对抗网络(adversarial network, GAN) (Goodfellow et al., 2014)等多种深度学习模型的出现,简化DL模型的编程是实现其广泛应用的关键。
在工业界和学术界的不断努力下,为了简化各种DL模型的应用,提出了TensorFlow (Abadi et al., 2016)、PyTorch (Paszke et al., 2019)、MXNet (Chen et al., 2015)和CNTK (Seide and Agarwal, 2016)等流行的DL编程框架。尽管上述DL编程框架的优点和缺点取决于它们设计中的折衷方案,但是互操作性对于减少冗余工程工作变得非常重要,因为它支持跨现有DL模型的新出现的DL模型。为了提高互操作性,ONNX (onn, [n.d.])被提议定义一种表示DL模型的开放源码格式,这种格式可以促进不同DL框架之间的模型转换。
与此同时,矩阵乘法等独特的计算特性激发了芯片架构师为提高效率而设计定制DL芯片的热情。互联网巨头(如谷歌TPU (Jouppi等,2017年),Hisilicon NPU (Liao等,2019年),苹果Bonic (Kingsley-Hughes, 2017年)),处理器供应商(如NNP (tur, [n.d]),英特尔NNP (NNP, [n.d])),服务提供商(如亚马逊推论(inf, [n.d]),阿里巴巴Hanguang (han, [n.d]),甚至初创公司(如Cambricon (Liu等,2016年),Graphcore (Jia等,为了提高DL模型的计算能力,他们投入了大量的人力和资本来开发DL芯片。一般来说,DL芯片的种类包括:1)通用芯片与软硬件协同设计;2)完全定制DL机型专用芯片;3)受生物脑科学启发的神经形态芯片。例如,通用芯片(如CPU、GPU)增加了AVX512矢量单元和张量核等特殊硬件组件来加速DL模型。而对于专用芯片,如谷歌张量处理单元(TPU),应用特定的集成电路(如矩阵乘法引擎和高带宽内存)被设计来提高性能和能源效率到极致。在可预见的未来,DL芯片的设计将变得更加多样化。
为了在不同的DL芯片上加速DL模型,有效地将计算映射到DL芯片是很重要的。在通用芯片上,高度优化的线性代数库(如基本线性代数子程序库(BLAS))(如MKL和cuBLAS)是DL模型高效计算的基础。以卷积运算为例,DL框架将卷积转换为矩阵乘法,然后调用BLAS库中的GEMM函数。此外,芯片厂商还为DL计算(例如,MKL-DNN和cuDNN)发布了特别优化的库,包括前向和后向卷积、池化、标准化和激活。还开发了更高级的工具来进一步加速DL操作。以TensorRT (ten, [n.d])为例,它支持图形优化(如层融合)和低比特量化,并具有大量高度优化的GPU内核。在专用的DL芯片上,供应商还提供了类似的库和工具链,以有效地执行DL模型。然而,依赖上述的库和工具在不同的DL芯片上映射DL模型的缺点是,它们通常落后于DL模型的快速开发,因此不能有效地利用DL芯片。
为了解决DL库和工具的缺陷,以及减轻手工优化每个DL芯片上的DL模型的负担,DL社区求助于领域特定的编译器技术来解决问题。很快地,一些流行的DL编译器被提了出来,例如来自工业界和学术界的TVM (Chen等人,2018b)、Tensor Comprehension (Vasilache等人,2018)、Glow (Rotem等人,2018)、nGraph (Cyphers等人,2018)和XLA (Leary和Wang, 2017)。DL编译器将DL框架中描述的模型定义作为输入,并在各种DL芯片上生成高效的代码实现作为输出。模型定义和特定代码实现之间的转换针对模型规范和硬件体系结构进行了高度优化。具体地说,DL编译器结合了面向DL的优化,比如层和操作符融合,这支持高效的代码生成。此外,现有的DL编译器也利用了来自通用编译器(如LLVM (Lattner和Adve, 2004))的成熟工具链,这些工具链提供了跨不同硬件架构的更好的可移植性。与传统编译器相似,DL编译器也采用了包括前端、中间表示(IR)和后端在内的分层设计。然而,DL编译器的惟一性在于多层IRs和DL特定优化的设计。
在这篇论文中,我们通过将编译器的设计分为前端、多级IR和后端,并着重介绍了IR的设计和优化方法,对现有的DL编译器进行了全面的研究。据我们所知,这是第一篇全面介绍DL编译器设计的论文。具体而言,本文的贡献如下:
我们从硬件支持、DL框架支持、代码生成和优化等方面对现有的DL编译器进行了全面的比较,为最终用户选择合适的DL编译器提供了指导。
我们剖析了现有DL编译器的总体设计,并详细分析了多级IR设计和编译器优化技术,如数据流底层优化、硬件内在映射、内存延迟隐藏和并行化。
我们为DL编译器的未来发展提供了一些见解,包括自动调优、多面体编译器、量化、可微编程和隐私保护,希望能促进DL编译器社区的研究。
本文其余部分组织如下。第2节介绍DL编译器的背景,包括DL框架、DL芯片以及硬件(FPGA)特定的DL编译器。第3节详细比较了现有的DL编译器。第4节描述了DL编译器的总体设计,重点介绍了IR和前端/后端优化。第五部分是全文的总结,并对未来的研究方向进行了展望。