利用MLIR实现矩阵乘法的高性能GPU码生成：一些早期结果
介绍了使用MLIR编译器基础设施在NVIDIA GPU上针对张量核生成代码的一些早期结果。当今高性能深度学习的最先进技术主要是由高度调优的库驱动的。这些库通常是由专业程序员手工优化与调优的，他们使用低级抽象，付出了巨大的努力。对于类似的硬件与未来的硬件，可能需要重复很多这样的工作。因此，这个过程不像LLVM这样的编译器基础结构那样是模块化的或可重用的。手动优化通常不使用标准的中间表示（IR），尽管所执行的优化可以编码为IR上的一系列转换步骤与自定义过程。手动调整也可能错过对只有自动代码生成，才能轻松到达的设计点的探索。直到最近引入MLIR（多级中间表示），IR基础设施才能够有效地解决特定领域库的自动生成问题。特别是，很难使用单个IR来表示与转换高、中、低级别的计算抽象。
通过在MLIR中进行适当的抽象，构建了一个实验性的降低流水线，该流水线能够在NVIDIA GPU上针对其张量核自动生成矩阵乘法的代码。在评估的一组问题大小上，初始性能结果表明，在基于NVIDIA的Ampere微体系结构的Geforce 3090 RTX上，F32与F16积累的性能分别为CuBLAS的95-119%与80-160%。相信，这些结果可以作为进一步研究与开发使用类似专业加速器的IR基础设施自动生成代码与库的动机。
1简介
深度学习与高性能人工智能在很大程度上依赖于高性能计算。计算机硬件与微体系结构、库、编译器、运行时与编程模型的创新不断满足相关的计算需求。目前，大量高性能深度学习由硬件供应商提供的高度调优的库提供支持：例如，CuDNN、CUBLAS与MKL。创建这些图书馆需要大量的努力与专业知识。这个开发过程可能必须在每一个主要的硬件或软件版本中重复，并且以有效的方式探索与优化的内容是有限的。
矩阵乘法（matmul）作为一种计算内核，是许多基于Transformer（如BERT）与一般高性能计算的深度学习框架的核心。它还可以作为一个极好的测试用例来评估可以实现的目标。虽然自动代码生成器的优势通常是优化内核的组成，而不是众所周知的单个内核，但无法为经过充分研究的内核生成好的代码，可能会成为让自动代码生成器处理所有峰值性能代码生成的整体故事的障碍。特别针对NVIDIA GPU上的张量内核，NVIDIA GPU是矩阵乘法累加（MMA）操作的专用单元，其吞吐量通常是普通CUDA内核的3-4倍。
一些工作集中在以张量核为目标的GPU GEMM上。通过在Julia中创建三级API来解决两种语言的问题，该API使用户能够编写高效的GEMM内核。主要关注点是开发一个足够灵活的API，以满足各种各样的应用程序，而不是使用具有多个抽象级别的统一IR基础设施。 采用多面体代码生成方法来生成Volta张量核的代码，使用调度树来表示计算，并使用ISL为其生成CUDA代码。可以为matmul与融合运算（如偏置加法与ReLU）生成代码，同时实现高达2.55倍的速度。这项工作是针对Volta的，包括一些特定于设备的专业化，以实现有竞争力的性能。Triton是一种用于神经网络计算的IR与优化编译器。该框架基于网格的概念，网格是一个静态形状的多维数组。编译器由一个python包公开，该包允许用户编写python代码，并自动生成高效的机器代码。这项工作同时支持CUDA与张量核，并取得了良好的性能。
在这里的方法是使用编译器中间表示（IR）基础设施来促进高性能代码与库的生成。以矩阵乘法为核进行实验，以NVIDIA的张量核为目标。MLIR是在这里使用的编译器基础设施，目的是使整个过程在很大程度上更加模块化、系统化与自动化。通过小步骤下译IR并应用正确的IR转换与优化集，可以实现与手工编写的库相当的性能，而无需实际手工编写任何代码。虽然之前的一项工作对高性能CPU的单核进行了类似的研究，但目标是专门的加速器。
这项工作的一些贡献包括：
1）在MLIR方言中引入Warp矩阵乘法累加（Warp Matrix Multiply Accumulate，WMMA）操作，并将其降低到LLVM/NVPTX后端；
2）演示GPU上的matmul如何作为MLIR转换与方言降低的序列系统地、渐进地生成代码；
3）构建一个针对张量核的端到端matmul代码生成管道，初步结果表明，所获得的性能与手动调优库不相上下，在某些情况下甚至高达1.60x。
如果存在从这些模型到MLIR的下译流程，那么方法是基于IR的方法，可以与不同的编程模型与语言一起工作。
2背景
2.1 MLIR
多级中间表示（MLIR）旨在提供可重用、可扩展的编译器基础设施，并降低构建特定领域编译器的成本。MLIR可用于实现多个目标，例如：
1）可以表示数据流图（如Tensorflow），包括动态形状、变量等；
2）可以用于以适合优化的形式表示机器学习操作的内核；
3）能够跨内核进行高性能计算风格的循环优化（融合、循环交换、平铺、展开与阻塞）；
4）可以表示特定于目标的操作，例如特定于加速器的高级操作；
5）可以表示不同抽象级别的内核（MLIR中的方言），这有助于执行单一级别无法实现的转换与优化。最近一些使用MLIR的工作包括。
MLIR结构由以下部件组成：
1）操作：这是MLIR中的基本语义单元，被称为操作。从指令到函数再到模块的所有内容都被建模为MLIR的操作。运算取零个或多个值，分别称为操作数与结果；
2）属性：它是结构化的编译时静态信息，例如整数常数值、字符串数据等。属性是类型化的，每个Op实例都有一个从字符串名称到属性值的开放键值字典；
3）区域与块：区域包含块列表，块包含可能进一步包含区域的操作列表。区域内的块构成控制流图（CFG）。每个块以终止器操作结束，终止器操作可以具有控制流可以传输到的后续块；
4）方言：它是一个唯一名称空间下的操作、属性与类型的逻辑分组。来自不同方言的行动可以在任何时间在IR的任何级别共存。方言允许扩展性并提供灵活性，有助于执行特定的优化与转换。仿射、GPU、LLVM与Linalg是一些重要的方言；
5）功能与模块：模块是一个Op，具有包含单个块的单个区域，并由不传输控制流的伪Op终止。函数是一个具有单个区域的Op，其自变量与函数自变量相对应；
在工作中使用的一些MLIR方言解释如下：
6）仿射方言：该方言使用多面体编译技术，使相关性分析与循环转换高效可靠。已经在仿射方言的层面上进行了大多数优化与转换；
7）GPU方言：GPU方言为通用GPU编程范式建模，类似于MLIR中的CUDA或OpenCL。它的目标是提供抽象来建模GPU特定的操作与属性。它在很大程度上意味着与供应商无关。可以找到一些附加信息与GPU方言文档；
8）NVVM方言：由于专注于张量核代码生成，使用并扩展了NVVM方言；
此方言提供直接映射到LLVM中的NVPTX后端的操作。
9）LLVM方言：代码生成的最后阶段涉及降低到LLVM IR，从那里LLVM后端控制并生成目标代码。为了对LLVM IR进行建模，使用了这种方言。这是MLIR中存在的最低抽象级别。
2.2 GPU背景
GPU是通用的大规模并行计算设备。内存与计算层次结构在优化任何应用程序并因此实现高性能方面发挥着重要作用。可以将GPU内存抽象为4级层次结构、全局内存、L2缓存、可配置的L1缓存/共享内存与寄存器。GPU上的处理器也可以抽象为两级层次结构，即流式多处理器（SM）与SM内部的计算核心。这些计算核心通常被称为CUDA核心。除了CUDA核，张量核等特殊单元也存在于计算层次结构中相同级别的较新GPU中。每个SM被进一步划分成处理块，这些处理块具有各自的翘曲调度器。GPU的编程模型也被构造为与存在的处理器层次结构相匹配。线程是GPU上的单个执行实体，可以与其他线程并行执行。这些线被分成32组，称为经线。warp在SM上的计算核心上以步调一致的方式执行。warp调度器选择准备执行的warp，并将其调度到计算核心。当一个warp遇到数据依赖项时，它会暂停，而warp调度程序会选择另一个准备执行的warp。根据SM上存在的处理块的数量，可以并行执行多个warp。因此，一般来说，更多的warp有助于实现：
1）warp级别的并行性
2）更好的延迟隐藏
3）更好地利用底层资源。
现在，这些warp被进一步分组到一个线程块中。可以有几个线程块可以在GPU上并行执行。线程块绑定到SM。它在SM的执行周期内不能更改SM，必须在同一个SM上完成执行，并在完成后释放分配给它的所有资源。同一warp中的线程可以使用warp级别的混洗指令交换数据。线程块中的所有线程都可以使用低延迟共享内存进行通信，而不同线程块中线程需要使用高延迟全局内存进行通信。同步基元存在于线程块与warp级别。同步将确保线程块或warp中的任何线程（取决于所使用的同步类型）都不会继续执行下一条指令，直到所有线程都到达同步点。在数据首先写入共享内存，然后由所有线程读取的情况下，使用共享内存时需要同步。在读取与写入共享内存缓存之前，必须同步所有线程，以确保正确性。
2.3张量核
张量核是NVIDIA GPU上的可编程矩阵乘法与累加（MMA）单元。在Volta体系结构中首次引入，它们也出现在图灵与安培身上。显著高于CUDA核心的吞吐量，使其非常适合加速深度学习工作负载。他们执行MMA运算，表示为，D=A*B+C，其中运算大小在图灵与伏特上为4×4×4，而在安培上为8×4×8。张量核执行像HMMA这样的翘曲同步指令来执行MMA操作。Warp同步意味着Warp中的所有线程协同执行这些特殊指令，以产生MMA操作的输出。由于张量核指令的这种warp同步性质，在编程张量核时，有必要在warp级别而不是线程级别编写或生成代码。张量核最初只支持FP16用于输入，而支持FP16或FP32用于累加与输出。但现在它们支持多种输入与输出格式，如TF32、BF16、INT8与INT4。TF32具有与FP32相同的范围与与FP16相同的精度，但用19位表示。它可以用于精度不太重要的地方。要使用此模式，输入必须在FP32中，它们将在内部转换为TF32，在FP32中将进行累加，输出也将在FP32产生。这在CUDA核心上提供了比普通FP32模式更快的速度。BF16提供与FP32相同的范围，精度低于FP16。张量核在BF16与FP16模式下提供相同的速度，但两者都比TF32快。整数类型旨在用于训练后量化。
当谈到可编程性时，有三种方法可以利用张量核：
1）使用像cuBLAS这样的高级库；
2）在CUDA中使用高级C++WMMA API的程序；
3）使用汇编级指令显式编程。
张量核的参数与性能特征，见表8-5。
表8-5 张量核的参数与性能特征
|
|
性能
|
共享内存组性能冲突
|
易用性
|
支持算子融合
|
|
高级库
|
最佳
|
最小
|
用函数调用访问
|
有限
|
|
WMMA API
|
多数情况较优
|
较高
|
要求良好的编程工作
|
好
|
|
汇编
|
所有情况较优
|
最小
|
要求大量良好的程序
|
好
|
虽然cuBLAS可以仅使用函数调用来使用，但使用其他两种方法需要大量的编程工作。WMMA API提供更大的矩阵运算（16×16×16，32×8×16），以及加载与存储操作数矩阵的实用函数。将这些API函数转换为特定于GPU微体系结构的汇编指令的任务也转移到NVIDIA的专有编译中。使用WMMA API加载的矩阵一旦被加载到寄存器中，就具有不透明的布局，即哪个线程持有加载的矩阵的哪个元素（线程-数据映射）是未知的。由于这种不透明的性质，在进行诸如偏置添加之类的操作的融合时，需要一些特定的步骤，这些步骤需要了解线程数据映射。显式使用汇编指令编程张量核更具挑战性，因为程序员必须处理寄存器中的线程数据映射以及共享内存与寄存器之间的数据移动等复杂问题。表1总结了这些方法。
LLVM中的NVPTX后端将WMMA API函数公开为内部函数。这使得通过MLIR对张量核进行编程成为可能。这些内部函数与WMMA API函数一一映射，并在编程与使用方面表现出相同的行为。
3设计
将介绍管道设计。管道基于一系列优化与转换，这些优化与转换构成了GPU上快速matmul的配方。使用的配方与之前的一些工作中强调的配方非常接近。它们中共同的部分是两级阻塞，以最大限度地提高内存层次结构不同级别的重用性，矩阵乘法到GPU设备的下译流程，如图8-17所示。
在工作之前，MLIR中提供了一些支持，已经在管道中重复使用了这些支持，但缺少一些核心组件。主要是，MLIR中没有使用WMMA API对张量核进行编程所需的操作，而是引入这些操作的人。在必要时对现有的MLIR基础设施进行更改与添加。
图8-17 矩阵乘法到目标GPU张量核的下译流程
图8-17展示了所采取的下译流程。尽管可以有不同的下译流程来实现相同的目标，但只要要生成的目标核是仿射的，就应该选择通过仿射方言的下译流程。这在许多方面都有帮助，如快速内存缓存的创建与放置、循环平铺、展开阻塞、向量化、并行循环的检测以及同步屏障的放置，代码如下：
//第8章/mlir_Matmul_thread.c
算法1：两级平铺张量核Matmul。
全局内存: A[M][K] B[K][N], C[M][N];
共享内存: a_smem[tbm][tbk], b_smem[tbk][tbn];
寄存器: c_reg[wm][wn], a_reg[wm], b_reg[wn];
for threadBlockK ← 0 to M step tbk do
__syncthreads();
所有线程加载tbm × tbk块形式A到a_smem；
所有线程加载tbk × tbn块形式B到b_smem；
Warp加载wm×wn块形式C到c_reg[wm][wn]；
__syncthreads();
// wmmaM、wmmaN与wmmaK表示WMMA固有大小;
for warpK ← 0 to tbk step wmmaK do
for warpM ← 0 to wm step wmmaM do
Warp将A形式a_smem的片段加载到a_reg[warpM];
for warpN ← 0 to wn step wmmaN do
Warp将B形式b_smem的片段加载到b_reg[warpM];
c_reg[warpM][warpN] += a_reg[warpM] × b_reg[warpN];
end
end
end
Warp将c_reg[wm][wn]存储到c中的相应块；
End
尽管为了简洁起见在算法中没有突出显示，但高性能只能使用一组更多的优化来实现，这些优化包括以下内容：
1）在共享内存缓存中进行填充以减少库冲突；
2）寄存器平铺或warp平铺；
3）加载-存储向量化；
4）全局存储器加载延迟隐藏。
现在，将详细描述下译管道，讨论如何启用主要优化。
3.1起点
代码生成方法的起点是一个高级操作，如lmhlo.dot或linalg.matmul，或者简单地说是一个IR，它是从面向用户的编程模型中生成的，目标是线性代数方言。在前一种情况下，可以将op降低到三环路仿射matmul，而在后一种情况中，可以生成三环路仿射matmul。代码如下：
//第8章/mlir_affine_memred.asm
affine.for %i = 0 to %M {
affine.for %j = 0 to %N {
affine.for %k = 0 to %K {
%a = affine.load %A[%i, %k]: memref<8192×8192×f16>
%b = affine.load %B[%k, %j]: memref<8192×8192×f16>
%c = affine.load %C[%i, %j]: memref<8192×8192×f32>
%aq = fpext %a: f16 to f32
%bq = fpext %b: f16 to f32
%q = mulf %aq, %bq: f32
%co = addf %c, %q: f32
affine.store %co, %C[%i, %j]: memref<8192×8192×f32>
}
}
}
3.3创建与放置共享内存缓存
平铺完成后，下一步是创建共享内存缓存，并将其放置在正确的循环深度。使用
affineDataCopyGenerate实用程序来生成矩阵A与B的副本。在这里采用的方法与以前的一些工作略有不同。只为矩阵A与B创建共享内存缓存。由于C每次warp只加载一次，因此将其从全局内存直接流式传输到寄存器中。将C从全局存储器流到共享存储器，然后从共享存储器流到寄存器。通过共享内存流式传输C的基本原理是防止对全局内存的随机访问，并可能促进全局内存中的联合访问，这可能更有效。推测，情况可能并不总是如此，尤其是对于大的问题大小，因为C网格每次翘曲只加载一次。
此外，这种方法还需要使用动态分配的共享内存，因为C的最佳网格大小可以很容易地耗尽某些设备上48KB的静态共享内存限制。因此，必须动态分配用于保存网格的缓存，并且必须重用该缓存来存储所有三个操作数的网格。MLIR目前不支持动态分配共享内存。因此，将自己限制在静态分配的共享内存中，以避免代码生成器中的特定复杂性，这可能不值得付出努力。即使不这样做，在大多数情况下，也已经接近手动调优的库。创建共享内存缓存是故事的一部分，而确保以最小的库冲突进行共享内存访问是另一部分。银行冲突会显著降低共享内存的吞吐量。
避免库冲突的一种通用技术是在前导维度中填充共享内存缓存。通过将affineDataCopyGenerate生成的共享内存缓存的leadingDimension更改为leadingDimition+PpaddingF actor来实现同样的效果。这样做将更改共享内存缓存的底层布局图，以考虑前导维度的变化，并且IR的其余部分不需要更改。同样值得注意的是，可以在这里尝试不同的填充因子，看看什么执行得最好，但填充因子必须是8的倍数，即F16元素的128位。这是因为WMMA API的校准要求。
3.4生成WMMA操作
既然已经拥有了所需的所有基本东西，就可以继续生成gpu.subgroup_mma操作了。WMMA操作有不同的大小，在这项工作中使用了16×16×16版本的操作。在这里生成的操作应该取代已经存在的标量操作，并且必须相应地调整相应循环的循环步骤，代码如下：
//第8章/mlir_affine_map_memred.asm
#map0 = affine_map<(d0) -> (d0)>
#map1 = affine_map<(d0) -> (d0 + 64)>
#map2 = affine_map<(d0) -> (d0 + 128)>
module {
// A与B的共享内存缓存
memref.global "private" @b_smem_global: memref<64×136×f16, 3>
memref.global "private" @a_smem_global: memref<128×72×f16, 3>
func @main() {
...
affine.for %i = 0 to 8192 step 128 {
affine.for %j = 0 to 8192 step 128 {
// 对共享内存缓存的引用。
%b_smem = memref.get_global @b_smem_global: memref<64×136×f16, 3>
%a_smem = memref.get_global @a_smem_global: memref<128×72×f16, 3>
// 主k-循环
affine.for %k = 0 to 8192 step 64 {
// 复制B的循环
affine.for %copykk = #map0(%k) to #map1(%k) {
affine.for %copyjj = #map0(%j) to #map2(%j) {
%11 = affine.load %B[%copykk, %copyjj]: memref<8192×8192×f16>
affine.store %11, %b_smem[%copykk - %k, %copyjj - %j]: memref<64×136×f16, 3>
}
}
// 复制A的循环
affine.for %copyii = #map0(%i) to #map2(%i) {
affine.for %copykk = #map0(%k) to #map1(%k) {
%11 = affine.load %A[%copyii, %copykk]: memref<8192×8192×f16>
affine.store %11, %a_smem[%copyii - %i, %copykk - %k]: memref<128×72×f16, 3>
}
}
affine.for %ii = 0 to 128 step 64 {
affine.for %jj = 0 to 128 step 32 {
affine.for %kk = 0 to 64 step 32 {
affine.for %kkk = 0 to 32 step 16 {
affine.for %iii = 0 to 64 step 16 {
affine.for %jjj = 0 to 32 step 16 {
...
%a = gpu.subgroup_mma_load_matrix %a_smem[%11, %12] {leadDimension = 72:
index}: memref<128×72×f16, 3> -> !gpu.mma_matrix<16×16×f16, "AOp">
%b = gpu.subgroup_mma_load_matrix %b_smem[%12, %14] {leadDimension = 136:
index}: memref<64×136×f16, 3> -> !gpu.mma_matrix<16×16×f16, "BOp">
%c = gpu.subgroup_mma_load_matrix %C[%16, %17] {leadDimension = 8192: index
}: memref<8192×8192×f32> -> !gpu.mma_matrix<16×16×f32, "COp">
%res = gpu.subgroup_mma_compute %a, %b, %c: !gpu.mma_matrix<16×16×f16, "AOp
">, !gpu.mma_matrix<16×16×f16, "BOp"> -> !gpu.mma_matrix<16×16×f32, "
COp">
gpu.subgroup_mma_store_matrix %res, %C[%16, %17] {leadDimension = 8192:
index}: !gpu.mma_matrix<16×16×f32, "COp">, memref<8192×8192×f32>
}
}
}
}
}
}
}
}
}
}
}
清单2：带有WMMA操作的平铺与填充仿射matmul
现在已经生成了WMMA操作，执行以下IR转换：
1）排列最外面的六个循环，从（i，j，k，ii，jj，kk）顺序到（i，j，ii，jj，k、kk）次序，将计算循环映射到GPU计算层次结构。此外，它还有助于将C上的不变加载存储操作移动到尽可能远的位置。
2）排列最里面的三个循环，从（i，j，k）到（k，i，j）。这代表了经向水平MMA操作，并增强了ILP。
3）完全展开最里面的三个环。
清单2显示了在创建WMMA操作之后得到的IR。应该注意到最里面的循环在这里被调整的步骤。该列表进一步显示了想要的排列中的循环嵌套。最外面的两个循环稍后将映射到栅格中的线程块，下面的两个环路将映射到warp。接下来的两个循环是对应于线块的k循环，然后是经线。展开后，进行了两次观察。
A）C矩阵上的运算现在变得独立于周围的两个循环，因此现在将C上的运算提升到最外层的k-循环。通过这种方式，防止在全局内存中重复加载与存储到C，并且只在线程块网格的处理开始与结束时执行这些操作
B）展开这些循环完全显示a与B上的所有负载。其中一些负载在k维上是相同的，通过应用CSE，可以完全消除冗余负载，并实现展开阻塞的效果，代码如下：
//第8章/mlir_thread_affine_map_memred.asm
...
#map0 = affine_map<(d0, d1) -> (d0 + d1)>
...
// Thread block ‘i‘ loop.
affine.for %i = 0 to 8192 step 128 {
// 线程块j循环
affine.for %j = 0 to 8192 step 128 {
%b_smem = memref.get_global @b_smem_global: memref<64×136×f16, 3>
%a_smem = memref.get_global @a_smem_global: memref<128×72×f16, 3>
// Warp ‘i‘ 循环
affine.for %ii = 0 to 128 step 64 {
// Warp ‘j‘ 循环
affine.for %jj = 0 to 128 step 32 {
// C上的提升加载
%11 = affine.apply #map0(%i, %ii)
%12 = affine.apply #map0(%j, %jj)
%c_reg_0 = gpu.subgroup_mma_load_matrix %C[%11, %12] {leadDimension = 8192: index}:
memref<8192×8192×f32> -> !gpu.mma_matrix<16×16×f32, "COp">
...
// 主k-循环，加载的C操作数为iter_args
%res:8 = affine.for %k = 0 to 8192 step 64 iter_args(%c_in_0 = %c_reg_0, %c_in_1 = %
c_reg_1...) -> (!gpu.mma_matrix<16×16×f32, "COp">, !gpu.mma_matrix<16×16×f32, "COp
">) {
...
%a = gpu.subgroup_mma_load_matrix %a_smem[%ii, %c_in_0] {leadDimension = 72: index}:
memref<128×72×f16, 3> -> !gpu.mma_matrix<16×16×f16, "AOp">
%b = gpu.subgroup_mma_load_matrix %b_smem[%c_in_0, %jj] {leadDimension = 136: index}:
memref<64×136×f16, 3> -> !gpu.mma_matrix<16×16×f16, "BOp">
%c_res = gpu.subgroup_mma_compute %a, %b, %c_in_0: !gpu.mma_matrix<16×16×f16, "AOp">,
!gpu.mma_matrix<16×16×f16, "BOp"> -> !gpu.mma_matrix<16×16×f32, "COp">
...
// 主k-产生当前迭代结果的循环
affine.yield %104, %107...: !gpu.mma_matrix<16×16×f32, "COp">, !gpu.mma_matrix<16
×16×f32, "COp">...
}
// C上的提升存储
gpu.subgroup_mma_store_matrix %res#0, %C[%11, %12] {leadDimension = 8192: index}: !gpu.
mma_matrix<16×16×f32, "COp">, memref<8192×8192×f32>
...
}
}
}
}
清单3：环路展开与恒载存储吊装后的仿射matmul。
上述优化之后的循环结构如清单3所示。必须注意的是，在C矩阵的不变加载-存储对移动后，循环结构发生了怎样的变化。第20行中的affine.for运算表示main-k循环，现在被修改为将加载的C操作数作为循环iter_args。这些操作数将用作此循环中发生的乘法的累加器。在每次迭代之后，这个循环都会产生累积的乘积，这些乘积作为iter_args传递给下一次迭代。这些iter_args驻留在寄存器中，并在主k循环的不同迭代中重复使用。
3.5全局内存加载延迟隐藏
随着gpu.subgroup_mma操作与上一节中的一些其他优化的引入，正在朝着最终IR中的结构迈进。专注于在仿射方言本身没有任何gpu特定信息的情况下进行尽可能多的优化。在目前的IR中，在加载A与B的共享内存网格之前，无法开始计算。就延迟而言，全局内存加载是最昂贵的操作之一，因此消除操作数上的长等待时间非常重要。通过在迭代0中取出A与B的副本，并在迭代n-1中进行计算，来分割主k循环或线程块k循环。
副本就放在k循环之前，计算就放在它之后。在该循环中执行的计算的索引也需要移动以向前移动一次迭代。因此，计算发生在共享内存中已经可用的数据上，下一次迭代的负载已经发布。这个阶段的IR的结构，代码如下：
//第8章/mlir_thread_affine_map_memred_copy_load.asm
#map4 = affine_map<(d0) -> (d0)>
#map5 = affine_map<(d0) -> (d0 + 128)>
#map6 = affine_map<(d0) -> (d0 + 64)>
// k循环的迭代0的仿射复制循环
affine.for %copyk = 0 to 64 {
affine.for %copyj = #map4(%j) to #map5(%j) {
%35 = affine.load %B[%copyk, %copyj]: memref<8192×8192×f16>
affine.store %35, %b_smem[%copyk, %copyj - %j]: memref<64×136×f16, 3>
}
}
affine.for %copyi = #map4(%i) to #map5(%i) {
affine.for %copyk = 0 to 64 {
%35 = affine.load %A[%copyi, %copyk]: memref<8192×8192×f16>
affine.store %35, %a_smem[%copyi - %i, %copyk]: memref<128×72×f16, 3>
}
}
// Main k-loop.
affine.for %k = 0 to 8128 step 64 {
// 复制k循环的迭代%k+1的循环
affine.for %copyk = #map6(%k) to #map5(%k) {
affine.for %copyj = #map4(%j) to #map5(%j) {
%36 = affine.load %B[%copyk, %copyj]: memref<8192×8192×f16>
affine.store %36, %b_smem[%copyk - %k - 64, %copyj - %j]: memref<64×136×f16, 3>
}
}
affine.for %copyi = #map4(%i) to #map5(%i) {
affine.for %copyk = #map6(%k) to #map5(%k) {
%36 = affine.load %A[%copyi, %copyk]: memref<8192×8192×f16>
affine.store %36, %a_smem[%copyi - %i, %copyk - %k - 64]: memref<128×72×f16, 3>
}
}
affine.for %kk= 0 to 64 step 32 {
...
}
}
// k循环的最后一次迭代的仿射计算循环
affine.for %arg4 = 8128 to 8192 step 64 {
...
}
清单4:WMMA仿射matmul与移位的k-循环。
虽然这为延迟隐藏奠定了基础，但要看到这一点，需要将共享内存中的存储与线程块k循环内复制循环的全局内存负载解耦。这对于优化的正确性与功能都是必需的。为此，将复制循环与延迟存储展开为外部k循环中的后续操作。将这种优化推迟到流水线中的另一个点，因为需要一些特定于GPU的信息来实现它。
3.6插入同步屏障
已经完成了IR中的大部分部件的生成，这可能需要同步屏障。共享内存缓存将由线程块中的所有线程读取与写入，因此在写入这些缓存之前与之后进行同步是至关重要的。一般来说，这个过程也可以使用基于内存的依赖性分析来实现自动化，需要复制循环的静态信息来设置这些同步屏障。
3.7全局到共享拷贝向量化
虽然延迟隐藏会起到一定的作用，但它无法使实际拷贝运行得更快。众所周知，向量加载存储指令比其标量对应指令执行得更好，因为它减少了内存事务的数量，并且通常会更好地利用可用带宽。
使用MLIRX中已经存在的向量化实用程序。将全局上实用程序称为共享内存副本。可以使用这个实用程序来尝试不同的向量宽度。尝试了32、64与128位宽的向量，并找出了128位宽向量的最佳工作方式。向量化复制循环，代码如下：
//第8章/mlir_thread_affine_map_memred_cast.asm
...
#map4 = affine_map<(d0) -> (d0)>
#map5 = affine_map<(d0) -> (d0 + 128)>
#map6 = affine_map<(d0) -> (d0 + 64)>
...
// 全局内存memrefs的强制转换操作
%a_cast = memref.vector_cast %A: memref<8192×8192×f16> to memref<8192×1024×vector<8×f16>>
%b_cast = memref.vector_cast %B: memref<8192×8192×f16> to memref<8192×1024×vector<8×f16>>
// 共享内存memref的强制转换操作
%b_smem_cast = memref.vector_cast %b_smem: memref<64×72×f16, 3> to memref<64×9×vector<8×f16>,
3>
%a_smem_cast = memref.vector_cast %a_smem: memref<128×72×f16, 3> to memref<128×9×vector<8×f16>,
3>
// 向量化复制循环
affine.for %copyk = #map6(%k) to #map5(%k) {
affine.for %copyj = #map4(%j) to #map5(%j) step 8 {
%135 = affine.load %b_cast[%copyk, %copyj floordiv 8]: memref<8192×1024×vector<8×f16>>
affine.store %135, %b_smem_cast[%copyk - %k - 64, (%copyj - %j) floordiv 8]: memref<64
×17×vector<8×f16>, 3>
}
}
affine.for %copyi = #map4(%i) to #map5(%i) {
affine.for %copyk = #map6(%k) to #map5(%k) step 8 {
%135 = affine.load %a_cast[%copyi, %copyk floordiv 8]: memref<8192×1024×vector<8×f16>>
affine.store %135, %a_smem_cast[%copyi - %i, (%copyk - %k) floordiv 8 - 8]: memref<128
×9×vector<8×f16>, 3>
}
}
清单5：向量化复制循环。
3.8提取并行循环。
这是在仿射方言中所做的最后一步。使用MLIR中的isLoopParallel实用程序来查找所有并行循环，然后使用affineParallelize将它们转换为并行循环。这些并行循环稍后被处理并映射到GPU处理器层次结构，而顺序循环是唯一保留在内核中的循环。
3.9映射到GPU计算层次
前一步是仿射方言中的最后一步，之后立即转换为SCF方言。从SCF方言开始，要做的第一件事就是将并行循环映射到GPU计算层次结构。MLIR中用于映射的现有实用程序与过程不支持将循环映射到单个warp，这在情况下是必需的。扩展了实用程序与Pass，以添加对matmul的支持。理想情况下，应该概括这一步骤中使用的过程与实用程序，以处理各种各样的循环嵌套，将其作为未来的工作。采取了所有必要的措施来确保联合的全局内存访问，这对于有效的带宽利用率与更快地从全局内存复制到共享内存至关重要。映射完成后，最外面的两个循环将转换为GPU启动操作，接下来的两个环路将映射到warp，其余的计算环路实际上是连续的，并保持原样。
3.10完成延迟隐藏
在将负载与存储解耦之前，延迟隐藏是不完整的。为了在不在代码中引入任何复杂性的情况下实现这一点，首先在线程块k循环内完全展开复制循环，然后延迟存储，使其在计算完成后发生。所采取的方法与所指出的方法非常相似。IR的一般结构，代码如下：
//第8章/mlir_gpu_launch_block_thread.asm
gpu.launch blocks(%blockIdX, %blockIdY, %blockIdX) in (%arg6 = %c64, %arg7 = %c64, %arg8 = %c1)
threads(%threadIdX, %threadIdY, %threadIdZ) in (%arg9 = %c256, %arg10 = %c1, %arg11 = %c1)
{
...
%c_reg_0 = gpu.subgroup_mma_load_matrix %C[%26, %27] {leadDimension = 8192: index}: memref
<8192×8192×f32> -> !gpu.mma_matrix<16×16×f32, "COp">
...
// k循环的迭代0的仿射复制循环
scf.for %copy = %c0 to %c4 step %c1 {
...
}
scf.for %copy = %c0 to %c4 step %c1 {
...
}
gpu.barrier
// Main k-loop
%res:8 = scf.for %k = %c0 to %c8128 step %c64 iter_args(%c_in_0 = %c_reg_0, %c_in_1 = %c_reg_1
...) -> (!gpu.mma_matrix<16×16×f32, "COp">, !gpu.mma_matrix<16×16×f32, "COp">...) {
gpu.barrier
// k循环迭代i+1的全局内存加载
%a_next_iter_0 = memref.load %a_cast[%74, %81]: memref<8192×1024×vector<8×f16>>
%b_next_iter_0 = memref.load %b_cast[%94, %101]: memref<8192×1024×vector<8×f16>>
...
scf.for %kk = %c0 to %c64 step %c32 iter_args(%arg16 = %c_in_0, %arg17 = %c_in_1) -> (!gpu.
mma_matrix<16×16×f32, "COp">, !gpu.mma_matrix<16×16×f32, "COp"> {
...
}
gpu.barrier
// k循环迭代i+1的共享内存存储
memref.store %b_next_iter_0, %b_smem_cast[%51, %68]: memref<64×17×vector<8×f16>, 3>
memref.store %a_next_iter_0, %a_smem_cast[%150, %167]: memref<128×9×vector<8×f16>, 3>
...
}
gpu.barrier
// k循环的迭代n-1的仿射计算循环
scf.for %arg14 = %c0 to %c64 step %c32 {
...
}
gpu.subgroup_mma_store_matrix %res#0, %C[%26, %27] {leadDimension = 8192: index}: !gpu.
mma_matrix<16×16×f32, "COp">, memref<8192×8192×f32>
...
}
清单6：Gobal内存加载延迟隐藏。
这是优化的终点，也是在SCF方言中的最后一步。
3.11把一切放在一起
由于上一步是优化方面的最后一步，因此现在要设置生成的IR以供执行。MLIR中已经存在设计良好的基础设施来实现这一点，MLIR中的现有设计允许在单个MLIR文件中表示要在加速器上执行的IR，如GPU。IR将有两个部分：在CPU上运行的主机端组件与在GPU上运行的设备端组件或内核。主机端组件调用设备端组件，可以等待其执行完成，也可以继续执行其他任务。主机与设备侧组件的下降路径略有不同：
1）主机端编译：主机端代码转换为std方言，然后转换为llvm方言。在转换为llvm方言的过程中，通过MLIR的CUDA运行库提供的瘦包装器接口，从GPU方言（如GPU.launch）的操作被下译为对CUDA驱动程序与CUDA运行时API的函数调用。然后将MLIR转换为LLVM IR，并生成目标代码。最后，可通过MLIR cpu runner（使用MLIR的基于LLVM Orc的JIT）执行IR。
将要链接的共享库作为参数，其中可以提供与CUDA驱动程序API相对应的库。
2）设备端编译：设备端代码也被转换为std方言，然后转换为llvm与nvvm方言的混合。这又被转换为LLVM IR，由LLVM中的NVPTX后端转换为PTX。然后使用NVIDIA的编译器将PTX转换为cubin（CUDA二进制格式）。NVIDIA的编译器通过MLIR的CUDA驱动程序API调用。MLIR中的gpu-to-cubin过程可以访问驱动程序API，并为执行PTX-to-cubin编译与集成。扩展了该过程以获取其他选项，如优化级别与每个线程的最大寄存器数，这是编译PTX-to-cubin时所需的。
执行这些最后步骤的基础设施已经存在于MLIR中。虽然评估使用了MLIR JIT，但也可以使用类似的设置执行提前编译。介绍的GPU方言操作，GPU.subgroup_mma_load_matrix、GPU.subgroup _mma_store_matrix与GPU.subroup_mma_compute，都是开源的，并上传到官方LLVM/MLIR存储库。
4评估
介绍了内核的性能，并将其与CuBLAS 11.2进行了比较。该评估是在基于NVIDIA Ampere的Geforce RTX 3090上执行的，该Geforce RTX 3090安装在x86-64系统上，该系统具有AMD Ryzen Threadipper 3970X CPU，运行Ubuntu 20.04 LTS。为所有实验设置了以下参数：
1）SM时钟设置为白皮书中提到的所有实验的提升频率，即1695 MHz。
2）将自己限制为静态分配的共享内存，该内存等于48 KB。
3）每个线程的最大寄存器数设置为255。
使用NVIDIA Nsight系统进行计时，并仅考虑内核运行时来计算所获得的TFLOP。这适用于内核以及CuBLAS。考虑线程块级别网格与曲速级别网格的不同组合，并报告性能最佳的版本。所报告的性能是经过十次运行的平均值。
最后，考虑形式为C = AB + C的矩阵（所有三个矩阵都存储在一个行主布局中）。使用WMMA内在的m16n16k16版本，并将自己限制在1024到16384之间的平方问题大小，步长为256。假设问题大小是线程块的倍数，这也是经线块的倍数。经线块又是WMMA固有值的倍数。
4.1.混合精度性能
将介绍自动生成的混合精度内核的性能。在F16中具有A、B的矩阵乘法与在F32中进行的乘积的累加被称为混合精度matmul。输出矩阵C也在F32中，如图8-18所示。
图8-18 正方形大小矩阵上的混合精度（fp16输入、fp32累加与输出）性能。
性能始终在CuBLAS 11.2的95-119%以内。将性能与器件的绝对峰值进行比较，维持了器件峰值的95.4%。图2显示了在Ampere RTX 3090上自动生成的内核的性能，在评估的尺寸上非常接近CuBLAS。在一些较小的尺寸上，表现优于CuBLAS。对于较小的尺寸，CuBLAS内核可能不会像对于较大的尺寸那样进行很好的调整。在较大尺寸上，MLIR生成的代码在cuBLAS性能的2-8%以内，较小的线程块大小（如64×64×64）在较小的问题大小上表现更好。
自动代码生成方法通过选择性地启用或禁用优化来研究单个优化的影响。以增量的方式展示了图3中前面讨论的每个优化的影响，从初始版本到完全优化版本。
4.2半精度性能
将介绍自动生成的半精度内核的性能。在这个版本的matmul中，所有三个矩阵A、B与C都在FP16中。产品的积累也在FP16中完成。Thi版本通常比F32版本快，但由于尾数与指数的表示更窄，因此容易出现不精确性，性能始终是cuBLAS 11.2的80-160%，如图8-19所示。
图8-19 在启用与禁用各种优化的情况下，M=N=K=8192的混合精度性能（fp16输入、fp32累加与输出）。
在Ampere RTX 3090上自动生成的内核的性能，如图8-20所示。cuBLAS在整个范围内具有不一致的性能，尤其是在W=8848以上的问题大小上。这表明cuBLAS并没有很好地针对所有问题大小进行调整。在分析cuBLAS内核时，观察到cuBLAS选择的线程块大小，实际上小于具有最佳性能的线程块，例如，对于W=11264， cuBLAS会选择128×128×32，128×256×32。当cuBLAS使用五个阶段时，有一个管道阶段来隐藏全局内存加载的延迟。对于cuBLAS来说，全局内存加载的停滞要多得多。这可能是次优延迟隐藏的结果。
图8-20 在正方形大小的矩阵上使用fp16（fp16中的输入、累积与输出）的性能
5.小结
介绍了针对NVIDIA张量核支持的专用matmul指令的自动代码生成的早期结果。这些初步结果表明，在许多情况下，自动代码生成器可以实现与手动调优库相比具有竞争力的性能。在NVIDIA Geforce 3090 RTX（基于NVIDIA Ampere架构）上的实验结果证明了所提出方法的有效性。这些结果只是作为设计健壮代码与库生成器的垫脚石，这些代码与库生成程序不仅能够优化单个内核，而且能够实现内核的组合与融合。尽管已经做出了许多努力来允许通过DSL编译器或图形重写器进行融合与代码生成，但仍然缺少一种基于统一IR基础设施的稳健方法。
参考文献链接
https://arxiv.org/pdf/2108.13191.pdf
