Graphene: An IR for Optimized Tensor Computations on GPUs
Graphene: An IR for Optimized Tensor Computations on GPUs
作者/机构: Bastian Hagedorn, Bin Fan, Hanfeng Chen, Cris Cecka, Michael Garland, Vinod Grover (均来自 NVIDIA)
A1 主要贡献
本文针对在现代 GPU 上表达优化张量计算的挑战，提出了一种名为 Graphene 的新型中间表示（IR）。
核心问题：
现代 GPU 在硬件层面提供了加速多维张量计算和数据移动的指令（如 Tensor Cores），但用于编程的软件（如 CUDA C++）仍然以一维内存中的扁平缓冲区为中心，缺乏对多维数据和线程的有效抽象。这使得即便是专家也很难编写和优化张量计算代码。现有的张量 IR 要么依赖于厂商库、要么对关键指令（如需要复杂数据到线程映射的指令）表达能力不足、要么依赖于复杂的内置编译器转换，难以扩展。
研究目标：
本文旨在引入 Graphene，一种用于在 GPU 上表示优化张量计算的 IR。Graphene 的目标是作为张量编译器和性能专家的一个底层目标语言，它比 CUDA C++/PTX 等提供同等控制水平的语言更贴近张量计算领域，同时又能明确地表示针对张量指令的高度优化 GPU 代码。
创新贡献：
本文做出了以下主要贡献：
1. 新颖的张量表示方法：引入了一种新的张量形状、布局和分块（tile）的表示方法。Graphene 的张量可以层次化地分解为分块（由更小的嵌套张量表示），并且能够表示复杂的非连续内存布局，这些对于支持 GPU 张量指令至关重要。
2. 逻辑线程组（Logical Thread Groups）：将 GPU 计算层次结构表示为一个处理元素张量。这种方法使得线程可以像数据张量一样被任意地重塑和分块，从而最大限度地减少了对内置计算层次结构的依赖。
3. 可分解的规约（Decomposable Specifications, Specs）：引入 Specs 作为集体计算和数据移动的抽象。Specs 将数据张量映射到线程张量，并能表示从设备级核函数到线程级可执行指令的各类计算。GPU 指令被暴露为描述其张量操作语义的原子规约（atomic specifications）。优化的 GPU 核函数通过将核函数级别的 Specs 层次化地分解为已知的原子 Specs 来表达。
4. 全面的实验评估：通过对深度学习中一些最重要的张量计算（包括 GEMM、多层感知机 MLP、Layernorm、LSTM 和融合多头注意力 FMHA）进行评估，证明了 Graphene 能够表达实现与现有库实现相同实践峰值性能所需的所有优化。在 Graphene 中表达的、超越库例程的融合核函数显著提升了 Transformer 网络的端到端推理性能，其性能与 cuBLAS(Lt)、cuDNN 和定制手写核函数相匹配或更优。
A3 关键观察与动机
优化 GPU 数据移动的挑战。考虑一个优化 GPU 数据移动的实现。NVIDIA 的 Ampere 架构 GPU 能够通过单一指令在内存层次结构中移动二维张量。具体来说，ldmatrix 指令使用一个 warp（32 个线程）从共享内存中移动多达四个 8×8 矩阵到该 warp 中线程的寄存器中。ldmatrix 在快速的库核函数中被大量使用，在 GEMM 核函数中若将其替换为等效但更简单的数据移动，会导致性能下降高达 17%。
ldmatrix 的严格数据到线程映射。ldmatrix 指令规定了严格的数据到线程的映射关系。图 1a 展示了每个线程必须在共享内存中访问的值，图 1b 展示了执行该指令后它在寄存器中接收到的值。从概念上讲，一个 warp 被分块为四个 8 线程组。每个 8 线程组被分配给共享内存中一个唯一的 8×8 分块。然后，每个线程访问共享内存中每个 8×8 分块的一行（图 1a），并相应地在寄存器中接收到每个 8×8 分块的两个相邻值，总共八个值（图 1b）。
现有编程模型的局限性。ldmatrix 仅在 PTX 中公开，需要如图 1c 所示的代码才能使用。这是当前表达优化 GPU 张量计算的最新技术水平，也是性能专家必须编写和编译器必须理解的代码类型。该指令的多维特性在图 1a 和 1b 中可以轻松可视化，但在代码中却被严重混淆了：在 CUDA C++ 中，将一个 warp 在概念上重塑为 2×2 个 8 线程组，必须表示为一组标量线程索引操作（第 1-3 行）。这些操作随后被用于访问一维共享内存缓冲区的标量索引（第 7,8 行）。由于 CUDA 和 PTX 处理地址空间的方式，我们还必须将对共享内存缓冲区的访问强制转换为一个有效的共享内存指针（第 9-13 行），然后调用 ldmatrix（第 16-19 行）。
现有 IR 的表达能力不足。由于缺乏多维张量抽象，这种优化张量数据移动的表示方式编写起来极其困难，也难以理解。据我们所知，现有的张量 IR 无法表示使用 ldmatrix 的数据移动，因为它们缺乏表达其所规定的数据到线程映射的能力。
Graphene 的表示方法。图 1d 展示了同样优化的数据移动在 Graphene 中是如何表示的。所需的数据张量和可用的 GPU 处理元素在前四行中声明。然后我们描述了一个数据移动（Move），它在每个线程块（#3）和每个 warp（#4）上执行，将一个 16×16 的共享内存张量（%1）移动到 2×4 的寄存器（%2）中（第 5 行）。注意寄存器是线程本地的，2×4 个值乘以 32 个线程总共对应 256 个值，足以容纳整个源共享内存缓冲区。花括号中的实现（第 7-19 行）明确应用了 ldmatrix 规定的映射（见图 1a 和 1b）。首先，我们将 warp 分块为 4 个逻辑线程组，每个组有 8 个线程，并将这些组重塑为 2×2 的形状（第 7-9 行）。然后，源张量被分块为四个 8×8 矩阵（第 12 行），每个线程组一个（第 13 行）。每个 8×8 分块再被分块为行（第 14 行），并分配给各个线程（第 15 行）。最后，目标张量也被分块（第 18 行），我们指定了另一个 Move 操作作用于先前定义的分块。这最后的 Move 匹配了 ldmatrix 指令的预定义语义，在 Graphene 中被认为是一个原子规约（atomic specification）（见 5.2 节），无需进一步实现。给定图 1d 所示的 IR，Graphene 会生成上文所示的 CUDA C++ 代码。
A2 方法细节
3 张量的未来形态
本节介绍 Graphene 的张量、它们的形状和布局。多维张量和将张量分解为分块（tile）在多面体编译中已有研究。然而，Graphene 的张量表示法特别适合表达优化的 GPU 计算，主要原因有二：1) 我们使用一种简洁的表示法来描述层次化分解的张量，其分块本身就是更小的嵌套张量。层次化分块对于将数据映射到 GPU 的多层计算和内存层次结构是必需的。2) 我们允许表达非平凡的形状，例如包含非连续元素的张量和具有超越标准“行/列主序”或“NHWC”布局的交错（swizzled）内存布局的张量。这类布局对性能至关重要，例如，当将中间张量存储到共享内存时。现有通常使用整数列表来指定张量形状和步幅的张量 IR 无法表达此类布局。Graphene 的张量表示法明确地捕捉了当今手写优化实现中使用的所有形状和布局。Graphene 的形状表示法受到 NVIDIA CuTe 编程模型【索引1，CUTLASS: Python API, Enhancements, and CUTLASS 3.0 Preview (Announcing the CuTe programming model)，2022，NVIDIA GTC Fall 2022】【索引24，CUTLASS，2023，GitHub】的启发，并建立在其形状代数【索引17，CUTLASS - CuTe documentation，2022，NVIDIA】之上。
3.1 在 Graphene 中表达张量
Graphene 的张量语法。图 2 展示了 Graphene 中张量的语法。我们暂时只关注数据张量。由该语法构成的表达式能够将前面章节讨论的张量表示为一等公民。在 Graphene 中，每个张量由一个名称、一个形状、一个元素类型和一个表示其存储在 GPU 内存层次结构中位置的标签组成。我们支持标准的 CUDA 内存区域：全局内存（片外）、共享内存（片上，由线程块内的线程共享）和寄存器（线程本地）。为简洁起见，在以下示例中我们省略了内存标签。一个形状由维度和步幅组成，两者都由整数元组表示。注意，我们在图 1 中使用了简化表示法，并省略了所有张量的步幅。例如，A:[16,16].fp16.SH 实际上表示为一个行主序张量 A:[(16,16):(16,1)].fp16.SH，其中冒号分隔了维度和步幅。Graphene 通过两种方式来表达高级张量：
高级张量的表达方式。
1. 层次化维度（Hierarchical Dimensions）: IntTuple 是递归定义的。这允许使用多个整数来定义单个维度的大小和步幅。这有两个目的：a. 它能够表达复杂的内存布局（3.2 节），b. 它能够表达含有非连续元素的张量（3.3 节）。
2. 分块（Tiles）: 张量的 ElementType 是递归定义的，可能是另一个嵌套的 Shape。我们用它来表示层次化分块的张量，其中外部（即左侧）形状表示分块的排列，而内部形状表示分块内元素的排列。
3.2 内存布局
标准内存布局。图 3 展示了将一个二维 4×8 张量在内存中布局的不同示例。图 3a 和 3b 展示了标准的列主序和行主序布局，它们通过相应的步幅指定。一个逻辑坐标 (i, j) 在物理一维内存中的位置是通过计算该坐标与张量步幅的点积得到的（如张量内灰色数字所示）。
现有 IR 的局限性。像这样用步幅来表示内存布局是现有张量 IR 的标准方式。然而，这种表示方法将可表达的内存布局限制为那些某一维度的所有元素严格出现在其他维度元素之前的布局。例如，在图 3b 所示的二维行主序布局中，我们固定一行，然后遍历所有列，再移动到下一行。仅有填充的布局（padded layouts），即某一维度的步幅超过前一维度的大小（例如，[(4,8):(9,1)]），在这种格式下也是可以表达的。
复杂布局的必要性。当需要将优化张量计算的中间结果临时存储到共享内存时，这种局限性尤其成问题。根据线程读写这些中间张量的方式，需要更复杂的布局才能尽可能高效地使用硬件。例如，GPU 上的共享内存是按 bank 组织的，每个 bank 一次只能服务一个线程。一旦多个线程试图访问存储在同一个 bank 中的不同值（即所谓的 bank 冲突），所有冲突线程的访问必须串行化，这会因内存延迟增加而导致流水线停顿，从而严重影响性能。实现峰值性能的优化核函数经常以比我们目前描述的简单布局更复杂的方式来布局张量。
Graphene 中的复杂布局。图 3c 和 3d 可视化了两种这样的复杂内存布局及其在 Graphene 中的表示。图 3c 展示了一个二维 4×8 张量，其第二维由一个整数元组而非单个整数表示。我们称之为层次化维度（hierarchical dimension）。层次化维度使我们能为每个维度指定多个步幅。在这种情况下，我们将两个相邻的列值在内存中连续布局，但之后在逻辑上先向下移动行，然后再布局接下来的两个相邻列值。图 3d 展示了一个类似但稍微更复杂的布局，它两次使用了层次化维度。
层次化维度的优势。关键在于层次化维度不会增加张量的秩（rank）。我们仍然可以用逻辑二维坐标访问图 3c 和 3d 中所示的张量，然后在内部计算相应的层次化坐标。这使得我们可以在张量分配时一次性优化和指定其布局。此后，无论张量在内存中如何布局，都可以保持使用相同的二维逻辑坐标来访问它。在 CUDA 中表达这类布局也是可能的，但需要复杂的索引算术，并且每次布局改变时，在整个核函数中每次访问张量时都必须进行调整。
3.3 张量分块
分块的重要性。为 GPU 表达优化的张量计算需要指定如何将多维张量映射到多层计算和内存层次结构。Graphene 表示这种映射的关键抽象是分块（tiles）。图 4 展示了一个基本的 4×8 张量（图 4a）以及三种不同的分块示例（b,c,d）。
常规连续分块。在 Graphene 中，分块被表示为嵌套的形状。一个分块张量的 ElementType 就是另一个描述分块的形状。图 4b 展示了图 4a 的一个分块版本，其中 B 描述了一个包含四个元素（以 2×2 排列的分块）的张量，每个元素是一个 2×4 的浮点元素分块。需要注意的是，按照惯例，所有形状的步幅都指定了内存中最内层标量类型元素之间的距离，以简化 Graphene 代码生成期间的索引计算。因此，图 4b 中外部（即左侧）形状的 (2, 16) 步幅指定了要移动到下一个行分块，必须跳过两个标量元素；要移动到下一个列分块，必须跳过 16 个元素。
指定分块大小。通常，分块大小是为每个维度指定一个整数，以说明每个分块在各维度包含多少元素。在 Graphene 中，分块大小是为每个维度使用一维张量来指定的。图 4b 中使用的分块大小 ([2:1],[4:1]) 被解释如下：第一维度中的两个逻辑上相邻的元素（[2:1]）和第二维度中的四个逻辑上相邻的元素（[4:1]）形成一个分块。按照惯例，我们省略单位步幅张量的步幅，可以如图 1 所示那样指定分块大小。请注意，这种指定分块大小的方式与待分块张量的内存布局无关。两个分块大小参数都指定了步幅为 1，意味着我们希望将逻辑上相邻的元素分组到同一个分块中，无论它们在物理内存中如何排列。最终得到的张量 B 的步幅取决于张量 A 的步幅，并会自动计算。
非连续分块。一旦我们需要指定在一个或多个维度上非连续的分块，使用张量作为分块大小参数就变得特别有用。例如，图 4c 展示了如何将张量 A 分块为 2×4 形状且在第一维度交错的分块。在可视化中，相同颜色的元素仍然属于同一个分块。将步幅从图 4b 中用于组合两个逻辑上相邻行的 [2:1] 增加到 [2:2]，现在创建的分块逻辑上包含了每隔一行的元素。结果张量 C 中的步幅也反映了这一变化。
多维非连续分块。最后，还可以使用具有层次化维度的一维张量作为分块大小参数，如图 4d 所示。这里，我们表示的分块同样包含每隔一行的元素，但现在列维度也是非连续的。具体来说，每个分块包含两个逻辑上相邻的列（[2:1]），并以步幅 4 重复两次（[2:4]），如颜色所示。请注意，我们已经在高度优化的 GPU 张量计算（如 ldmatrix 示例，见图 1）中看到过类似的二维非连续分块。
3.4 参数化形状和部分分块
参数化形状。尽管本文中的所有示例都显示了具体的整数值作为维度和步幅，Graphene 能够表示具有参数化（即符号化）形状的张量，例如 [M,N].fp32。这对于为具有动态形状的神经网络生成核函数尤为重要。参数化形状在代码生成期间会产生额外的核函数参数。使用参数化形状的索引表达式会通过代数简化规则进行简化，例如，若 $M < 256$，则 $(M \pmod{256}) \to M$。
部分分块的处理。当分块大小不能整除输入张量的维度时，会导致一个或多个所谓的部分分块（例如，用大小为 128 的分块去划分 [1023].fp32）。为了表示涉及部分分块的实现，Graphene 采用了 CuTe 的方法，对涉及的形状进行过近似（overapproximating）【索引18，CUTLASS - CuTe Predication，2022，NVIDIA】。随后对可能存在部分分块的张量的访问必须加上谓词（predicated），以防止越界访问。
4 逻辑线程组
不同线程排布的需求。高效的张量计算在核函数的不同阶段需要不同的线程排布。例如，使用 ldmatrix 进行数据移动时的线程排布与使用 Tensor Core mma.m8n8k4 指令【索引20，PTX ISA - SM70 mma-884-f16，2022，NVIDIA】计算矩阵乘法所需的排布显著不同。ldmatrix 指令，如第 2 节所讨论的，是由一个 warp 内的 8 个线程组执行的。然而，mma.m8n8k4 指令是由一个 warp 内另一组特定的八个非连续线程（称为 quad-pairs）执行的。在没有 Graphene 的情况下，必须通过一组标量索引计算来建立所有线程排布，并仔细协调线程如何映射到数据张量（例如，见图 1c）。
Graphene 的解决方案：线程即张量。在 Graphene 中，我们将 GPU 计算层次结构表示为一个张量。这种表示允许像处理数据一样操纵线程。对线程张量进行分块和重塑，使得线程排布可以表示为逻辑线程组。我们的方法有两大优势：
1. 显式形状（Explicit shapes）：逻辑线程组将线程的排布明确表示为一个多维分块张量，而不是使用多个标量线程索引计算。高性能核函数通常在单个核函数内需要不同的排布，导致大量此类索引转换。在 Graphene 中，我们只需使用不同分块方式的线程张量，并在代码生成时自动生成所需的标量索引表达式。
2. 灵活的层次结构（Flexible hierarchies）：特定的线程组仅在针对特定架构时才需要。例如，quad-pairs 是在 Volta 架构中引入的，而在更新的架构中又消失了。为所有架构（包括未来的架构）表示高性能代码是理想的，但又不希望为特定的计算层次结构添加内置支持。逻辑线程组使得任意线程排布的表达成为可能。
线程张量的语法。表示线程的张量语法与图 2 中表示数据张量的语法略有不同。线程张量的 ScalarType 是 thread 或 block，呼应了 Graphene 代码生成目标语言 CUDA C++ 中的两个基本层次结构。用于数据张量的 Memory 标签对于表示线程张量是不需要的，因此被省略。按照惯例，为了在视觉上区分线程张量和数据张量，我们使用 % 作为数据张量名称的前缀，使用 # 作为线程张量名称的前缀。
示例：ldmatrix 的线程排布。图 5 展示了我们如何在 Graphene 中表示第 2 节讨论的 ldmatrix 线程排布。灰色框中显示了在 CUDA C++ 中表示等效线程排布的相应标量线程索引表达式。我们从一个表示 warp 的一维 32 个连续线程的张量开始（图 5a）。这个 warp 被分块，就像我们上一节讨论的分块数据张量一样，分为四个 8 线程组，如图 5b 所示。之后，我们通过在分块张量的最外层（深度 0）应用 reshape 函数，将这四个分块重新排列成一个二维 2×2 的形状（图 5c）。在 Graphene 中，分块的线程张量允许其嵌套分块具有不同的秩。例如这里，分块是二维排列的，而分块内的线程是一维的。相应的索引表达式在 CUDA C++ 代码生成期间会自动计算。
示例：Volta 的 quad-pairs。图 6 展示了如何在 Graphene 中表示执行 Volta 的 mma.m8n8k4 指令所需的 quad-pairs。一个 quad-pair 是一个由 8 个线程组成的组，由两个特定的 quad（一组 4 个连续线程）构成。例如，第一个 quad-pair 由线程 0-3 和线程 16-19 组成。这种排布是由硬件规定的【索引20，PTX ISA - SM70 mma-884-f16，2022，NVIDIA】，在使用 Volta 的 Tensor Cores 时必须严格遵守，以避免未定义行为。在 Graphene 中，这种排布通过使用一个非连续张量 [(4,2):(1,16)] 对 warp 进行分块来表示，这精确地描述了 quad-pairs 的形状和布局。
5 规约与分解
本节我们讨论 Graphene 如何使用数据和线程张量来表示优化的张量计算。在本节余下部分，当我们使用“计算”一词时，它既指对张量的操作，也指数据移动。
5.1 使用规约表达计算
规约（Specifications, specs）的概念。对张量的计算由所谓的规约（Specifications，简称 specs）来表示。这个概念受到了 Fireiron【索引9，Fireiron: A Data-Movement-Aware Scheduling Language for GPUs，2020，PACT】的启发。其核心思想是，specs 封装了一个自包含的计算块，比如一个设备级的矩阵乘法核函数或一个 warp 级的数据移动。图 7 展示了在 Graphene 中使用 specs 的语法。一个 spec 包含了它的输入和输出张量，以及一个描述执行此计算可用线程的执行配置。可选的分解（decomposition）描述了此计算是如何实现的，可能包含简单的控制流或其他嵌套的 specs。
层次化分解。例如，我们通常从一个描述核函数级计算的 spec 开始。然后，我们通过将其分解为更细粒度的 specs（通常作用于数据和线程张量的分块）来逐步描述其实现，直到只剩下我们知道如何为其生成代码的 specs。这些剩下的 specs 被称为原子规约（atomic specifications），因为它们不需要进一步分解。例如，一个 GEMM 核函数可以被分解为“更小”的计算构建块，例如，作用于分块的线程块级矩阵乘法，以及这些分块在内存层次结构各级别之间的数据移动。图 1d 展示了一个具体例子。我们声明了一个从共享内存到寄存器的 warp 级数据移动（第 5 行），并逐步将其分解为另一个代表 ldmatrix 指令的嵌套数据移动 spec（第 19 行）。
5.2 内置规约
内置规约与原子规约。Graphene 提供了一套内置规约。内置规约描述一种特定类型的计算，而原子规约是内置规约的具体实例，描述由 GPU 指令实现的计算。例如，Graphene 提供了一个内置的 Move spec 来明确表示数据移动，同时还提供了一组与不同数据移动指令（如 ldmatrix）相关联的预定义原子 Move。
内置规约的种类。表 1 展示了 Graphene 的完整内置规约集。除了 Move，Graphene 还提供了其他内置规约，它们的原子版本映射到 ISA 中暴露的不同类型的指令。MatMul 表示类矩阵乘法的计算，原子 MatMul 映射到标量和向量化的融合乘加指令以及 Tensor Core 指令。Reduction 和 Unary/BinaryPointwise 表示符合预期的计算。Shfl 用于表示数据移动，但不是在内存层次结构级别之间（如 Move 所表示的），而是在特定的线程组之间。原子 Shfl 映射到 warp 级的 shfl.sync PTX 指令。最后，Allocate 用于在另一个 spec 的实现中引入新的临时数据张量，Init 用于将标量值统一赋给一个张量。
原子规约示例。表 2 展示了 Graphene 中原子规约的例子。在代码生成期间，每个没有分解的 spec 都会与目标架构的预定义原子规约集进行匹配。例如，每当我们遇到一个由单个线程执行的、从全局内存到寄存器的单个标量浮点值移动（表 2，第 1 行），我们就会发出 ld.global.u32 PTX 指令，该指令精确地实现了这个 Move。当我们看到一个从全局内存到寄存器的八个连续 fp16 值的 Move 时，我们发出向量化的 ld.global.v4.u32（第二行）。
张量指令的原子规约。像 ldmatrix 这样的张量指令不再由单个线程执行，而是由线程组协作执行。这些指令也不再操作于标量或一维向量，而是操作于多维张量。Graphene 的原子规约明确地捕捉了所需的线程排布和张量形状。例如，表 2 展示了两个代表 Tensor Core mma 指令的原子规约。它们需要不同的线程排布以及二维（分块的）输入和输出张量，所有这些都在我们的 IR 中明确指定。
5.3 表示融合核函数
通用规约（Generic Spec）。Graphene 能够表示 GPU 上的各种张量计算，包括那些不能由 MatMul 等内置规约之一表示的计算。例如，融合核函数在单个核函数中实现多个张量操作，例如 GEMM 后跟逐点操作。在现代深度学习网络中，融合核函数被大量使用以获得最佳性能。为了表示融合计算，我们使用一个通用的 Spec。通用规约描述了所需的输入和输出张量以及执行此计算的参与线程。该规约所代表的计算完全由其分解方式定义。
5.4 示例：一个简单的 GEMM 核函数
Graphene IR 的生成。图 8 展示了一个实现矩阵乘法的核函数级 spec 的最简单但完整的分解。由于其冗长和冗余的形状注解，Graphene IR 并不适合直接编写。目前，我们使用 Python API 生成 Graphene IR。未来，我们设想将 Graphene 集成到现有的深度学习编译器中，如 XLA【索引8，TensorFlow XLA，2017，Google】或 Triton【索引25，Triton: an intermediate language and compiler for tiled neural network computations，2019，MAPL@PLDI】，在这些编译器中，它可以作为 CUDA C++ 和 PTX 的替代目标语言。
简单 GEMM 核函数分解。Graphene 代码首先描述输入和输出张量以及执行此计算可用的块和线程（第 1-5 行）。最外层的 spec（第 6 行）代表 CUDA C++ 核函数。Graphene 还提供了基本的控制流语句，包括循环和 if 语句，以及其他不操作张量的表达式，如同步或屏障。在这种情况下，我们使用一个简单的三层嵌套 for 循环（第 9-11 行）来迭代每个线程计算的标量输出元素。我们为线程块进行分块（第 12-18 行），然后立即再次为线程进行分块（第 20-26 行）。最后，我们指定了每个线程的顺序标量计算（第 34 行）。这个 spec 不需要分解，因为它将匹配预定义的原子 hfma spec（比较内部 MatMul spec 的张量类型与表 2 中所示的原子 hfma spec）。更优化的 GEMM 实现会描述多个数据移动，并针对向量化和 Tensor Core 指令。
5.5 代码生成
代码生成过程。由于 Graphene IR 精确地描述了张量计算的实现，生成 CUDA C++ 代码就简化为将 IR 打印成有效的 CUDA C++。Graphene IR 可能包含 specs、张量操作或像循环和条件这样的控制流，以及其他不涉及张量的表达式。控制流语句、同步和屏障被发出为有效的 CUDA C++ 语法。一个没有实现的 spec 会与原子规约集进行匹配，我们会发出对相关指令的调用，如图 8 所示。对于有分解的 specs，我们递归地发出它们的实现；对于张量操作，我们构建抽象语法树（ASTs）并将它们编译成线程索引和缓冲区访问表达式。生成的索引会进行算术简化。
A4 实验环境
- 硬件配置:
- GPU: 实验在两款 GPU 上进行：一块 NVIDIA V100 (Volta 架构, SM70) 和一块 NVIDIA RTX A6000 (Ampere 架构, SM86)。
- 时钟频率: 使用 NVIDIA Nsight-Compute 性能分析器时，GPU 时钟被自动锁定到基础频率。
- 软件配置:
- CUDA: CUDA 11.7。
- 驱动: 510.68.02 版本。
- 依赖库:
- cuBLAS(Lt): 11.10 版本。
- 性能分析器: NVIDIA Nsight-Compute (2021.3.1.0 版本)。
- 模型与数据集:
- 实验中的所有张量计算默认使用 FP16 数据类型，并采用 FP32 进行 Tensor Core 累加。
- 算子层面: 评估了深度学习中的关键张量计算，包括通用矩阵乘法 (GEMM)、多层感知机 (MLP)、长短期记忆网络 (LSTM)、层归一化 (Layernorm) 和融合多头注意力 (Fused Multi-Head Attention, FMHA)。
- 端到端层面: 将 Graphene 生成的 FMHA 核函数作为自定义算子注入到多个 Huggingface Transformer 网络中进行推理性能测试。
A4 实验结果
实验旨在回答两个问题：1) Graphene 是否能表示性能与不同架构上的库实现相媲美的核函数？2) Graphene 是否能与超越库函数的手写融合核函数竞争？
假设 A: Graphene 能表示与高性能库实现竞争的核函数
- 实验内容 (GEMM): 将 Graphene 生成的优化 GEMM 核函数与 NVIDIA cuBLAS 库进行性能比较。实验在 Volta 和 Ampere 架构上进行，确保使用与 cuBLAS 相同的分块大小，并精确测量 GPU SM 的平均利用率。
- 实验结果 (GEMM): Graphene 生成的核函数在两个架构上都精确地匹配了 cuBLAS 的性能。通过分析器数据（图 9），可以看出这些核函数都是计算密集型的，Tensor Cores 以最大容量运行。
- 实验内容 (Fused GEMM): 比较了 Graphene 为 GEMM 加上融合的逐点操作（如加偏置或 ReLU 激活）生成的核函数与 cuBLASLt 提供的融合核函数的性能。
- 实验结果 (Fused GEMM): Graphene 生成的核函数再次在两个架构上都精确匹配了高度优化的库实现性能（图 10）。
- 分析结论: 假设 A 得到证实。Graphene 能够表达实现当今 GPU 上最高 GEMM 性能所需的所有优化。
假设 B: Graphene 能为重要的深度学习张量计算生成有竞争力的融合核函数
- 实验内容 (MLP): 将多个 MLP 层（GEMM + 加偏置 + ReLU）融合成单个 Graphene 核函数，并与调用多次单层 cuBLASLt 实现的累积性能进行比较。
- 实验结果 (MLP): 在特定问题规模下，通过避免中间张量往返于全局内存，Graphene 的融合核函数比 cuBLASLt 取得了高达 2.39 倍的性能提升（图 11）。
- 实验内容 (LSTM): 将一个简化的 LSTM 单元（两个独立的 GEMM、一个加法和两个逐点操作）融合成单个 Graphene 核函数。与两种基线进行比较：1) 使用 cuBLAS/cuDNN 的完全非融合实现；2) 使用 cuBLASLt 的部分融合实现。
- 实验结果 (LSTM): Graphene 的完全融合核函数比非融合基线在 Volta 和 Ampere 上分别取得了 1.75 倍和 1.82 倍的显著加速（图 12），因为它实现了超越现有库能力的额外融合。
- 实验内容 (Layernorm): 将 Graphene 生成的 Layernorm 核函数（仅包含逐点和归约计算）与多种 PyTorch 实现（Eager、JIT、内置融合算子）以及 NVIDIA Apex 提供的融合核函数进行比较。
- 实验结果 (Layernorm): Graphene 的性能与目前已知的最佳实现（NVIDIA Apex）相匹配（图 13），证明了其在非 GEMM 计算上的高性能表达能力。
- 实验内容 (FMHA): 针对 MLPerf BERT 推理的问题规模，实现了融合多头注意力（FMHA）核函数。与非融合的 cuBLAS 基线以及 NVIDIA 用于 MLPerf 提交的手写融合核函数进行比较。
- 实验结果 (FMHA): Graphene 能够生成最先进的融合核函数，并且由于优化的共享内存布局，甚至比 MLPerf 使用的核函数取得了微小的性能提升，相比非融合基线加速高达 2.92 倍（图 14）。
- 实验内容 (端到端 Transformer 推理): 将 Graphene 生成的 FMHA 核函数作为自定义算子注入到五个不同的 Huggingface Transformer 网络中。
- 实验结果 (端到端 Transformer 推理): 与常规 PyTorch 推理相比，使用 Graphene 的自定义核函数实现了高达 59% 的端到端性能提升，且加速比与网络中 FMHA 的出现频率相关（图 15）。
- 分析结论: 假设 B 得到证实。Graphene 能够为重要的深度学习计算生成匹配或超越最先进水平的融合核函数，并在实际应用中带来显著收益。
A5 结论
本文介绍了 Graphene，一种用于优化 GPU 张量计算的中间表示（IR）。Graphene 主要解决了表示问题：高度优化的 GPU 核函数必须使用 CUDA/PTX 编写，而这并非适合张量计算的 IR。通过引入一种更贴近张量计算领域的 IR，Graphene 为机器学习编译器和性能专家提供了一个替代的目标语言。
Graphene 的核心思想是将数据和线程都表示为一等公民的可分解张量。其 "specs"（规约）是统一的概念，用于表示从核函数到可执行指令的计算和数据移动。高性能 GPU 代码通过将核函数级别的计算分解为线程级别的可执行指令来表示。快速的张量指令被作为原子规约暴露出来，在代码生成阶段，Graphene 会生成使用内联 PTX 汇编的 CUDA C++ 代码。
实验结果表明，Graphene 能够达到与手动调优的库实现相媲美的性能，因此能够表示当今已知的最快核函数。此外，Graphene 还能表示尚不存在库例程的融合计算，并能匹配或超越手动开发的核函数。当这些生成的核函数部署到真实的深度学习网络中时，取得了显著的加速效果。
未来展望：Graphene 为新颖的机器学习编译器研究奠定了基础，包括系统性地推导优化的张量计算，以及生成高性能、针对特定架构的 GPU 核函数。
💬 评论讨论
欢迎在这里分享您的想法和见解！