[MAPL@PLDI'19] Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations 阅读笔记

发表于 2026-02-07 分类于 mlsys 阅读次数：本文字数： 6.1k 阅读时长 ≈ 22 分钟

#1. 引言

深度神经网络 (DNNs) 的近期复兴在很大程度上得益于可编程并行计算设备的普及. 特别是, 多核架构 (例如 GPU) 性能的持续改进发挥了基础性作用, 使研究人员和工程师能够探索种类日益增多、规模越来越大的模型, 并使用越来越多的数据. 这一努力得到了一系列供应商库 (cuBLAS、cuDNN) 的支持, 这些库旨在尽快将最新的硬件创新带给从业者. 不幸的是, 这些库仅支持有限的一组 Tensor 操作, 将创新原语的实现留给了专家.

这一观察推动了针对深度神经网络 (DNNs) 的多种领域特定语言 (DSLs) 的开发, 这些语言基于 polyhedral machinery (例如 Tensor Comprehension) 和/或循环合成技术 (例如, Halide、TVM 和 PlaidML). 然而, 虽然这些系统在特定问题类别 (如深度可分离卷积, 例如 MobileNet) 上通常表现良好, 但在实践中它们往往比供应商库慢得多 (参见图 1), 并且缺乏实现结构化稀疏模式的表达能力, 而这些模式不能直接使用嵌套循环中的仿射数组索引来指定.

alt text

图 1. $C = AB^T$ 不同实现与 Roofline 模型的性能对比 (NVIDIA GeForce GTX1070), 其中 $A \in \mathbb{R}^{1760 \times 1760}, B \in \mathbb{R}^{N \times 1760}$ , N 用于控制算术强度。

这些问题通常通过使用 Micro-Kernel (即手工编写的 Tile 级原语) 来解决–但这种方法需要大量的人工工作, 并且缺乏可移植性. 虽然最近已经提出了几种支持 Tiling 抽象的高级编程语言 (Cutlass, TiDA), 但 底层的编译器后端仍然缺乏对 Tile 级操作和优化的支持. 为此, 我们提出了 Triton, 一种开源的中级语言和编译器, 用于指定和编译 Tile 程序为高效的 GPU 代码.

本文的主要贡献总结如下:

Triton-C: 一种类似 C 的语言, 用于通过参数化 Tile 变量表达 Tensor 程序. 该语言旨在为现有的 DNN 转编译器 (例如 PlaidML、Tensor Comprehensions) 和熟悉 CUDA 的程序员提供稳定的接口. 列表 1 显示了与简单矩阵乘法任务相关的 Triton-C 源代码.
Triton-IR: 一种基于 LLVM 的中级表示 (IR), 它提供了一个适合 Tile 级程序分析、转换和优化的环境. 列表 5 显示了 ReLU 函数的 Triton-IR 代码. 在这里, Triton-IR 程序在解析过程中直接从 Triton-C 构建, 但未来也可以探索从嵌入式 DSL 或更高层次的 DNN 编译器 (例如 TVM) 自动生成.
Triton-JIT: 一种即时编译器 (JIT) 和代码生成后端, 用于将 Triton-IR 程序编译为高效的 LLVM 位码. 这包括 (1) 一套面向 tile 级别、与机器无关的优化过程, 旨在独立于任何编译目标; (2) 一套针对 Tile 级别的、与机器相关的编译过程, 用于生成高效的 GPU 准备好的 LLVM-IR; (3) 一个自动调优器, 用于优化与上述编译过程相关的任何元参数.
数值实验: 对 Triton 的数值评估, 展示了其能力在于:
1. 在循环和 Transformer 神经网络中生成与 cuBLAS 相当且比其他 DSL 快 3 倍的矩阵乘法实现
2. 在不损失性能的情况下重新实现 cuDNN 的密集卷积算法
3. 为诸如移位卷积等新型研究思想创建高效的实现.

alt text

图 2. Triton 概述

表1. 在 Triton-C 中计算 $C = A \times B^T$

// Tile shapes are parametric and can be optimized by compilation backends
const tunable int TM = {16, 32, 64, 128};
const tunable int TN = {16, 32, 64, 128};
const tunable int TK = {8, 16};

// C = A * B.T
kernel void matmul_nt(float* a, float* b, float* c, int M, int N, int K)
{
    // 1D tile of indices
    int rm[TM] = get_global_range(0);
    int rn[TN] = get_global_range(1);
    int rk[TK] = 0 ... TK;

    // 2D tile of accumulators
    float C[TM, TN] = 0;

    // 2D tile of pointers
    float* pa[TM, TK] = a + rm[:, newaxis] + rk * M;
    float* pb[TN, TK] = b + rn[:, newaxis] + rk * K;

    for (int k = K; k >= 0; k -= TK) {
        bool check_k[TK] = rk < k;
        bool check_a[TM, TK] = (rm < M)[:, newaxis] && check_k;
        bool check_b[TN, TK] = (rn < N)[:, newaxis] && check_k;

        // Load tile operands
        float A[TM, TK] = check_a ? *pa : 0;
        float B[TN, TK] = check_b ? *pb : 0;

        // Accumulate
        C += dot(A, trans(B));

        // Update pointers
        pa = pa + TK * M;
        pb = pb + TK * N;
    }

    // Write-back accumulators
    float* pc[TM, TN] = c + rm[:, newaxis] + rn * M;
    bool check_c[TM, TN] = (rm < M)[:, newaxis] && (rn < N);
    @check_c *pc = C;
}

#2. 相关工作

深度学习框架和库的存在对于新型神经网络架构和算法的出现至关重要. 然而, 尽管线性代数编译器的分析和经验启发式方法取得了进步, 这些软件仍然不可避免地依赖于手工优化的子程序 (例如 cuBLAS 和 cuDNN). 这导致了各种用于 DNN 的 DSL 和编译器的开发, 通常基于三种不同的方法:

XLA 和 Glow 使用 Tensor 级别的 IR 将 Tensor 程序转换为预定义的 LLVM-IR 和 CUDA-C 操作模板 (例如, Tensor 收缩、逐元素操作等), 通过模式匹配的方法.
Tensor Comprehension 和 Diesel 使用 Polyhedral Machinery 参数化和自动化地将一个或多个 DNN 层编译成 LLVM-IR 和 CUDA-C 程序.
Halide 和 TVM 使用 循环合成 将 Tensor 计算转换为 loop nests, 同时支持用户微调调度 (尽管可能是参数化的).

相比之下, Triton 依赖于在传统编译管道中添加基于 Tile 级别的操作和优化. 这种方法提供了 (1) 比 XLA 和 Glow 更高的灵活性; (2) 与 TC 和 Diesel 相反, 支持非仿射 Tensor 索引; (3) 自动推断可能的执行调度, 否则必须手动指定给 Halide 或 TVM. Triton 的好处是以增加编程工作量为代价的–请参阅表 2, 其中包含这些 DSL 中矩阵乘法的实现.

表2. 在TF, PlaidML, TC 和 TVM 中计算 $C = A \times B^T$

C = tf.matmul(A, tf.transpose(B))   // TF
C[i, j: I, J] = +(A[i, k] * B[j, k]) ; // PlaidML
C(i, j) +=! A(i, k) * B(j, k) // TC
tvm.sum(A[i, k] * B[j, k], axis=k) // TVM

#3. Triton-C

Triton-C 的目的是为现有/未来的 DNN 转译器, 以及熟悉低级 GPU 编程的程序员提供一个稳定的前端. 在本节中, 我们描述了 Triton-C 的 CUDA 类似语法 (第 3.1 节)、其 Numpy 类似语义 (第 3.2 节) 及其单程序多数据流 (SPMD) 编程模型 (第 3.3 节).

#3.1 语法

Triton-C 的语法基于 ANSI C (更具体地说是 CUDA-C) 的语法, 但为了适应下两个小节中描述的语义和编程模型, 进行了修改和扩展 (参见列表 3). 这些更改可分为以下几类:

Tile 声明: 我们添加了用于声明多维数组的特殊语法 (例如, int tile[16, 16]), 以强调其与 ANSI C 中嵌套数组 (例如 int tile[16][16]) 的语义差异. Tile 形状必须是常量, 但也可以使用关键字使其参数化. 一维整数 Tile 可以使用省略号进行初始化 (例如, int range[8] = 0...8).

内置函数: 常见的 C 语法用于元素级数组操作 (如+、*等) 仍然保留, 但添加了多种内置函数 (如 get_global_range) 以支持 tile 语义 (第 3.2.1 节) 和 SPMD 编程模型.

广播: N 维 Tile 可以使用 newaxis 关键字沿着任何特定轴进行广播 (例如 int broadcast[8, 8] = range[:, newaxis] 用来堆叠列). 注意, 切片 tile 以检索标量或子数组是禁止的.

谓词: 在 Tile 操作中的基本控制流 (第 4.3 节) 是通过使用带 @ 前缀的谓词语句实现的.

表3. Triton-C 的语法扩展:

// Broadcasting semantics
slice       : ':'   | 'newaxis'
slice_list  : slice | slice_list ',' slice
slice_expr  : postfix_expr | expr '[' slice_list ']'
// Range initialization
constant_range : expr '...' expr
// Intrinsics
global_range    : 'get_global_range'  '(' constant      ')'
dot             : 'dot'               '(' expr ',' expr ')'
trans           : 'trans'             '(' expr ',' expr ')'
intrinsic_expr  : global_range | dot | trans
// Predication
predicate_expr : '@' expr
// Tile extensions for abstract declarators
abstract_decl : abstract_decl | '[' constant_list ']'
// Extensions of C expressions
expr          : expr | constant_range | slice_expr | intrinsic_expr
// Extensions of C specifiers
storage_spec  : storage_spec  | 'kernel'
type_spec     : type_spec     | 'tunable'
// Extensions of C statements
statement     : statement | predicate_expr statement

#3.2 语义

#3.2.1 Tile 语义

Triton-C 中内置 Tile 类型和操作 (即 Tile 语义) 的存在提供了两个主要优势.

通过隐藏 Intra-Tile Memory Coalecsing、缓存管理和特化硬件利用等重要的硬件细节来简化 Tensor 程序的结构.
为编译器自动执行这些优化打开了大门, 如第 5 节所述.

#3.2.2 广播语义

Triton-C 中的 tile 具有强类型特性, 某些指令静态地要求其操作数必须满足严格的形状约束. 例如, 除非先进行适当的广播, 否则标量不能与数组相加. 广播语义 提供了一套规则来隐式执行这些转换 (示例参见列表 4):

填充: 最短操操作数的形状用全 1 进行左填充, 直到两个操作数具有相同的维数.
广播: 两个作数会根据需要被复制多次, 直到它们的形状相同; 如果无法完成, 则会发出错误.

表 4. 广播语义的实际应用

int a[16], b[32, 16], c[16, 1];
// a is first reshaped to [1, 16]
// and then broadcast to [32, 16]
int x_1[32, 16] = a[newaxis, :] + b;
// Same as above but implicitly
int x_2[32, 16] = a + b;
// a is first reshaped to [1, 16]
// a is broadcast to [16, 16]
// c is broadcast to [16, 16]
int y[16, 16] = a + c;

#3.3 编程模型

在 GPU 上执行 CUDA代码是由 SPMD 编程模型支持的, 其中每个 Kernel 都与 launch grid 中一个可识别的 thread-block 相关联. Triton 编程模型类似, 但每个 Kernel 都是单线程的 – 尽管自动并行化 – 并与一组全局范围相关联, 这些范围在不同的实例之间有所不同 (见图 3). 这种方法导致 Kernel 更简单, 其中不包含 CUDA 类似的并发原语 (Shared Memory Synchronization、线程间通信等).

可以使用 get_global_range(axis) 内置函数查询与 Kernel 相关联的全局范围, 以便创建 Pointer 的 Tile, 如列表 1 所示.

alt text

图 3. CUDA 编程模型与 Triton 编程模型之间的差异

#4. Triton IR

Triton-IR 是一种基于 LLVM 的中间表示 (IR), 其目的是提供一个适合于 tile 级别程序分析、转换和优化的环境. 在这项工作中, Triton-IR 程序是在解析期间从 Triton-C 中直接构建的, 尽管它们也可以在将来直接从更高级别的 DSL 生成.

Triton-IR 程序和 LLVM-IR 程序具有相同的高级结构 (如第 4.1 节所述), 但前者还包括一些用于块级数据流 (第 4.2 节) 和控制流 (第 4.3 节) 分析的扩展. 这些新颖的扩展对于执行第 5 节中概述的优化以及安全地访问任意形状的 Tensor 至关重要, 如第 6 节所示.

#4.1 结构

#4.1.1 模块

在最高层次上, Triton-IR 程序由一个或多个称为模块的编译基本单元组成. 这些模块相互独立地编译, 最终由一个链接器聚合, 其作用是解析前向声明并适当地合并全局定义.

每个模块本身由函数、全局变量、常量和其他杂项符号组成 (例如, 元数据、函数属性).

#4.1.2 函数

Triton-IR 函数定义由返回类型、名称和可能为空的参数列表组成. 如果需要, 可以添加额外的可见性、对齐和链接规范符. 函数属性 (如内联提示) 和参数属性 (如只读、别名提示) 也可以指定, 允许编译器后端通过例如更好地利用只读内存缓存等方式执行更激进的优化.

该头后面是一个由基本块列表组成的数据体, 这些基本块之间的相互依赖关系形成了函数的控制流图 (CFG).

#4.1.3 基本块

基本块是直线代码序列, 其末尾只能包含所谓的终止指令 (即分支、返回).

Triton-IR 使用静态单赋值 (SSA) 形式, 这意味着每个基本块中的每个变量必须 (1) 只被赋值一次, 并且 (2) 在使用之前必须被定义. 通过这种方式, 每个基本块隐式地定义了一个数据流图 (DFG), 其不同路径对应于程序 SSA 表示中的 def-use 链. 这种形式可以直接从抽象语法树 (AST) 中创建.

#4.2 Tile 级别 DFA 支持

#4.2.1 类型

多维 Tile 是 Triton-IR 中数据流分析的核心, 可以使用与 LLVM-IR 中向量声明相似的语法进行声明. 例如, i32<8, 8> 表示对应于 8x8 32 位整数 Tile. 请注意, Triton-IR 中没有 tunable 关键字, 因此必须在生成程序之前求解参数化形状值. 在我们的案例中, 这是通过 Triton-JIT 的自动调优器 (第 5.3 节) 完成的.

#4.2.2 指令

Triton-IR 引入了一套 retiling 指令, 其目的是支持第 3.2.2 节中描述的广播语义:

reshape 指令使用其输入参数的数据创建指定形状的 Tile. 这特别适用于通过用 1 填充其输入形状来重新解释变量为高维数组, 为隐式或显式广播做准备.
broadcast 指令通过在其输入参数上沿大小为 1 的维度进行必要的复制来创建指定形状的 Tile --如图 4 所示.

常规的标量指令 (如 cmp, getelementptr, add, load) 被保留并扩展, 表达对 Tile 操作数的逐元素操作. 最后, Triton-IR 还暴露了用于转置 (trans) 和矩阵乘法 (dot) 的专用算术指令.

alt text

图 4. broadcast<3, 3> 指令

#4.3 对 Tile 级别控制流分析的支持

Triton-IR 中存在 Tile 级别操作所带来的一个问题是 Tile 内部发散控制流的不可表达性. 例如, 程序可能需要部分保护 Tile 级别 Load 以防止内存访问违规, 但这是无法通过分支实现的, 因为 Tile 元素不能单独访问.

我们提出通过使用 Predicated-SSA (PSSA) 形式和 ψ-函数来解决这个问题. 这需要在 Triton-IR 中添加两个指令类 (参见列表 6):

cmpp 指令与普通的比较(cmp)指令类似, 不同之处在于它们返回两个相反的谓词, 而不是一个.
psi 指令合并来自不同谓词指令流的指令.

; pt[i, j], pf[i, j] = (true, false) if x[i, j] < 5
; pt[i, j], pf[i, j] = (false, true) if x[i, j] >= 5
%pt, %pf = icmpp slt %x, 5
@%pt %x1 = add %y, 1
@%pf %x2 = sub %y, 1
; merge values from different predicates
%x = psi i32 <8, 8> [%pt, %x1], [%pf, %x2]
%z = mul i32 <8, 8> %x, 2

#5. Triton-JIT 编译器

Triton-JIT 的目标是将 Triton-IR 程序简化并编译成高效的机器码, 通过一系列与机器无关 (第 5.1 节) 和与机器相关的 (第 5.2 节) 的 passes, 并由自动调优引擎 (第 5.3 节) 支持.

#5.1 与机器无关的 passes

#5.1.1 预取

循环内部的 Tile 级内存操作可能会存在问题, 因为它们可能会引入严重的延迟, 如果没有足够的独立指令则无法隐藏. 然而, 通过检测循环并在必要时添加适当的预取代码, 可以直接在 Triton-IR 中缓解这个问题 (参见列表 7).

表 7. 自动预取

B0:
  %p0 = getelementptr %1, %2
B1:
  %p = phi [%p0, B0], [%p1, B1]
  %x = load %p
  ; increment pointer
  %p1 = getelementptr %p

B0:
  %p0 = getelementptr %1, %2
  %x0 = load %p0
B1:
  %p = phi [%p0, B0], [%p1, B1]
  %x = phi [%x0, B0], [%x1, B1]
  ; increment pointer
  %p1 = getelementptr %p, %3
  ; prefetching
  %x1 = load %p

#5.1.2 Tile 级窥孔优化

Triton-IR 中存在 Tile 级别的操作, 为窥视优化器提供了新的机遇. 例如, 可以使用恒等式 $X = (X^T)^T$ 简化任何 Tile X 的转置链. 我们相信, 未来还可以利用对角 Tile 等其他代数性质做优化.

#5.2 与机器相关的优化

我们现在介绍一套针对遵循图 5 所示高级模型的机器的优化过程. 具体而言, Triton-JIT 执行的优化包括: (1) Hierarchical Tiling, (2) Memory Coalecsing, (3) Shared Memory Allocation, (4) Shared Memory Synchronization.

#5.2.1 Hierarchical Tiling

嵌套 Tile 策略 (见图 5) 旨在将 Tile 分解为 Micro-Tile, 最终分解为 Nano-Tile, 以尽可能紧密地适应机器的计算能力和内存层次结构. 虽然这种技术在自动调优框架中常规使用, 但 Triton-IR 的结构使其能够自动枚举和优化任何可表达程序的有效的嵌套 Tile 配置 (且无需多边形机制).

alt text

图 5. Triton-IR 机器模型中的 Hierarchical Tiling

#5.2.2 Memory Coalecsing

当相邻线程同时访问邻近的内存位置时, 内存访问被认为是合并的. 这很重要, 因为内存通常是从 DRAM 中以大块的形式检索的.

由于 Triton-IR 程序是单线程的并且自动并行化的, 我们的编译器后端能够在每个微块内部内部排序线程, 以便在可能的情况下避免非合并的内存访问. 这种策略减少了加载一个块列所需的内存事务数量 (见图 6).

alt text

图 6. Uncoalesced 和 coalesced 的 DRAM 访问. 不同线程以不同颜色显示.

#5.2.3 Shared Memory Allocation

对于具有高算术强度的块级操作 (例如, dot), 可以将它们的操作数临时存储在快速共享内存中. Shared Memory Allocation 阶段的目的在于确定何时以及何地将块存储到该空间中. 如图 7 所示, 这可以通过首先计算每个感兴趣变量的活动范围, 然后使用 Algorithms for Compile-time Memory Optimization 文献中提出的线性时间存储分配算法来完成.

alt text

图 7. Shared Memory Allocation

#5.2.4 Shared Memory Synchronization

在我们的机器模型中, 对共享内存的读取和写入是异步的. Shared Memory Synchronization 阶段的目的是自动在生成的 GPU 源代码中插入 Barrier, 以保持程序的正确性. 这是通过使用前向数据流分析来检测读后写 (RAW) 和写后读 (WAR) 风险, 并使用以下数据流方程式来完成的:

$\mathrm{in}_s^{(\mathrm{RAW})} = \bigcup_{p \in \mathrm{pred}(s)} \mathrm{out}_p^{(\mathrm{RAW})} \\ \mathrm{in}_s^{(\mathrm{WAR})} = \bigcup_{p \in \mathrm{pred}(s)} \mathrm{out}_p^{(\mathrm{WAR})} \\ \mathrm{out}_s^{(\mathrm{RAW})} = \begin{cases} \varnothing, & \textbf{if } \mathrm{in}_s^{(\mathrm{RAW})} \cap \mathrm{read}(s) \neq \varnothing \\ \mathrm{in}_s^{(\mathrm{RAW})} \cup \mathrm{write}(s), & \textbf{otherwise} \end{cases} \\ \mathrm{out}_s^{(\mathrm{WAR})} = \begin{cases} \varnothing, & \textbf{if } \mathrm{in}_s^{(\mathrm{WAR})} \cap \mathrm{write}(s) \neq \varnothing \\ \mathrm{in}_s^{(\mathrm{WAR})} \cup \mathrm{read}(s), & \textbf{otherwise} \end{cases}$

#5.3 自动调优器

传统自动调优器通常依赖于手工编写的参数化代码模板, 以在预定义的工作负载上实现良好性能. 相比之下, Triton-JIT 可以通过简单地将与每个优化过程相关的元参数连接起来, 直接从 Triton-IR 程序中提取优化空间.

在本工作中, 仅考虑了 Hierarchical Tiling 过程, 导致每个 Tile 每个维度最多有 3 个 Tile 参数. 然后使用穷举搜索优化这些参数, 搜索范围为:

Tile 大小: 32 到 128 之间的 2 的幂次方;
Micro-Tile 大小: 8 到 32 之间的 2 的幂次方;
Nano-Tile 大小: 1 到 4 之间的 2 的幂次方.

未来可以使用更好的自动调优方法.

#6. 数值实验

在本节中, 我们评估了 Triton 在各种深度学习文献中的工作负载上的性能. 我们使用 NVIDIA GeForce GTX1070, 并将我们的系统与最新的供应商库 (cuBLAS 10.0, cuDNN 7.0) 以及相关的编译器技术 (AutoTVM, TC, PlaidML) 进行了比较. 在适用的情况下, 我们根据官方文档指南, 针对每个单独的问题大小对这些 DSL 进行自动调优.

#6.1 矩阵乘法

矩阵乘法任务的形式: $A = D \times W (D \in R, W \in R)$ 是神经网络计算的核心. 这里我们考虑了从 Recurrent (DeepSpeech2) 到 Transformer 的各种任务; 我们在图 8 中报告了它们的性能.

Triton 和 cuBLAS 通常性能相当, 并在某些任务上达到了设备峰值性能的 90%以上.
cuBLAS 在较浅的 Transformer 神经网络中仍然比 Triton 更快, 这得益于使用了一种 3D 算法, 该算法将深度归约分解为独立块, 以在 M 和 N 过小时提供更多的并行性.
否则, 现有的 DSL 比我们的解决方案慢 2-3 倍
当输入形状是 32 的倍数时, TVM 慢 <2 倍.

alt text

图 8. 矩阵乘法的性能

#6.2 卷积

CNN 是一类重要的机器学习模型, 需要得到 DSLs 和编译器的良好支持. 它们基于卷积层 (图 9a), 其作为矩阵乘法 (图 9b) 的实现是利用专用 Tensor 处理硬件的必要条件–然而现有的 DSLs 并未支持这一点. 在这里, 我们基准测试了 Triton 对 cuDNN 的 IMPLICIT_GEMM 算法 (第 6.2.1 节) 的重新实现, 并提供了第一个用于移位卷积的融合 Kernel (第 6.2.2 节). 我们使用指针增量查找表实现了这些例程, 如列表 8 所示.

alt text

图 9. 密集和移位卷积层 (a) 被视为矩阵乘法 (b)

#6.2.1 密集卷积

本节考虑的卷积层来自深度学习文献, 如下表所示.

	H	W	C	B	K	R	S	应用场景
Task 1	112	112	64	4	128	3	3	ResNet
Task 2	56	56	128	4	256	3	3	ResNet
Task 3	28	28	256	4	512	3	3	ResNet
Task 4	14	14	512	4	512	3	3	ResNet
Task 5	7	7	512	4	512	3	3	ResNet
Task 6	161	700	1	8	64	5	5	DeepSpeech2
Task 7	79	341	32	8	32	5	10	DeepSpeech2

如图 10 所示, Triton 在 ResNet 的 IMPLICIT_GEMM 的 cuDNN 实现上表现更优. 这可能是因为 cuDNN 也维护了更好的 3 x 3 卷积算法 (即 Winograd), 对于重要性较低的核心留下的工程优化空间较小. 当快速算法不可用时 (例如 DeepSpeech2), cuDNN 和 Triton 表现相当.

alt text

图 10. 隐式矩阵乘法性能

#6.2.2 移位卷积

最后, 我们考虑表 1 中 Task1-5 的一种实现方式为移位卷积–这是一种针对 CNN 的新方法 (见图 9a). 我们将 Triton 中融合的移位-卷积模块实现 (见清单 8) 与依赖手写移位核心和单独调用 cuBLAS 的简单实现进行了比较. 我们还报告了在不移位时 (即 1×1 卷积) 可达到的最大性能. 如图 11 所示, 我们的 Triton 实现几乎完全隐藏了移位的成本.

alt text

图 11. Triton 中移位卷积的性能

表 8. Triton-C 中的移位卷积

const tunable int TM = {16, 32, 64, 128};
const tunable int TN = {16, 32, 64, 128};
const tunable int TK = {8};
__constant__ int* delta = alloc_const int[512];

for (int c = 0; c < C; c++)
    delta[c] = c * H * W + shift_h[c] * W + shift_w[c];

void shift_conv(restrict read_only float* a,
                restrict read_only float* b, float* c,
                int M, int N, int K) {
    int rxa[TM] = get_global_range[TM](0);
    int ryb[TN] = get_global_range[TN](1);
    int rka[TK] = 0 ... TK;
    int rkb[TK] = 0 ... TK;
    float C[TM, TN] = 0;
    float* pxa[TM, TK] = a + rxa[:, newaxis];
    float* pb[TN, TK] = b + ryb[:, newaxis] + rkb * N;
    __constant__ int* pd[TK] = delta + rka;

    for (int k = K; k > 0; k = k - TK) {
        int delta[TK] = *pd;
        float* pa[TM, TK] = pxa + delta[newaxis, :];
        float a[TM, TK] = *pa;
        float b[TN, TK] = *pb;
        C = dot(a, trans(b), C);
        pb = pb + TK * N;
        pd = pd + TK;
    }

    int rxc[TM] = get_global_range[TM](0);
    int ryc[TN] = get_global_range[TN](1);
    float* pc[TM, TN] = c + rxc[:, newaxis] + ryc * M;
    bool checkc0[TM] = rxc < M;
    bool checkc1[TN] = ryc < N;
    bool checkc[TM, TN] = checkc0[:, newaxis] && checkc1;
    @checkc * pc = C;
}

#7. 结论

在本文中, 我们介绍了 Triton, 一种开源的语言和编译器, 用于表达和编译 tile 化神经网络计算为高效的机器代码. 我们展示了在 LLVM-IR 中仅添加少量数据流和控制流扩展即可启用各种 tile 级优化过程, 这些过程共同带来了与供应商库相当的性能. 我们还提出了 Triton-C, 一种高级语言, 我们能够在其中简洁地实现针对 CNN 的新型神经网络架构的高效Kernel.

未来工作的方向包括对 Tensor Core 的支持、量化 Kernel 的实现以及集成到更高级别的 DSL 中.

#1. 引言

#2. 相关工作

#3. Triton-C

#3.1 语法

#3.2 语义

#3.2.1 Tile 语义

#3.2.2 广播语义

#3.3 编程模型

#4. Triton IR

#4.1 结构

#4.1.1 模块

#4.1.2 函数

#4.1.3 基本块

#4.2 Tile 级别 DFA 支持

#4.2.1 类型

#4.2.2 指令

#4.3 对 Tile 级别控制流分析的支持

#5. Triton-JIT 编译器

#5.1 与机器无关的 passes

#5.1.1 预取

#5.1.2 Tile 级窥孔优化

#5.2 与机器相关的优化

#5.2.1 Hierarchical Tiling

#5.2.2 Memory Coalecsing

#5.2.3 Shared Memory Allocation

#5.2.4 Shared Memory Synchronization

#5.3 自动调优器

#6. 数值实验

#6.1 矩阵乘法

#6.2 卷积

#6.2.1 密集卷积

#6.2.2 移位卷积

#7. 结论

#参考资料