【论文解读】基于MLIR的矩阵乘法高性能GPU代码生成：一些早期结果

0x0. 前言

本文是对 https://arxiv.org/abs/2108.13191 这篇论文进行解读，学习一下如何基于MLIR编译器基础设施生成高效的GPU代码。本文的阅读的先后顺序分别为：

1. 标题
2. 摘要
3. 引言
4. 结论
5. 背景
6. 设计
7. 实验
8. 评论

这篇论文是MLIR支持Tensor Core工作对应的论文，这篇论文涉及到的代码在llvm-project中已经开源。

0x1. 标题

本文题目为基于MLIR的矩阵乘法高性能GPU代码生成：一些早期结果。这说明论文可能还会继续完善，也许是实验或部分还要补充吧。作者团队是来自PolyMage Labs以及印度理工学院的。

0x2. 摘要

这篇文章介绍了使用MLIR编译器基础架构针对NVIDIA GPU上的Tensor Core生成代码的一些结果。当前高性能深度学习的最新技术主要由高度调优的库驱动。这些库通常由专业的程序员在low-level的级别进行手工优化和调优，并付出了很大的努力。对于类似的硬件或者将来可能出现的新硬件，可能需要重复很多这样的工作以及努力。因此，这个过程不像LLVM这样的编译器基础设施那样模块化以及可重用性很强。手工优化通常不使用IR，尽管这些优化可以被编码为一系列在IR上定义的pass。手工调优也可能会错过只有通过自动代码生成才可以轻松实现的一些优化点。本文认为，在引入MLIR之前，IR基础设施并不能有效地解决自动生成特定领域库的问题。特别是，很难使用单个IR来表示和转换高，中，低级别的抽象。

通过MLIR中的适当抽象，我们构建了一个实验性的递降(lowering)管道(pipline)，该管道可以自动生成基于GPU的Tensor core硬件的矩阵乘法的代码。在我们的实验中，初始性能结果表明，可以在NVIDIA的Ampere架构上对FP32和FP16分别达到CuBLAS性能的95-119%和80-160%的性能（显卡为Geforce 3090 RTX）。我们相信，这些结果可以作为使用IR基础设施进一步研究和开发为类似的专业加速器自动生成代码和库的动力。

0x3. 引言

深度学习和人工智能通常严重依赖于高性能计算。计算机硬件和微架构，库，编译器，运行时和编程模型的创新不断满足相关的计算需求。目前，大量高性能深度学习应用由硬件厂商提供的高度优化的库所支持，如CuDNN，CUBLAS和MKL（现在应该更名为oneDNN了）等。创建这些库需要大量的努力和专业知识，且这个开发过程可能必须在每种硬件或者软件版本中重复，并且可以有效地探索和优化的内容都是有限的。

矩阵乘法计算Kernel是许多基于Transformer（如Bert）架构的核心。它还可以作为了一个良好的测试样例来衡量可以实现的目标。虽然自动代码生成器的优势通常是优化Kernel的组合形式，而不是单个Kernel，但无法为研究充分的Kernel自动生成接近硬件峰值性能的代码难以让自动代码生成整个故事自洽。在本报告中，我们专门针对NVIDIA GPU Tensor Core，这是用于矩阵乘累加（MMA）操作的专用单元，其吞吐量通常是普通CUDA核心的3-4倍。

最近已经有一些工作聚焦于GPU Tensor Core上的 GEMM。Faingnaert 等人尝试通过在Julia创建一个三层API来解决这两种语言的问题，使得用户可以编写高效的GEMM Kernel。它们的主要关注点是开发足够灵活的API来满足各种应用程序的需求，而不是使用具有多个抽象级别的统一IR基础架构。Bhaskaracharya等人使用多面体代码生成的方法为Volta Tensor Core生成代码，它们使用调度树来表示计算并使用ISL[27]并为其生成CUDA代码。它们可以为MatMul和融合操作（如BiasAdd+ReLU）生成代码并实现高达2.55倍加速。这项工作是针对Volta的，包括一些特定于硬件的优化以实现高性能。Tillet等提出了Triton，一种IR和神经网络计算优化编译器。该框架基于tile的概念，tile是一个静态的多维数组。Triton编译器被暴露为一个Python包，它允许用户编写Python代码，编译器将自动生成高效的机器码。这项工作同时支持CUDA和Tensor Core并取得了很好的性能。

本文的方法使用编译器中间表示（IR）基础设施来做高性能代码库生成。这里使用矩阵乘法Kernel进行实验，以NVIDIA Tensor Core为目标后端。MLIR是我们在这里使用的编译器基础设施，其目标是在很大程度上使整个过程更加模块化，系统化和自动化。我们证明，通过逐级递降IR并应用正确的IR转换和优化，我们实现了和手写库相当的性能，而无需实际手动编写任何代码。虽然之前的工作对CPU单核的高性能实现进行了类似的研究，但我们这里的目标是专用的加速器。

本文贡献：

• 在 MLIR Dialect中引入 Warp Matrix Multiply Accumulate (WMMA) [13] Operation，并将它们递降到 LLVM/NVPTX 后端。
• 演示如何将 GPU 上的 matmul 系统地和渐进地生成为一系列 MLIR 变换和dialect loweing pass的代码。
• 构建针对Tensor Core的端到端matmul代码生成管道，初步结果表明，获得的性能与手动优化库的性能相当，在某些情况下加速达到1.60倍。

如果存在从此类模型到 MLIR 的递降(lowering)路径，我们这个基于 IR 的方法可以与不同的编程模型和语言一起使用。

本节强化了摘要，先列举了一系列和GPU Tensor Core GEMM的相关工作，这些工作有手工开发库的，也有Triton这种基于编译器的。然后作者引出本文的思路是基于MLIR这个基础设施来探索一下生成高性能的GPU Tensor Core GEMM代码，并列出了本文的贡献。（可以看出这篇论文是MLIR的一个应用，偏工程方向。

0x4. 结论

我们展示了针对NVIDIA Tensor Core支持的专用MatMul指令做自动代码生成的早期结果。这些初步结果表明，在许多情况下，自动代码生成器可以实现和手工优化库想媲美的性能。在NVIDIA Geforce 3090 PTX（基于NVIDIA Ampere架构）上的实验结果证明了本文方法的有效性。本文的研究只是设计鲁棒的代码库生成器的奠基石，它们不仅可以优化单个kernel，还可以实现kernel的组合和融合。这是一个众所周知的优化库有局限性的领域。虽然人们已经在通过DSL编译器或者图重写器来实现融合及代码生成方面付出了很多努力，但仍然缺少基于一个统一的IR基础设施的鲁棒性方法。

我总感觉这个结论怪怪的，似乎没说完。作者可能是想说，本文基于MLIR的方法让这种基于统一的IR基础设施针对特定加速器进行代码生成和优化成为了可能，并且因为接入了MLIR基础设施图重写也更加方便。

0x5. 背景

0x5.1 MLIR

MLIR的介绍这里就不多说了，我之前解读过MLIR的论文，感兴趣可以看看：MLIR：摩尔定律终结的编译器基础结构 论文解读 。

这个工作和MLIR的几个Dialect是有关的，这里再简要介绍一下这几种Dialect。

• Affine Dialect：这种Dialect使用来自多面体编译的技术使依赖分析和循环转换高效可靠。 我们已经在Affine Dialect级别进行了大部分优化和转换。
• GPU Dialect：MLIR中的GPU Dialect模拟了类似于CUDA或OpenCL的通用GPU编程范式。 它的目标是提供抽象来模拟 GPU 特定的操作和属性。它在很大程度上意味着与供应商无关。 一些附加信息可以在 [11, 12] 和 GPU Dialect文档 [16] 中找到。
• NNVM Dialect：由于我们专注于Tensor Core代码生成，我们使用和扩展 NVVM Dialect。 这种Dialect提供了直接映射到 LLVM 中的 NVPTX 后端的操作。
• LLVM Dialect：代码生成的最后阶段涉及递降到 LLVM IR，LLVM 后端从这里控制并生成目标代码。 为了对 LLVM IR 进行建模，使用了这种Dialect。 这是 MLIR 中存在的最低抽象级别。

0x5.2 GPU背景

GPU是通用的大规模并行计算设备。内存和计算层次结构在优化任何应用程序从而实现高性能方面发挥着重要作用。我们可以将GPU内存抽象为4级层次结构，global memory，L2-cache，可配置的L1-cache(shared memrory)，和寄存器。GPU上的处理器也可以抽象为两级层次结构，即流式多处理器（SM）和SM内的计算核心。计算核心通常也被叫作CUDA Cores。除了CUDA cores之外，tensor cores这种特殊单元也在较新的GPU中出现在和CUDA cores同一级别的计算层次结构中。每个SM被进一步划分为具有各自warp调度器的处理块。GPU的编程模型的结构也和当前的处理器层次结构相匹配。线程是GPU上可以和其它线程并行执行的单个执行实体。这些线程以32个为一组，叫作warp。warp在SM的计算核心上以锁步的方式执行。warp调度器选择一个准备好执行的warp并将其派发到compute cores。当一个warp碰到数据依赖时它会停止，并且warp调度器会选择另一个准备好执行的warp。

根据SM上要处理的block的数量，可能会并行执行多个warp。因此，一般而言，更多的wrap有助于实现：（i）warp级别的并行。(ii) 更好的延迟隐藏 （iii）更好的利用底层资源。现在，这些 warp 被进一步分组到一个线程块中。可以有多个线程块在 GPU 上并行执行。一个线程块会绑定到一个SM。它在执行的生命周期中不能更改SM，必须在同一个SM上完成执行，并在完成时释放分配给它的所有资源。同一个warp中的线程可以使用warp级别的shuffle指令交换数据。同一个线程块中的所有线程都可以使用低延迟的shared memory进行通信，不同线程块中的线程需要使用高延迟的global memoey进行通信。 同步源语存在于线程块和warp级别。根据所使用的同步类型，同步将确保线程块或warp中的任何线程都不会继续执行下一条指令，直到所有线程都到达同步点。在数据首先写入shared memory然后由所有线程读取的情况下，使用同步是必要的。在读取和写入shared memory缓冲区之前，所有线程必须同步，以确保正确性。

这段话是NVIDIA相关博客的缝合，对CUDA编程模型，执行模型以及内存模型进行了简要概述。

0x5. 3 Tensor Cores

Tensor Cores是NVIDIA GPU上的可编程矩阵乘法累加（MMA）单元。首先在Volta架构中引入，它们也出现在Turiong和Ampere架构上。显著高于CUDA cores的吞吐量使其非常适合加速深度学习的工作。它们执行表示为$D=A\ast B+C$的MMA操作，其中操作的尺寸在Turing和Volta架构上是$4\times 4\times 4$，而在Ampere上为 8 × 4 × 8 8 \times 4\times 8