标题:打包还是不打包:一种通用的打包分析与变换
作者/机构:Caio Salvador Rohwedder, Nathan Henderson, João P. L. De Carvalho, Yufei Chen, and José Nelson Amaral
单位:University of Alberta, Edmonton, Canada
本文针对目前编译器中 Packing 优化的局限性——即主要局限于手工优化的通用矩阵乘法(GEMM)实现或自动调优技术,而通用循环优化器(如 Polly、Pluto)要么仅通过模式匹配将其应用于 GEMM,要么完全不支持——提出了一个通用的解决方案。
核心问题与研究目标:
现有的循环优化器,如 TVM、Halide、MLIR 的 Affine 方言以及多面体框架(Pluto、Polly),普遍缺乏一个基于成本效益分析的、模块化的、可应用于通用计算的 Packing 传递。Polly 仅在成功模式匹配为矩阵乘法时才可能应用 Packing,且其模式匹配能力有限【索引6,KernelFaRer: Replacing Native-Code Idioms with High-Performance Library Calls+2021+ACM Transactions on Architecture and Code Optimization】。其他方法如 pragma 指令【索引31,Autotuning PolyBench Benchmarks with LLVM Clang/Polly Loop Optimization Pragmas Using Bayesian Optimization+2020+IEEE/ACM Performance Modeling, Benchmarking and Simulation of High Performance Computer Systems (PMBS)】则需要手动添加。因此,Packing 的应用被限制在特定循环顺序的手工编码或自动调优中。本文旨在打破这一局限,提出一个能够将 Packing 推广到通用循环嵌套的编译器框架,并采用分析模型而非试错调优来决策。
主要贡献:
本文的主要贡献在于提出了一种名为 GPAT (Generalized Packing Analysis and Transformation) 的通用打包分析与变换方法,其核心创新点如下:
1. GPAT 分析与变换框架:这是一个模块化的编译时 Packing 分析与变换方法。它利用一个分析模型来决策何时 Packing 能够通过数据布局变换减少 TLB 未命中并促进向量化,从而带来性能收益。(详见第 3 节)
2. 开源实现:在 MLIR 的 Affine 方言中实现了 GPAT,并作为一个开源工件提供。(详见第 3.5 节)
3. 全面的评估:通过在 Polybench 基准测试套件上的评估,证明了 GPAT 的有效性。
* GPAT 的启发式方法能够选择出色的 Packing 组合,相较于 Polly 和 Pluto 循环优化器取得了显著的加速。
* GPAT 与分块(tiling)变换是正交的,无论循环嵌套是否经过分块(以及采用何种分块策略和分块因子),GPAT 都能应用 Packing 并提升性能。(详见第 4 节)
许多计算任务会访问非连续的内存数据,例如大矩阵中的子矩阵,这会导致空间局部性差。除了缓存性能不佳外,这种访问模式还可能给虚拟内存系统带来压力,因为它要求在单个循环嵌套内翻译属于许多不同虚拟页的虚拟地址。地址翻译通过使用 TLB(一个地址翻译的缓存)来提高效率。然而,TLB 条目的数量有限,可能不足以翻译循环嵌套中访问的所有地址。因此,生成能避免或减少缓存和 TLB 未命中的代码对于获得高性能至关重要。Packing 是一种可以提高缓存和 TLB 利用率的技术。
Packing 包含从一个n维张量中复制一个子张量到内存中。子张量是张量元素的一个块,通常在内存中不连续。在 Packing 复制过程中,子张量元素的顺序可以被重排,使得连续访问的元素在内存中变得连续。对非连续的子张量进行 Packing 会带来三个重要的好处。首先,访问打包后的子张量时,缓存自冲突未命中会减少,因为连续的元素不太可能映射到同一个缓存组【索引16,The Cache Performance and Optimizations of Blocked Algorithms+1991+ACM SIGPLAN Notices】。这个好处是早期提出 Packing 的工作的主要动机,因为那时的缓存关联度较低。在这些早期工作中,Packing 被称为数据复制(data copying),因为不需要改变数据布局就能获得这种效果【索引28,To Copy or Not to Copy: A Compile-Time Technique for Assessing When Data Copying Should be Used to Eliminate Cache Conflicts+1993+ACM/IEEE Supercomputing Conference】。其次,随着问题规模的增长,目前更重要的是,打包后的子张量的地址翻译需要更少的 TLB 条目。第三,更好的向量化。通过在 Packing 复制中采用数据布局变换,Packing 可以使编译器使用向量指令来访问打包后的子张量,因为其元素被重排为它们被访问的顺序。如果元素以其在原始张量中的存储顺序被复制到打包的子张量中,向量化的效率会低得多。由于重排后的元素在内存中遵循访问顺序,流预取也可能变得更有效。
Goto 等人【索引9,Anatomy of High-Performance Matrix Multiplication+2008+ACM Trans. Math. Software】描述了 GEMM 的高性能 CPU 实现,其中 Packing 是一个关键步骤。Packing 被用来减少子矩阵所需的 TLB 条目数量,以使 TLB 不成为计算的限制因素。在他们的工作中,Packing 还改变了打包子矩阵的数据布局,以便连续的操作访问内存中的连续数据。结果是,由于空间局部性的增加,数据更容易加载到寄存器中。许多先进的线性代数库,如 Eigen【索引12,Eigen v3+2010】,OpenBLAS【索引32,Model-driven Level 3 BLAS Performance Optimization on Loongson 3A Processor+2012+IEEE International Conference on Parallel and Distributed Systems (ICPADS 2012)】和 BLIS【索引29,BLIS: A Framework for Rapidly Instantiating BLAS Functionality+2015+ACM Trans. Math. Software】在它们的 GEMM 实现中都遵循了 Goto 等人描述的策略。除了一些可以将 Packing 作为用户指令应用或作为自动调优策略一部分的工作外【索引1,High Performance Code Generation in MLIR: An Early Case Study with GEMM+2020+arXiv:2003.00532】【索引30,LoopStack: a Lightweight Tensor Algebra Compiler Stack+2020+arXiv:2205.00618】【索引31,Autotuning PolyBench Benchmarks with LLVM Clang/Polly Loop Optimization Pragmas Using Bayesian Optimization+2020+IEEE/ACM Performance Modeling, Benchmarking and Simulation of High Performance Computer Systems (PMBS)】,这些库实现的 GEMM 是目前 Packing 使用的范围。
尽管 Packing 提供了前面提到的好处,但决定打包哪些张量以及在循环嵌套的哪个点放置张量的 Packing 是困难的。此外,如果计算修改了打包的子张量,那么它需要被解包——复制回原始张量。解包会产生额外的复制开销。复制操作的开销可能会超过 Packing 的好处。因此,只有当 Packing 的性能增益大于这种开销时,才应应用 Packing。正如 Lam 等人【索引16,The Cache Performance and Optimizations of Blocked Algorithms+1991+ACM SIGPLAN Notices】和 Temam 等人【索引28,To Copy or Not to Copy: A Compile-Time Technique for Assessing When Data Copying Should be Used to Eliminate Cache Conflicts+1993+ACM/IEEE Supercomputing Conference】所指出的,一个无限制的、打包所有张量的 Packing 算法的复制开销很容易超过其收益。
GPAT 概览。本文的主要贡献是 GPAT,一种推广了 Packing 思想的 Packing 分析和代码变换。与现有的 Packing 方法相比,GPAT 对分块策略和循环顺序的变化更具鲁棒性。事实上,对 GPAT 而言,循环分块是一个可选的前置变换:任何深度至少为三的循环嵌套结构都可能包含有效的 Packing 机会。尽管如此,之前的循环分块变换可能会为 GPAT 引入额外的 Packing 机会以供识别。
中间表示 (IR) 假设。GPAT 的介绍假设了一个中间表示(IR),其中 for 循环被表示为一个操作,该操作编码了归纳变量(IV)、上界和下界、增量步长以及循环体。for 循环操作的循环体可以包含其他操作,包括其他 for 循环。在这个 IR 中,for 循环是控制流图中的一个单入口单出口(SESE)区域。
张量与访存。在这个 IR 中,一块连续的内存区域可以表示为一个 n 维张量。通过用一个索引元组的元素分别索引张量的每个维度来访问张量元素。张量元素的地址由张量基地址和每个维度的索引与步长(stride)的线性组合之和给出。索引元组的每个元素都可以是循环归纳变量、常数和程序中变量的仿射表达式。张量的形状由其每个维度的长度定义。
核心定义。
* 张量的 footprint 是该张量在内存中占用的字节数;它由每个维度的长度与元素数据类型大小(以字节为单位)的乘积计算得出。
* 张量 T 在 for 循环 L 中的 working set 是由在 L 内部访问的 T 的元素集合构成的子张量。
Packing 步骤。对一个目标循环 L 中的张量 T 进行 Packing 包含三个步骤:(i) 在 L 的 SESE 区域的入口块之前立即插入一个 Packing 循环。这个 Packing 循环创建 T',一个包含 T 在 L 中的工作集副本的张量,并可能改变 T 在 T' 中的数据布局。(ii) 在 L 中将所有对 T 的内存引用替换为对 T' 的引用。(iii) 如果对 T' 存在写操作,则在 L 的出口块之后立即插入一个解包(unpacking)循环。解包将 T' 的元素复制回 T 中各自的位置。一个循环-张量对 (L, T) 代表一个 Packing 候选者。
数据布局变换。Packing 循环可以交换 T' 的索引元组的元素,以改变元素被复制到 T' 中的顺序。这种数据布局的改变用一个大小为 n 的置换向量表示,其中 n 是 T 的维度。单位置换向量 [0, 1, 2, ..., n-1] 表示没有数据布局变化。每个元素 p 对应于张量 T 的第 p 个维度。交换单位向量的元素会创建一个改变数据布局的置换。
静态控制部分 (SCoP) 与静态形状。与广泛采用的 IR 实践一致,GPAT 的分析要求 for 循环是程序的静态控制部分(SCoPs)【索引15,SCoP Detection: A Fast Algorithm for Industrial Compilers+2016+International Workshop on Polyhedral Compilation Techniques (IMPACT 2016)】,并且张量的形状是静态已知的。其思想是,一个早期的传递可以将具有未知张量形状的计算转换为固定大小的张量分块——随着在固定大小操作符上进行计算的硬件加速器的普及,这种做法变得更加重要。在这些要求下,一个张量在循环中的工作集形状也是静态已知的。
示例代码。本节使用列表 1 中的伪 IR 代码表示的三维张量收缩来举例说明 GPAT。在这个未分块的示例中,输入张量是 $A_{80 \times 100 \times 50}$,$B_{100 \times 80 \times 60}$ 和 $C_{50 \times 60}$,其中 C 已初始化为零。
1 // A_80x100x50
2 for(i=0; i<50; i++) // A_80x100x1
3 for(j=0; j<60; j++) // A_80x100x1
4 for(k=0; k<80; k++) // A_1x100x1
5 for(l=0; l<100; l++) // A_1x1x1
6 a = load A[k][l][i]
7 b = load B[l][k][j]
8 prod = mul a, b
9 c = load C[i][j]
10 sum = add c, prod
11 store sum, C[i][j]
列表 1. 3D 张量收缩: $C_{i,j} = \sum_{k} \sum_{l} A_{k,l,i} \times B_{l,k,j}$。右侧显示了 A 的工作集在循环嵌套外和每个循环内的形状。
Packing 变换示例。列表 2 展示了对列表 1 应用 Packing 变换后的代码。循环通过其归纳变量的名称来引用,例如 forj 是列表 1 中第 3 行的循环。列表 2 突出了应用于张量 A、目标为 fori 的 Packing 循环,以及对打包后的子张量 A' 的加载。这次 Packing 通过在第 6 行将 A' 的索引元组从 [m][n][0] 重排为 [0][m][n] 来改变数据布局,这对应于置换向量 [2,0,1](从单位向量 [0,1,2] 变换而来)。由于 A' 的最外层维度长度为 1,这个变化是微不足道的。结果,A' 的形状为 (1, 80, 100)。A 的工作集形状在列表 1 的每个 for 循环旁显示,类似地,A' 的工作集形状在列表 2 中显示。
1 for(i=0; i<50; i++)
2 A’ = alloc(1x80x100) // A’_1x80x100
3 for(m=0; m<80; m++)
4 for(n=0; n<100; n++)
5 tmp = load A[m][n][i]
6 store tmp, A’[0][m][n]
7 for(j=0; j<60; j++) // A’_1x80x100
8 for(k=0; k<80; k++) // A’_1x1x100
9 for(l=0; l<100; l++) // A’_1x1x1
10 a = load A’[0][k][l]
11 b = load B[l][k][j]
12 prod = mul a, b
13 c = load C[i][j]
14 sum = add c, prod
15 store sum, C[i][j]
列表 2. 应用于 3D 张量收缩的 Packing。右侧显示了 A' 的工作集。
输入与输出。GPAT 的分析接收一个循环嵌套作为输入。在一个给定的循环嵌套中,将包含的 for 循环数量乘以访问的张量数量,得到可能的 Packing 候选者的总数。调优策略必须考虑所有候选者集合的子集,不包括冗余的子集。如果两个候选者在不同的目标循环中打包同一个张量,但一个循环包含另一个循环,则该候选者子集是冗余的。与调优相反,GPAT 的分析静态地确定其输出:一个有利可图的 Packing 候选者选择及其各自的数据布局变化(如果有)。为此,GPAT 使用以下特定于体系结构的参数:(i) L1、L2 和 L3 缓存的大小(KiB);(ii) L1 数据 TLB 条目的数量;以及 (iii) 这些条目所寻址的页大小(KiB)。
四个分析阶段。GPAT 的分析有四个阶段。第一和第二阶段是过滤阶段,它们缩小了 Packing 候选者的范围,只将表现出数据重用并能保持缓存驻留的候选者转发到后续阶段。第三阶段,目标实现,转发那些至少满足两个 Packing 目标之一的剩余候选者:(i) 最内层访问步长缩减和 (ii) TLB 未命中缩减。这些目标防止 GPAT 选择那些 Packing 收益被 Packing 开销所抵消的候选者。虽然最终池中的所有候选者都满足分析目标,但某些候选者子集可能包含冗余候选者。因此,最后的选择阶段定义了一个成本效益函数,并从第三阶段的候选者中贪婪地选择一个最大化成本效益的非冗余候选者子集。
阶段 1. 数据重用过滤器。此阶段会淘汰一个候选者 (L, T),如果由 Packing 循环创建的子张量 T' 在 L 的迭代中没有被重用;这是克服 Packing 循环开销的一个要求。如果一个访问对于循环 L 的归纳变量是不变的,那么在 L 的每次迭代中都会访问相同的元素集合,这些元素被重用。此过滤器检查循环 L 中与张量 T 相关的所有内存指令是否对 L 的归纳变量(IV)是不变的。由于一个归纳变量 i2 可能依赖于另一个归纳变量 i1(如果定义 i2 的循环在其上界或下界表达式中使用了 i1),因此需要进行流分析来确定循环不变性。这种依赖关系是可传递的:依赖于 i2 的归纳变量 i3 也依赖于 i1。必须对 L 的每个子循环的归纳变量检查这种依赖性——这在分块循环中很常见。因此,一个张量 T 的子张量 T' 在目标循环 L 的迭代中被重用,当且仅当每个访问 T' 的指令的索引元组既不依赖于 L 的归纳变量,也不依赖于任何依赖于 L 的归纳变量的归纳变量。在列表 1 中,有四个表现出重用的候选者 (L, T):{(fori, A), (fori, B), (fori, C), (forj, C)}。在候选者 (fori, A)(在列表 2 中被打包)中,A' 的所有元素在 forj 的每次迭代中都被重用,因为列表 1 第 6 行对 A 的加载对归纳变量 j 是不变的。此阶段是保守的,并且对重用的概念很简单;某些可能打包有利的情况可能会被过滤掉。然而,模块化的设计允许未来集成更复杂的重用分析。
阶段 2. 缓存驻留过滤器。GPAT 分析的第二步是基于缓存驻留性来过滤 Packing 候选者 (L, T)。如果打包后的张量 T' 在 L 中被重用,但缓存中没有足够的可用空间,T' 可能会在 L 的迭代之间被驱逐,导致缓存未命中并增加对 T' 的访问延迟。为避免这种情况,第二个过滤器会淘汰那些打包后的张量无法在使用期间保持缓存驻留的候选者。现代 CPU 有多级缓存。一个打包后的张量 T' 可能在一级缓存中保持驻留,而在另一级被驱逐。GPAT 为其输入循环嵌套设置一个目标缓存级别,打包后的张量应在该级别保持驻留。目标缓存级别被选为不能存储循环嵌套中访问的所有张量总 footprint 的最大缓存。例如,列表 1 中张量收缩的总 footprint 是 A、B 和 C 的 footprint 之和。最后,如果总 footprint 能放入 L1,则所有数据都可以在计算期间保留在 L1 中。在这种情况下,TLB 未命中不太可能成为性能瓶颈,因此不应用 Packing。如果目标缓存足够大,可以容纳:(i) T' 的 footprint 和 (ii) 目标循环 L 一次迭代中所有其他张量工作集 footprint 的两倍,那么候选者的打包张量 T' 就可以保持缓存驻留。对 T' 同一元素的两次访问之间的重用距离是一次 L 迭代中的访问次数。在完成这个距离内的访问次数后,T' 的元素会被带回缓存。这个距离表明,在缓存中有空间容纳 T' 和 L 一次迭代中所有其他张量的工作集将确保 T' 的驻留性。然而,与 T' 的工作集不同,其他张量的工作集可能会在 L 的迭代之间发生变化。对 T' 的访问也可能与其他张量的访问交错。因此,在最近最少使用(LRU)缓存策略中,T' 的元素可能是最近最少使用的。因此,其他所有张量工作集 footprint 的两倍确保了 T' 能够保持驻留,因为这不仅为第 i 次迭代提供了足够的空间,也为第 (i+1) 次迭代提供了空间。GPAT 推广了 Mitchell 等人【索引21,Quantifying the Multi-Level Nature of Tiling Interactions+1998+International Journal of Parallel Programming】关于确保分块 GEMM 驻留性的思想。循环 L 可能有多个直接子循环和条件语句。张量访问可能属于由条件语句选择的互斥代码块。对于这些情况,GPAT 的分析可能会高估确保驻留性所需的缓存大小。然而,即使在此阶段保守,GPAT 对于使用 Polymer 分块的 2mm 基准测试也显示出显著的加速,该基准测试具有互斥的 if 条件(见第 4.3.1 节)。假设列表 2 针对 L2 缓存。打包张量 A' 的元素具有相当于目标循环 fori 一次迭代中访问次数的重用距离。在一次迭代中,会访问 A' 的所有元素,以及形状为 $B_{100 \times 80 \times 1}$ 和 $C_{1 \times 1}$ 的工作集。为了确保 A' 在 forj 的整个迭代过程中都驻留在 L2 中,L2 应该有空间容纳 A' 以及 B 和 C 上述工作集 footprint 的两倍。这就是两次 forj 迭代的 footprint。
阶段 3. 目标实现。在过滤阶段之后,GPAT 的分析会确定哪些通过了第二阶段的候选者满足以下段落中详述的两个目标中的至少一个。目标实现表明选择一个候选者将有利于性能。在验证目标实现之前,GPAT 的分析会检查每个候选者是否有机会进行数据布局更改,以最小化对打包张量的连续访问步长。访问的步长最小化数据布局更改通过以下方式获得:(i) 列出创建访问索引元组中使用的归纳变量的循环深度;(ii) 为索引元组的每个元素保存最大深度;(iii) 置换该索引元组的元素,使循环嵌套深度较低的元素先于深度较高的元素。这个置换可以应用于 T' 的 Packing 循环,从而应用于对 T' 的所有访问,以改变其数据布局。表示此数据布局更改的置换向量会为该候选者保存下来。如果不需要置换,或者最小化步长的置换顺序对于所有对 T' 的访问不相同,则置换向量为单位向量。列表 2 中 A' 的置换在第 3.3 节中讨论。
最内层访问步长缩减。减少最内层循环迭代之间的步长可以改善缓存局部性并可能启用向量化。因此,如果 T' 的数据布局更改减少了最内层循环两次连续迭代中对 T' 的访问步长,则该候选者满足此目标。非最内层循环迭代之间的步长缩减主要减少了访问 T' 所需的 TLB 条目数量,并在下一个目标中考虑。此外,为了满足此目标,候选者要求最内层循环访问至少相当于两个缓存行大小的打包张量元素。这个经验驱动的标准是为了避免保留那些性能收益被 Packing 开销抵消的候选者。在运行示例中,应用于列表 2 的 Packing 候选者通过置换向量 [2, 0, 1] 更改了 A' 的数据布局。与列表 1 中最内层循环 forl 迭代间对 A 的 50 步长访问相比,A' 的数据布局更改确保了在列表 2 中 forl 的连续迭代间访问具有单位步长。
TLB 未命中缩减。如果在一个循环内访问数据所需的条目数超过 TLB 容量,该循环可能会遭受 TLB 容量未命中。一个打包后的子张量,可能带有数据布局更改,需要更少的 TLB 条目来解析地址翻译。如果一个候选者的打包子张量将所需的 L1 数据 TLB(dTLB)条目数减少到 L1 dTLB 容量以下或等于容量,则此目标得以实现。
1 def improvesTLB(L, T, PV, TLBEntries):
2 packedTShape = perm(PV, wSS(T,L))
3 for l in {L, child loops of L}:
4 Packing = 0, NoPacking = 0
5 for t in {tensors accessed in l}:
6 Entries = estTLB(wSS(t,l), shape(t))
7 NoPacking += Entries
8 if t == T:
9 Packing += estTLB(perm(PV,wSS(t,l)), packedTShape)
10 else:
11 Packing += Entries
12 if NoPacking > TLBEntries and Packing <= TLBEntries:
13 return true
14 return false
列表 3. 检查一个 Packing 候选者是否达到减少 TLB 未命中目标的函数。
此阶段使用三个函数:(i) wSS(t, l) 返回张量 t 在循环 l 中的工作集形状;(ii) perm(PV, s) 将置换向量 PV 应用于形状 s;(iii) estTLB(WSS, S) 估计在形状为 S 的张量中寻址形状为 WSS 的工作集所需的 TLB 条目数。使用这些函数,列表 3 描述了一个布尔函数来确定目标实现。对于一个具有置换向量 PV 的候选者 (L, T),在一个拥有 TLBEntries 条目的 L1 dTLB 中,对于 L 及其每个子循环,列表 3 估计了访问张量工作集所需的 TLB 条目数——分别对应打包和不打包该候选者的情况。变量 NoPacking 和 Packing 存储了每个循环的 TLB 条目估计值。只有访问 T 所需的 TLB 条目数受 Packing (L, T) 的影响。对于所有其他张量,这个数字在第 6 行通过估计访问 t 在 l 中的工作集所需的条目数来计算,给定 t 的形状是 shape(t)。第 9 行近似计算了如果 T 在 L 中的工作集被打包到 T' 中,访问 T 所需的 TLB 条目数。在这一行中,estTLB 的两个形状参数都通过候选者的 PV 进行了置换。第二个参数是 T' 的形状而不是 T 的形状。例如,在列表 1 中,A 在 forl 中的工作集形状为 $A_{1 \times 100 \times 1}$,而在列表 2 中,A' 在同一个循环中的工作集形状为 $A'_{1 \times 1 \times 100}$。回想一下,A 的形状是 $A_{80 \times 100 \times 50}$,并且在列表 2 中打包的候选者是 (A, fori)。假设一个 L1 dTLB 包含 64 个条目,为简单起见,假设每个条目可以寻址 50 个 A 的元素。在这种设置下,访问 forl 中 A 的工作集需要 100 个 TLB 条目,而通过将 A 打包成 A',同样的工作集只需要两个条目。这个例子会在列表 3 的 improvesTLB 函数中运行,当 L 是 fori 且 T 是 A 时。其他张量的 TLB 条目也需要被估计,以检查第 12 行的 if 条件。
阶段 4. 贪婪选择。为了确定选择哪些 Packing 候选者,GPAT 的最后阶段对每个满足第三阶段中至少一个目标的候选者进行成本效益分析。一个候选者的收益通过打包 T' 和重用 T' 所减少的 TLB 条目数量来量化。收益是列表 3 中第 12 行评估为 true 的每个循环的 TLB 条目减少量与相应循环执行次数的线性组合。此外,成本通过 T' 的 footprint 来估计——如果 T' 被写入,则为 T' 的两倍 footprint。GPAT 将每个候选者的效益与成本之比作为性能提升的代理,提供了一种从最有益到最无益对候选者进行排序的方法。如果两个候选者具有相同的成本效益比,GPAT 通过比较目标循环的深度来解决排序问题;较小的深度会使 T' 在更多的计算中持续存在,并表明 T' 中元素的重用率更高。一旦排序完成,GPAT 会遍历列表,贪婪地选择候选者。要被选中,一个候选者必须:(i) 相对于已提交的候选者集合不冗余;(ii) 当与已提交的候选者集合一起考虑时,必须继续满足第三阶段中的任一目标。为了验证 (ii),在选择一个候选者之前,会根据已提交候选者的信息重新运行 improvesTLB 函数。
MLIR 框架。MLIR【索引18,MLIR: Scaling Compiler Infrastructure for Domain Specific Computation+2021+IEEE International Symposium on Code Generation and Optimization (CGO 2021)】是 LLVM 编译器基础设施项目【索引17,LLVM: A Compilation Framework for Lifelong Program Analysis & Transformation+2004+ACM/IEEE International Symposium on Code Generation and Optimization (CGO 2004)】的最新成员。其目标是扩展 LLVM 中可能的表示级别,并成为高级抽象的通用框架。MLIR 是 IR 方言的组合,而不是像 LLVM 那样的通用 IR。Affine 是这些方言之一,提供了一种适合多面体分析的 IR,因为它保留了 C 等语言中可用的高级循环嵌套结构。
实现细节。GPAT 分析描述中使用的大多数高级概念都存在于 Affine 方言中。张量由 memref 表示,通过索引映射访问。for 循环是一等操作。由于 GPAT 的实现建立在 Affine 的数据复制生成转换传递之上,计算张量形状和 footprint 以及张量在循环中的工作集形状的实用函数已经定义好了。与 Affine 的循环分块类似,GPAT 的 Affine 实现作为一个模块化传递在基础设施内可用,并将作为工件提供。
-march=native -flto,启用了 fast-math,并假设指针不别名(no-alias)。该实验旨在评估 GPAT 选择 Packing 候选组合的有效性,并分析其对 TLB 和向量化的影响。
2mm 基准上,比较 GPAT、单个 Packing 候选以及 Polymer 本身的性能。分析结论:GPAT 能够有效地选择一组 Packing 候选,实现最佳性能。在小分块尺寸下的减速归因于 LLVM 的循环展开导致的寄存器压力。
gemm 基准测试 (Fig 1b):
gemm 上进行评估。gemm 生成了具有单位步长访问的优化代码。因此,GPAT 的决策仅基于减少 TLB 未命中。在 GPAT 决定应用 Packing 的分块大小上,其性能优于 Polymer。分析结论:即使在访存模式已经很好的情况下,GPAT 仍能通过优化 TLB 利用率来提升性能。
gemm (BLIS 循环顺序) (Fig 1c):
gemm 采用 BLIS 框架建议的循环顺序,并评估 GPAT 的效果。分析结论:GPAT 对先前的循环变换策略具有鲁棒性(正交性),能适应不同的循环结构并找到优化机会。
TLB 未命中与向量化分析 (Fig 2, Fig 3):
perf 工具收集 2mm 实验中的硬件性能计数器。
该实验将 GPAT 与其他循环优化方法(Polly, Clang -O3)进行比较,评估其在更广泛基准上的通用性。
gemm, 2mm, 3mm 基准上,比较 GPAT+Polymer, Polly, 以及 Polymer 相对于 Clang -O3 的加速比。2mm 和 3mm 上,GPAT 的性能始终优于 Polly,因为它能应用 Polly 不支持的数据布局变换。在 gemm 上,GPAT 通过减少 TLB 需求也超越了 Polly。在 3mm 上,GPAT 结合了对 gemm 和 2mm 的优化策略,取得了最高的整体加速。分析结论:GPAT 与先进的多面体优化器(如 Polymer/Pluto)结合使用时,能进一步提升性能,效果优于 Polly。
基于 Affine 分块引擎 (Fig 5):
2mm, 3mm, doitgen, gramschmidt, trmm)和运行示例(contract3D)中发现了 Packing 机会。gramschmidt,证明了其通用性。contract3D 和 gramschmidt 上,GPAT 也能独立工作并带来性能提升。trmm 和 doitgen 中发现了机会,而这些机会在 Polymer 优化后并不存在,表明 GPAT 的有效性取决于输入的循环结构,而非特定的前置优化。contract3D 和 gramschmidt 上 GPAT 优于 Polly。在 trmm 上,虽然 GPAT 提升了 Affine 的性能,但未能超越 Polly 更全面的优化策略。
本文提出了 GPAT,一个模块化的编译器分析与代码变换框架,用于决策在何处以及如何应用 Packing 优化。GPAT 使用一个分析模型,该模型考虑了 TLB 利用率和数据布局变换,来判断应用 Packing 是否能提升性能。尽管 MLIR 中的抽象为实现 GPAT 提供了便利,但 GPAT 的理念可以被集成到其他生产级编译器(如 LLVM)中,以自动优化 GEMM 之外的更广泛的计算任务,同时保持与分块策略的正交性。本文的实验证明,GPAT 不仅自身能在某些基准上超越现有的循环优化器,还能与它们协同工作,进一步提升整体性能。
工件内容。本工件提供了一个 Docker 镜像,其中包含了执行论文中两个实验所需的所有二进制文件和指令。此外,工件还包含了源代码、脚本、基准测试,以及论文中展示的实验数据和图表。
$ sudo apt install linux-tools-common linux-tools-generic linux-tools-$(uname -r) $ sudo sh -c 'echo 1 > /proc/sys/kernel/perf_event_paranoid'
$ sudo cpupower frequency-set --governor performance $ sudo cpupower frequency-set -u 3GHz $ sudo cpupower frequency-set -d 3GHz
$ docker load --input docker-packing-artifact.tar.gz $ docker create --privileged -it --name artifact packing-artifact $ docker start artifact
\1 标志是为了在实验中使用 perf。此外,perf 需要在容器中安装与您的系统匹配的版本:
$ docker exec -u 0 -it artifact bash $ apt update && apt install -y linux-tools-$(uname -r) $ exit
交互式运行容器。在 docker 容器已经启动的情况下,使用以下命令以交互方式运行它。
$ docker exec -it artifact bash
配置参数。首先通过修改容器内的 spec.file 来设置目标 CPU 的体系结构特定参数。
$ vim $HOME/scripts/experiments/spec.file
spec.file 中输入的参数被 Packing 分析和 LLVM 的 Polly 使用。它们还定义了在 Polybench 评估中作为 Affine Tiling 输入的缓存参数。在主机上,可以通过运行以下命令找到 x86 机器的缓存和 TLB 信息:
$ lscpu -C $ sudo apt install -y x86info && x86info -c
$ mkdir $HOME/replica $ $HOME/scripts/experiments/auto-eval.sh $HOME/replica
最后,在主机上,通过运行以下命令来检索结果:
$ ID="$(docker ps -aqf 'name=^artifact$')"
$ docker cp ${ID}:/home/packing/replica/.
$ BENCHMARK="gemm"
$ DATASET_SIZE="LARGE"
$ cd $HOME/scripts/experiments/packing-selection-evaluation/
$ OUTPUT_DIR="$HOME/output/output-${BENCHMARK}-${DATASET_SIZE}"
$ mkdir -p ${OUTPUT_DIR}/graphs
$ ./generate-files.sh -D ${DATASET_SIZE} -B ${BENCHMARK} ${OUTPUT_DIR}
$ ./run.sh -D ${DATASET_SIZE} ${OUTPUT_DIR}/executables ${OUTPUT_DIR}
$ ./parse-log.py ${OUTPUT_DIR}/output.log ${OUTPUT_DIR}/graphs ${BENCHMARK}
要运行 2mm 或 BLIS 循环顺序的 gemm,只需将 BENCHMARK 变量更改为 "2mm" 或 "gemm-blis"。
运行 Polybench 评估 (Polymer):
$ TILING="Polymer"
$ DATASET_SIZE="LARGE"
$ cd $HOME/scripts/experiments/polybench-evaluation/
$ OUTPUT_DIR="$HOME/output/output-polybench-${TILING}-${DATASET_SIZE}"
$ mkdir -p ${OUTPUT_DIR}/logs ${OUTPUT_DIR}/graphs
$ ./generate-files.sh -D ${DATASET_SIZE} -T ${TILING} ${OUTPUT_DIR}
$ ./run.sh -D LARGE ${OUTPUT_DIR} ${OUTPUT_DIR}/logs
$ ./parse-log.py ${OUTPUT_DIR}/logs ${OUTPUT_DIR}/graphs ${TILING}
要使用 Affine Tiling 引擎运行,只需将 TILING 变量更改为 "AffineTiling"。
所有脚本都有一个通过 -h 标志访问的帮助信息。run.sh 脚本支持 -p 标志(收集 perf 事件)和 -r NUMBER 标志(指定运行次数)。DATASET_SIZE 变量可以更改为 Polybench 支持的其他大小。
结果预期。如果使用类似的环境,结果趋势应与本文中呈现的一致。然而,不同的 CPU 可能具有不同的特性,从而影响最终结果。例如,缓存和 TLB 大小的巨大变化,或对额外指令的支持可能会影响最终结果。
定制选项。脚本允许选择实验中使用的 Polybench 数据集大小。此外,在实验的 spec.file 中,有三个默认为 false 的变量可以设置为 true:
* POLLY_ENABLE_PATTERN_MATCHING: 在 Polly 中启用基于模式匹配的优化。
* LLVM_DISABLE_VECTORIZATION: 在 LLVM 中禁用向量化传递。
* LLVM_DISABLE_UNROLLING: 在 LLVM 中禁用展开传递。
为了定制 Polybench 评估工作流,脚本位于 ~/scripts 的各个文件夹中,提供了对编译流程的更多细节。
sudo 权限。trmm 生成不正确的代码,从而导致失败。run.sh 脚本中的 perf 事件计数器名称。