
C++
Triton主要包含两方面的意义:一方面,它能够全自动地处理GPU编程中的复杂问题,使开发者可以更加专注于算子逻辑的实现;另一方面,为了确保其具有广泛的适用性,Triton不会在流式多处理器(SM)之间自动进行任务调度,而是将这一决策权留给开发者。这与传统的CUDA编程模型有着本质区别。在CUDA中,采用的是单指令多线程(SIMT)的执行模式,而Triton则引入了一种新的编程范式——SPMB(Single Program Multi Block,这是我自己创造的一个术语)。这种范式更类似于使用numpy时的感觉,开发者可以直接描述张量之间的操作。不过,与numpy不同的是,在Triton中,开发者需要明确指出每个Block应完成哪些工作,也就是如何在各个SM之间分配任务,但完全无需考虑线程级别的细节。以两个矩阵相加为例,具体的kernel写法可以参考官方文档。在高性能算子的开发过程中,理想的做法是从较高的抽象层次开始描述问题,然后逐步降低抽象层次(lowering),即将模糊、笼统的描述转化为具体、清晰的操作步骤。每次经过一次lowering后,都可以对代码进行相应的优化,直到最终降至汇编语言为止。Triton的实现正是基于这样的思路。在ttir层面上,提供了最高层次的抽象,定义了基础的数据类型和一系列基本算子,包括常见的数学运算、reduce操作、矩阵点积计算以及简单的内存操作等。而在ttgir层面上,则实现了针对GPU硬件特性的抽象,统一支持异步数据拷贝、共享内存(SharedMemory)操作、矩阵乘加(MMA)操作等功能。通过结合ttir和ttgir,我们可以方便地描述一些优化算法,例如合并内存访问请求、避免共享内存中的bank冲突、在共享内存与寄存器之间实现双缓冲机制、优化计算单元的流水线结构等。ttgir会被进一步降低到llvmir,并生成最终的汇编代码。有关lowering过程的具体细节,请参阅本章第三节Triton编译过程分析(尚未完成)。相比之下,
C++的抽象层次非常低。如果直接用
C++编写kernel代码,上述提到的优化策略往往难以显式表达出来,而且针对不同的输入张量尺寸调整代码也显得十分困难。而这恰恰是Triton的优势所在。理论上,使用Triton编写的kernel可以无缝移植到任何受支持的后端硬件上,并且能够在这些硬件上保持高性能表现。