英伟达 CUDA 太难：OpenAI 出手要取代它，新语言性能相当但编程更简单

用 CUDA 为 GPU 编程实在太难了。

为了让没有 CUDA 编程经验的人写出和专家效率相当的 GPU 代码，现在 OpenAI 推出了一种新的语言和编译器 ——Triton。

它的难度比 CUDA 低，但是性能却可与之相媲美。

OpenAI 声称：

Triton 只要 25 行代码，就能在 FP16 矩阵乘法 shang 上达到与 cuBLAS 相当的性能。

OpenAI 的研究人员已经使用 Triton，来生成比同等 Torch 效率高出 1 倍的内核。

Triton 项目的负责人 Philippe Tillet 说：“我们的目标是使 Triton 成为深度学习 CUDA 的可行替代方案。”

25 行代码实现最佳性能

Triton 起源于 Tillet 在 2019 年学术会议 MLPF 上的一篇论文，当时他还是哈佛大学的一名研究生。

Tillet 解决的问题是如何开发一种 cuDNN 更具表现力的语言，既能够处理神经网络中涉及的矩阵的各种操作，同时兼具可移植性且以及和 cuDNN 相媲美的性能。

现代 GPU 大致分为三个主要组件 ——DRAM、SRAM、ALU，对这些资源进行调度管理十分复杂，即便是熟悉 CUDA 的程序员。

Triton 可以将这些优化过程完全自动化，让开发者可以更好地专注于并行代码的高级逻辑。

以矩阵乘法为例，能够为逐元素运算和归约编写融合内核很重要，但考虑到神经网络中矩阵乘法任务的重要性，这还不够。

Triton 非常适合这些应用，只需约 25 行 Python 代码即可实现最佳性能。

而另一方面，在 CUDA 中实现类似的过程需要花费更多的精力，甚至可能会降低性能。

手写矩阵乘法内核的一个重要优点是它们可以根据需要进行定制，以适应其输入和输出的融合变换。

如果没有 Triton，对于没有特殊 GPU 编程经验的开发者来说，矩阵乘法内核的修改是非常困难的。

Triton 背后的原理

Triton 的良好性能，来自于以 Triton-IR 为中心的模块化系统架构，这是一种基于 LLVM 的中间表示。

@triton.jit decorator 通过遍历提供 Python 函数的抽象语法树（AST），产生的 Triton-IR 使用通用 SSA 构建算法上的动态。

生成的 IR 代码随后由编译器后端进行简化、优化和自动并行化，然后转换为高质量的 LLVM-IR（最终转换为 PTX）。

研究人员发现，数据可以通过查看计算密集型块级操作（例如 tl.dot）的操作数自动存储到共享内存中，并使用标准活性分析技术进行分配/同步。

另一方面，Triton 程序可以通过同时执行不同的内核实例跨 SM 进行高效和自动并行化，以及通过分析每个块级操作的迭代空间，并在不同的 SIMD 中进行充分分区将 SM 内单元并行化。

目前 Triton 仅适用于英伟达 GPU，但官方表示 AMD GPU 以及 CPU 的版本正在开发中。

开源地址：

https://github.com/openai/triton

论文：

https://dl.acm.org/doi/abs/10.1145/3315508.3329973