过去十年，深度神经网络（DNN）已经成为最重要的机器学习模型之一，在计算机视觉和计算神经科学等许多领域创建SOTA实现。DNN模型的强大之处在于它的层次结构，导致计算量巨大，同时也产生大量高度并行化的工作，特别适用于多核和众核处理器。深度学习领域的新研究思路往往结合原生框架算子实现。虽然这种方法很方便，但它需要创建或移动许多临时张量，这可能会导致神经网络的性能损失。编写专用的GPU内核可能会解决这个问题，但GPU编程确实是一项相当复杂的工作。DNN的计算潜力与GPU编程难度之间的紧张关系长期存在。NVIDIA于2007年发布了CUDA的初始版本。CUDA平台是一个软件层，允许用户直接访问GPU的虚拟指令集和并行计算单元来执行计算内核。近年来，几乎所有主流的深度学习框架都是基于CUDA加速的，NVIDIA也一直在完善CUDA工具包，但对于普通开发者来说，CUDA仍然“没那么好用”。今天，OpenAI正式推出了Triton1.0，这是一种类Python的开源编程语言。即使没有CUDA经验的研究人员也可以高效地编写GPU代码。例如，它可以用不到25行代码编写出与cuBLAS性能相匹配的FP16矩阵乘法内核，这是许多专业GPU程序员目前还无法做到的。此外，OpenAI的研究人员已经使用Triton成功生成了比同类PyTorch实现效率高2倍的内核。代码地址：https://github.com/openai/tritonTriton最初的想法来自现任OpenAI科学家PhilippeTillet2019年在哈佛大学攻读研究生时发表的一篇论文，他当时的导师是H.T.Kung和大卫考克斯。论文链接：http://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdfTillet希望解决的问题是创建一个比特定供应商更好的库库，例如Nvidia的CUDA库，它处理涉及神经网络中矩阵的各种操作，其可移植性和性能可与cuDNN或类似的供应商库相媲美。该团队表示：“直接使用CUDA对GPU进行编程太难了。例如，由于GPU编程的复杂性，为GPU编写原生内核或函数将出奇地困难。”Facebook人工智能研究中心科学家SoumithChintala也在推特上表达了我对Triton的期待：新发布的Triton可以为一些核心神经网络任务（例如矩阵乘法）提供显着的易用性优势。“我们的目标是使其成为CUDA深度学习的可行替代方案，”Triton项目负责人PhilippeTillet说。GPU编程挑战现代GPU的架构可大致分为三个主要组件：DRAM、SRAM和ALU。在优化CUDA代码时，必须考虑每个组件：来自DRAM的内存传输必须合并为大型事务，以利用现代内存接口的总线宽度；在重新使用数据之前，必须手动管理存储到SRAM的操作，以最大限度地减少检索时共享内存库的冲突；必须在流处理器(SM)内或跨流处理器(SM)仔细划分和调度计算，以促进指令/线程级并行性和专用算术逻辑单元(ALU)的利用。GPU基础设施。多种因素使得GPU编程变得异常困难，即使对于具有多年经验的CUDA程序员也是如此。Triton的目标是将这些优化自动化，让开发者更专注于并行代码的高层逻辑。为了通用性，Triton不会自动调度跨流处理器的工作，而是将一些重要的算法考虑（例如平铺、SM之间的同步）留给开发者来决定。CUDA与Triton编译器优化比较。编程模型在所有可用的特定领域语言和JIT编译器中，Triton可能与Numba最相似：内核被定义为经过修饰的Python函数，并在实例网格上以不同的program_ids同时启动。但差异值得注意：如下面的代码片段所示，Triton通过对块进行操作来演示实例内并行性，其中块是维度为2的幂的数组，而不是单指令多线程(SIMT)执行模型。通过这种方式，Triton有效地抽象出了CUDA线程块内与并发相关的所有问题（例如内存合并、共享内存同步/冲突、TensorCore调度）。Triton中的矢量加法。虽然这可能无助于令人尴尬的并行（即逐元素）计算，但它可以简化更复杂的GPU程序的开发。例如，在融合softmax内核的情况下，对于每个输入张量X∈R^M×N，每个实例对给定输入张量的不同行进行归一化。这种并行化策略的标准CUDA实现可能很难编写，需要线程之间的显式同步，因为这种策略同时减少了X的同一行。Triton在很大程度上消除了这种复杂性，每个内核实例加载感兴趣的行并使用类似NumPy的原语按顺序对其进行规范化。importtritonimporttriton.languageastl@triton.jitdefsoftmax（Y，stride_ym，stride_yn，X，stride_xm，stride_xn，M，N）：#rowindexm=tl.program_id（0）#colindices#thisspecifickernelonlyworksformatricestthat#havelessthanBLOCK_SIZEcolumnsBLOCK_SIZE=1024n=tl.arange（0，BLOCK_SIZE)#thememoryaddressofalltheelements#thatwewanttoloadcanbecomputedasfollowsX=X+m*stride_xm+n*stride_xn#loadinputdata;padout-of-boundselementswith0x=tl.load(X,mask=n
