AutoConfig: 面向深度学习编译优化的自动配置机制

深度学习在处理图像[1]、语音[2]、文本[3]等各个场景中得到了越来越广泛的应用. 为了取得更优异的表现性能, 深度学习模型的设计日益精妙和复杂. 举例而言, 早期经典的图像分类模型AlexNet[1]有600万参数, 如今流行的大语言模型ChatGPT[4]的参数量已经达到1750亿, 参数量的飞跃式增长为深度学习模型在硬件设备中高效落地带来重要挑战.

目前支持深度学习模型落地主要有两类技术. 一类是通过开发高性能算子库来加速模型[5−7]. 这些算子库通常使用编译器内建函数或汇编级别指令来手动实现算子的核心逻辑, 可以在特定的计算场景中提供充分的算子优化机会. 然而要在大部分计算场景中都获得较好的性能, 需要在算子库中为不同规模的算子输入重复编写程序逻辑, 这个开发流程费时费力. 而且程序逻辑中对优化算法和优化参数的选择也与硬件平台强相关, 这使得一套程序逻辑无法被不同的硬件平台直接复用, 因此这类技术依赖于手动调优, 通常不具有跨平台的普适性.

另一类技术是构建支持深度学习模型的编译优化框架. 这些框架的优化流程与处理高级语言程序的传统编译器(如LLVM[8])有共通之处, 它们将深度学习模型的落地视作一个编译优化的过程: 首先采用某种高层抽象的语言来表示模型, 然后将高层抽象逐步下降到硬件级别中间表示, 同时在下降的不同阶段结合多种策略对模型进行优化. 与调用高性能算子库相比, 深度学习编译优化技术能够从全局的角度发掘模型的优化空间, 发挥跨硬件平台做自动调优的潜力. 然而现有深度学习编译优化框架的自动调优方式通常基于搜索策略. 当面对程序逻辑复杂且具有大量调优参数的算法时, 庞大的搜索空间会导致调优过程的开销变得难以承受, 而且参数选取的过程也不具备可解释性.

针对上述问题, 本文提出了AutoConfig, 一种面向深度学习编译优化的自动配置机制. AutoConfig由重写模式(rewrite pattern)和优化分析模型两个部分组成. 重写模式提供可配置的代码生成策略, 以驱动优化分析模型进行参数选取和算法选择. 优化分析模型可以根据硬件平台信息确定调优参数范围, 并使用基准程序从计算、访存和特殊指令开销等多个方面量化模型中的权重系数, 以指导优化算法的选择. 通过综合考虑计算负载特点和硬件特征等要素, AutoConfig能够完成调优参数选取、优化算法选择和自动代码生成, 只需一次优化实现就能适配多种硬件平台. 本文将AutoConfig集成在深度学习编译器Buddy Compiler[9]中, 并针对实际的深度学习计算负载场景, 在多种SIMD平台上对AutoConfig的参数配置及优化算法的选择进行探究, 并将基于优化分析模型的代码生成策略与TVM[10]的自动调优策略进行跨平台的性能比较. 实验结果验证了AutoConfig作为全新的深度学习编译优化开发范式的实用性和有效性.

本文第 1 节介绍相关工作. 第2节介绍基础知识, 包括深度学习计算负载及其优化方法、深度学习编译优化的基本流程和AutoConfig的实现所依托的编译基础设施. 第3节介绍AutoConfig的设计思路、模块组成、使用方式和生态集成. 第4节介绍AutoConfig中的优化分析模型, 并以矩阵乘法和卷积为例, 针对不同的优化算法进行建模和分析. 第5节介绍AutoConfig中的静态信息提取和动态开销测量策略. 第6节呈现实验设计、结果和分析. 第7节给出结论.

在深度学习模型诞生的早期, 人们通过开发高性能算子库来加速模型的推理过程. cuDNN[5]是一个由NVIDIA开发的深度学习加速库, 它提供了一系列深度学习计算负载实现, 采用启发式方法为负载实现选择最佳的算法, 并使用CUDA技术来加速, 可以在NVIDIA GPU上实现高效的深度学习推理. oneDNN[6]是由Intel开发的用于深度神经网络的数学内核库, 它可以发挥Intel处理器体系结构和多核处理器的优势, 实现深度学习计算负载的高性能计算. MIOpen[7]是由AMD开发的数学和科学计算库, 专为深度学习在AMD设备上的并行计算而设计. 高性能算子库的开发通常针对某类特殊硬件, 实现上一般依赖手工调优. 这种调优方式尽管在某些具体场景能取得良好的效果, 但是不具备跨平台迁移的能力.

近年来, 随着摩尔定律趋于终结, 为了应对深度学习领域对高算力的需求, 越来越多计算场景引入了FPGA、CGRA、TPU[11]等深度神经网络加速器进行加速. 为了将深度学习模型的计算模式高效映射到多样的硬件设备上, 基于深度学习编译优化的框架设计成为了研究重点. TVM是目前工业界中较为常用的深度学习编译框架, 它将模型优化的过程视为对张量程序的变换, 可借助autoTVM[12]、Ansor[13]等基于代价模型和启发式算法的搜索策略优化算子执行逻辑. 但是在面对程序逻辑复杂、调优参数众多的算法时, 庞大的搜索空间会使调优开销变得不可承受, 且调优方式也不具备理论上的可解释性. 除了以搜索策略为主的编译框架, 还有其他的工作如ROLLER[14]和SparTA[15]. ROLLER可基于硬件平台完成张量程序的分块调优, 并采用递归算法实现代码生成, SparTA则提出了一种充分利用稀疏性的端到端模型优化通路. 这两种框架通过针对某种特定优化方式或利用数据特性来从新的视角发掘优化的机会, 但对于其他通常情况则存在应用的局限.

本节首先以卷积和矩阵乘法为例介绍深度学习模型计算负载及其优化方法, 然后概述深度学习编译优化的基本流程, 最后介绍AutoConfig的实现所依托的编译基础设施MLIR[16].

近年来, 随着大语言模型的迅速发展, Transformer[17]已成为自然语言处理领域中极具影响力的经典模型. 注意力机制是 Transformer 中的关键模块, 它能够有效处理长序列和目标序列中的依赖关系. 该模块通过计算序列中不同位置之间的相似度, 并为每个位置分配一个权重系数, 从而实现对不同位置的加权关注. 在具体实现中, 需要首先计算查询向量和键向量之间的相似度, 然后将相似度与值向量相乘并加权求和, 最终得到加权后的值向量. 因为此类计算模式在注意力机制中广泛存在, 所以优化矩阵乘法是高效执行 Transformer 模型的关键[18].

矩阵乘法的一种在算法级别的优化方法是Strassen方法[19] , 它通过减少乘法操作的次数来降低计算的理论时间复杂度. 在工业界中, 矩阵乘法的优化策略主要集中在提高计算效率和减少资源消耗上, 包括利用多线程和多核处理器做并行计算、将大矩阵划分成小块做分块计算、优化矩阵在内存中的排布来减少缓存未命中次数等高性能优化方法.

卷积神经网络是在计算机视觉领域中广泛使用的模型. 它通过卷积模块来提取图像的局部特征. 卷积层是卷积模块的核心, 它通过在输入数据上滑动一个固定大小的卷积核来提取图像的局部信息. 卷积层中可以包含多个卷积核, 每个卷积核可以提取不同的特征. 在深度学习模型中, 卷积还有空洞卷积、深度可分离卷积等不同的变种, 它们根据输入的张量表示也可分为“NHWC”和“NCHW”等不同的数据排列方式. 不同的卷积类型、输入尺寸和卷积核大小、数据排列方式会影响计算负载优化的策略.

卷积的优化方法有很多种, 包括从矩阵乘法的视角看待卷积[20]并应用Strassen方法做优化、基于FFT和Winagrad的变换方法[21]和Im2Col算法[22]等. 其中Im2Col算法是目前最通用的优化算法之一, 其主要思想是先将卷积核重叠部分展开和变换输入矩阵, 然后利用高性能的矩阵乘法实现来加速卷积运算, 如图1(a)所示. 需要注意的是变换的开销由于涉及到访存, 因此不可忽略.

图 1 卷积优化算法图示    

在执行卷积和矩阵乘法等计算负载时, 可以通过改变原始算法的操作顺序来换取更好的访存模式, 以为向量化提供充分的空间. 本文采用了一种基于广播操作的优化算法(Broadcast算法), 它的计算逻辑是从一个矩阵中顺序取出向量, 从另一个矩阵中顺序取出标量并使用广播操作构造出向量, 从而达到向量化计算的效果, 如图1(b)所示. 在具体实现时, Broadcast算法会将需要跨行的矩阵计算转换为访存模式更友好的多次迭代单行计算, 并专注于对内层循环做向量化. 本节的后续内容将基于该算法和其他算法采用优化分析模型进行建模、比较和验证.

深度学习编译优化的基本流程如图2(a) 所示. 首先, 来自深度学习模型中的计算负载会通过框架前端引入成为各种中间表示. 接着, 在中间表示的各个层级会进行编译优化, 优化后的结果随后被送到编译器和工具链中完成硬件代码生成. 在这个过程中, 中间表示层级是优化和调优的关键部分[23].

图 2 AutoConfig的概述、设计、实现与集成    

现有的优化和调优策略一般分为编译器的自动调优和调用手动优化高性能库两个大类, 二者的特性决定了它们各自适合不同的场景. 自动调优在循环优化上有出色的效果, 在同类硬件上可以复用基础设施避免重复开发, 但是由于缺乏可解释性, 这种优化和调优方法在新型硬件架构上需要准确的代价模型和耗时的搜索调优过程才能发挥作用. 手动优化能够最大化地利用硬件特性, 可以达到极致性能表现, 但往往需要使用平台特定的编程模型或接口, 即使同类硬件也需要重复开发优化, 造成大量的工程和维护开销.

MLIR[16]是一种现代编译基础设施, 它致力于提供模块化、可复用、可扩展的多层中间表示. 目前, MLIR已经逐渐形成一个编译生态, 尤其是在深度学习编译优化领域, 许多前端框架和后端硬件提供了对MLIR的编译支持[24−26]. MLIR的特点在于其所抽象出的多层中间表示. 每一层中间表示称作方言(Dialect), 这是MLIR扩展机制的核心概念. MLIR提供一系列内置的核心方言, 包括高层级面向领域计算的方言(TOSA Dialect, Linalg Dialect). 循环变换级别的方言(Transform Dialect, Affine Dialect, SCF Dialect). 硬件抽象级别的方言(MemRef Dialect, Vector Dialect). 特定后端级别的方言(LLVM Dialect, GPU Dialect). 在每一种方言内可以定义操作(operation)、类型(type)、属性(attribute), 来提供特定抽象表示的语义.

MLIR提供了强大的扩展机制和自动化能力, 用户可以使用声明式的方法灵活地扩展出自定义方言. 所有的内置方言均可与自定义方言兼容, 并且可以复用所有的内置编译下降通路与优化. 这种扩展能力还使得MLIR支持与其他已有的编译器对接, 便于支持更多的后端和硬件平台.

本节概述AutoConfig机制的设计思路、模块组成、使用方式和生态集成.

深度学习编译器间的差异主要体现在优化方式和调优策略的设计方面. 如图2(a)所示, 目前领域内广泛使用如下两种优化和调优策略: (1)计算与调度分离的编译优化方式, 结合编译器的自动调优机制; (2)图级别编译器生成特定优化算子, 结合手动调优的高性能算子库. 前者通过提供优化与调优编程接口, 将调优的复杂度交给编译器的使用者; 后者通过使用硬件平台特定的编程模型来编写高性能库, 将调优的复杂度交给高性能库的编写者.

本文提出的AutoConfig是一种新的面向深度学习计算负载的编译优化和调优方法, 其核心设计思路是编写可配置的代码生成重写模式, 结合优化分析模型自动配置调优过程, 将调优的复杂度交给编译器的设计者和开发者. 这种设计思路可以利用编译器桥接软硬件的天然优势, 获取计算负载的语义信息和计算平台硬件信息, 实现可配置的编译优化和调优机制. 如图2(a)所示, AutoConfig采用优化分析模型自动配置调优策略, 并选择合适的优化算法进行代码生成, 有效地权衡了多种优化算法和调优策略. 相比自动调优和手动优化, AutoConfig的设计能够达成一次优化实现适配多种硬件平台的目标, 并且具有性能可解释、可迁移的特性.

可配置的代码生成重写模式是AutoConfig优化方法的基础. 通过采用粗颗粒度重写的优化策略, AutoConfig可以在保障优化通用性和高性能的同时, 将优化算法实现为可配置的编译优化. 这种策略集成到编译工具链中, 使得不同目标硬件平台的代码生成可以进行灵活的配置. 此外, 重写模式的配置过程依赖于优化分析模型, 它可以在计算负载的编译时评估计算开销, 并根据此评估结果自动选择并配置最优的代码生成策略.

优化分析模型是AutoConfig调优方法的基础. 该模型对优化算法的总体计算开销进行建模, 将其分解为计算、访存和特殊指令的开销总和. 编译优化的开发者在实现代码生成重写模式时, 也需配套构建参数化的优化分析模型, 并根据硬件信息确定可配置参数取值范围. 在编译过程中, 优化分析模型通过具体的参数配置和动态测量的程序特性来确定总开销. 各优化算法对应的预测总开销是选择代码生成策略的关键依据, 以确保最终的编译优化结果能够最小化计算开销.

值得强调的是, AutoConfig中的各模块构成了可以服务于任意编译优化和调优的基础设施, 并具有良好的可扩展性. 在当前领域特定计算和体系结构发展的黄金时代, 各领域计算负载到多种硬件平台的编译和优化需求日益增长. AutoConfig中的硬件信息收集策略可以针对特定硬件进行定制, 用户可根据需要扩展默认提供的动态测量程序集合, 以及划定给定程序的子集. 无论是优化分析模型的定义还是配置规则的注册, AutoConfig都提供了通用的基础设施, 以满足灵活多变的调优需求. AutoConfig的设计能够帮助开发者解耦编译优化的开发、分析、调优的过程, 并为更精细的分工提供了可能性. 这种设计使得开发者可以专注于特定子模块进行深入开发, 并通过复用现有组件降低开发成本.

AutoConfig的架构由驱动器、可配置的代码生成重写模式和优化分析模型这3大核心组件构成, 它们共同协作为编译优化算法服务. 其中, 优化分析模型抽象出优化算法的特征, 而代码生成重写模式负责表达算法的优化逻辑. 如图2(b)所示, 分析模型与重写模式之间存在直接的映射关系, 由AutoConfig驱动器维护这一映射, 并确保其无缝集成到深度学习编译工具链中.

AutoConfig驱动器是核心控制模块, 负责调用各模块接口进行协同工作, 并且对接深度学习编译工具链. 当深度学习编译工具执行搭载AutoConfig的编译优化Pass时, AutoConfig驱动器在所有已注册的分析模型中筛选出总开销最低的一个. 随后, 它利用选定的分析模型和相对应的重写模式来生成优化代码.

优化分析模型融合了静态信息提取与动态开销测量这两大环节, 并结合编译优化开发者对优化算法的特征分析, 最终形成一个综合的评估函数. 模型的分析是一个动静结合的过程: 静态信息主要来源于对优化算法及硬件平台的特性提取; 动态信息来自于算法特征在目标硬件平台上的执行开销. 如图2(b)所示, 对优化算法的静态分析用来确定计算次数、访存频率和特殊指令的发生次数, 而硬件信息的静态提取用于明确配置项取值范围. 动态开销测量模块根据这一配置范围获取具体测量结果作为各部分开销的权重. 硬件信息的精确提取和分析对于明确可配置的范围至关重要, 这不仅可以减少动态测量的开销, 也能够高效地选择配置项. 总的来说, 优化分析模型的设计与评估是AutoConfig准确性的关键, 本文第4节将详尽阐述以向量计算尺寸为调优点的优化分析模型构建过程.

代码生成重写模式是对编译优化算法逻辑的直接实现. 在代码生成方面, AutoConfig采用MLIR作为中间表示, 并利用MLIR基础设施提供的各层级操作接口实现重写模式. AutoConfig要求在这些重写模式中嵌入可配置的参数, 它们的设置依据来自前述优化分析模型的结果. AutoConfig的设计思路强调将这些可配置参数与特定的重写模式相绑定, 而非与整体的编译优化Pass绑定, 这样的设计提升了调优的精细度和灵活性, 使得每一次优化都可以根据特定的需求和条件进行个性化调整与配置.

AutoConfig旨在服务于编译优化的开发者, 提供了一系列接口以辅助优化和调优的开发. 区别于使用固定策略开发编译优化, AutoConfig要求开发者在创建重写模式时明确定义可配置项及其可能的取值范围. 此外, 用户需要针对优化算法构造分析规则和相应的动态测量程序集合. 这些分析规则、动态测量程序集合和可配置项取值范围共同构成优化分析模型. 针对每种计算负载的不同优化方式, 需要独立定义相应的分析模型和重写模式, AutoConfig驱动器负责将它们配对并完成注册. 集成到MLIR优化Pass中的AutoConfig驱动器, 在触发时根据计算负载和硬件平台信息评估每一组分析模型的开销. 随后AutoConfig开始执行自动配置, 它通过整合计算负载的尺寸信息、分析模型和硬件信息, 全面评估所有可能的算法与配置选择, 从而选择出最佳策略进行优化代码生成.

如算法1所示, 本文以可配置参数vs为例展示针对计算负载op的两种编译优化算法的自动配置和代码生成过程. 首先, 为可配置参数vs设置可选的配置数值. 这里可以采用统一的配置集合以适应所有硬件后端, 或者根据各种硬件特性定制专门的配置集合. 例如, 对于x86的AVX512硬件平台, 可选配置集合可以是{64, 128, 256}, 而对于Arm Neon硬件平台, 则可以是{16, 32, 64}. 在选择了特定的配置集合后, 可以配置动态测量程序子集以适应所需的开销函数, 以此来降低配置成本.

算法1. AutoConfig 机制的使用.

进一步地, 开发者需针对不同优化算法定义相应的重写模式和分析模型. 这些重写模式以MLIR为中间表示, 利用其多层级的MLIR方言来丰富语义表达. 优化分析模型结合分析规则、动态测量程序集合以及可配置参数的取值范围, 共同定义了编译优化的配置空间以供编译优化过程中的搜索和决策使用. 在用户接口方面, 本文提出的方法使用rewritePattern函数实现可配置的重写模式, 并使用lambda表达式来约定优化分析模型. 其中, rewritePattern函数接受目标操作以及可配置参数, 并在函数体内使用各个操作的构造函数实现编译优化的代码生成逻辑. 同时, 为了构建分析模型, lambda表达式通过其捕获列表传递动态测量程序集合, 通过参数列表传递可选配置数值集合. 该模型在lambda函数体内部完成构建, 详细的构建过程将在本文的第4节中展开讨论.

在定义了多种优化算法的重写模式和分析模型之后, 开发者需要初始化一个AutoConfig对象, 这一对象作为驱动器, 负责将每对重写模式和分析模型进行注册. AutoConfig对象的populate函数负责为每个重写模式选择最优的配置项, 然后对不同的重写模式进行比较和评估, 最终选择出执行开销最小的重写模式和最佳配置项, 将其整合进编译优化Pass及其工具链.

本文将AutoConfig集成到了Buddy Compiler中, 展示了AutoConfig对深度学习编译器的适配能力及其在优化深度学习计算任务方面的调优能力. Buddy Compiler是一个基于MLIR的特定领域编译器, 它致力于打造面向深度学习领域的软硬件协同编译生态. 如图2(b)所示, Buddy Compiler中的多级中间表示编译工具buddy-opt会调用AutoConfig的各个模块对关键的深度学习操作进行分析和配置, 随后选择适当的代码生成策略进行编译优化, 并将代码编译下降后翻译至LLVM IR, 最终生成针对目标硬件的汇编代码. Buddy Compiler支持面向CPU SIMD/Vector、GPU、加速器的代码生成, 通过集成AutoConfig的优化与调优方法, 它能够实现面向多种硬件后端的性能迁移和优化复用, 从而打造出高效通用的编译通路.

除了在Buddy Compiler中的集成之外, AutoConfig更重要的贡献是为MLIR生态提供了新的调优方法和思路. MLIR是深度学习编译器的主流生态之一, MLIR上游提供PDL和Transform方言作为基础对高层中间表示进行调度和优化, 例如循环分块、顺序调整、向量化等. 基于MLIR的深度学习编译器IREE除了采用上游提供的基础方言之外还通过自定义方言进行调度, 并且使用强大的运行时环境实现多种平台的部署. 不同于使用特定方言进行调优和调度, AutoConfig的创新点在于使用多层级的方言(例如Vector方言, Affine方言等)编写可配置的优化算法重写模式, 并为优化算法建立分析模型, 搭配硬件信息收集策略提供的开销函数, 可以在编译时确定优化重写模式的配置参数, 从而实现面向特定计算负载和硬件的代码生成. 同时, AutoConfig与MLIR生态中的其他调优方式并不冲突, 使用上游调度方言或者自定义方言的方式也可以集成AutoConfig使得其调优过程具备可解释性和可预测性.

AutoConfig的可解释性源于其优化分析模型. 该分析模型可以对优化算法的计算特征进行建模, 从而近似预测其执行开销. 本节提出一种向量尺寸可配置的优化分析模型建模方法, 可以在编译时根据高层中间表示中计算负载的输入尺寸和SIMD硬件平台特征来预测计算负载优化算法的期望开销, 从而有效指导AutoConfig进行代码生成. 值得注意的是, AutoConfig并不耦合与某个特定的优化分析模型, 而是开放出可注册的接口, 用户可以按需注册自己的优化分析模型.

在深度学习编译优化中, 对优化算法的选择以及算法内的参数配置都会对性能产生影响. 这些影响主要体现在计算开销、访存开销和特殊指令执行开销3个方面. 一般来说, 优化算法与朴素算法相比具有更低的计算开销, 但是由于对算法的优化通常会引入额外的内存排布操作和特殊指令, 这会带来额外的访存开销和特殊指令执行开销. 设Cost_All 为一个优化算法的理论总开销, 则有:

 (1)    

其中, Cost_Arith 表示计算开销, 在SIMD浮点计算场景中主要指融合乘加操作(FMA)的执行开销. Cost_Memory 表示访存开销, 主要指访问内存、缓存和寄存器组等存储模块所需的开销. Cost_SpecIns 则涵盖了除了计算指令和访存指令以外其他特殊指令的执行开销. 不同的优化算法通常会引入不同的特殊指令, 比如Broadcast向量化算法会引入广播操作指令, 针对卷积优化的Im2Col算法会引入维度变换指令等. 朴素算法一般不引入特殊指令, 它的理论总开销只由计算开销和访存开销构成.

为了衡量这些开销, 设优化算法中包含的浮点运算指令个数(floating point operations)为FPO , 内存访问指令个数(data movement)为DM , 所包含的特殊指令种数为M, 且第i条特殊指令的个数为Num_SpecInsi, 则上述公式可进一步描述为:

 (2)    

其中, λ_Arith、λ_Memory 和λ_SpecIns 分别表示计算开销、访存开销和特殊指令开销的重要性系数. 以计算开销为例, 由于它的大小与浮点运算指令个数成正比, 因此λ_Arith 的含义就是单位运算指令开销. 重要性系数与优化算法的参数配置和执行平台紧密相关, 优化算法的参数调优越有效, 越能充分发挥执行平台的硬件特点, 则重要性系数越小. 在实际优化中, 该系数将通过第5.2节的动态测量流程确定.

基于上述对一个优化算法理论总开销的建模方式, 进一步提出优化加速比的概念, 以综合比较不同优化算法所带来的性能差异, 并作为AutoConfig选择最优算法的依据. 优化加速比的定义是: 相对于某个基准算法, 在相同的编译优化流程和硬件执行环境下, 优化算法的理论总开销与基准算法的比值. 具体而言, 对于优化算法A和优化算法B而言, A相较于B的优化加速比可表示为:

 (3)    

在优化分析时, 还可以在计算、访存、特殊指令等不同方面定义局部优化加速比SpeedU_pArith 、SpeedU_pMemory 和SpeedU_pSpecIns , 它们能够帮助分析不同算法在不同方面的优劣情况. 以计算方面为例, 当A 和B 两种算法的计算开销重要性系数保持一致时, 其局部优化加速比可表示为:

 (4)    

此时通过分别统计两种算法的浮点运算指令个数, 就可以求得两者在计算方面的局部优化加速比. 第4.2节会借助该指标来探究不同算法在计算、访存方面的优化表现.

矩阵乘法是一种典型的深度学习计算负载. 设两个输入矩阵分别为M×K 的矩阵A和K×N 的矩阵B, 输出矩阵为M×N 的矩阵C, 下面基于优化分析模型, 对朴素的标量算法和向量化的Broadcast算法做建模分析.

标量算法采用3层嵌套循环完成计算逻辑, 总的循环次数为MNK . 在每一次循环中, 会进行一次FMA操作, 因此其浮点运算指令个数为MNK . 此外在每一次循环中, 默认会涉及到对矩阵A、B和C各一次的读操作和对矩阵C的一次写操作, 但是在实际实现时可以把对矩阵C的读和写操作提取到长度为K𝐾的内层循环外面, 因此内存访问指令个数为2MN+2MNK . 标量算法没有特殊指令执行开销.

基于 Broadcast 算法的向量化实现可通过参数vs𝑣𝑠来调节向量化长度, 以对长度为N 的最内层循环做向量化. 在 Broadcast 算法中, 总的循环次数为, 因此浮点运算指令个数为, 每个FMA操作同时对vs 个元素进行乘加运算. 在进入最内层循环之前, 需要使用一种标量广播的特殊指令将矩阵A的单个元素广播成长度为 𝑣𝑠的向量数组, 该指令共需调用MK 次. 在每次运算的过程中, 分别需要对矩阵B和C进行一次长度为𝑣𝑠的向量读操作, 将其与广播后的向量数组进行乘加运算, 并在最内层循环结束后将结果向量写回矩阵C. 因此该算法的内存访问指令个数为.

设将向量化长度设置为vs 时, 执行一条向量FMA指令的计算开销为Cost_FMA[vs] , 以向量数组的形式访问内存的访存开销为Cost_Memory[vs] , 将标量操作视为 𝑣𝑠=1时的向量操作, 则对标量算法与向量算法的特性分析结果如表1所示. 从中可以看出, 向量算法通过引入了特殊的广播指令并改变了原先的访存模式来为向量化提供优化空间. 一般来说 𝑣𝑠越大, 则向量化越充分. 但是当 𝑣𝑠过大时, 一方面负载的输入尺寸规模可能较小, 可向量化的空间不足. 另一方面向量化所带来的计算便利可能难以弥补寄存器溢出的开销. 因此硬件平台信息会对向量化长度的配置形成约束, 而基于优化分析模型也容易探究不同向量化长度对算法在计算、访存等方面上的影响.

表 1 标量算法与向量算法的特性分析    

深度学习模型里的卷积层通常处理的是四维张量, 这些张量会带有batch、channel和feature等维度. 设N 表示batch, C表示channel, F 表示feature, H_i×W_i 、H_k×W_k 和H_o×W_o 分别表示张量二维切片中的输入、卷积核和输出所对应的尺寸, 下面以数据排布形式为NCHW_FCHW的卷积层为例(即输入张量的数据排布为NCHW, 卷积核张量的数据排布为FCHW, 输出张量的数据排布为NFHW), 对 Im2Col 与 Broadcast两种算法做优化分析.

Im2Col算法首先通过Im2Col策略, 对卷积层的输入进行数据重排, 并对卷积核进行维度变换, 然后通过矩阵乘法操作得到结果矩阵张量, 最后将该张量再进行维度变换恢复成卷积层的输出. 由于维度变换直接对张量进行原地操作, 不涉及浮点计算和内存访问, 因此主要关注该算法的数据重排开销和矩阵乘法操作开销. 在数据重排的过程中, 需要进行NCH_kW_kH_oW_o 次迭代, 每次迭代会读取输入中的一个元素, 再存储到矩阵B中, 该部分的访存开销为2NCH_kW_kH_oW_o, 不涉及计算开销. 在矩阵乘法操作中, 开销等价于N 次固定规模矩阵乘法的开销. 设每次矩阵乘法的两个输入矩阵和一个输出矩阵分别为A、B、C, 则A矩阵的行数A_row=F , B矩阵的行数B_col=H_oW_o , A矩阵的列数A_col=CH_kW_k . 若采用针对矩阵乘法的Broadcast向量算法进行计算, 则一次向量化矩阵乘法的迭代次数为, 浮点运算指令个数为, 内存访问指令个数为, 共需要用到FCH_kW_k 条标量广播的特殊指令. 综合Im2Col策略和所有矩阵乘法操作, Im2Col 算法的总浮点运算指令个数为, 总内存访问指令个数为, 总共需要用到NFCH_kW_k 条标量广播的特殊指令.

若直接采用Broadcast算法, 在朴素算法的基础上对内层循环W_o 进行向量化, 则迭代次数为. 总浮点运算指令个数等于迭代次数, 内存访问指令个数为, 标量广播的特殊指令使用次数为𝑁𝐹𝐶𝐻_𝑜𝐻_𝑘𝑊_𝑘.

对Im2Col算法和Broadcast算法的特性分析如表2所示. 可以看出, 当输入宽度W_o 较大时, Broadcast算法在访存上更占优势. 而当输入高度H_o较大时, Im2Col算法的优化性能更好. 这是因为Broadcast算法直接对W_o 所对应的内层循环进行向量化, 而Im2Col算法会通过将卷积层转换成矩阵乘法操作来对 𝐶𝐻_𝑘𝑊_𝑘所对应的内层循环进行向量化. 所以在同样的向量化长度下, 虽然两种优化算法所需的浮点运算指令个数相同, 但是Im2Col算法需要花费额外的访存开销进行数据重排, 与此同时Broadcast算法需要花费较多的特殊指令执行开销. 在这种情况下便需要优化分析模型基于输入尺寸对算法的选择做出合理的判断.

表 2 Im2Col 与 Broadcast算法的特性分析    

优化分析模型是AutoConfig可解释性的来源, 其准确性由静态信息的提取和动态开销的测量共同决定. 静态信息提取阶段收集计算负载的尺寸信息和硬件平台数据, 旨在确定可配置参数的范围, 可以看作是对配置空间的精化和剪枝, 从而降低调优开销. 动态开销测量阶段通过在指定优化平台上测量示例程序, 以最终确定对应计算、访存和特殊指令的权重系数, 为优化分析模型提供量化数据来选择合适的优化算法, 进而完成代码生成. 为了聚焦于AutoConfig的设计思想, 本文只使用单指令多数据流SIMD作为示例进行论述.

静态信息提取根据计算负载的优化分析模型的需求进行采集, 结合这些硬件信息以及低层编译器的代码生成策略可以分析出加速指令的使用和额外的开销. 面向 CPU 的优化需要利用SIMD的加速能力并且使用 Cache降低访存开销, 因此本文展示的静态信息提取主要包括收集SIMD寄存器堆的尺寸和 Cache的信息. 配合生成的LLVM代码和汇编代码进行静态分析.

SIMD计算单元的关键信息包括SIMD寄存器长度和SIMD寄存器堆的容量. SIMD寄存器的长度标志着一条SIMD指令能够操作的元素个数. SIMD寄存器堆的容量标志着一个SIMD中间表示语句能够使用的最大非溢出尺寸. 使用这两个硬件信息能够结合编译器使用的SIMD抽象尺寸分析出计算的内存溢出次数.

编译策略选择的主要的挑战在于向量中间表示参数和实际机器关键信息之间存在差距. MLIR和LLVM IR的向量抽象提供了灵活的语义, 并且根据后端硬件平台生成SIMD代码, 这意味着中间表示层面有更强的向量语义的表达能力来承载计算负载的编译优化, 同时能够避免在多个后端重复实现优化算法. 然而, 物理机器上的SIMD寄存器的长度和数量是固定的, 需要生成多条SIMD指令来弥合向量中间表示到SIMD器件的鸿沟. 一旦中间表示层面使用的向量抽象尺寸超过了SIMD寄存器堆的容量, SIMD寄存器数量不足以完成所有的计算, 此时汇编代码生成器就会将生成访存指令, 将寄存器中的元素溢出到内存, 腾挪出空闲寄存器完成计算, 当溢出的元素被后续计算使用时, 再将这部分元素从内存中取回到SIMD寄存器, 这就造成了访存操作的额外开销.

在向量中间表示的转换过程中, 为了兼容多种后端硬件平台, 向量抽象的长度和类型不会发生改变. 图3展示了使用融合乘加操作(FMA)给出的具体分析示例. 示例函数由MLIR编写, 接受3个内存抽象的中间表示(memref)作为参数, 从这3个内存抽象中加在向量中间表示, 对3个向量进行融合乘加运算, 并将输出的向量存储到目标内存中间表示中. 示例使用两个不同长度的向量表示, 长度分别为128和1024, 向量中的每个元素为32位浮点数. 在从MLIR编译至LLVM IR之前, 在MLIR层级还需要将不同方言的操作改写为低层级方言中的等价操作这个过程向量表示的尺寸不会改变. 最终, 所有低层级方言中的操作都被翻译为LLVM IR. 图3(a2)和图3(b2)描述了融合乘加函数的LLVM IR代码. 图3高亮了MLIR和LLVM IR中的向量类型. 在这个编译过程中, 中间表示的转换和翻译会重写各层级的方言和操作, 但这些向量抽象的长度和类型从未改变.

图 3 不同向量配置下的融合乘加计算负载与SIMD寄存器的使用情况    

但是, 与MLIR和LLVM IR所抽象的向量中间表示参数不同, 汇编代码级别的参数与实际机器关键信息强相关. 此处选择AVX512扩展来展示编译结果. AVX512的寄存器堆包括32个 zmm SIMD寄存器, 每个寄存器的宽度为512位. 图3(a3)展示了向量抽象长度为128的融合乘加函数的汇编代码. 汇编代码清楚地表现出了SIMD指令的加载、计算和存储的过程. 加载和存储操作使用 vmovups 指令, 融合乘加操作使用vfmadd213ps指令. 在汇编代码中, 每个vfmadd213ps指令需要两个zmm寄存器. 所以向量抽象类型为vector<128xf32> 的融合乘加操作需要16个zmm寄存器(4096/512×2)就可以完成计算. 当向量类型为vector<1024xf32> 时, 理论上需要128个zmm寄存器, SIMD寄存器堆但只有32个可用. 图3(b3)的汇编代码展示了寄存器首先用于进行加载操作, 直至32个zmm寄存器用尽. 此后指令vfmadd213ps进行融合乘加的计算. 由于没有空闲寄存器可用, 需要将计算结果溢出到内存中, 将寄存器腾挪出来负责后续计算. 最后进行存储操作时, 溢出的计算结果被重新加载回寄存器然后存储到目标缓冲区. 因此, 由于汇编代码和硬件平台指令集绑定的, 不合理的向量化配置会降低执行性能.

综上所述, 选择向量配置实际是在权衡迭代开销和溢出开销. 如果使用保守的向量配置策略, 寄存器堆可以承载所有的计算负载需求, 但是需要多次迭代才能完成所有的计算, 即增加了迭代的开销. 如果使用激进的向量配置策略, 寄存器堆需要溢出到内存才可以承载所有的计算负载, 这增加了访存开销, 但是可以用更少次数的迭代完成全部计算, 节省了循环迭代的开销. 这种权衡关系可以指导参数配置范围的选择.

优化算法的访存模型对实际性能有重要的影响, 因此编译优化生成代码对Cache的利用十分关键. 本节对L1数据Cache进行信息收集和分析, 主要关注Cache容量和Cache块尺寸. 这些信息可以辅助编译器制定循环分块或者循环顺序的策略, 尽可能利用Cache降低访存开销.

以实验平台Intel(R) Xeon(R) Gold 5218R CPU为例, L1 数据Cache的容量为 32 KB, 而每个Cache块的大小为 64 字节. 这个 Cache块的大小可以看作与 vector<16xf32> 向量抽象尺寸相当, 这与AVX512的寄存器长度恰好一致. 当使用vector<16xf32> 的整数倍作为向量抽象时, 数据之间并不会存在Cache命中或未命中的相互影响. 正如 第5.1.1节所提及的, 本文所选择的配置项均为16的整数倍. 因此, 对于本文这种以向量化为核心的优化策略, Cache的状态不会对优化分析模型的准确性产生显著的影响. 对于访存密集型的优化任务, Cache相关的可配置项是非常重要的, 并应当纳入优化分析模型. 针对这种优化任务, Cache容量的抽象表示和Cache的结构都是必须进行详细提取和分析的要素. 以该实验平台为例, 访存操作一共可以在L1数据Cache中存放vector<8192xf32>, 以及对应的二维向量表示, 例如vector<8x1024xf32>、vector<4x512xf32> 等. L1数据Cache采用8路组相连的设计, 每一路64个Cache块. 因此, 在优化过程中还需要关注访存操作的抽象长度是否会超出每一路的Cache块数量. 为了达到高效的访存模式, 高层代码优化策略可以根据以上分析采用特定的循环分块和循环顺序策略.

动态开销测量针对优化所在平台确定计算、访存和特殊指令开销的权重系数, 并传入优化分析模型来得到最终的执行时间表达式. 为此, AutoConfig维护了一个动态测量的程序集合, 其中的每个实例程序都体现了某一种程序特性. 该集合允许用户进行删改和添加, 使得用户可以根据其优化需求来调整程序集合. AutoConfig提供的集合能够覆盖大部分与性能有关的MLIR操作, 并且这些动态测量的示例程序都源自实际的测试和优化片段. 这些示例程序的种类和迭代次数都是可配置的, 用户可据此定制合适的示例, 确保得到的测量值既准确又可靠, 从而避免随机误差的影响.

本节以Broadcast向量化算法的开销表示为例阐述动态测量流程, 涉及的测量开销如表3所示. 其中包括访存开销Cost_Memory[vs] 、SIMD计算开销Cost_FMA[vs] 和表示广播指令的特定操作开销Cost_Broadcast[vs] , 它们分别对应了优化分析模型的3项权重系数λ_Arith 、λ_Memory和λ_SpecIns . 具体来说, 在SIMD浮点计算场景中, λ_Arith 是单条FMA指令的执行时间, λ_Memory 则是单条访存指令的执行时间. 而λ_SpecIns , 在对Broadcast向量化算法的分析中, 等价于单条广播操作指令的执行时间, 因为该优化算法涉及的特殊指令只有广播指令. 其他优化算法中特殊指令也同理.

表 3 Broadcast向量化算法的开销表示    

为了量化这些开销, 本节使用动态测量程序集合中的特定子集在多个目标硬件平台上进行了实验. 其中访存开销的测量选择内存拷贝示例程序, 其中包括使用向量方言的内存加载和存储指令. FMA开销的测量选择了FMA示例程序, 该程序结合了访存指令和融合乘加指令. 表示广播指令的特殊指令开销的测量选择了BroadcastFMA的示例程序, 其中包括访存指令、广播指令和融合乘加指令. 以上示例程序均使用2的20次方作为计算负载尺寸, 计算负载迭代10000次, 通过测量总的执行时间来确定Cost_Memory[vs] 、Cost_FMA[vs] 和Cost_Broadcast[vs] 的量化数值.

表4 是在AVX512平台下的实验测量结果, 从中可以看出, 访存指令开销Cost_Memory[vs] 平均为 FMA指令开销Cost_FMA[vs] 和表示广播指令的特定操作开销Cost_Broadcast[vs] 的20倍左右. 在同样的硬件平台下, 这里得出的开销比值与权重系数λ_Arith 、λ_Memory 和λ_SpecIns 之间的比值是一致的, 因此这些数据可以为优化分析模型提供权重数值, 进而分析出整体算法的性能表现, 并选择合适的算法和配置进行代码生成.

表 4 开销表示的测量结果 (ms/每10000次迭代)    

本节面向实际的深度学习计算负载场景, 在AVX512和ARM Neon两种SIMD平台上, 对AutoConfig的参数配置及优化算法的选择进行探究, 并将基于分析模型的代码生成策略与TVM的自动调优策略进行跨平台的性能比较. 其中, 在第6.2节会探究向量化长度参数的选取对优化性能的影响, 指出基于不同硬件平台的特性来确定参数配置的意义. 在第6.3节会利用硬件信息和动态开销数据来计算不同算法的预测加速比, 以验证AutoConfig的算法选择是否准确. 在第6.4节将AutoConfig与基于自动调优策略的TVM进行跨平台的代码生成性能比较. AutoConfig基于Buddy Compiler实现, 在代码生成阶段使用了开源编译器框架LLVM, 版本为17.0.0; 在对比实验中采用的TVM版本为v0.14.0.

本文选取了一个经典的卷积神经网络EfficientNet (模型文件来自https://coral.ai/models/image-classification/), 并针对其中的关键深度学习计算负载即不同输入尺寸的卷积层进行优化. 卷积神经网络是以卷积层作为主要模块的深度学习模型. 对于卷积神经网络而言, 靠近输入部分的卷积层其功能主要在于从输入中提取初步特征, 因此卷积的输入尺寸较大而通道数较少. 处于模型中间部分或靠近输出部分的卷积层其功能主要在于从上一层输出的特征中进一步提取深层信息, 因此卷积的输入尺寸较小而通道数较多.

EfficientNet是经典的卷积神经网络模型, 本节基于EfficientNet完成实验探究. 表5对该模型文件中与计算和访存相关的粗粒度计算负载进行了统计, 可以看出卷积(conv_2d_nhwc_hwcf)的计算占所有计算负载的大多数(38/64), 可见通过从模型中的不同部分选取了尺寸多样的卷积作为优化对象具有一定的代表性, 能够反映优化算法的选择对整体性能的影响.

表 5 EfficientNet高层计算负载统计    

为了验证硬件信息对优化算法参数选取的影响, 在AVX512和ARM Neon平台中分别以{16,32,64,128,256}{16,32,64,128,256}的数值设置向量化长度, 并在不同的卷积优化场景中计算优化加速比, 随后对于这两个平台分别选取平均实际加速比最大的向量化长度设置, 将相关数据记录在表6和表7中. 为了方便对比, 也将适合AVX512的向量化长度配置应用于ARM Neon平台中进行测试, 将相关数据记录在表8中. 图4是面向ARM Neon平台时, 选取不同向量化参数vs的Broadcast算法优化卷积的执行时间比较.

表 6 基于AVX512平台, 向量化长度为64时的最佳算法选择探究    

表 7 基于ARM Neon平台, 向量化长度为16时的最佳算法选择探究    

一方面, 从图4中可以看出, 在ARM Neon平台上选择16作为向量化长度能够获得更小的执行时间和更准确的预测结果, 这说明向量化长度参数的选取会对优化性能产生关键影响. 另一方面, 根据表6, 在AVX512平台上选择64作为向量化长度是更优的参数配置, 但是根据表7和表8, 该参数对ARM Neon平台来说并不是最优的. 这说明采用较大的向量化长度不一定在每个平台上都能得到较好的执行性能, 还需要考虑硬件平台里SIMD寄存器个数和Cache容量的限制. 实践中发现, 在面向AVX512的硬件后端时, {32,64,128} 是合适的参数项配置范围. 而在面向ARM Neon的硬件后端时, 参数选择 {16,32,64}可以得到理想的调优效果, 这反映了参数的可取范围受硬件信息的约束(举例而言, 不同平台有不同的指令向量长度. AVX512的指令向量长度为512位, 而ARM Neon的指令向量长度只有128位). 由此可见, 合理的参数配置能够提高算法选择的精确度, 同时提升代码的执行性能.

表 8 基于ARM Neon平台, 向量化长度为64时的最佳算法选择探究    

图 4 ARM Neon平台下不同向量化参数vs的Broadcast算法优化卷积的执行时间(ms)    

为了验证AutoConfig进行算法选择的有效性, 首先基于硬件信息和动态开销求出Im2Col相对于Broadcast的预测加速比. 预测加速比大于1表示在优化分析模型中Im2Col的执行性能优于Broadcast, 此时AutoConfig会选择Im2Col算法实施代码生成策略. 反之选择Broadcast算法实施代码生成策略. 然后在不同的硬件平台上基于不同的配置参数, 分别显性指定Im2Col和Broadcast算法进行代码生成, 测量两种算法所得到优化代码各自的实际执行时间, 计算出Im2Col相对于Broadcast的实际加速比, 最后在不同的卷积优化场景中AutoConfig的预测加速比进行验证, 实验结果如表6、表7和表8所示. 图5是在AVX512平台下, 确定最佳参数配置vs=64时的最佳算法选择结果. 图5中计算得出了Im2Col和Broadcast算法的相对执行时间占比, 一个算法对应的柱形部分越短, 则该算法的性能与另一个算法相比越好, 此时选择该算法是最佳的策略.

图 5 不同算法优化卷积的相对执行时间占比    

根据图5, 结合表6–表8对算法选择结果的正确性验证, 可以看出AutoConfig在多种硬件平台和不同向量化长度配置下, 能够有效地完成代码生成策略中的算法选择, 这体现在对于大部分卷积优化场景, AutoConfig能够预测得到实际性能更好的优化算法. 尽管如此, AutoConfig的分析模型也有其局限性, 这会导致它在某些情况下(如处理表8中后两个卷积尺寸时)的算法选择结果存在偏差. 具体而言, 在第4.1节的建模中笼统地将各种访存操作带来的开销视为同一类开销Cost_Memory , 但实际上, Cost_Memory 还应包括Cache未命中的代价、从内存中预取数据的代价、数据从寄存器中溢出的代价等. 在实际的访存中预取数据的操作是穿插在运算操作之间来隐藏延迟的, 它与读取Cache操作的频率共同反映了算法对内存和缓存的利用程度, 而寄存器溢出的代价则反映了算法对寄存器组的使用情况. 由此可见, 提高优化分析模型对硬件行为的抽象程度和描述精度是实现更加实用的自动配置机制的关键.

TVM采用搜索的方式进行自动调优, 搜索的过程中需要通过枚举优化参数并实际执行程序来确定最佳的优化策略. 与TVM相比, AutoConfig进行配置的所需开销主要体现在动态开销的测量上. 基于静态硬件信息和动态开销测量的结果, AutoConfig能够直接借助代价模型完成参数配置和算法选择, 实现代码生成的过程. 针对深度学习计算负载场景, 本节将AutoConfig与基于自动调优策略的TVM进行跨平台的性能比较.

在AVX512平台和ARM Neon平台上, 对不同尺寸的矩阵乘法(Matmul)和卷积(Conv)实施调优的实验结果如表9和表10所示. 可以看出AutoConfig和TVM的优化相比于原始程序取得了较明显的性能提升, 且AutoConfig的优化加速比和所生成优化代码的绝对执行时间与TVM的优化在一个数量级内, 具有可比性.

表 9 AVX512平台下AutoConfig与其他调优机制的性能对比    

进一步的, 为了评估AutoConfig相比TVM在调优时的性能优势, 本文定义单位加速比所需调优开销, 表示优化算法在与标量算法相比实现性能提升时, 平均每单位倍数的性能提升所需的调优开销, 该指标直观反映了调优行为对生成代码加速的贡献力度. 基于表10使用单位加速比所需调优开销, 对AutoConfig和TVM在卷积优化场景进行描述, 实验结果如图6所示. 结果表明, 与TVM的自动调优方式相比, AutoConfig无需通过耗时较长的搜索方式确定最优参数, 只需要完成硬件信息提取和动态开销测量, 即可依托代价模型实现可自动配置的代码生成, 且由生成的代码可达到与自动调优相似的执行性能.

表 10 ARM Neon平台下AutoConfig与其他调优机制的性能对比  

图 6 单位平台加速比所需调优开销的跨平台比较    

本文提出了AutoConfig, 一种面向深度学习编译优化的自动配置机制. 针对不同的深度学习计算负载和硬件平台, 该机制构建了具备可配置的代码生成重写模式和可解释性的优化分析模型, 通过分析静态提取的信息和动态测量的开销来确定最佳的参数配置与优化算法, 从而进行代码生成. 本文还将AutoConfig集成到深度学习编译器Buddy Compiler中, 旨在达成一次优化实现适配多种硬件平台的目标. 通过对深度学习模型中的卷积和矩阵乘法的优化实验, 本文验证了AutoConfig自动配置优化和调优方法能够解决当前深度学习编译器调优开销大, 优化效果可解释性差的问题. AutoConfig生成的优化代码可以达到与编译器自动调优相似的性能表现, 同时避免了重复实现和反复调优.

本文期望AutoConfig成为自动配置编译优化的基础设施, 并拥有开放的使用模式. 本文阐述的优化分析模型和动静融合的调优策略并非与AutoConfig的基础设施紧密耦合. 用户可以提供不同的优化分析模型和调优机制来驱动AutoConfig, 从而实现更加高效的编译优化. 例如, 用户可以针对其他深度学习场景进行代码生成, 可以采用更多样的优化算法进行比较和调优, 也可以从硬件平台收集更细颗粒度的信息并建立理论模型来提高性能预测的准确性, 还可以集成Amdahl模型、Roofline模型等不同的分析模型进行针对性的调优等. 因此, 本文更长远的意义在于提供了全新的编译优化开发范式, 将编译优化解耦为开发、分析和调优的过程, 为编译优化提供了新的研究方向.