简介¶

动机¶

在过去十年中，深度神经网络(DNNs)已成为机器学习(ML)模型中重要的一类，能够在从自然语言处理[SUTSKEVER2014]到计算机视觉[REDMON2016]再到计算神经科学[LEE2017]等多个领域实现最先进的性能。这些模型的优势在于其层次结构，由一系列参数化(如卷积)和非参数化(如修正线性)层组成。这种模式虽然以计算成本高昂著称，但也产生了大量高度可并行化的工作，特别适合多核和众核处理器。

因此，图形处理单元(GPU)已成为探索和/或部署该领域新研究理念的经济实惠且易于获取的资源。这一趋势因多个通用GPU计算框架(如CUDA和OpenCL)的发布而加速，这些框架使得开发高性能程序变得更加容易。然而，针对局部性和并行性优化GPU仍然极具挑战性，特别是对于那些无法通过组合现有优化原语有效实现的计算。更复杂的是，GPU架构也在快速演进和专业化，NVIDIA(以及最近的AMD)微架构中添加张量核心就是明证。

This tension between the computational opportunities offered by DNNs and the practical difficulty of GPU programming has created substantial academic and industrial interest for Domain-Specific Languages (DSLs) and compilers. Regrettably, these systems – whether they be based on polyhedral machinery (e.g., Tiramisu [BAGHDADI2021], Tensor Comprehensions [VASILACHE2018]) or scheduling languages (e.g., Halide [JRK2013], TVM [CHEN2018]) – remain less flexible and (for the same algorithm) markedly slower than the best handwritten compute kernels available in libraries like cuBLAS, cuDNN or TensorRT.

该项目的主要前提如下：基于分块算法的编程范式[LAM1991]可以促进神经网络高性能计算内核的构建。我们特别重新审视了传统的GPU"单程序多数据"(SPMD [AUGUIN1983])执行模型，并提出了一种变体，其中程序（而非线程）被分块。例如，在矩阵乘法的情况下，CUDA和Triton的区别如下：

CUDA编程模型

(标量程序，块状线程)

Triton编程模型

(块状程序，标量线程)

#pragma parallel
for(int m = 0; m < M; m++)
#pragma parallel
for(int n = 0; n < N; n++){
  float acc = 0;
  for(int k = 0; k < K; k++)
    acc += A[m, k] * B[k, n];

  C[m, n] = acc;
}

#pragma parallel
for(int m = 0; m < M; m += MB)
#pragma parallel
for(int n = 0; n < N; n += NB){
  float acc[MB, NB] = 0;
  for(int k = 0; k < K; k += KB)
    acc +=  A[m:m+MB, k:k+KB]
          @ B[k:k+KB, n:n+NB];
  C[m:m+MB, n:n+NB] = acc;
}

这种方法的一个关键优势在于，它形成了块结构化的迭代空间，相比现有领域特定语言(DSL)，在实现稀疏操作时为程序员提供了更大的灵活性，同时允许编译器针对数据局部性和并行性对程序进行激进优化。

挑战¶

我们提出的范式面临的主要挑战是工作调度问题，即如何划分每个程序实例的工作以在现代GPU上高效执行。为了解决这个问题，Triton编译器大量使用了块级数据流分析技术，这是一种基于目标程序控制流和数据流结构静态调度迭代块的方法。最终实现的系统表现惊人地出色：我们的编译器能够自动应用各种有趣的优化（例如自动合并、线程重排、预取、自动向量化、张量核心感知指令选择、共享内存分配/同步、异步拷贝调度）。当然，实现这些并非易事；本指南的目的之一就是让您了解其工作原理。