在CUDA的天下，OpenAI开源GPU编程语言Triton，同时支持N卡和A卡

机器之心报道
编辑：蛋酱、陈萍
OpenAI开源了全新的GPU编程语言Triton ，它能成为CUDA的替代品吗？
过去十年中，深度神经网络(DNN)已成为最重要的机器学习模型之一，创造了从自然语言处理到计算机视觉、计算神经科学等许多领域的SOTA实现。 DNN模型的优势来自于它的层次结构，这一特征导致其计算量巨大，但也会产生大量高度并行化的工作，特别适合多核和众核处理器。
深度学习领域的新研究思路往往是结合原生框架operator来实现的，这种方法虽然方便，但需要创建或移动许多临时张量，因此可能会造成神经网络的性能损失。编写专门的GPU内核或许可以解决这个问题，但GPU编程的确是一件相当复杂的事。
DNN计算潜力与GPU编程困难之间的矛盾由来已久。英伟达在2007年发布了CUDA的初始版本， CUDA平台是一个软件层，使用者可以直接访问GPU的虚拟指令集和并行计算单元，用于执行计算内核。近年来，主流深度学习框架几乎都是基于CUDA进行加速，英伟达也一直在完善CUDA工具包，但对于一般的开发者来说， CUDA还是「不那么容易上手」。

文章图片
今天， OpenAI正式推出Triton1.0 ，这是一种类Python的开源编程语言。即使没有CUDA经验的研究人员，也能够高效编写GPU代码。例如，它可以用不到25行代码写出与cuBLAS性能相匹配的FP16矩阵乘法内核，后者是许多专业的GPU编程者尚且无法做到的。此外， OpenAI的研究者已经使用Triton成功生成了比PyTorch同类实现效率高2倍的内核。
代码地址：https://github.com/openai/triton
Triton的最初想法来源于现任OpenAI科学家的PhilippeTillet2019年在哈佛大学攻读研究生学位时发表的一篇论文，当时他的导师是H.T.Kung和DavidCox 。

文章图片
论文链接：http://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf
Tillet希望解决的问题是打造一种比英伟达的CUDA等特定供应商库更好用的库，能够处理神经网络中涉及矩阵的各种操作，具备可移植性，且性能可与cuDNN或类似的供应商库相媲美。团队表示：「直接用CUDA进行GPU编程太难了，比如为GPU编写原生内核或函数这件事，会因为GPU编程的复杂性而出奇困难。」
FacebookAI研究中心科学家SoumithChintala也在推特上表达了自己对Triton的期待：

文章图片
新发布的Triton可以为一些核心的神经网络任务（例如矩阵乘法）提供显著的易用性优势。「我们的目标是使其成为深度学习CUDA的可行替代方案，」PhilippeTillet作为Triton项目负责人如此表示。
GPU编程面临的挑战
现代GPU的架构大致可以分为三个主要组件：DRAM、SRAM和ALU 。优化CUDA代码时，必须考虑到每一个组件：
来自DRAM的内存传输必须合并进大型事务，以利用现代内存接口的总线位宽；必须在数据重新使用之前手动存储到SRAM中，并进行管理以最大限度地减少检索时共享内存库冲突；计算必须在流处理器（SM）内部或之间细致分区和调度，以促进指令/线程级的并行以及专用算术逻辑单元（ALU）的利用。
文章图片
GPU基础架构。
种种因素导致GPU编程难度骤增，即使对于具有多年经验的CUDA程序员也是如此。 Triton的目的是将这些优化过程自动化，以此让开发人员更专注于并行代码的高级逻辑。出于对泛用能力的考量， Triton不会自动调度跨流处理器的工作，而是将一些重要的算法考虑因素（例如tiling、SM间同步）留给开发者自行决定。