并行计算任务的划分

在 Triton 编程中，任务的划分被认为是最为关键也最难优化的部分。在 CUDA 编程中，这两者都需要由程序员亲自规划，给编程带来了很大的阻碍。在此，有必要形式化地研究一下这个概念的含义。

并行计算的划分

在 GPGPU 中，一个任务可以由多种并行方式的组合进行完成：

• 线程并行（thread-level parallelism）：GPGPU 相比于 CPU 或者是 NPU 的最突出特点在于允许数量巨大的逻辑线程同时执行，提高计算的吞吐量，因此 GPGPU 编程又称为大规模并行（massively parallel）。GPGPU 的硬件设计也得益于丰富的线程并行，通过在多个并行的线程间切换来掩盖计算单元的延迟。因此，将任务拆解为线程并行通常是 CUDA 编程的第一课。
• 指令并行（instruction-level parallelism）：和 CPU 中类似的，GPU 中没有数据依赖的多条指令间可以多发射（但没有 CPU 中复杂的重排）来并行计算，而一条指令中也可以利用 SIMD 技术加速。指令并行同样是并行，可以用于达到在线程数相同的情况下，增加算法的并行数量。
• 迭代次数：由于线程的启动是有开销的，为了均摊线程启动的开销，会主动地选择更多的迭代次数采取线程粗化（thread coarsening）来选择在一个线程中处理多次任务。
• 串行：任何形式的并行都是有限度的，在超出并行的范围后，就只能采取串行的方法进行处理，例如 CUDA 中如果指定的线程数多于实际并行的能力，则多出的部分需要串行执行。

因此，一个并行计算任务需要划分为以下几个维度：

• 线程并行：需要指定一个线程以及对应的线程块（block）和线程束（warp）分别负责哪部分的任务。
• 指令并行：需要指定线程中一次计算处理哪部分的任务。
• 迭代次数：需要指定线程中要重复几次计算。
• 串行：前面分配完了，剩下的都只能串行处理，一般不分析。

上述的划分在实际的硬件中会受到诸多的限制，因此是程序参数调整的要点。例如，当程序所用到的数据可以复用时，通常数据的复用都伴随着迭代次数的增加。而当任务处理时没有充分应用到计算资源时，可以尝试增加线程并行和指令并行来提高硬件的利用率。在 Triton 编程中，任务划分无法自动推倒，需要由程序员自己进行任务的划分，也可以利用自动调参的接口来调整划分的参数。

示例

以基本的一维向量加为例，假设核函数的功能就是简单地从全局内存读取，计算后再写回全局内存，不使用共享内存，则进行编程前需要规划：

• 线程并行：向量加任务计算密度很低应该是完全的内存读写瓶颈，需要做好全局内存的访问合并。因此，一个线程束中的线程应该访问连续的内存，即连续的32个线程应该分到连续的32个任务。
• 指令并行：由于计算密度很低，可以不做指令并行。
• 迭代次数：如果发现每个线程的启动开销较大，可以每个线程多算几次。 如果选择迭代次数为1，则可能线程启动的影响较大；如果选择一个线程的连续几次迭代计算连续的几次任务，则全局内存的访问无法完全合并。

对于神经网络中常用的矩阵乘法而言，任务的划分更是重要。这里假设矩阵乘法的算法为一个线程块合作地读取矩阵的一块（tile）到共享内存，对这一块做完矩阵乘后，再将结果写回全局内存。矩阵计算以英伟达常用的C(m x n) = A(m x k) @ B(k * n)表示。这个任务的一种基本的划分方法为：

• 线程并行：把 m 和 n 分为若干块，每个线程块负责一块的计算。
• 指令并行：由于矩阵读取的延迟很高，可以做软件上的流水编排，在计算的同时读取后续需要的数据。
• 迭代次数：当 k 的维度大于一次计算的大小时，每次迭代沿着 k 的方向迭代，累加计算的结果。 当只划分 m 和 n 的情况下线程并行度还不够时，有把 k 维度在线程块间划分和线程块内划分的方法，称为 split-k 和 sliced-k。参考CUTLASS 的文档。 在上述算法中不难发现，全局内存的读写量和每一块的大小负相关。这说明我们应该尽可能地给每一块分配更多的共享内存，来服用矩阵数据。然而，GPGPU 中总的共享内存是有限的，当一个线程块用了过多的共享内存时，线程并行的程度就受限，无法掩盖访存延迟，需要在指令并行的维度上做流水编排。因此，矩阵惩罚最优的参数以及划分方式需要基于程序员的经验与实际运行状况的分析来确定。