To Pack or Not to Pack: A Generalized Packing Analysis and Transformation

标题:打包还是不打包:一种通用的打包分析与变换
作者/机构:Caio Salvador Rohwedder, Nathan Henderson, João P. L. De Carvalho, Yufei Chen, and José Nelson Amaral
单位:University of Alberta, Edmonton, Canada

A1 主要贡献

本文针对目前编译器中 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 节)

A3 背景知识

许多计算任务会访问非连续的内存数据,例如大矩阵中的子矩阵,这会导致空间局部性差。除了缓存性能不佳外,这种访问模式还可能给虚拟内存系统带来压力,因为它要求在单个循环嵌套内翻译属于许多不同虚拟页的虚拟地址。地址翻译通过使用 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 算法的复制开销很容易超过其收益。

A2 方法细节

3. Packing 优化分析

GPAT 概览。本文的主要贡献是 GPAT,一种推广了 Packing 思想的 Packing 分析和代码变换。与现有的 Packing 方法相比,GPAT 对分块策略和循环顺序的变化更具鲁棒性。事实上,对 GPAT 而言,循环分块是一个可选的前置变换:任何深度至少为三的循环嵌套结构都可能包含有效的 Packing 机会。尽管如此,之前的循环分块变换可能会为 GPAT 引入额外的 Packing 机会以供识别。

3.1 初步定义和符号

中间表示 (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',一个包含 TL 中的工作集副本的张量,并可能改变 TT' 中的数据布局。(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 个维度。交换单位向量的元素会创建一个改变数据布局的置换。

3.2 分析约束

静态控制部分 (SCoP) 与静态形状。与广泛采用的 IR 实践一致,GPAT 的分析要求 for 循环是程序的静态控制部分(SCoPs)【索引15,SCoP Detection: A Fast Algorithm for Industrial Compilers+2016+International Workshop on Polyhedral Compilation Techniques (IMPACT 2016)】,并且张量的形状是静态已知的。其思想是,一个早期的传递可以将具有未知张量形状的计算转换为固定大小的张量分块——随着在固定大小操作符上进行计算的硬件加速器的普及,这种做法变得更加重要。在这些要求下,一个张量在循环中的工作集形状也是静态已知的。

3.3 运行示例

示例代码。本节使用列表 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' 的工作集。

3.4 分析概览

输入与输出。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 条目数——分别对应打包和不打包该候选者的情况。变量 NoPackingPacking 存储了每个循环的 TLB 条目估计值。只有访问 T 所需的 TLB 条目数受 Packing (L, T) 的影响。对于所有其他张量,这个数字在第 6 行通过估计访问 t 在 l 中的工作集所需的条目数来计算,给定 t 的形状是 shape(t)。第 9 行近似计算了如果 TL 中的工作集被打包到 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 函数中运行,当 LforiT 是 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 函数。

3.5 GPAT 在 MLIR Affine 方言中的实现

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 实现作为一个模块化传递在基础设施内可用,并将作为工件提供。

A4 实验环境与结果

实验环境

  • 硬件配置
  • CPU: Intel i5 8500,频率锁定在 3.0 GHz。
  • 内存: 32 GiB DDR4。
  • 软件配置
    • 操作系统: Ubuntu 20.04.1 (Kernel 5.15.0-46)。
  • 代码实现: 实验基准测试用 C 语言编写,使用了 PolyBench/C 基准测试套件。
  • 依赖库与工具链:
  • 使用 Polygeist 将 C 程序编译为 MLIR SCF 和 Affine 方言代码。
  • 对比的优化器包括 Polly、Pluto(通过 MLIR 工具 Polymer 使用)、MLIR Affine 方言的分块传递。
  • 使用 Google Benchmark 和 Polybench/C 内置计时工具进行性能测量。
    * 编译标志:-march=native -flto,启用了 fast-math,并假设指针不别名(no-alias)。
  • Polly 的基于模式匹配的优化被禁用。
  • 数据集
  • 使用了 PolyBench/C 的 "large" 数据集。
  • 实验设置
  • 所有实验均为单线程执行。
  • GPAT 使用了体系结构相关的参数,如缓存大小、TLB 条目数等。
  • 分块实验中,Polymer 接受了 8 到 512 的一系列分块因子,Affine 分块传递则根据目标缓存大小自动选择因子。

实验结果

1. Packing 选择评估(回答问题 1, 3, 5)

该实验旨在评估 GPAT 选择 Packing 候选组合的有效性,并分析其对 TLB 和向量化的影响。

  • 2mm 基准测试 (Fig 1a)
    • 实验内容:在 Polymer 优化(包括部分循环融合)后的 2mm 基准上,比较 GPAT、单个 Packing 候选以及 Polymer 本身的性能。
  • 实验结果:GPAT 在 Polymer 融合后的循环嵌套中发现了数据布局变换的机会。根据分块大小,GPAT 选择了 0 到 4 个 Packing 候选,在大多数分块大小下,其性能超越了任何单个 Packing 候选以及 Polymer。
  • 分析结论:GPAT 能够有效地选择一组 Packing 候选,实现最佳性能。在小分块尺寸下的减速归因于 LLVM 的循环展开导致的寄存器压力。

  • gemm 基准测试 (Fig 1b)

    • 实验内容:在 Polymer 优化的 gemm 上进行评估。
    • 实验结果:Polymer 已为 gemm 生成了具有单位步长访问的优化代码。因此,GPAT 的决策仅基于减少 TLB 未命中。在 GPAT 决定应用 Packing 的分块大小上,其性能优于 Polymer。
  • 分析结论:即使在访存模式已经很好的情况下,GPAT 仍能通过优化 TLB 利用率来提升性能。

  • gemm (BLIS 循环顺序) (Fig 1c)

    • 实验内容:对 gemm 采用 BLIS 框架建议的循环顺序,并评估 GPAT 的效果。
  • 实验结果:这种不同的循环顺序为 GPAT 提供了数据布局变换的机会,通过减少内层循环的访存步长,GPAT 在所有分块大小上都获得了持续的性能提升。
  • 分析结论:GPAT 对先前的循环变换策略具有鲁棒性(正交性),能适应不同的循环结构并找到优化机会。

  • TLB 未命中与向量化分析 (Fig 2, Fig 3)

    • 实验内容:使用 Linux perf 工具收集 2mm 实验中的硬件性能计数器。
  • 实验结果:Fig 2 显示,GPAT 选择的 Packing 候选显著减少了 L1 数据 TLB 的加载未命中。Fig 3 显示,通过数据布局变换,GPAT 促进了 CPU 向量指令的使用,从而减少了总指令数。
  • 分析结论:实验数据证实了 GPAT 的两个主要优化目标:减少 TLB 压力和通过数据布局变换改善向量化。

图1. Packing 选择评估图表
图2. 2mm 的 L1 dTLB 加载未命中次数
图3. 2mm 的指令计数

2. PolyBench 评估(回答问题 2, 3, 4)

该实验将 GPAT 与其他循环优化方法(Polly, Clang -O3)进行比较,评估其在更广泛基准上的通用性。

  • 基于 Polymer 分块引擎 (Fig 4)
    • 实验内容:在 gemm, 2mm, 3mm 基准上,比较 GPAT+Polymer, Polly, 以及 Polymer 相对于 Clang -O3 的加速比。
    • 实验结果:GPAT 在这三个基准上均识别出优化机会。在 2mm3mm 上,GPAT 的性能始终优于 Polly,因为它能应用 Polly 不支持的数据布局变换。在 gemm 上,GPAT 通过减少 TLB 需求也超越了 Polly。在 3mm 上,GPAT 结合了对 gemm2mm 的优化策略,取得了最高的整体加速。
  • 分析结论:GPAT 与先进的多面体优化器(如 Polymer/Pluto)结合使用时,能进一步提升性能,效果优于 Polly。

  • 基于 Affine 分块引擎 (Fig 5)

  • 实验内容:将 GPAT 应用于仅由 MLIR Affine 传递进行分块的基准上,并与 Polly 和 Clang -O3 对比。
    • 实验结果:GPAT 在五个 PolyBench 基准(2mm, 3mm, doitgen, gramschmidt, trmm)和运行示例(contract3D)中发现了 Packing 机会。
  • 分析结论
    * 通用性:GPAT 成功优化了非 GEMM 类的 gramschmidt,证明了其通用性。
    * 与分块的正交性:即使在 Affine 分块失败的 contract3Dgramschmidt 上,GPAT 也能独立工作并带来性能提升。
    * 对输入代码的鲁棒性:GPAT 在 Affine 分块后的 trmmdoitgen 中发现了机会,而这些机会在 Polymer 优化后并不存在,表明 GPAT 的有效性取决于输入的循环结构,而非特定的前置优化。
    * 性能对比:在 contract3Dgramschmidt 上 GPAT 优于 Polly。在 trmm 上,虽然 GPAT 提升了 Affine 的性能,但未能超越 Polly 更全面的优化策略。

图4. 基于 Polymer 分块引擎的 Polybench 评估
图5. 基于 Affine 分块引擎的 Polybench 评估。基准的基线时间(ms)按图中顺序为:1673, 3121, 482, 93, 13919, 和 2966

A5 结论

本文提出了 GPAT,一个模块化的编译器分析与代码变换框架,用于决策在何处以及如何应用 Packing 优化。GPAT 使用一个分析模型,该模型考虑了 TLB 利用率和数据布局变换,来判断应用 Packing 是否能提升性能。尽管 MLIR 中的抽象为实现 GPAT 提供了便利,但 GPAT 的理念可以被集成到其他生产级编译器(如 LLVM)中,以自动优化 GEMM 之外的更广泛的计算任务,同时保持与分块策略的正交性。本文的实验证明,GPAT 不仅自身能在某些基准上超越现有的循环优化器,还能与它们协同工作,进一步提升整体性能。

A6 附录

A.1 摘要

工件内容。本工件提供了一个 Docker 镜像,其中包含了执行论文中两个实验所需的所有二进制文件和指令。此外,工件还包含了源代码、脚本、基准测试,以及论文中展示的实验数据和图表。

A.2 工件清单(元信息)

  • 算法: 一种通用的 Packing 分析与代码变换。
  • 程序: 包括 Polybench/C 4.2 和运行示例(列表 1)的源代码。
  • 编译: 包括 LLVM 14.0.1。
  • 变换: 包括第 4.1 和 4.2 节中指定版本的 Polygeist 和 Polymer。
  • 二进制文件: 在 Docker 镜像中提供。
  • 运行环境: 在 Ubuntu 20.04.1 (Kernel 5.15.0-46) 上使用 Docker 测试。需要 root 权限以使用 perf 进行性能分析。
  • 硬件: 在 Intel x86-64 CPU 上测试。需要 sudo 权限以锁定 CPU 频率。
  • 指标: 执行时间,Linux perf 计数器(L1, L2, L3 缓存未命中;L1, L2 TLB 未命中)。
  • 输出: 实验的执行日志和相关图表。预期结果也已包含。
  • 实验: Packing 选择评估(第 4.3 节)和 Polybench 评估(第 4.4 节)。
  • 所需磁盘空间 (约): 基础 Docker 容器 4.1 GiB。源代码、脚本和论文结果额外需要 1.4 GiB。
  • 完成实验所需时间 (约): 3-4 天,但高度依赖于 CPU 和锁定的频率。
  • 是否公开: 是。
  • 代码许可证: Apache License v2.0 with LLVM Exceptions。
  • 存档 (提供 DOI): 10.5281/zenodo.7517506

A.3 描述

  • A.3.1 如何访问。工件已存档:10.5281/zenodo.7517506,所需空间约 5.5 GiB【索引26,Artifact of "To Pack or Not to Pack: A Generalized Packing Analysis and Transformation"+2023+Zenodo】。
  • A.3.2 硬件依赖。需要能够访问硬件计数器(用于缓存未命中、TLB 未命中和指令执行)的 CPU。
  • A.3.3 软件依赖。可用的 Docker 安装和 Linux perf。

A.4 安装

  • A.4.1 主机上的 Perf。要启用 perf,请安装它并允许非特权用户查看 perf 事件计数器:
$ sudo apt install linux-tools-common linux-tools-generic linux-tools-$(uname -r)
$ sudo sh -c 'echo 1 > /proc/sys/kernel/perf_event_paranoid'
  • A.4.2 CPU 伸缩。为获得更一致和可靠的结果,请锁定 CPU 频率。本文实验中,CPU 被锁定在其基础频率(3 GHz),这因 CPU 而异。
$ sudo cpupower frequency-set --governor performance
$ sudo cpupower frequency-set -u 3GHz
$ sudo cpupower frequency-set -d 3GHz
  • A.4.3 Docker 容器。下载 Docker 镜像,导航到包含下载文件的目录以加载镜像,然后创建并启动一个容器。
$ 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

A.5 实验工作流

交互式运行容器。在 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
  • A.5.1 复现实验。要完全复现本文中获得结果的方法,请在容器中运行以下命令:
$ mkdir $HOME/replica
$ $HOME/scripts/experiments/auto-eval.sh $HOME/replica

最后,在主机上,通过运行以下命令来检索结果:

$ ID="$(docker ps -aqf 'name=^artifact$')"
$ docker cp ${ID}:/home/packing/replica/.
  • A.5.2 详细工作流。以下指令提供了一种更模块化的方式来运行相同的实验,也允许在没有 perf 的情况下运行实验,因此执行时不需要 root 权限。
    运行 Packing 选择评估 (gemm)
$ 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 支持的其他大小。

A.6 评估和预期结果

结果预期。如果使用类似的环境,结果趋势应与本文中呈现的一致。然而,不同的 CPU 可能具有不同的特性,从而影响最终结果。例如,缓存和 TLB 大小的巨大变化,或对额外指令的支持可能会影响最终结果。

A.7 实验定制

定制选项。脚本允许选择实验中使用的 Polybench 数据集大小。此外,在实验的 spec.file 中,有三个默认为 false 的变量可以设置为 true:
* POLLY_ENABLE_PATTERN_MATCHING: 在 Polly 中启用基于模式匹配的优化。
* LLVM_DISABLE_VECTORIZATION: 在 LLVM 中禁用向量化传递。
* LLVM_DISABLE_UNROLLING: 在 LLVM 中禁用展开传递。
为了定制 Polybench 评估工作流,脚本位于 ~/scripts 的各个文件夹中,提供了对编译流程的更多细节。

A.8 注意事项

  • 根据主机的 Docker 配置,可能需要 sudo 权限。
  • 某些分块大小会导致 MLIR 为 trmm 生成不正确的代码,从而导致失败。
  • 在使用 Affine 分块引擎的 Polybench 评估中,摘要图(图 5)对硬件参数变化敏感,可能无法正确渲染。
  • 非 Ubuntu 或非 Intel CPU 的主机可能需要更改 run.sh 脚本中的 perf 事件计数器名称。

方法细节中的引用汇总

  • [1] Uday Bondhugula. 2020. High Performance Code Generation in MLIR: An Early Case Study with GEMM.
    • 引用位置: 引言 (第1节, 第2段); 背景知识 (第2节, 第3段); 相关工作 (第5节, 第6段)
  • 原文描述:
  • 引言中提到,Bondhugula 证明了要让传输操作模拟 Packing,必须遵循一个无法泛化的预定义策略。
  • 背景知识中提到,一些工作(包括此文)可以将 Packing 作为用户指令应用或作为自动调优策略的一部分。
  • 相关工作中详细描述了 Bondhugula 在 MLIR 中尝试用编译器生成的代码匹配 GEMM 库性能的开创性工作,指出其方法依赖于为 GEMM 手动调整的 DCG 传递,而 GPAT 则是完全与计算无关的。
  • [2] Uday Bondhugula, et al. 2008. A Practical Automatic Polyhedral Parallelizer and Locality Optimizer.
    • 引用位置: 引言 (第1节, 第2段); 主要贡献 (第1节, 第3段); GPAT 评估 (第4节, 第1段)
  • 原文描述: 提及 Pluto 作为一个多面体框架,通常不应用 Packing 优化。GPAT 的评估显示,在 Pluto 的输出上应用 GPAT 能获得可观的加速。
  • [6] João P. L. De Carvalho, et al. 2021. KernelFaRer: Replacing Native-Code Idioms with High-Performance Library Calls.
    • 引用位置: 引言 (第1节, 第2段)
  • 原文描述: 指出 Polly 的模式匹配能力有限,仅限于特定写法的矩阵乘法代码。
  • [9] Kazushige Goto and Robert A. van de Geijn. 2008. Anatomy of High-Performance Matrix Multiplication.
    • 引用位置: 引言 (第1节, 第1段); 引言 (第1节, 第3段); 背景知识 (第2节, 第3段); 评估 (第4.3.3节, 第1段)
  • 原文描述: 描述了 GEMM 内核依赖于分块和 Packing 两个关键变换。GPAT 的分析灵感来源于 Goto 等人描述的用于优化 GEMM 的 Packing 方法。背景部分详细说明了 Goto 的工作如何使用 Packing 减少 TLB 压力和改变数据布局。实验部分提及 BLIS 框架的循环顺序基于 Goto 的工作。
  • [11] Tobias Grosser, et al. 2012. Polly - Performing Polyhedral Optimizations on a Low-Level Intermediate Representation.
    • 引用位置: 引言 (第1节, 第2段); 主要贡献 (第1节, 第3段); GPAT 评估 (第4节, 第1段)
  • 原文描述: 提及 Polly 作为一个循环优化器,仅在模式匹配为矩阵乘法时才可能应用 Packing。GPAT 的评估表明其性能优于 Polly。
  • [12] Gaël Guennebaud, et al. 2010. Eigen v3.
    • 引用位置: 引言 (第1节, 第1段); 背景知识 (第2节, 第3段); 相关工作 (第5节, 第5段)
  • 原文描述: 列举 Eigen 作为一个高效的线性代数库,其效率关键在于复杂的 GEMM 实现。背景部分提到 Eigen 遵循 Goto 的策略。相关工作部分将其列为使用 Packing 的 BLAS 库之一。
  • [15] Aditya Kumar and Sebastian Pop. 2016. SCoP Detection: A Fast Algorithm for Industrial Compilers.
    • 引用位置: 方法细节 (第3.2节, 第1段)
  • 原文描述: 说明 GPAT 的分析要求 for 循环是程序的静态控制部分(SCoPs)。
  • [16] Monica D. Lam, et al. 1991. The Cache Performance and Optimizations of Blocked Algorithms.
    • 引用位置: 背景知识 (第2节, 第2段); 背景知识 (第2节, 第4段); 相关工作 (第5节, 第1段)
  • 原文描述: 背景部分引用其观点,即 Packing 减少缓存自冲突未命中,并指出其警告:无限制的 Packing 其复制开销可能超过收益。相关工作部分提到 Lam 等人研究了数据复制(一种不改变数据布局的 Packing)以减少缓存干扰。
  • [17] Chris Lattner and Vikram Adve. 2004. LLVM: A Compilation Framework for Lifelong Program Analysis & Transformation.
    • 引用位置: 方法细节 (第3.5节, 第1段)
  • 原文描述: 提及 MLIR 是 LLVM 编译器基础设施项目的一部分。
  • [18] Chris Lattner, et al. 2021. MLIR: Scaling Compiler Infrastructure for Domain Specific Computation.
    • 引用位置: 引言 (第1节, 第2段); 方法细节 (第3.5节, 第1段)
  • 原文描述: 提及 MLIR 的 Affine 方言缺乏通用的 Packing 传递。方法部分介绍了 MLIR 作为 GPAT 实现的基础设施。
  • [21] Nicholas Mitchell, et al. 1998. Quantifying the Multi-Level Nature of Tiling Interactions.
    • 引用位置: 方法细节 (第3.4节, 第3段)
  • 原文描述: 说明 GPAT 推广了 Mitchell 等人关于确保分块 GEMM 缓存驻留性的思想。
  • [22] William S. Moses, et al. 2021. Polygeist: Raising C to Polyhedral MLIR.
    • 引用位置: 实验环境 (第4.1节, 第2段); 实验方法 (第4.2节, 第1段)
  • 原文描述: 提及使用 Polygeist 作为 C/C++ 前端,将 C 程序生成为 MLIR Affine 方言代码。
  • [24] Louis-Noel Pouchet and Tomofumi Yuki. 2016. Polybench: The polyhedral benchmark suite.
    • 引用位置: 主要贡献 (第1节, 第3段); 实验环境 (第4.1节, 第1段)
  • 原文描述: 指出 GPAT 的评估使用了 Polybench 基准测试套件。
  • [26] Caio Salvador Rohwedder, et al. 2023. Artifact of "To Pack or Not to Pack: A Generalized Packing Analysis and Transformation".
    • 引用位置: 附录 (第A.3.1节, 第1段)
  • 原文描述: 提供工件的 Zenodo 存档链接和所需空间信息。
  • [28] Olivier Temam, et al. 1993. To Copy or Not to Copy: A Compile-Time Technique for Assessing When Data Copying Should be Used to Eliminate Cache Conflicts.
    • 引用位置: 背景知识 (第2节, 第2段); 背景知识 (第2节, 第4段); 相关工作 (第5节, 第1段)
  • 原文描述: 背景部分引用其观点,即在早期工作中 Packing 被称为数据复制,并提到其警告:复制开销可能超过收益。相关工作部分提到 Temam 等人开创了确定何时何地应用数据复制的编译时技术。
  • [29] Field G. Van Zee and Robert A. van de Geijn. 2015. BLIS: A Framework for Rapidly Instantiating BLAS Functionality.
    • 引用位置: 引言 (第1节, 第1段); 背景知识 (第2节, 第3段); 评估 (第4.3.3节, 第1段); 相关工作 (第5节, 第5段)
  • 原文描述: 列举 BLIS 作为一个高效的线性代数库。背景部分提到 BLIS 遵循 Goto 的策略。评估部分提到使用 BLIS 建议的循环顺序进行实验。相关工作部分将其列为使用 Packing 的 BLAS 库之一。
  • [30] Bram Wasti, et al. 2020. LoopStack: a Lightweight Tensor Algebra Compiler Stack.
    • 引用位置: 背景知识 (第2节, 第3段); 相关工作 (第5节, 第6段)
  • 原文描述: 背景知识中提到,一些工作可以将 Packing 作为用户指令或自动调优的一部分。相关工作中提到,其他代码生成方法要么不将 Packing 作为编译器传递实现,要么只考虑 GEMM 的 Packing。
  • [31] Xingfu Wu, et al. 2020. Autotuning PolyBench Benchmarks with LLVM Clang/Polly Loop Optimization Pragmas Using Bayesian Optimization.
    • 引用位置: 引言 (第1节, 第2段); 背景知识 (第2节, 第3段)
  • 原文描述: 提及可以通过 pragma 指令应用 Packing,但这需要手动添加。
  • [32] Zhang Xianyi, et al. 2012. Model-driven Level 3 BLAS Performance Optimization on Loongson 3A Processor.
    • 引用位置: 引言 (第1节, 第1段); 背景知识 (第2节, 第3段); 相关工作 (第5节, 第5段)
  • 原文描述: 列举 OpenBLAS 作为一个高效的线性代数库。背景部分提到 OpenBLAS 遵循 Goto 的策略。相关工作部分将其列为使用 Packing 的 BLAS 库之一。