相关工作¶
初看之下，Triton 可能只是又一个用于 DNN 的 DSL。本节旨在介绍 Triton 的背景，并强调它与该领域的两种主要方法的区别：多面体编译和调度语言。
多面体编译¶
传统编译器通常依赖于中间表示，例如 LLVM-IR [LATTNER2004]，它们使用（非）条件分支来编码控制流信息。这种相对低级的格式使得难以静态分析输入程序的运行时行为（例如，缓存未命中），也难以通过使用分块（tiling）[WOLFE1989]、融合（fusion）[DARTE1999] 和交换（interchange）[ALLEN1984] 来自动优化循环。为了解决这个问题，多面体编译器 [ANCOURT1991] 依赖于具有静态可预测控制流的程序表示，从而实现针对数据局部性和并行性的激进编译时程序转换。尽管许多用于 DNN 的语言和编译器（例如 Tiramisu [BAGHDADI2021]、Tensor Comprehensions [VASILACHE2018]、Diesel [ELANGO2018] 以及 MLIR 中的 Affine 方言 [LATTNER2019]）都采用了这种策略，但它也伴随着一些局限性，本节稍后将对此进行描述。
程序表示¶
多面体编译是一个广阔的研究领域。本节仅概述该主题最基本的方面，但对底层坚实数学基础感兴趣的读者可以参考关于线性和整数规划的大量文献。
for(int i = 0; i < 3; i++)
for(int j = i; j < 5; j++)
A[i][j] = 0;
|
|
多面体编译器关注一类通常称为 静态控制部分 (SCoP) 的程序，i.e.，连续语句的最大集合，其中条件和循环边界是围绕循环索引和全局不变参数的仿射函数。如上所示，这种格式的程序总是产生由仿射不等式界定的迭代域，即多面体。这些多面体也可以代数定义；对于上面的例子，
在 \(\mathcal{P}\) 中的每个点 \((i, j)\) 都代表一个多面体语句，即一种程序语句，它 (1) 不会引起控制流副作用 (例如，for、if、break)，且 (2) 在数组访问中只包含循环索引和全局参数的仿射函数。为了便于别名分析，数组访问也通过所谓的访问函数进行数学抽象。换句话说，A[i][j] 简单地表示为 A[f(i,j)]，其中访问函数 \(f\) 定义为
注意，SCoP (Static Control Part) 的迭代域并不指定其语句的执行顺序。事实上，这个迭代域可以按许多不同的合法顺序遍历，即调度。形式上，一个调度被定义为循环索引 \(\mathbf{x}\) 和全局不变参数 \(\mathbf{g}\) 的一个 p 维仿射变换 \(\Theta\)
其中 \(\Theta_S(\mathbf{x})\) 是一个 p 维向量，表示在遍历围绕 \(S\) 的循环嵌套时，增长最慢到最快 (从左到右) 的索引。对于上面所示的代码，C 语言循环嵌套定义的原始调度可以通过以下方式获取
其中 \(i\) 和 \(j\) 分别是嵌套中增长最慢和最快的循环索引。如果 \(T_S\) 是一个向量 (或张量)，则称 \(\Theta_S\) 是一维的 (或多维的)。
优点¶
适用于多面体编译的程序可以进行激进的变换和优化。这些变换大多数实际上归结为生成调度和迭代域，以支持循环变换，从而促进并行性以及空间/时间数据局部性 (例如，融合、交换、分块、并行化)。
多面体编译器还可以自动进行复杂的验证过程，以确保在整个优化阶段保留其输入程序的语义。注意，多面体优化器与更标准的优化技术并非不兼容。事实上，这些系统通常作为一组 LLVM Pass 实现，可以在更传统的编译技术之前运行 [GROSSER2012]。
总而言之，在适用时，多面体机制极其强大。已证明它可以支持大多数常见的循环变换，并且在密集矩阵乘法方面确实取得了与最先进的 GPU 库相当的性能 [ELANGO2018]。此外，它也是全自动的，除了 C 语言风格的源代码外，无需程序员提供任何提示。
局限性¶
不幸的是，多面体编译器存在两个主要局限性，阻碍了其成为神经网络代码生成的通用方法。
首先，可能的程序变换集合 \(\Omega = \{ \Theta_S ~|~ S \in \text{program} \}\) 很大，并且随着程序中语句数量及其迭代域大小的增加而增长。验证每次变换的合法性可能还需要求解复杂的整数线性规划问题，使得多面体编译计算成本非常高昂。更糟的是，硬件属性 (例如，缓存大小、SM 数量) 和上下文特性 (例如，输入张量形状) 也必须由该框架考虑，这会导致昂贵的自动调优过程 [SATO2019]。
其次，多面体框架的适用范围不广；SCoPs 相对常见 [GIRBAL2006] 但要求循环边界和数组下标是循环索引的仿射函数，这通常只发生在规则的密集计算中。因此，该框架仍需成功应用于稀疏 (甚至结构化稀疏) 神经网络，而这些网络的重要性在过去几年中迅速提升。
另一方面，本论文所倡导的基于块的程序表示形式，适用范围更广，并且可以使用标准数据流分析实现接近峰值性能。
调度语言¶
关注点分离 (Separation of concerns) [DIJKSTRA82] 是计算机科学中一个众所周知的经典设计原则：程序应该被分解为模块化的抽象层，这些层将算法的语义与其实现的细节分离开来。诸如 Halide 和 TVM 之类的系统将这一理念向前推进了一步，并通过使用调度语言在语法层面强制实现这种分离。这种方法的好处在矩阵乘法案例中尤其明显，如下所示，算法的定义 (第 1-7 行) 与其实现 (第 8-16 行) 完全分离，这意味着两者可以独立维护、优化和分发。
1// algorithm
2Var x("x"), y("y");
3Func matmul("matmul");
4RDom k(0, matrix_size);
5RVar ki;
6matmul(x, y) = 0.0f;
7matmul(x, y) += A(k, y) * B(x, k);
8// schedule
9Var xi("xi"), xo("xo"), yo("yo"), yi("yo"), yii("yii"), xii("xii");
10matmul.vectorize(x, 8);
11matmul.update(0)
12 .split(x, x, xi, block_size).split(xi, xi, xii, 8)
13 .split(y, y, yi, block_size).split(yi, yi, yii, 4)
14 .split(k, k, ki, block_size)
15 .reorder(xii, yii, xi, ki, yi, k, x, y)
16 .parallel(y).vectorize(xii).unroll(xi).unroll(yii);
然而，生成的代码可能不完全可移植，因为调度有时依赖于并不广泛可用的执行模型 (例如 SPMD) 或硬件 intrinsic (例如矩阵乘累加指令)。这个问题可以通过自动调度机制来缓解 [MULLAPUDI2016]。
优点¶
这种方法的主要优点是它允许程序员只写一次算法，然后分别专注于性能优化。这使得手动指定那些多面体编译器无法使用静态数据流分析自动发现的优化成为可能。
毫无疑问，调度语言是神经网络代码生成中最流行的方法之一。用于此目的最流行的系统可能是 TVM，它在广泛的平台上提供了良好的性能，并且内置了自动调度机制。
局限性¶
这种开发的便捷性是有代价的。首先，遵循这种范例的现有系统在现代硬件上适用时 (例如 V100/A100 Tensor Core 且分块大小相同时) 往往比 Triton 慢很多。我相信这不是调度语言的一个根本性问题——从可以通过更多努力来解决的意义上讲——但这可能意味着这些系统更难于设计和实现。更重要的是，现有的调度语言生成的循环，其边界和步长不能依赖于周围的循环索引，否则至少会对可能的调度施加严格限制——甚至可能完全破坏系统。这对于迭代空间可能不规则的稀疏计算来说是个问题。
for(int i = 0; i < 4; i++)
for(int j = 0; j < 4; j++)
float acc = 0;
for(int k = 0; k < K[i]; k++)
acc += A[i][col[i, k]] * B[k][j]
C[i][j] = acc;
|
|
另一方面，我们通过这项工作倡导的基于块的程序表示形式允许块结构的迭代空间，并允许程序员根据需要手动处理负载均衡。
参考文献¶
Lattner et al., “LLVM: a compilation framework for lifelong program analysis transformation”, CGO 2004
Wolfe, “More Iteration Space Tiling”, SC 1989
Darte, “On the Complexity of Loop Fusion”, PACT 1999
Allen et al., “Automatic Loop Interchange”, SIGPLAN Notices 1984
Ancourt et al., “Scanning Polyhedra with DO Loops”, PPoPP 1991
Baghdadi et al., “Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code”, CGO 2021
Vasilache et al., “Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions”, ArXiV 2018
Elango et al. “Diesel: DSL for Linear Algebra and Neural Net Computations on GPUs”, MAPL 2018
Lattner et al., “MLIR Primer: A Compiler Infrastructure for the End of Moore’s Law”, Arxiv 2019
Grosser et al., “Polly - Performing Polyhedral Optimizations on a Low-Level Intermediate Representation”, Parallel Processing Letters 2012
Sato et al., “An Autotuning Framework for Scalable Execution of Tiled Code via Iterative Polyhedral Compilation”, TACO 2019
Girbal et al., “Semi-Automatic Composition of Loop Transformations for Deep Parallelism and Memory Hierarchies”, International Journal of Parallel Programming 2006
Dijkstra et al., “On the role of scientific thought”, Selected writings on computing: a personal perspective 1982
Mullapudi et al., “Automatically scheduling halide image processing pipelines”, TOG 2016