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

OpenAI 开源了新的GPU 编程语言Triton。它能成为CUDA的替代品吗？

在过去的十年中，深度神经网络(DNN) 已成为最重要的机器学习模型之一，在从自然语言处理到计算机视觉、计算神经科学等许多领域创造了最先进的实现。 DNN模型的优势来自于它的层次结构，导致计算量很大，但也有大量高度并行化的工作，特别适合多核和众核处理器。

深度学习领域的新研究思路往往是与原生框架算子结合实现的。虽然这种方法很方便，但它需要创建或移动许多临时张量，这可能会导致神经网络的性能损失。编写专门的GPU 内核可能会解决这个问题，但GPU 编程相当复杂。

DNN 的计算潜力与GPU 编程难度之间的紧张关系由来已久。 NVIDIA于2007年发布了CUDA的初始版本。CUDA平台是一个软件层，允许用户直接访问GPU的虚拟指令集和并行计算单元以执行计算内核。近年来主流的深度学习框架几乎都是基于CUDA进行加速，NVIDIA也一直在完善CUDA工具包。不过，对于普通开发者来说，CUDA 还是“没那么容易上手”。

今天，OpenAI 正式推出类似Python 的开源编程语言Triton 1.0。即使没有CUDA 经验的研究人员也可以高效地编写GPU 代码。例如，它可以用不到25行代码编写出与cuBLAS性能相匹配的FP16矩阵乘法内核，这是许多专业GPU程序员还无法实现的。此外，OpenAI 研究人员已使用Triton 成功生成内核，其效率比PyTorch 中的类似实现高2 倍。

代码地址：https://github.com/openai/triton

论文链接：http://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf

Tillet 希望解决的问题是创建一个比特定于供应商的库（例如Nvidia 的CUDA）更好的库，可以处理涉及神经网络中矩阵的各种操作，可移植，并且具有与cuDNN 或类似提供商相当的性能。相当于商业图书馆。该团队表示：“直接使用CUDA 进行GPU 编程太困难了。例如，由于GPU 编程的复杂性，为GPU 编写原生内核或函数将会出奇地困难。”

Facebook AI 研究中心科学家Soumith Chintala 也在Twitter 上表达了对Triton 的期待：

新发布的Triton可以为矩阵乘法等一些核心神经网络任务提供显着的易用性优势。 Triton 项目负责人Philippe Tillet 表示：“我们的目标是使其成为CUDA 深度学习的可行替代方案。”

GPU 编程挑战

现代GPU的架构大致可分为三个主要组件：DRAM、SRAM和ALU。优化CUDA 代码时，必须考虑每个组件：

来自DRAM 的内存传输必须合并为大型事务，以利用现代内存接口的总线宽度；数据在重用之前必须手动存储到SRAM 中，并进行管理以最大程度地减少检索期间的共享内存组冲突；计算必须在流处理器(SM) 内部或之间进行仔细的分区和调度，以促进指令/线程级并行性和专用算术逻辑单元(ALU) 的利用。 GPU基础设施。

各种因素使得GPU编程变得极其困难，即使对于具有多年经验的CUDA程序员来说也是如此。 Triton 旨在自动化这些优化过程，使开发人员能够更多地关注并行代码的高级逻辑。出于一般目的的考虑，Triton 不会自动跨流处理器调度工作，而是将一些重要的算法考虑因素（例如平铺、SM 间同步）留给开发人员自行决定。

CUDA 与Triton 编译器优化比较。

编程模型

在所有可用的特定领域语言和JIT 编译器中，Triton 可能与Numba 最相似：内核被定义为修饰的Python 函数，并在实例网格上与不同的program_id 同时启动。但区别值得注意：如下面的代码片段所示，Triton 通过对块进行操作来显示实例内并行性，其中块是一个维度为2 的幂的数组，而不是单指令多线程（SIMT）执行。模型。通过这种方式，Triton 有效地抽象了CUDA 线程块内与并发相关的所有问题（例如内存合并、共享内存同步/冲突、张量核心调度）。

Triton 中的矢量加法。

虽然这对于令人尴尬的并行（即逐元素）计算可能没有帮助，但它可以简化更复杂的GPU 程序的开发。例如，在融合softmax 内核的情况下，对于每个输入张量XR^MN，每个实例对给定输入张量的不同行进行归一化。此并行化策略的标准CUDA 实现可能很难编写，需要线程之间的显式同步，因为此策略同时减少X 的同一行。 Triton 在很大程度上消除了这种复杂性，每个内核实例加载感兴趣的行并使用类似NumPy 的原语按顺序对它们进行规范化。

import tritonimport triton.language as tl@triton.jitdef softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N): # 行索引m=tl.program_id(0) # col 索引# 仅此特定内核有效对于# 列少于BLOCK_SIZE 的矩阵BLOCK_SIZE=1024 n=tl.arange(0, BLOCK_SIZE) # 我们要加载的所有元素的内存地址可以计算如下X=X + m * stride_xm + n * stride_xn # 加载输入数据；用0 填充越界元素x=tl.load(X, mask=n N, other=-float('inf')) # 计算数值稳定的softmax z=x - tl .max(x, axis=0) num=tl.exp(z) denom=tl.sum(num, axis=0) y=num/denom # 写回Y Y=Y + m * stride_ym + n * stride_yn tl .store(Y, y, mask=n N)import torch# 分配输入/输出张量X=torch.normal(0, 1, size=(583, 931), device='cuda')Y=torch.empty_like(X )# SPMD启动gridgrid=( X.shape[0], )# 入队GPU kernelsoftmax[grid](Y, Y.stride(0), Y.stride(1), X, X.stride(0), X.stride(1), X. shape[0], X.shape[1]) 在Triton 中融合softmax

Triton JIT 将X 和Y 视为指针而不是张量。最重要的是，softmax 的这种特殊实现保留了行。这与PyTorch 的内部CUDA 代码不同，后者使用临时内存，使其更通用，但速度明显慢（见下图）。

Softmax 的A100 性能，M=4096。

Torch (v1.9) JIT 的较低性能凸显了从高级张量操作序列自动生成CUDA 代码的难度。

@torch.jit.scriptdef softmax(x): x_max=x.max(dim=1)[0] z=x - x_max[: None] 分子=torch.exp(x) 分母=numerator.sum(dim=1) 返回分子/分母[: None]合并softmax和Torch JIT

矩阵乘法

能够为逐元素运算和归约运算编写融合内核很重要，但考虑到矩阵乘法在神经网络中的重要性，这还不够。事实证明，Triton 在这些方面表现非常出色，仅用大约25 行Python 代码就实现了最佳性能。相比之下，CUDA 的效率并不高。

Triton 中的矩阵乘法。

手写矩阵乘法内核的一个重要优点是，它们可以根据需要进行定制，以适应其输入（例如切片）和输出（例如Leaky ReLU）的融合变换。如果像Triton 这样的系统不存在，那么对于没有出色GPU 编程专业知识的开发人员来说，矩阵乘法内核将很难进行彻底修改。

高层系统架构

Triton的良好性能得益于以Triton-IR为中心的模块化系统架构。 Triton-IR是基于LLVM的中间表示，其中多维值块是最重要的。

Triton 的高层架构。

@triton.jit 装饰器的工作原理是遍历Python 函数提供的抽象语法树(AST)，该函数使用常见的SSA 构造算法动态生成Triton-IR。然后，生成的IR 代码经过编译器后端的简化、优化和自动并行化，然后转换为高质量的LLVM-IR，最后转换为PTX，以便在最新的NVIDIA GPU 上执行。目前Triton 不支持CPU 和AMD GPU，但该团队表示对两者的支持正在进行中。

编译后端

研究人员发现，通过Triton-IR 使用块程序表示可以让编译器自动执行各种重要的程序优化。例如，通过查看计算密集型块级操作（例如tl.dot）的操作数，可以将数据自动存储到共享内存中，并使用标准活性分析技术进行分配和同步。

Triton 编译器通过分析计算密集型操作中使用的块变量的活动范围来分配共享内存。

此外，Triton 通过同时执行不同的内核实例来实现SM 之间的高效和自动并行化，以及通过分析每个块级操作的迭代空间并将其完全划分为不同的SIMD 单元来实现SM 内的高效和自动并行化。如下图：