(2021-3-1) The Deep Learning Compiler: A Comprehensive Survey (深度学习编译器:全面调查)

作者: Mingzhen Li; Yi Liu; Xiaoyan Liu; Qingxiao Sun; Xin You; Hailong Yang; Zhongzhi Luan; Lin Gan; Guangwen Yang; Depei Qian;
期刊: IEEE Transactions on Parallel and Distributed Systems (发表日期: 2021-3-1)
期刊分区: ㅤㅤ ㅤㅤIF 5.6 ㅤㅤ ㅤㅤ ㅤㅤ ㅤㅤ ㅤㅤ ㅤㅤSCI Q1 ㅤㅤ ㅤㅤ ㅤㅤ ㅤㅤ ㅤㅤ ㅤㅤ中科院 工程技术3区 ㅤㅤ ㅤㅤ
本地链接: Li 等 - 2021 - The Deep Learning Compiler A Comprehensive Survey.pdf
DOI: 10.1109/TPDS.2020.3030548
摘要: The difficulty of deploying various deep learning (DL) models on diverse DL hardware has boosted the research and development of DL compilers in the community. Several DL compilers have been proposed from both industry and academia such as Tensorflow XLA and TVM. Similarly, the DL compilers take the DL models described in different DL frameworks as input, and then generate optimized codes for diverse DL hardware as output. However, none of the existing survey has analyzed the unique design architecture of the DL compilers comprehensively. In this paper, we perform a comprehensive survey of existing DL compilers by dissecting the commonly adopted design in details, with emphasis on the DL oriented multi-level IRs, and frontend/backend optimizations. Specifically, we provide a comprehensive comparison among existing DL compilers from various aspects. In addition, we present detailed analysis on the design of multi-level IRs and illustrate the commonly adopted optimization techniques. Finally, several insights are highlighted as the potential research directions of DL compiler. This is the first survey paper focusing on the design architecture of DL compilers, which we hope can pave the road for future research towards DL compiler.

此文为机器辅助翻译,仅供个人学习使用,如有翻译不当之处欢迎指正

1 ABSTRACT

在各种深度学习(DL)硬件上部署不同的深度学习模型存在困难,这推动了学术界和工业界对深度学习编译器的研发。诸如TensorFlow XLA和TVM等多个深度学习编译器已被提出。同样,深度学习编译器将不同深度学习框架中描述的深度学习模型作为输入,然后为各种深度学习硬件生成优化代码作为输出。然而,现有调查均未全面分析深度学习编译器独特的设计架构。在本文中,我们通过详细剖析常用的设计,对现有的深度学习编译器进行了全面的综述,重点关注面向深度学习的多级中间表示(IR)以及前端/后端优化。我们对多级IR的设计进行了详细分析,并阐述了常用的优化技术。最后,我们强调了几个具有潜力的研究方向,作为深度学习编译器未来的研究方向。这是第一篇聚焦于深度学习编译器设计架构的综述论文,我们希望它能为深度学习编译器的未来研究铺平道路。

2 INTRODUCTION

深度学习(DL)的发展对各个科学领域产生了深远影响。它不仅在自然语言处理(NLP)[64]和计算机视觉(CV)[26]等人工智能领域展现出非凡价值,还在电子商务[36]、智慧城市[68]和药物研发[15]等更广泛的应用中取得了巨大成功。随着卷积神经网络(CNN)[54]、循环神经网络(RNN)[80]、长短期记忆网络(LSTM)[38]和生成对抗网络(GAN)[29]等多种功能强大的深度学习模型的出现,简化各种深度学习模型的编程对于实现其广泛应用至关重要。

在工业界和学术界的持续努力下,为了简化各种深度学习模型的实现过程,已经提出了多个广受欢迎的深度学习框架,如TensorFlow[1]、PyTorch[75]、MXNet[16]和CNTK[81]。尽管上述深度学习框架因其设计上的权衡各有优劣,但在支持现有深度学习模型与新兴深度学习模型交互时,互操作性对于减少冗余的工程工作变得十分重要。为了实现互操作性,开放神经网络交换格式(ONNX)[66]应运而生,它定义了一种统一的格式来表示深度学习模型,以便于在不同的深度学习框架之间进行模型转换。

与此同时,诸如矩阵乘法等独特的计算特性激发了芯片架构师设计定制化深度学习加速器以提高效率的热情。互联网巨头(如谷歌的张量处理单元(TPU)[44]、海思的神经网络处理器(NPU)[56]、苹果的仿生芯片(Bonic)[49])、处理器供应商(如英伟达的图灵(Turing)架构[72]、英特尔的神经网络处理器(NNP)[41])、服务提供商(如亚马逊的Inferentia[8]、阿里巴巴的含光芯片[7]),甚至初创企业(如寒武纪[57]、Graphcore[43])都在投入大量人力和资金开发深度学习芯片,以提升深度学习模型的性能。 一般来说,深度学习硬件可分为以下几类:1)采用软硬件协同设计的通用硬件;2)为深度学习模型完全定制的专用硬件;3)受生物脑科学启发的神经形态硬件。例如,通用硬件(如CPU、GPU)增加了特殊的硬件组件,如AVX512向量单元和张量核心,以加速深度学习模型的运行。而对于像谷歌TPU这样的专用硬件,其专用集成电路(如矩阵乘法引擎和高带宽内存)被设计用于将性能和能效提升到极致。在可预见的未来,深度学习硬件的设计将变得更加多样化。

为适应硬件的多样性,将计算高效地映射到深度学习硬件上至关重要。在通用硬件上,高度优化的线性代数库,如基本线性代数子程序(BLAS)库(例如英特尔数学核心函数库(MKL)和CUDA基本线性代数子程序库(cuBLAS)),是深度学习模型高效计算的基础。以卷积运算为例,深度学习框架将卷积转换为矩阵乘法,然后调用BLAS库中的通用矩阵乘法(GEMM)函数。此外,硬件供应商发布了专门为深度学习计算优化的库(例如英特尔深度学习库(MKL-DNN)和CUDA深度神经网络库(cuDNN)),涵盖前向和反向卷积、池化、归一化和激活等操作。还有更先进的工具被开发出来,以进一步加速深度学习操作。例如,TensorRT[73]支持图优化(如层融合)和低比特量化,拥有大量经过高度优化的GPU内核。在专用深度学习硬件上,也提供类似的库[43, 57]。然而,依赖这些库的缺点是,它们通常跟不上深度学习模型的快速发展,因此无法高效利用深度学习芯片。

为了解决深度学习库和工具的缺陷,同时减轻在每种深度学习硬件上手动优化深度学习模型的负担,深度学习领域开始借助特定领域的编译器来解决问题。很快,学术界和工业界便提出了几种广受欢迎的深度学习编译器,如TVM[17]、Tensor Comprehension[91]、Glow[79]、nGraph[21]和XLA[53]。深度学习编译器将深度学习框架中描述的模型定义作为输入,为各种深度学习硬件生成高效的代码实现作为输出。模型定义与具体代码实现之间的转换针对模型规格和硬件架构进行了高度优化。具体而言,它们融入了面向深度学习的优化方法,如层融合和操作符融合,这使得高效的代码生成成为可能。此外,现有的深度学习编译器还利用了通用编译器的成熟工具链(例如LLVM[51]),这为跨多种硬件架构提供了更好的可移植性。与传统编译器类似,深度学习编译器也采用分层设计,包括前端、中间表示(IR)和后端。然而,深度学习编译器的独特之处在于其多级IR的设计和深度学习特定的优化。

在本文中,我们通过将编译器设计剖析为前端、多级中间表示(IR)和后端,对现有的深度学习编译器进行了全面的综述,特别强调了IR设计和优化方法。据我们所知,这是第一篇对深度学习编译器设计进行全面综述的论文。具体而言,本文做出了以下贡献:

  • 我们剖析了现有深度学习编译器常用的设计架构,并对关键设计组件进行了详细分析,这些组件包括多级中间表示(IR)、前端优化(涵盖节点级、块级和数据流级优化)以及后端优化(包括硬件特定优化、自动调优和优化的内核库)。
  • 我们从多个方面对现有深度学习编译器进行了全面的分类,这与本综述中描述的关键组件相对应。这种分类的目的是为从业者在根据自身需求选择深度学习编译器时提供指导,同时为研究人员提供深度学习编译器的全面总结。
  • 我们对深度学习编译器在卷积神经网络(CNN)模型上的性能进行了量化比较,这些模型包括完整模型和轻量级模型。我们比较了端到端性能以及每层(由于卷积层主导推理时间,因此重点比较卷积层)的性能,以展示优化的有效性。评估脚本和结果已开源,可供参考 。
  • 我们强调了深度学习编译器未来发展的几个方向,包括动态形状和预处理/后处理、先进的自动调优、多面体模型、子图分区、量化、统一优化、可微编程和隐私保护,我们希望这些能推动深度学习编译器领域的研究。

本文的其余部分结构如下。第2节介绍深度学习编译器的背景,包括深度学习框架、深度学习硬件,以及硬件(FPGA)特定的深度学习代码生成器。第3节描述深度学习编译器的常见设计架构。第4节讨论深度学习编译器的关键组件,包括多级中间表示(IR)、前端优化和后端优化。第5节给出全面的分类体系。第6节进行定量的性能比较。第7节强调深度学习编译器研究的未来方向。

3 背景

3.1 深度学习框架

在本节中,我们将概述流行的深度学习框架。这里的讨论可能并不详尽,但旨在为深度学习从业者提供指导。图 1 展示了深度学习框架的全景,包括当前流行的框架、历史框架以及 ONNX 支持的框架。

TensorFlow:在所有深度学习框架中,TensorFlow 对语言接口的支持最为全面,包括 C++、Python、Java、Go、R 和 Haskell。TensorFlow 采用带有受限控制边的基本算子数据流图来表示可微程序 [78]。TensorFlow Lite 专为移动和嵌入式深度学习而设计,并提供了 Android 神经网络 API。为降低使用 TensorFlow 的复杂性,谷歌采用 Keras 作为 TensorFlow 核心的前端。此外,TensorFlow 中的即时执行模式(eager - mode)采用了类似于 PyTorch 的方法,以更好地支持动态计算图。

Keras:Keras [19] 是一个用于快速构建深度学习模型的高级神经网络库,用纯 Python 编写。虽然它本身不是一个深度学习框架,但 Keras 提供了一个高级 API,可以与 TensorFlow、MXNet、Theano 和 CNTK 集成。使用 Keras,深度学习开发者只需编写几行代码就能构建一个神经网络。此外,Keras 还可以与其他常见的深度学习包(如 scikit - learn)集成。然而,由于过度封装,Keras 的灵活性不足,这使得添加算子或获取底层数据信息变得过于困难。

PyTorch:Facebook 用 Python 重写了基于 Lua 的深度学习框架 Torch,并在张量级别重构了所有模块,从而推出了 PyTorch。作为最受欢迎的动态框架,PyTorch 在 Python 中嵌入了用于构建动态数据流图的原语,其控制流在 Python 解释器中执行。PyTorch 1.0 集成了 PyTorch 0.4 和 Caffe2 的代码库,创建了一个统一的框架。这使得 PyTorch 能够吸收 Caffe2 的优势,以支持高效的图执行和移动部署。FastAI [39] 是基于 PyTorch 上层封装的高级 API 层,它充分借鉴了 Keras,以简化 PyTorch 的使用。

Caffe/Caffe2:Caffe [42] 是由加州大学伯克利分校为深度学习和图像分类而设计的。Caffe 拥有命令行、Python 和 MATLAB API。Caffe 的简洁性使其源代码易于扩展,适合开发者进行深入分析。因此,Caffe 主要定位于研究领域,从一开始就广受欢迎。Caffe2 建立在原始 Caffe 项目的基础上,其代码结构与 TensorFlow 类似,但 API 更轻量级,并且更易于访问计算图中的中间结果。

MXNet:MXNet 支持多种语言 API,包括 Python、C++、R、Scala、Julia、Matlab 和 JavaScript。它旨在实现可扩展性,并从降低数据加载和 I/O 复杂性的角度进行设计 [16]。MXNet 提供了不同的编程范式:类似于 Caffe 和 TensorFlow 的声明式编程,以及类似于 PyTorch 的命令式编程。2017 年 12 月,亚马逊和微软基于 MXNet 联合发布了 Gluon [69],这是一个类似于 Keras 和 FastAI 的高级接口。Gluon 既支持灵活的动态图,也支持高效的静态图。

CNTK:CNTK 可以通过 Python、C++ 和 C# API,或者其自己的脚本语言(即 BrainScript)来使用。CNTK 旨在易于使用,并为生产环境中的大规模数据做好了准备 [37]。然而,CNTK 目前尚不支持 ARM 架构,这限制了它在移动设备上的使用。它使用类似于 TensorFlow 和 Caffe 的静态计算图,在这种图中,深度学习模型被视为通过有向图的一系列计算步骤。

PaddlePaddle:PaddlePaddle [11] 最初的设计与 Caffe 类似,每个模型都可以表示为一组层。然而,PaddlePaddle v2 参考 TensorFlow 采用了算子的概念,将层分解为更细粒度的算子,从而支持更复杂的深度学习模型。并且 PaddlePaddle Fluid 与 PyTorch 类似,因为它提供了自己的解释器,以避免 Python 解释器的性能限制。

ONNX:开放神经网络交换格式(ONNX)[66] 定义了一个可扩展的计算图模型,因此不同深度学习框架构建的计算图可以轻松转换为 ONNX 格式。有了 ONNX,在不同深度学习框架之间转换模型变得更加容易。例如,它允许开发者构建一个 MXNet 模型,然后使用 PyTorch 进行推理。如图 1 所示,ONNX 已经集成到 PyTorch、MXNet、PaddlePaddle 等框架中。对于一些尚未直接支持的深度学习框架(如 TensorFlow 和 Keras),ONNX 添加了转换器。

![[Pasted image 20250513103958.png]]

历史框架:由于深度学习领域的快速发展,许多历史深度学习框架不再活跃。例如,PyTorch 已经取代了 Torch [20]。作为最古老的深度学习框架之一,Theano [86] 不再维护。Deeplearning4J [85] 是一个基于 Java 和 Scala 的分布式深度学习框架,但由于缺乏庞大的开发者社区而变得不活跃。Chainer [87] 曾是动态计算图的首选框架,但现在已被具有类似功能的 MXNet、PyTorch 和 TensorFlow 所取代。

以前的工作 [10, 25, 35, 70, 82, 100] 已经在不同应用(如计算机视觉和图像分类)和不同硬件(如 CPU、GPU 和 TPU)上比较了深度学习框架的性能。有关每个深度学习框架的详细信息,读者可以参考 [37]。与这些工作不同,本综述侧重于深度学习编译器的研究,深度学习编译器提供了一种更通用的方法,可以在各种硬件上高效执行各种深度学习模型。

3.2 深度学习硬件

深度学习硬件根据通用性可分为三类:1)通过硬件和软件优化能够支持深度学习工作负载的通用硬件;2)通过完全定制的电路设计专注于加速深度学习工作负载的专用硬件;3)模仿人类大脑功能的神经形态硬件。

通用硬件:深度学习模型最具代表性的通用硬件是图形处理单元(GPU),它通过多核架构实现了高并行性。例如,英伟达的 GPU 自 Volta 架构起就引入了张量核心。张量核心可以并行加速混合精度矩阵乘累加计算,这在深度学习模型的训练和推理过程中都有广泛应用。与硬件协同优化的是,英伟达还推出了高度优化的深度学习库和工具,如 cuDNN [18] 和 TensorRT [73],以进一步加速深度学习模型的计算。

专用硬件:专用硬件是为深度学习计算完全定制的,旨在将性能和能源效率提升到极致。深度学习应用和算法的迅速扩展促使许多初创企业开发专用深度学习硬件(例如 Graphcore GC2、寒武纪 MLU270)。此外,传统硬件公司(如英特尔的 NNP、高通的 Cloud AI 100)和云服务提供商(如谷歌的 TPU、亚马逊的 Inferentia 和阿里巴巴的含光芯片)也在这一领域进行了投资。最著名的专用深度学习硬件是谷歌的 TPU 系列。一个 TPU 包含矩阵乘法单元(MXU)、统一缓冲区(UB)和激活单元(AU),由主机处理器通过 CISC 指令驱动。MXU 主要由脉动阵列组成,针对矩阵乘法的功耗和面积效率进行了优化。与 CPU 和 GPU 相比,TPU 仍然是可编程的,但它以矩阵作为基本单元,而不是向量或标量。最近,亚马逊的 Inferentia 也引起了人们的关注。这款芯片有四个专为张量级操作设计的神经核心,并且具有较大的片上缓存,可避免频繁访问主内存。

神经形态硬件:神经形态芯片使用电子技术模拟生物大脑。这类芯片的代表性产品有 IBM 的 TrueNorth 和英特尔的 Loihi。神经形态芯片(例如 TrueNorth)的人工神经元之间具有很高的连接性。神经形态芯片还复制了类似于脑组织的结构:神经元可以同时存储和处理数据。传统芯片将处理器和内存分布在不同位置,而神经形态芯片通常有许多微处理器,每个微处理器都有少量的本地内存。与 TrueNorth 相比,Loihi 具有更接近大脑的学习能力。Loihi 引入了脉冲时间依赖的突触可塑性模型(STDP),这是一种根据突触前和突触后脉冲的相对时间来调节突触强度的机制。然而,神经形态芯片距离大规模商业生产还很远。尽管如此,在计算机科学领域,神经形态芯片可以帮助捕捉常规深度学习模型忽略的快速终身学习过程;在神经学领域,它们有助于弄清楚大脑的各个部分是如何协同工作以产生思想、情感甚至意识的。

3.3 特定硬件的深度学习代码生成器

现场可编程门阵列(FPGA)是可重新编程的集成电路,包含一组可编程逻辑块。程序员可以在制造后对其进行配置。除了可重新编程的特性外,FPGA 的低功耗和高性能特性使其在通信、医疗、图像处理和 ASIC 原型设计等众多领域得到广泛应用。在深度学习领域,高性能的 CPU 和 GPU 具有高度的可编程性,但功耗较大;而高效节能的 ASIC 则专为固定应用而设计。然而,FPGA 可以弥合 CPU / GPU 和 ASIC 之间的差距,这使得 FPGA 成为深度学习领域一个具有吸引力的平台。

高级综合(HLS)编程模型使 FPGA 程序员能够使用 C 和 C++ 等高级语言方便地生成有效的硬件设计。它避免了编写大量的 Verilog 或 VHDL 描述,降低了编程门槛,缩短了设计周期。赛灵思的 Vivado HLS 和英特尔的 FPGA SDK for OpenCL 是两个针对各自 FPGA 的流行 HLS 工具。然而,即使使用 HLS,将深度学习模型映射到 FPGA 仍然是一项复杂的工作,原因如下:1)深度学习模型通常是用深度学习框架的语言描述的,而不是纯粹的 C / C++ 代码;2)深度学习特定的信息和优化很难被利用。

针对 FPGA 的特定硬件代码生成器将深度学习模型或其领域特定语言(DSL)作为输入,进行特定领域(关于 FPGA 和深度学习)的优化和映射,然后生成 HLS 或 Verilog / VHDL 代码,最终生成比特流。根据基于 FPGA 的加速器生成的架构,它们可以分为两类:处理器架构和流架构 [93]。

处理器架构与通用处理器有相似之处。这种架构的 FPGA 加速器通常由几个处理单元(PU)组成,而处理单元又由片上缓冲区和多个较小的处理引擎(PE)组成。它通常有一个虚拟指令集(ISA),硬件的控制和执行的调度由软件决定。此外,静态调度方法避免了冯・诺依曼执行的开销(包括指令获取和解码)。硬件模板是一种通用的基本实现,具有可配置的参数。针对这种架构的深度学习代码生成器采用硬件模板自动生成加速器设计。通过模板的可配置参数,代码生成器实现了可扩展性和灵活性 [104]。可扩展性意味着代码生成器可以为从高性能到高效节能的 FPGA 生成设计;灵活性意味着代码生成器可以为具有不同层类型和参数的各种深度学习模型生成设计。PU 的数量和每个 PU 中 PE 的数量是重要的模板参数。此外,分块大小和批处理大小也是将深度学习模型映射到 PU 和 PE 的重要调度参数。所有这些参数通常通过使用各种策略进行设计空间探索来确定,例如结合性能模型和自动调优。DNN Weaver [83]、AngelEye [33]、ALAMO [63]、FP - DNN [32]、SysArrayAccel [101] 是针对处理器架构的典型 FPGA 深度学习代码生成器。此外,PU 和 PE 通常负责粗粒度的基本操作,如矩阵向量乘法、矩阵矩阵乘法、池化和一些逐元素操作。这些基本操作的优化主要由并行性和数据重用之间的权衡来指导,这与一般的优化类似。

流架构与流水线有相似之处。这种架构的 FPGA 加速器由多个不同的硬件块组成,对于输入的深度学习模型的每一层,它几乎都有一个硬件块。对于深度学习模型的输入数据,这种加速器按照层的顺序通过不同的硬件块处理数据。此外,通过流式输入数据,所有硬件块可以以流水线方式充分利用。然而,流架构通常基于一个初始假设,即目标 FPGA 上的片上内存和计算资源足以容纳深度学习模型,这给部署具有复杂层的深度模型带来了障碍。针对这种架构的深度学习代码生成器可以通过利用 FPGA 的可重构性或采用动态控制流来解决这个问题。单个块的进一步优化与处理器架构的基本操作类似。fpgaConvNet [92]、DeepBurning [98]、Haddoc2 [2] 和 AutoCodeGen [59] 是典型的相应深度学习代码生成器。

有关将深度学习模型映射到 FPGA 的特定编译技术的详细综述,读者可以参考 [34, 93, 104]。与 [34, 93, 104] 不同,本综述侧重于可以应用于更广泛深度学习硬件的通用深度学习编译技术,而不限于 FPGA。

![[Pasted image 20250513104327.png]]

4 深度学习编译器的常见设计架构

深度学习编译器的常见设计架构主要包含两部分:编译器前端和编译器后端,如图2所示。中间表示(IR)贯穿于前端和后端。一般来说,IR是程序的一种抽象,用于程序优化。具体而言,在深度学习编译器中,DL模型会被转换为多级IR,其中高级IR位于前端,低级IR位于后端。编译器前端基于高级IR进行与硬件无关的转换和优化。编译器后端基于低级IR进行特定硬件的优化、代码生成和编译。需要注意的是,本综述侧重于深度学习编译器的设计原理。关于深度学习编译器的功能和实验比较,读者可参考[55, 102]。

高级中间表示(IR),也被称为图IR,用于表示计算过程和控制流,并且与硬件无关。高级IR的设计难点在于对计算和控制流的抽象能力,这种能力能够捕捉并表达各种深度学习(DL)模型。高级IR的目标是建立算子和数据之间的控制流以及依赖关系,同时为图级别的优化提供接口。它还包含丰富的编译语义信息,并为自定义算子提供可扩展性。关于高级IR的详细讨论将在4.1节中给出。

低级中间表示(IR) 旨在针对不同的硬件目标进行特定硬件的优化和代码生成。因此,低级IR的粒度要足够精细,以反映硬件特性并表示特定硬件的优化。它还应允许在编译器后端使用成熟的第三方工具链,例如Halide [77]、多面体模型 [31] 和LLVM [51]。关于低级IR的详细讨论将在4.2节中给出。

前端将现有深度学习框架中的深度学习模型作为输入,然后将模型转换为计算图表示(例如图IR)。前端需要实现各种格式转换,以支持不同框架中的多样化格式。计算图优化结合了通用编译器的优化技术和深度学习特定的优化方法,这些优化可以减少冗余并提高图IR的效率。这类优化可分为节点级(例如空操作消除和零维张量消除)、块级(例如代数简化、算子融合和算子下沉)和数据流级(例如公共子表达式消除、死代码消除、静态内存规划和布局转换)。经过前端处理后,生成优化后的计算图并传递给后端。关于前端的详细讨论将在4.3节中给出。

后端将高级中间表示(IR)转换为低级中间表示,并执行特定于硬件的优化。一方面,它可以直接将高级中间表示转换为诸如LLVM中间表示(LLVM IR)之类的第三方工具链,以利用现有的基础设施进行通用优化和代码生成。另一方面,它可以利用深度学习(DL)模型和硬件特性的先验知识,通过自定义的编译过程来实现更高效的代码生成。常见的特定于硬件的优化包括硬件固有映射、内存分配与获取、内存延迟隐藏、并行化以及面向循环的优化。为了在大型优化空间中确定最优的参数设置,现有深度学习编译器广泛采用了两种方法,例如自动调度(如多面体模型)和自动调优(如自动TVM)。经过优化的低级中间表示使用即时编译(JIT)或提前编译(AOT)技术,为不同的硬件目标生成代码。关于后端的详细讨论将在第4.4节中给出。

5 深度学习编译器的关键组件

5.1 高级中间表示(IR)

为了克服传统编译器中所采用的中间表示(IR)的局限性,这种局限性会限制深度学习(DL)模型中所使用的复杂计算的表达,现有的深度学习编译器利用高级中间表示(也称为图 IR)并采用特殊设计,以实现高效的代码优化。为了更好地理解深度学习编译器中使用的图 IR,我们将如下描述图 IR 的表示和实现。

5.1.1 Graph IR的表示形式。

图中间表示的表示形式会影响图中间表示的表达能力,同时也决定了深度学习编译器分析图中间表示的方式。

基于有向无环图的中间表示——基于有向无环图的中间表示是编译器构建计算图的最传统方式之一,其节点和边被组织成有向无环图(DAG)。在深度学习编译器[17, 21, 53, 79, 91]中,有向无环图的节点表示深度学习的原子算子(如卷积、池化等),边表示张量。并且该图是无环的,没有循环,这与通用编译器[51, 52]的数据依赖图(DDG)不同。借助有向无环图计算图,深度学习编译器可以分析各种算子之间的关系和依赖关系,并利用它们来指导优化。对于数据依赖图已经有许多优化方法,比如公共子表达式消除(CSE)和死代码消除(DCE)。通过将深度学习的领域知识与这些算法相结合,可以对有向无环图计算图应用进一步的优化,这将在4.3节中详细阐述。 基于有向无环图的中间表示因其简单性而便于编程和编译,但它存在一些缺陷,例如由于缺少计算作用域的定义而导致语义模糊。

基于let绑定的中间表示——let绑定是一种通过为某些函数提供具有受限作用域的let表达式来解决语义模糊问题的方法,许多高级编程语言(如JavaScript[30]、F#[76]和Scheme[3])都使用这种方法。当使用let关键字定义一个表达式时,会生成一个let节点,然后它指向表达式中的算子和变量,而不是像基于有向无环图(DAG)的方式那样仅仅在变量之间建立计算关系。在基于DAG的编译器中,当一个进程需要获取一个表达式的返回值时,它首先访问相应的节点并搜索相关节点,这也被称为递归下降技术。相比之下,基于let绑定的编译器会计算出let表达式中变量的所有结果,并构建一个变量映射。当需要特定结果时,编译器会查找这个映射来确定表达式的结果。在深度学习编译器中,TVM的Relay中间表示(IR)同时采用基于DAG的中间表示和基于let绑定的中间表示,以兼具两者的优势。

表示张量计算——不同的图中间表示(IR)有不同的方式来表示对张量的计算。各种深度学习框架的算子会根据这些特定的表示方式被转换为图中间表示。并且定制的算子也需要按照这种表示方式进行编程。张量计算的表示可以分为以下三类。

  1. 基于函数的:基于函数的表示方式仅仅提供封装好的算子,Glow、nGraph 和 XLA 采用了这种方式。以高级优化器(HLO,即 XLA 的中间表示)为例,它由符号编程中的一组函数组成,并且其中大多数函数没有副作用。指令被组织为三个层级,包括 HloModule(整个程序)、HloComputaion(一个函数)和 HloInstruction(一项操作)。XLA 使用 HLO 中间表示来同时表示图中间表示和操作中间表示,这样 HLO 的操作范围涵盖从数据流级别到算子级别。
  2. lambda 表达式:lambda 表达式是一种索引公式表达式,通过变量绑定和替换来描述计算过程。使用 lambda 表达式,程序员无需实现新函数就能快速定义一项计算。TVM 使用基于 lambda 表达式的张量表达式来表示张量计算。在 TVM 中,张量表达式中的计算算子是由输出张量的形状以及计算规则的 lambda 表达式来定义的。
  3. 爱因斯坦记号(Einstein notation):爱因斯坦记号,也被称为求和约定,是一种用于表示求和的记号。它在编程方面的简洁性优于 lambda 表达式。以张量计算(TC)为例,临时变量的索引无需定义。基于爱因斯坦记号,中间表示(IR)能够根据未定义变量的出现情况推算出实际的表达式。在爱因斯坦记号中,算子需要满足结合律和交换律。这一限制确保了归约算子可以按照任意顺序执行,为进一步的并行化提供了可能。
5.1.2 Graph IR的实现。

深度学习编译器中图中间表示的实现完成了对数据和操作的管理。

数据表示——深度学习编译器中的数据(例如输入、权重和中间数据)通常以张量的形式组织,张量也被称为多维数组。深度学习编译器可以通过内存指针直接表示张量数据,或者以占位符的更灵活方式来表示。占位符包含张量每个维度的大小。另外,张量的维度大小可以标记为未知。为了进行优化,深度学习编译器需要数据布局信息。此外,迭代器的边界应该根据占位符来推断。

  1. 占位符:占位符在符号编程中广泛使用(例如,Lisp[65]、TensorFlow[1])。占位符简单来说是一个带有明确形状信息(例如,每个维度的大小)的变量,并且它会在计算的后期阶段填充值。它允许程序员在描述操作和构建计算图时无需关注确切的数据元素,这有助于在深度学习编译器中把计算定义和确切的执行分离开来。此外,程序员使用占位符可以方便地改变输入/输出以及其他相应中间数据的形状,而无需更改计算定义。
  2. 未知(动态)形状表示:在声明占位符时,通常支持未知的维度大小。例如,TVM 使用“Any”来表示未知维度(例如,张量⟨(Any, 3), fp32⟩);XLA 使用“None”来达到相同的目的(例如,tf.placeholder(“float”, [None, 3]));nGraph 使用其“PartialShape”类。未知形状表示对于支持动态模型是必要的。然而,为了全面支持动态模型,边界推断和维度检查应该放宽。此外,还应实现额外的机制来确保内存的有效性。
  3. 数据布局:数据布局描述了张量在内存中的组织方式,通常是从逻辑索引到内存索引的一种映射。数据布局通常包括维度的顺序(例如,NCHW和NHWC)、分块、填充、步幅等。TVM和Glow将数据布局表示为算子参数,并且在计算和优化时需要此类信息。然而,将数据布局信息与算子相结合,而非与张量相结合,能够为特定算子实现直观的操作,并减少编译开销。XLA将数据布局表示为与其后端硬件相关的约束条件。Relay和MLIR打算将数据布局信息添加到它们的张量类型系统中。
  4. 边界推断:在深度学习编译器中编译深度学习模型时,边界推断用于确定迭代器的边界。尽管深度学习编译器中的张量表示便于描述输入和输出,但它给推断迭代器边界带来了特殊挑战。边界推断通常根据计算图和已知的占位符以递归或迭代的方式进行。例如,在TVM中,迭代器形成一个有向无环超图,图中的每个节点表示一个迭代器,每条超边表示两个或多个迭代器之间的关系(例如,拆分、融合或重基)。一旦根据占位符的形状确定了根迭代器的边界,就可以根据这些关系递归地推断出其他迭代器。

支持的算子——深度学习编译器所支持的算子负责表示深度学习的工作负载,它们是计算图的节点。这些算子通常包括代数算子(例如加法、乘法、指数运算和topK运算)、神经网络算子(例如卷积和池化)、张量算子(例如重塑、调整大小和复制)、广播和归约算子(例如最小值和求最小值索引),以及控制流算子(例如条件判断和循环)。在此,我们选取了三个在不同深度学习编译器中频繁使用的代表性算子进行说明。此外,我们还将讨论定制算子的情况。

  1. 广播:广播算子可以复制数据,并生成形状兼容的新数据。如果没有广播算子,输入张量的形状限制会更多。例如,对于加法算子,期望输入张量具有相同的形状。像XLA和Relay等一些编译器通过提供广播算子来放宽这种限制。例如,XLA允许通过复制向量,使其形状与矩阵匹配,从而实现矩阵和向量之间的按元素相加。
  2. 控制流:在表示复杂且灵活的模型时,控制流是必需的。诸如循环神经网络(RNN)和强化学习(RL)等模型依赖于递归关系和数据相关的条件执行 [103],这就需要控制流。如果深度学习编译器的图中间表示不支持控制流,这些模型就必须依赖于宿主语言(比如 Python 中的 if 和 while 语句)的控制流支持,或者进行静态展开,这会降低计算效率。Relay 注意到,任意的控制流都可以通过递归和模式来实现,这在函数式编程中已得到验证 [78]。因此,它提供了 if 算子和递归函数来实现控制流。相反,XLA 通过特殊的高级优化器(HLO)算子(如 while 和条件判断算子)来表示控制流。
  3. 求导:算子 Op 的求导算子将 Op 的输出梯度和输入数据作为其输入,然后计算 Op 的梯度。尽管一些深度学习编译器(如 TVM 和张量计算(TC))支持自动求导 [88],但当应用链式法则时,它们需要高级中间表示中所有算子的导数。TVM 正致力于提供代数算子和神经网络算子的求导算子。程序员可以使用这些求导算子来构建定制算子的导数。相反,PlaidML 可以自动生成求导算子,即使是对于定制算子也是如此。值得注意的是,不支持求导算子的深度学习编译器无法提供模型训练的能力。
  4. 定制算子:它允许程序员为特定目的定义自己的算子。对定制算子提供支持可提高深度学习编译器的可扩展性。例如,在 Glow 中定义新算子时,程序员需要实现其逻辑和节点封装。此外,如果有必要,还需要额外的工作,比如进行算子层级的转换、生成操作中间表示以及生成指令。而 TVM 和 TC 除了描述计算实现外,所需的编程工作较少。具体来说,TVM 的用户只需描述计算过程和调度,并声明输入 / 输出张量的形状。此外,定制算子通过钩子函数集成 Python 函数,这进一步减轻了程序员的负担。
5.1.3 讨论

几乎所有的深度学习编译器都有其独特的高级中间表示。然而,它们有着相似的设计理念,比如使用有向无环图(DAG)和基于let绑定的方式来构建计算图。此外,它们通常为程序员提供了便捷的方式来表示张量计算。在高级中间表示中设计的数据和算子足够灵活且具有可扩展性,能够支持各种不同的深度学习模型。更重要的是,高级中间表示与硬件无关,因此可以应用于不同的硬件后端。

5.2 Low-level IR

5.2.1 低级中间表示的实现。

与高级中间表示相比,低级中间表示以更细粒度的方式描述深度学习模型的计算过程,它通过提供调整计算和内存访问的接口,实现了与目标相关的优化。在本节中,我们将低级中间表示的常见实现分为三类:基于Halide的中间表示、基于多面体的中间表示以及其他独特的中间表示。

基于 Halide 的中间表示 ——Halide 最初是为了实现图像处理的并行化而提出的,并且已被证明在深度学习编译器(如 TVM)中具有扩展性和高效性。Halide 的基本理念是将计算和调度分离。采用 Halide 的编译器不是直接给出特定的方案,而是尝试各种可能的调度并选择最佳的方案。在 Halide 中,内存引用和循环嵌套的边界被限制在与轴对齐的有界框内。因此,Halide 无法表达具有复杂模式(如非矩形)的计算。幸运的是,深度学习中的计算相当规则,可以用 Halide 完美地表达。此外,Halide 可以轻松地对这些边界进行参数化,并将其暴露给调优机制。当应用于深度学习编译器的后端时,Halide 的原始中间表示需要进行修改。例如,Halide 的输入形状是无限的,而深度学习编译器需要知道数据的确切形状,以便将算子映射到硬件指令。一些编译器,如张量计算(TC),要求数据具有固定的大小,以确保张量数据具有更好的时间局部性。

TVM 通过后续的努力将 Halide 中间表示改进为独立的符号中间表示。它消除了对 LLVM 的依赖,并对项目模块的结构和 Halide 的中间表示设计进行了重构,以追求更好的组织方式,同时提高图中间表示和 Python 等前端语言的可访问性。可重用性也得到了提高,通过实现运行时调度机制可以方便地添加定制算子。TVM 将变量定义从字符串匹配简化为指针匹配,确保每个变量只有一个定义位置(静态单赋值,SSA)[22]。

基于多面体的中间表示 —— 多面体模型是深度学习编译器中采用的一项重要技术。它使用线性规划、仿射变换和其他数学方法来优化具有静态控制流的边界和分支的基于循环的代码。与 Halide 不同,在多面体模型中,内存引用和循环嵌套的边界可以是任意形状的多面体。这种灵活性使得多面体模型在通用编译器中得到了广泛应用。然而,这种灵活性也阻碍了与调优机制的集成。尽管如此,由于能够处理深度嵌套的循环,许多深度学习编译器,如张量计算(TC)和 PlaidML(作为 nGraph 的后端),都采用多面体模型作为它们的低级中间表示。基于多面体的中间表示使得应用各种多面体变换(如融合、分块、下沉和映射)变得容易,这些变换包括与设备相关和与设备无关的优化。基于多面体的编译器借鉴了许多工具链,如 isl [96]、Omega [48]、PIP [23]、Polylib [60] 和 PPL [9] 。

TC 在低级中间表示方面有其独特的设计,它结合了 Halide 和多面体模型。它使用基于 Halide 的中间表示来表示计算,并采用基于多面体的中间表示来表示循环结构。TC 通过抽象实例给出详细的表达式,并引入特定的节点类型。简而言之,TC 使用域节点来指定索引变量的范围,使用上下文节点来描述与硬件相关的新迭代变量。它使用带节点来确定迭代的顺序。过滤器节点表示与语句实例结合的迭代器。“集合(Set)” 和 “序列(Sequence)” 是用于指定过滤器执行类型(并行和串行执行)的关键字。此外,TC 使用扩展节点来描述代码生成所需的其他指令,例如内存移动指令。

PlaidML 使用基于多面体的中间表示(称为 Stripe)来表示张量操作。它通过将并行多面体块的嵌套扩展到多个级别来创建可并行化代码的层次结构。此外,它允许将嵌套的多面体分配到嵌套的内存单元中,为使计算与内存层次结构相匹配提供了一种方法。在 Stripe 中,硬件配置与内核代码无关。Stripe 中的标签(在其他编译器中称为 “遍(passes)”)不会改变内核结构,但会为优化遍提供关于硬件目标的额外信息。Stripe 将深度学习算子分割成适合本地硬件资源的小块。

其他独特的中间表示 —— 有些深度学习编译器实现了定制的低级中间表示,不使用 Halide 和多面体模型。基于这些定制的低级中间表示,它们应用特定于硬件的优化,并将其降低为 LLVM 中间表示。

Glow 中的低级中间表示是一种基于指令的表达式,对通过地址引用的张量进行操作 [79]。Glow 低级中间表示中有两种基于指令的函数:声明和程序。第一种函数声明在程序生命周期内存在的常量内存区域的数量(例如,输入、权重、偏置)。第二种函数是局部分配区域的列表,包括函数(例如,卷积和池化)和临时变量。指令可以在全局内存区域或局部分配区域上运行。此外,每个操作数都用限定符进行注释:“@in” 表示操作数从缓冲区读取;“@out” 表示操作数写入缓冲区;“@inout” 表示操作数对缓冲区进行读写。这些指令和操作数限定符有助于 Glow 确定何时可以执行特定的内存优化。

MLIR 深受 LLVM 的影响,并且它是比 LLVM 更纯粹的编译器基础设施。MLIR 重用了 LLVM 中的许多思想和接口,并且位于模型表示和代码生成之间。MLIR 具有灵活的类型系统,并允许多个抽象级别,它引入了方言来表示这些多个抽象级别。每个方言由一组定义的不可变操作组成。MLIR 当前的方言包括 TensorFlow 中间表示、XLA 高级优化器(HLO)中间表示、实验性多面体中间表示、LLVM 中间表示和 TensorFlow Lite。它还支持方言之间的灵活转换。此外,MLIR 可以创建新的方言以连接到新的低级编译器,这为硬件开发人员和编译器研究人员铺平了道路。

XLA 的 HLO 中间表示可以被视为高级中间表示和低级中间表示,因为 HLO 足够细粒度,可以表示特定于硬件的信息。此外,HLO 支持特定于硬件的优化,并且可以用于生成 LLVM 中间表示。

5.2.2 基于低级中间表示的代码生成

大多数深度学习编译器采用的低级中间表示最终可以降低为LLVM中间表示,并受益于LLVM成熟的优化器和代码生成器。此外,LLVM可以从头开始为专用加速器显式设计定制指令集。然而,传统编译器在直接转换为LLVM中间表示时可能会生成质量较差的代码。为了避免这种情况,深度学习编译器采用两种方法来实现与硬件相关的优化:1)在LLVM的上层中间表示(如基于Halide的中间表示和基于多面体的中间表示)中执行特定于目标的循环变换;2)为优化遍提供关于硬件目标的额外信息。大多数深度学习编译器会同时采用这两种方法,但侧重点有所不同。一般来说,更偏向于前端用户的深度学习编译器(如TC、TVM、XLA和nGraph)可能更关注方法1,而更倾向于后端开发人员的深度学习编译器(如Glow、PlaidML和MLIR)可能更关注方法2。

深度学习编译器中的编译方案主要可分为两类:即时编译(JIT)和提前编译(AOT)。对于即时编译器而言,它能够动态生成可执行代码,并且可以利用更好的运行时信息来优化代码。提前编译器会先生成所有可执行二进制文件,然后再执行它们。因此,与即时编译相比,提前编译在静态分析方面的范围更广。此外,提前编译方法可以应用于嵌入式平台的交叉编译器(例如C-GOOD [46]),还能够在远程机器(如TVM RPC)和定制加速器上执行代码。

5.2.3 讨论

在深度学习编译器中,低级中间表示是深度学习模型的一种细粒度表示,它反映了深度学习模型在不同硬件上的详细实现。低级中间表示包括基于Halide的中间表示、基于多面体的中间表示以及其他独特的中间表示。尽管它们在设计上有所不同,但它们都利用成熟的编译器工具链和基础设施,来提供针对特定硬件优化和代码生成的定制接口。低级中间表示的设计也会影响新的深度学习加速器的设计(例如,TVM的Halide中间表示与Inferentia,以及XLA的HLO中间表示与张量处理单元(TPU))。

5.3 前端优化

在构建计算图之后,前端会应用图级别的优化。许多优化在图级别更容易被识别和执行,因为计算图提供了计算的全局视图。这些优化仅应用于计算图,而不是后端的具体实现。因此,它们与硬件无关,并且可以应用于各种后端目标。

![[Pasted image 20250513114408.png]]

前端优化通常由“遍(passes)”来定义,并且可以通过遍历计算图的节点并执行图变换来应用。前端提供了以下方法:1)从计算图中捕获特定特征;2)重写计算图以进行优化。除了预定义的“遍”之外,开发人员还可以在前端定义自定义的“遍”。一旦深度学习模型被导入并转换为计算图,大多数深度学习编译器就能够确定每个操作的输入张量和输出张量的形状。这一特性使得深度学习编译器能够根据形状信息来执行优化。图3展示了一个使用TensorFlow XLA进行计算图优化的示例。

在本节中,我们将前端优化分为三类:1)节点级优化;2)块级(窥孔,局部)优化;3)数据流级(全局)优化。

5.3.1 节点级优化

计算图的节点足够粗糙,以便在单个节点内部进行优化。节点级优化包括节点消除(即消除不必要的节点)和节点替换(即用成本更低的其他节点来替换节点)。

在通用编译器中,空操作消除(Nop Elimination)会移除那些只占据少量空间但不指定任何操作的空操作指令。在深度学习编译器中,空操作消除用于消除那些缺少足够输入的操作。例如,只有一个输入张量的求和节点可以被消除,填充宽度为零的填充节点也可以被消除。

零维张量消除用于移除那些输入为零维张量的不必要操作。假设 A 是一个零维张量,B 是一个常量张量,那么 A 和 B 的求和操作节点可以被已有的常量节点 B 替代,且不影响正确性。假设 C 是一个三维张量,但其中一个维度的形状为零,比如 {0, 2, 3},那么 C 没有元素,其求最小值 / 最大值(argmin/argmax)操作节点可以被消除。

5.3.2 块级优化

代数简化 —— 代数简化优化包括:1)代数识别;2)强度削弱,即我们可以用成本更低的算子来替代成本较高的算子;3)常量折叠,即我们可以用常量表达式的值来替代常量表达式。这类优化会考虑一系列节点,然后利用不同类型节点的交换律、结合律和分配律来简化计算。

除了典型的算子(如加法、乘法等)外,代数简化也可应用于深度学习特定的算子(如重塑、转置和池化)。这些算子可以重新排序,有时还能被消除,从而减少冗余并提高效率。以下我们举例说明可应用代数简化的常见情况:1)计算顺序优化,在这种情况下,优化过程会根据特定特征查找并移除重塑 / 转置操作。以矩阵乘法(通用矩阵乘法,GEMM)为例,有两个矩阵(如矩阵 A 和矩阵 B),两个矩阵都进行转置(分别得到 AT 和 BT ),然后将 AT 和 BT 相乘。然而,实现通用矩阵乘法更高效的方法是交换参数 A 和 B 的顺序,将它们相乘,然后对通用矩阵乘法的输出进行转置,这样就把两次转置减少为一次;2)节点组合优化,在这种情况下,优化会将多个连续的转置节点合并为单个节点,消除恒等转置节点,并在转置节点实际上不移动数据时将其优化为重塑节点;3)均值归约(ReduceMean)节点优化,在这种情况下,如果归约算子的输入是四维的且最后两个维度要进行归约,优化过程会用平均池化(AvgPool)节点来替代均值归约节点(例如在 Glow 中)。

算子融合 —— 算子融合是深度学习编译器必不可少的优化手段。它能更好地实现计算共享,消除中间分配,通过合并循环嵌套来促进进一步优化, 还能减少启动和同步开销 [91]。在 TVM 中,算子被分为四类:内射型、归约型、复杂输出可融合型和不透明型。在定义算子时,会确定其对应的类别。针对上述类别,TVM 设计了跨算子的融合规则。在张量计算(TC)中,融合是基于自动多面体变换以不同方式执行的。然而,如何识别和融合更复杂的图模式,例如带有多个广播和归约节点的模块,仍然是一个问题。最近的研究工作 [61, 62] 试图解决这个问题,并提出了一个框架来探索和优化激进的融合方案。该框架不仅支持逐元素节点和归约节点,还支持其他具有复杂依赖关系的计算 / 内存密集型节点。

算子下沉 —— 这种优化会将诸如转置之类的操作下沉到批量归一化、修正线性单元(ReLU)、 sigmoid 函数和通道混洗等操作之下。通过这种优化,许多相似的操作会被移动得更接近彼此,从而为代数简化创造更多机会。

5.3.3 数据流级优化

公共子表达式消除(CSE)—— 如果表达式 E 的值在之前已经计算过,并且自上次计算后 E 的值没有改变,那么表达式 E 就是一个公共子表达式 [6]。在这种情况下,E 的值只需计算一次,后续就可以使用已经计算好的值,从而避免在其他地方重新计算。深度学习编译器会在整个计算图中搜索公共子表达式,并将后续出现的公共子表达式替换为之前计算的结果。

死代码消除(DCE)—— 如果一组代码的计算结果或副作用未被使用,那么这组代码就是死代码。死代码消除优化会移除这些死代码。死代码通常不是由程序员编写的,而是由其他图优化产生的。因此,死代码消除和公共子表达式消除通常在其他图优化之后应用。其他优化,如死存储消除(DSE,即移除对永远不会被使用的张量的存储操作),也属于死代码消除的范畴。

静态内存规划 —— 执行静态内存规划优化是为了尽可能地复用内存缓冲区。通常有两种方法:原地内存共享和标准内存共享。原地内存共享为一个操作的输入和输出使用相同的内存,并且只在计算前分配一份内存。标准内存共享在不重叠的情况下复用之前操作的内存。静态内存规划是离线完成的,这使得可以应用更复杂的规划算法。最近的一项工作 [4] 首先设计并执行了内存感知调度,以最小化边缘设备上的峰值激活内存占用,这为内存受限设备上的内存规划提出了新的研究方向。

布局转换 —— 布局转换试图找到最佳的数据布局来存储计算图中的张量,然后将布局转换节点插入到图中。请注意,实际的转换并不是在这里执行的,而是会在编译器后端评估计算图时执行。

实际上,同一操作在不同数据布局下的性能是不同的,而且在不同硬件上的最佳布局也各不相同。例如,在 GPU 上以 NCHW 格式进行的操作通常运行得更快,因此在 GPU 上转换为 NCHW 格式是高效的(例如 TensorFlow)。一些深度学习编译器依赖特定于硬件的库来实现更高的性能,而这些库可能需要特定的布局。此外,一些深度学习加速器更喜欢更复杂的布局(例如分块布局)。另外,边缘设备通常配备异构计算单元,不同的单元可能需要不同的数据布局以实现更好的利用率,因此布局转换需要仔细考虑。因此,编译器需要提供一种方法来在各种硬件之间执行布局转换。

不仅张量的数据布局会对最终性能产生重要影响,而且转换操作也会带来显著的开销,因为它们也会消耗内存和计算资源。

最近一项基于 TVM 针对 CPU 的工作 [58],首先在计算图中将所有卷积操作的布局更改为 NCHW [x] c,其中 c 表示通道 C 的分割子维度,x 表示子维度的分割大小。然后,在进行特定于硬件的优化时,会在提供硬件细节(如缓存行大小、向量化单元大小和内存访问模式)的情况下,通过自动调优来全局探索所有 x 参数。

5.3.4 讨论

前端是深度学习编译器中最重要的组件之一,它负责将深度学习模型转换为高级中间表示(如计算图),并基于高级中间表示进行与硬件无关的优化。尽管不同深度学习编译器的前端在高级中间表示的数据表示和算子定义方面可能有所不同,但与硬件无关的优化主要集中在三个级别:节点级、块级和数据流级。每个级别的优化方法都利用了深度学习特定的以及通用的编译优化技术,这些技术在计算图级别减少了计算冗余,提高了深度学习模型的性能。

5.4 后端优化

深度学习编译器的后端通常包含各种特定于硬件的优化、自动调优技术以及经过优化的内核库。特定于硬件的优化能够针对不同的硬件目标生成高效的代码。而自动调优在编译器后端至关重要,它可以减少为得出最优参数配置所需的手动工作。此外,高度优化的内核库也广泛应用于通用处理器和其他定制的深度学习加速器。

![[Pasted image 20250513114914.png]]

5.4.1 特定于硬件的优化

特定于硬件的优化,也称为与目标相关的优化,用于针对特定硬件获取高性能的代码。应用后端优化的一种方法是将低级中间表示转换为 LLVM 中间表示,以便利用 LLVM 的基础设施生成经过优化的 CPU/GPU 代码。另一种方法是利用深度学习领域的知识设计定制的优化方案,从而更高效地利用目标硬件。由于特定于硬件的优化是为特定硬件量身定制的,无法在本文中详尽阐述,我们介绍现有深度学习编译器中广泛采用的五种方法。这些特定于硬件的优化概述如图 4 所示,详细描述如下:

硬件内在映射 - 硬件内在映射可以将一组特定的低级中间表示指令转换为在硬件上已高度优化的内核。在 TVM 中,硬件内在映射通过可扩展的张量化方法实现,该方法可以声明硬件内在的行为以及内在映射的降低规则。这种方法使编译器后端能够将硬件实现以及高度优化的手工微内核应用于特定的操作模式,从而显著提升性能。而 Glow 支持诸如量化之类的硬件内在映射。它可以估计神经网络每个阶段可能的数值范围,并支持基于剖析引导的优化来自动执行量化。此外,Halide/TVM 将特定的中间表示模式映射到每种架构上的单指令多数据(SIMD)操作码,以避免 LLVM 中间表示映射在遇到向量模式时的低效率问题。

内存分配和获取 - 内存分配是代码生成中的另一个挑战,特别是对于 GPU 和定制的加速器。例如,GPU 主要包含共享内存空间(访问延迟较低但内存大小有限)和本地内存空间(访问延迟较高但容量较大)。这种内存层次结构需要高效的内存分配和获取技术来提高数据局部性。为实现这一优化,TVM 引入了内存作用域的调度概念。内存作用域调度原语可以将计算阶段标记为共享或线程本地。对于标记为共享的计算阶段,TVM 生成带有共享内存分配以及协同数据获取的代码,它会在适当的代码位置插入内存屏障以保证正确性。此外,张量计算(TC)也通过扩展 PPCG [97] 编译器提供类似的功能(称为内存提升)。然而,TC 仅支持有限的预定义规则。特别地,TVM 通过内存作用域调度原语在加速器中实现特殊的缓冲。

内存延迟隐藏 - 内存延迟隐藏也是后端使用的一项重要技术,通过重新排序执行流水线来实现。由于大多数深度学习编译器支持在 CPU 和 GPU 上的并行化,内存延迟隐藏可以自然地通过硬件实现(例如,GPU 上的线程束上下文切换)。但对于类似张量处理单元(TPU)的具有解耦访问执行(DAE)架构的加速器,后端需要执行调度和细粒度同步以获得正确且高效的代码。为实现更好的性能并减轻编程负担,TVM 引入了虚拟线程调度原语,它使用户能够在虚拟化的多线程架构上指定数据并行性。然后,TVM 通过插入必要的内存屏障来降低这些虚拟并行线程,并将来自这些线程的操作交织到单个指令流中,从而形成每个线程更好的执行流水线以隐藏内存访问延迟。

面向循环的优化 - 面向循环的优化也应用于后端,为目标硬件生成高效的代码。由于 Halide 和 LLVM [51](与多面体方法集成)已经纳入了此类优化技术,一些深度学习编译器在后端利用了 Halide 和 LLVM。面向循环的优化中应用的关键技术包括循环融合、滑动窗口、分块、循环重排序和循环展开。

  1. 循环融合:循环融合是一种循环优化技术,可以融合具有相同边界的循环以实现更好的数据重用。对于像 PlaidML、TVM、TC 和 XLA 这样的编译器,这种优化通过 Halide 调度或多面体方法执行,而 Glow 通过其算子堆叠来应用循环融合。
  2. 滑动窗口:滑动窗口是 Halide 采用的一种循环优化技术。其核心概念是在需要时计算值,并即时存储这些值以便数据重用,直到不再需要这些值。由于滑动窗口交织两个循环的计算并使其串行化,这是在并行性和数据重用之间的一种权衡。
  3. 分块:分块将循环分割成几个小块,因此循环被分为遍历小块的外层循环和在小块内迭代的内层循环。这种转换通过使小块适合硬件缓存来实现更好的数据局部性。由于小块的大小是特定于硬件的,许多深度学习编译器通过自动调优来确定分块模式和大小。
  4. 循环重排序:循环重排序(也称为循环置换)改变嵌套循环中迭代的顺序,这可以优化内存访问,从而提高空间局部性。它特定于数据布局和硬件特性。然而,当迭代顺序存在依赖关系时,执行循环重排序是不安全的。
  5. 循环展开:循环展开可以将特定的循环展开为固定数量的循环体副本,这使编译器能够应用激进的指令级并行性。通常,循环展开与循环分割结合应用,即首先将循环分割为两个嵌套循环,然后完全展开内层循环。

并行化 - 由于现代处理器通常支持多线程和单指令多数据(SIMD)并行性,编译器后端需要利用并行性来最大限度地提高硬件利用率以实现高性能。Halide 使用一种称为并行的调度原语来指定用于线程级并行化的循环并行维度,并通过将标记为并行的循环维度与块和线程注释进行映射来支持 GPU 并行化。并且它将大小为 n 的循环替换为 n 宽的向量语句,该语句可以通过硬件内在映射映射到特定于硬件的 SIMD 操作码。Stripe 开发了一种多面体模型的变体,称为嵌套多面体模型,它引入并行多面体块作为迭代的基本执行元素。经过这种扩展,嵌套多面体模型可以检测分块和跨步级别之间的层次并行性。此外,一些深度学习编译器依赖于手工编写的库,如 Glow 或硬件供应商提供的优化数学库(将在 4.4.3 节讨论)。同时,Glow 将向量化任务卸载到 LLVM,因为当提供张量维度和循环迭代次数的信息时,LLVM 自动向量化器能很好地工作。然而,完全由编译器后端利用并行性可以应用更多深度学习模型的特定领域知识,从而以更多的工程工作量为代价实现更高的性能。

5.4.2 自动调优

由于在特定硬件优化中参数调优的搜索空间巨大,因此有必要利用自动调优来确定最优的参数配置。在本次调研所研究的深度学习编译器中,TVM、TC 和 XLA 都支持自动调优。一般来说,自动调优的实现包含四个关键组件,例如参数化、成本模型、搜索技术和加速。

参数化 - 1)数据与目标:数据参数描述数据的规格,例如输入形状。目标参数描述在优化调度和代码生成过程中需要考虑的特定硬件特性和约束。例如,对于 GPU 目标,需要指定共享内存和寄存器大小等硬件参数。2)优化选项:优化选项包括优化调度和相应的参数,例如面向循环的优化和分块大小。在 TVM 中,预定义的和用户自定义的调度以及参数都会被考虑在内。而 TC 和 XLA 更倾向于对优化进行参数化,这些参数与性能密切相关,并且后续可以低成本地进行更改。例如,小批次维度是通常在 CUDA 中映射到网格维度的参数之一,并且可以在自动调优过程中进行优化。

成本模型 - 以下是在自动调优中应用的不同成本模型的比较。1)黑盒模型:该模型只考虑最终的执行时间,而不考虑编译任务的特性。构建黑盒模型比较容易,但在没有任务特性指导的情况下,很容易导致更高的开销和不太理想的解决方案。TC 采用这种模型。2)基于机器学习的成本模型:基于机器学习的成本模型是一种使用机器学习方法来预测性能的统计方法。它使模型能够在探索新配置时进行更新,这有助于实现更高的预测精度。TVM 和 XLA 采用这种模型,例如,分别采用梯度树提升模型(GBDT)和前馈神经网络(FNN)。3)预定义成本模型:基于预定义成本模型的方法期望建立一个基于编译任务特性的完美模型,并能够评估任务的整体性能。与基于机器学习的模型相比,预定义模型在应用时产生的计算开销较小,但在每个新的深度学习模型和硬件上重新构建模型需要大量的工程工作。

搜索技术 - 1)初始化和搜索空间确定:初始选项可以随机设置,也可以基于已知的配置来设置,例如用户给出的配置或历史最优配置。就搜索空间而言,需要在自动调优之前指定。TVM 允许开发人员利用他们特定领域的知识来指定搜索空间,并根据计算描述为每个硬件目标提供自动搜索空间提取。相比之下,TC 依赖于编译缓存和预定义规则。2)遗传算法(GA)[28]:遗传算法将每个调优参数视为基因,将每个配置视为一个候选方案。新的候选方案通过交叉、变异和根据适应度值进行选择来迭代生成,这是一种受自然选择过程启发的元启发式算法。最后,得出最优的候选方案。交叉、变异和选择的比率用于控制探索和利用之间的权衡。TC 在其自动调优技术中采用遗传算法。3)模拟退火算法(SA)[12]:模拟退火算法也是一种受退火启发的元启发式算法。它允许以递减的概率接受较差的解决方案,这可以找到近似的全局最优解,并在固定数量的迭代中避免陷入精确的局部最优解。TVM 在其自动调优技术中采用模拟退火算法。4)强化学习(RL):强化学习通过在探索和利用之间进行权衡来学习,以在给定环境中最大化奖励。Chameleon[5](基于 TVM 构建)在其自动调优技术中采用强化学习。

加速 - 1)并行化:加速自动调优的一个方向是并行化。考虑到遗传算法需要评估每一代中的所有候选方案,TC 提出了一种多线程、多 GPU 策略。首先,它将候选配置入队,并在多个 CPU 线程上进行编译。生成的代码在 GPU 上并行评估,每个候选方案都有其在父代选择步骤中使用的适应度。在完成整个评估后,生成新的候选方案,新的编译任务入队,等待在 CPU 上进行编译。同样,TVM 支持交叉编译和远程过程调用(RPC),允许用户在本地机器上编译,并在多个目标上运行具有不同自动调优配置的程序。2)配置重用:加速自动调优的另一个方向是重用先前的自动调优配置。TC 通过编译缓存存储与给定配置相对应的已知最快的生成代码版本。在编译过程中,每次内核优化之前都会查询缓存,如果缓存未命中,则触发自动调优。同样,TVM 生成一个日志文件,存储所有调度算子的最优配置,并在编译过程中查询日志文件以获取最佳配置。值得一提的是,TVM 对 Halide 中间表示中的每个算子(例如二维卷积 conv2d)进行自动调优,因此会分别为每个算子确定最优配置。

5.4.3 优化的内核库

有几种经过高度优化的内核库被广泛用于在各种硬件上加速深度学习的训练和推理。英特尔的 DNNL(前身为 MKL-DNN)、英伟达的 cuDNN 以及 AMD 的 MIOpen 都是被广泛使用的库。计算密集型原语(例如,卷积、通用矩阵乘法和循环神经网络)以及受内存带宽限制的原语(例如,批量归一化、池化和混洗)都根据硬件特性(例如,AVX-512 指令集架构、张量核心)进行了高度优化。并且支持可定制的数据布局,以便轻松集成到深度学习应用程序中,并避免频繁的数据布局转换。此外,还支持低精度训练和推理,包括单精度浮点数(FP32)、半精度浮点数(FP16)、8 位整数(INT8)以及非 IEEE 浮点格式的脑浮点数(bfloat16)[45] 。其他定制的深度学习加速器也维护着它们特定的内核库 [43, 57]。

现有的深度学习编译器,如 TVM、nGraph 和 TC,在代码生成过程中可以生成对这些库的函数调用。然而,如果深度学习编译器需要利用现有的优化内核库,它们首先需要将数据布局和融合风格转换为内核库中预定义的类型。这种转换可能会破坏最优的控制流。此外,深度学习编译器将内核库视为黑盒。因此,在调用内核库时,它们无法对跨算子进行优化(例如,算子融合)。总之,当计算可以由特定的高度优化原语满足时,使用优化的内核库可以显著提高性能,否则可能会受到进一步优化的限制,且性能不太理想。

5.4.4 讨论

后端负责基于低级中间表示的裸机优化和代码生成。尽管由于各种低级中间表示,后端的设计可能会有所不同,但它们的优化可以分为特定于硬件的优化、自动调优技术和优化的内核库。这些优化可以单独执行,也可以组合执行,通过利用硬件 / 软件特性来实现更好的数据局部性和并行性。最终,深度学习模型的高级中间表示被转换为在不同硬件上的高效代码实现。

6 深度学习编译器的分类

本次调研中研究的深度学习编译器包括 TVM、nGraph、TC、Glow 和 XLA。我们选择这些编译器是因为它们知名度高、维护良好,最重要的是,被广泛使用。因此,我们能够从工业界和学术界找到足够多的论文、文档和讨论,以便深入研究它们的设计和实现。表 1 从前端、后端、中间表示(IR)和优化这四个角度展示了所选深度学习编译器的分类,这与本次调研中描述的关键组件相对应。

具体来说,我们尽我们所知提供了关于这些编译器的更多信息。我们不仅会说明一个编译器是否支持特定功能,还会描述如何通过其编程接口来使用该功能。此外,我们还会描述特定功能的开发状态,以及某些编译器不支持特定功能的原因。

这种分类的目的是为从业者在考虑自身需求时提供选择深度学习编译器的指导方针,同时为研究人员全面总结深度学习编译器的相关内容。

在表 1 中,我们展示了每个深度学习编译器的特性,包括开发者、编程语言、对 ONNX / 框架的支持、训练支持和前端的量化支持,并且在后端类别中展示了编译方法和支持的设备。总结这些特性是因为它们对深度学习编译器在特定场景下的使用有很大影响。基于这些特性,从业者或研究人员可以轻松决定他们想要使用哪个深度学习编译器。

表 1 与图 2 一起,可作为本次调研的系统总结。通过它们,读者可以识别每个编译器所支持的特性以及每个编译器的关键组件。更详细的信息将在以下部分呈现。

![[Pasted image 20250513120845.png]]

7 评估

7.1 实验设置

我们的实验在两台配备 GPU 的机器上进行,硬件配置如表 2 所示。我们在 CPU 和 GPU 上评估 TVM(v0.6.0)、nGraph(0.29.0-rc.0)、TC(提交版本 fd01443)、Glow(提交版本 7e68188)和 XLA(TensorFlow 2.2.0)的性能。我们选择 19 个 ONNX 格式的神经网络模型作为我们的数据集,这些模型是从 Torchvison2 模型库和 GluonCV3 模型库转换而来的。这些模型包括功能完整的模型,如残差网络(ResNet)、密集连接卷积网络(DenseNet)和 VGG 系列,以及轻量级模型,如 MobileNet 和 MNASNet 系列。为了导入 ONNX 模型,如表 1 所示,我们使用 TVM 内置的 tvm.relay.frontend.from_onnx 接口、nGraph 的 ngraph-onnx Python 包、Glow 内置的 ONNXModelLoader 以及 XLA 的 tensorflow-onnx Python 包。值得注意的是,TC 不支持 ONNX,所以我们只在后续的逐层性能比较中对其进行评估。每个模型执行 15 次,我们报告每个编译器最后 10 次执行的平均执行时间,因为我们将前 5 次执行视为预热,以消除即时编译(JIT)的开销。

![[Pasted image 20250513120908.png]]

7.2 端到端性能比较

如图 5 所示,我们比较 TVM、nGraph、Glow 和 XLA 的端到端推理性能。我们在 CPU(Broadwell 和 Skylake)和 GPU(V100 和 2080Ti)上评估这些编译器。请注意,我们在此省略了对 TC 的比较。因为 TC 更像是一个内核库,而不是功能齐全的深度学习编译器,并且它要求用户使用其爱因斯坦概念手动实现模型的所有层,这使得进行公平比较需要大量的工程工作。另一个原因是 TC 仅支持在 GPU 上运行,因此我们无法获取其在 CPU 上的性能结果。然而,对于详细的比较(图 6 和图 8),我们仍然在 TC 中实现了几个残差网络(ResNet)和 MobileNetV2 模型。总之,我们从以下几个方面比较和分析性能结果。

![[Pasted image 20250513120922.png]]

兼容性 - 尽管nGraph和XLA声称支持ONNX,但仍然存在兼容性问题。1)由于张量具有动态形状,nGraph无法运行DenseNet121、VGG16/19和MNASNet0_5/1_0模型。作为替代,我们从ONNX模型库中选取相应模型来替换DenseNet121、VGG16/19模型,而MNASNet0_5/1_0模型无法获取。此外,当我们将PlaidML设置为nGraph在GPU上的后端时,所有MobileNet模型都无法运行。这是因为PlaidML无法处理不同深度学习框架中算子定义不一致的问题。2)XLA可以运行所有选定的模型,然而其性能相当低。因此,我们将选定的ONNX模型替换为来自TensorFlow Hub的保存模型,而MNASNet0_5/1_0模型不可用。使用来自TensorFlow Hub的模型时,XLA的速度提高了两个数量级,并且XLA的性能与其他编译器相比具有竞争力。

性能 - 从图5中,我们对性能有以下几点观察:

  1. 在CPU上,Glow的性能比其他编译器差。这是因为Glow不支持线程并行化,所以它无法充分利用多核CPU。而TVM、nGraph和XLA能够利用所有的CPU核心。
  2. 对于功能完整的模型(ResNet、DenseNet和VGG系列)以及轻量级模型(MobileNet和MNASNet系列),XLA的端到端推理性能相似。此外,它在CPU和GPU上的推理性能几乎相同。众所周知,XLA嵌入在TensorFlow框架中。与TVM、nGraph和Glow相比,TensorFlow包含一个复杂的运行时,这给XLA带来了不可忽视的开销。另外,如果我们增大批量大小(在我们的评估中默认设置为1)并关注深度学习编译器的吞吐量,那么XLA的开销在高吞吐量的情况下可以被忽略。
  3. 一般来说,在 CPU 上,由于上述 Glow 和 XLA 的局限性,TVM 和 nGraph 在所有模型上都比其他深度学习编译器性能更好。TVM 在功能完整的模型上与 nGraph 性能相当,而在轻量级模型上优于 nGraph。nGraph 依赖于 DNNL(前身为 MKL-DNN)库进行加速。因此,nGraph 可以将经过优化的子图卸载到 DNNL 上,并受益于 DNNL 针对英特尔 CPU 定制的细粒度指令级即时编译优化。
  4. 经过调优的 TVM(经过 200 次试验调优)在所有模型上,无论是在 CPU 还是 GPU 上,几乎都能实现最佳性能,尤其是在轻量级模型(MobileNet、MNASNet 系列)上。根据我们的研究,这是因为这些模型内部的经典算子调度已由 TVM 开发者精心设计,并且 TVM 的 tophub 中提供了默认参数。默认的调度和参数能帮助 TVM 实现与其他深度学习编译器相近的性能。此外,经过调优的 TVM 和未调优的 TVM 在 CPU 上的性能差异可以忽略不计,但在 GPU 上差异却相当显著(平均加速 41.26 倍)。这是因为 GPU 的线程和内存层次结构比 CPU 更复杂,因此为了充分利用其计算能力,GPU 需要更细粒度的调度(例如,TVM 中的分块、分割和重排序)。因此,确定 GPU 上的最优调度参数至关重要,而自动调优在这方面展现出了其有效性。

7.3 逐层性能比较

为了进一步比较深度学习编译器后端优化的能力,我们在V100 GPU和Broadwell CPU上(由于Glow缺乏多线程支持,所以在CPU上采用单线程)评估了ResNet50和MobileNetV2_1.0的逐层(卷积层,因为卷积层主导着推理时间)性能。

方法 - 为了测量各个层的执行时间,我们考虑深度学习编译器、硬件(CPU/GPU)和卷积神经网络(CNN)模型,采用不同的方法。具体来说:1)在TVM中,我们复用自动调优的日志来提取内核形状和最优调度。然后我们重建各个卷积层,并使用时间评估器进行评估。2)我们通过Glow的跟踪文件提取执行时间。3)在TC中,我们测量手写内核的执行时间。4)对于nGraph,我们利用时间线来测量在CPU上的执行时间。然而,其PlaidML后端(通过OpenCL提供GPU支持)不支持时间线。此外,没有可用的方法来分析OpenCL中的命令队列。因此,我们将nGraph在GPU上的逐层性能分析留到未来的工作中。4)对于XLA,我们利用内置的tf.profiler.experimental方法来评估CPU性能,并使用英伟达的DLProf [71]工具包来评估GPU性能。

![[Pasted image 20250513121533.png]]

![[Pasted image 20250513121546.png]]

性能 - 从图 6、7、8、9 中,我们对性能有以下几点观察:

  1. nGraph 在 CPU 上的卷积层性能更好,这得益于硬件(英特尔 CPU)和软件(编译器、库和运行时)的协同设计。而在这些编译器中,TVM 在 GPU 上的性能更好。在 MobileNetV2_1.0 模型上,TVM 的性能不稳定,尤其是在 conv1 层。这是因为自动调优过程会受到同一机器上其他进程的影响,因此往往会得出不准确甚至是负面的调度参数。
  2. TC 允许用户使用爱因斯坦概念来定义张量计算内核(例如卷积),而无需指定输入 / 输出张量的形状(例如内核大小)。然后对内核进行自动调优,并存储在其编译缓存中,以加速进一步的自动调优和编译。然而,在我们的评估中,我们发现 TC 的性能在很大程度上依赖于最初编译的内核。以 MobileNetV2_1.0 为例,如果我们用 c1 层来初始化自动调优,那么 c1 层可以表现良好。但随着网络层数变深(远离 c1 层),后续的 c∗_b∗_∗层的运行速度会慢很多。为了获得一致的性能,我们需要分别对每个内核进行调优。
  3. 在优化 MobileNetV2_1.0 的 1×1 卷积(例如 b∗_线性层)以及 ResNet50 的深度可分离卷积(例如 c∗_b∗_2 层)方面,Glow 落后于其他编译器。在 GPU 和 CPU 上计算这些卷积所需的时间都更长。我们注意到,在 Glow 上,卷积通常会与其他层(例如 ReLU、批量归一化)融合,这可能就是其性能比其他编译器低的原因。此外,在 CPU 上,MobileNetV2_1.0 末端的卷积所需时间比开头的卷积短很多。根据跟踪日志,我们注意到这些卷积通过 CPUConvDKKC8 优化([79])得到了加速,该优化对具有特定模式的卷积应用了分块、布局变换和向量化。
  4. 对于 XLA,它可以自动编译(XlaCompile)来自 TensorFlow 的符合条件的子图,并用生成的二进制文件(XlaRun)替换这些子图。此外,卷积层可能会与其他内核聚类,因此它们的性能不容易单独衡量。因此,我们统计了聚类和未聚类的卷积,数据见表 3。需要注意的是,TensorFlow 中的 MobileNetV2_1.0 模型在开头和结尾层与 ONNX 模型略有不同,但线性瓶颈层是相同的。此外,如果一个卷积要进行聚类,在_XlaCompile 完成之前最多可以测量两次。因此,在图 6 中有五个极值(对应于 MobileNetV2_1.0 中的 5 个聚类卷积)。实际上,只有聚类的内核会被 XLA 优化,而未聚类的内核则由 TensorFlow 优化。因此,无法测量由 XLA 优化的独立卷积层的执行时间。因此,我们决定在图 7 - 9 中不包含 XLA 的性能数据。

![[Pasted image 20250513121611.png]]

7.4 讨论

通过上述对深度学习编译器的定量性能比较,我们可以深入分析结合了前端(图级别)和后端(算子级别)优化的粗粒度端到端性能,以及经过后端优化的卷积的细粒度逐层性能。然而,要准确衡量不同深度学习编译器所采用优化的有效性,仍然存在一些尚未解决的挑战。在我们的评估过程中,一个特别的困难在于,在现有的深度学习编译器中,前端优化和后端优化通常紧密耦合,原因如下:1)前端优化通常会影响一系列算子。因此,作为后端优化输入的已优化算子在不同的编译器中会有所不同;2)这些优化往往是协同设计的,以便进一步挖掘性能潜力(例如,XLA 中的聚类以及更高级的优化 [58, 61])。因此,要单独评估和比较不同深度学习编译器的特定优化,即使不是不可能,也是非常困难的。

为了解决这个问题,我们一直在致力于为现有的深度学习编译器构建一个通用的基准测试框架,以衡量逐层性能。其基本思路是提取目标层的必要结构和参数(我们将其称为模型片段),并将这些层重建为特定深度学习编译器可接受的输入,这样编译器就可以忠实地应用相应的前端和后端优化。然后,我们可以衡量这些已优化模型片段的性能,以了解深度学习编译器在感兴趣的层上的有效性。使用模型片段的基准测试框架可扩展到感兴趣的定制层(例如,融合层)。有了这样的基准测试框架,我们可以得出每个深度学习编译器的粗粒度(例如,端到端)和细粒度(例如,逐层)性能指标,从而在感兴趣的层面上比较不同深度学习编译器优化的有效性。目前,我们已经成功地从最先进的卷积神经网络模型中提取目标层进行了实验,比如 ResNet50 的瓶颈层以及 MobileNetV2_1.0 的线性瓶颈层。我们的基准测试框架仍在快速开发中,我们希望能尽快将其提供给社区使用。

8 结论与未来方向

在本次调研中,我们针对现有深度学习编译器的设计原则进行了全面分析。首先,我们深入研究了现有深度学习编译器中采用的通用架构,包括多层中间表示、前端和后端。我们详细介绍了每个组件的设计理念和参考实现,重点关注深度学习编译器特有的中间表示和优化方法。我们总结了本次调研的发现,并突出了深度学习编译器的未来发展方向,具体如下:

动态形状与预处理 / 后处理:动态模型在深度学习领域越来越受欢迎,其输入形状甚至模型本身在执行过程中可能会发生变化。特别是在自然语言处理领域,模型可能会接受各种形状的输入,这对深度学习编译器来说具有挑战性,因为数据的形状在运行时才能确定。现有的深度学习编译器需要更多的研究工作,以便有效地支持新兴动态模型的动态形状。此外,随着未来深度学习模型变得更加复杂,其整个控制流不可避免地可能会包含复杂的预处理 / 后处理程序。目前,大多数深度学习编译器使用 Python 作为编程语言,当预处理 / 后处理由 Python 解释器执行时,可能会成为性能瓶颈。现有深度学习编译器尚未考虑到这种潜在的性能瓶颈。在深度学习编译器中支持整个控制流,能够对深度学习模型以及预处理 / 后处理进行表示和优化,这为模型部署中的性能加速开辟了新的机会。

高级自动调优:现有的自动调优技术专注于单个算子的优化。然而,局部最优的组合并不一定能得到全局最优解。例如,应用于不同数据布局的两个相邻算子可以一起进行调优,且不会在中间引入额外的内存转换。此外,随着边缘计算的兴起,执行时间不再是深度学习编译器唯一的优化目标。在自动调优中还应考虑新的优化目标,如内存占用和能耗。特别是对于基于机器学习的自动调优技术,有几个值得进一步探索的方向。首先,机器学习技术除了应用于成本模型外,还可以应用于自动调优的其他阶段。例如,在选择编译器选项和优化调度的阶段,机器学习技术可用于直接预测可能性,并开发算法来确定最终的配置。其次,基于机器学习的自动调优技术可以基于领域知识进行改进。例如,在自动调优技术中融入特征工程(选择特征来表示程序)[99] 可能是实现更好调优结果的一个潜在方向。

多面体模型:在深度学习编译器的设计中,将多面体模型与自动调优技术相结合以提高效率,是一个很有前景的研究方向。一方面,自动调优可以通过重用先前的配置来最小化多面体即时编译的开销。另一方面,多面体模型可用于执行自动调度,这可以减少自动调优的搜索空间。在深度学习编译器中应用多面体模型的另一个挑战是支持稀疏张量。一般来说,像 CSF [84] 这样的稀疏张量格式使用索引数组(例如,a [b [i]])来表示循环索引,这不再是线性的。这种间接索引寻址会导致非仿射的下标表达式和循环边界,从而阻碍了多面体模型的循环优化 [14, 90]。幸运的是,多面体领域在支持稀疏张量方面取得了进展 [94, 95],整合多面体模型的最新进展可以为深度学习编译器带来更多提升性能的机会。

子图划分:支持子图划分的深度学习编译器可以将计算图划分为几个子图,并且可以以不同的方式处理这些子图。子图划分为深度学习编译器带来了更多的研究机会。首先,它为集成图优化库提供了可能性。以 nGraph 和 DNNL 为例,DNNL 是一个具有图优化功能的深度学习库,它利用了大量经过高度优化的内核。nGraph 与 DNNL 的集成使 DNNL 能够加速 nGraph 生成的子图的执行。其次,它为异构和并行执行提供了可能性。一旦计算图被划分为子图,不同子图的执行可以同时分配给异构硬件目标。以边缘设备为例,其计算单元可能由 ARM CPU、Mail GPU、DSP 甚至 NPU 组成。从深度学习编译器生成能够有效利用所有计算单元的子图,可以显著加速深度学习任务的执行。

量化:传统的深度学习框架中应用的量化策略基于一组固定的方案和数据类型,对于在不同硬件上运行的代码几乎没有定制化。而在深度学习编译器中支持量化,可以利用编译过程中的优化机会,推导出更高效的量化策略。例如,Relay [78] 提供了一个量化重写流程,可以针对各种方案自动生成量化代码。为了支持量化,深度学习编译器需要解决几个挑战。第一个挑战是如何在不耗费大量工程工作的情况下实现新的量化算子。AWS 的尝试指出了一个可能的方向,即使用方言的概念在基本算子的基础上实现新算子,这样图级和算子级的优化可以被复用。第二个挑战是编译过程中量化与其他优化之间的交互。例如,确定量化的合适阶段以及与算子融合等优化的协作,需要未来的研究探索。

统一优化:尽管现有的深度学习编译器在计算图优化和特定硬件优化方面采用了类似的设计,但每个编译器在某些方面都有其自身的优势。目前缺乏一种方法来共享最先进的优化,以及对现有编译器中新兴硬件目标的支持。我们主张统一现有深度学习编译器的优化,以便复用每个深度学习编译器中采用的最佳实践。此外,统一深度学习编译器的优化可以汇聚强大的力量,影响通用和专用深度学习加速器的设计,并为深度学习编译器和硬件的高效协同设计提供环境。目前,Google 的 MLIR 是朝着这个方向的一个很有前景的举措。它提供了多层中间表示的基础设施,并包含中间表示规范和工具包,以在各个层级的中间表示之间进行转换。它还提供了灵活的方言,使得每个深度学习编译器都可以为高级和低级中间表示构建其定制的方言。通过方言之间的转换,一个深度学习编译器的优化可以被另一个编译器复用。然而,方言的转换需要进一步的研究工作,以减少对精细设计的依赖。

可微编程 :可微编程是一种编程范式,在这种范式下,程序可以进行全面的求导。以可微编程范式编写的算法能够自动进行求导,这对深度学习领域很有吸引力。许多编译器项目都采用了可微编程,比如 Myia [89]、Flux [40] 和 Julia [13]。遗憾的是,现有深度学习编译器对可微编程的支持很少。

对现有深度学习编译器来说,支持可微编程极具挑战性。这些困难不仅来自数据结构方面,还来自语言语义方面。例如,要实现从 Julia 到 XLA 的 HLO IR(高级中间表示)的转换,其中一个挑战 [24] 在于,Julia 使用的命令式语言与 XLA 使用的符号语言之间的控制流是不同的。为了高效使用 HLO IR,编译器还需要为 Julia 提供操作抽象,以便支持 XLA 的特定语义,比如映射规约(MapReduce)和广播。此外,Julia 和 XLA 在求导语义上的差异,也需要对编译器设计进行重大修改。

隐私保护:在边云系统中,深度学习模型通常被分成两部分,各部分模型分别在边缘设备和云服务上运行,这样可以提供更好的响应延迟,并减少通信带宽的消耗。然而,边云系统的缺点之一是用户隐私变得很容易受到威胁。原因是攻击者可以截获从边缘设备发送到云端的中间结果,然后利用这些中间结果训练另一个模型,从而泄露与原始用户任务相关的隐私信息。

为了在边云系统中保护隐私,现有方法 [27, 67, 74] 建议向中间结果中添加具有特殊统计属性的噪声,这样可以在不严重降低用户任务准确性的情况下,降低攻击者任务的准确性。然而,困难在于确定应该插入噪声的层,要确定最优的插入层非常耗费精力。上述困难为深度学习编译器支持隐私保护提供了很好的机会,因为编译器保存了深度学习模型的丰富信息,这些信息可以自动指导跨层的噪声插入。

训练支持 : 一般来说,当前深度学习编译器对模型训练的支持远远不够。如表 1 所示,nGraph 仅支持在英特尔 NNP-T 加速器上进行训练,TC 仅支持单个内核的自动求导,Glow 对有限的模型提供实验性的训练支持,TVM 的训练支持还在开发中,而 XLA 依赖于 TensorFlow 的训练支持。总的来说,当前深度学习编译器主要专注于有效弥合将深度学习模型部署到不同硬件上的差距,因此它们将推理作为主要的优化目标。然而,扩展深度学习编译器对模型训练的支持能力将带来大量的研究机会,比如梯度算子的优化和高阶自动求导。

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐