英伟达CUDA太难!OpenAI出手要取代它,新语言性能相当但编程更简单
晓查 发自 凹非寺
量子位 报道 | 公众号 QbitAI
用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代码即可实现最佳性能。
@triton.jit
defmatmul(A, B, C, M, N, K, stride_am, stride_ak,
stride_bk, stride_bn, stride_cm, stride_cn,
**META):
# extract metaparameters
BLOCK_M, GROUP_M = META['BLOCK_M'], META['GROUP_M']
BLOCK_N = META['BLOCK_N']
BLOCK_K = META['BLOCK_K']
# programs are grouped together to improve L2 hit rate
_pid_m = tl.program_id(0)
_pid_n = tl.program_id(1)
pid_m = _pid_m // GROUP_M
pid_n = (_pid_n * GROUP_M) + (_pid_m % GROUP_M)
# rm (resp. rn) denotes a range of indices
# for rows (resp. col) of C
rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
# rk denotes a range of indices for columns
# (resp. rows) of A (resp. B)
rk = tl.arange(0, BLOCK_K)
# the memory addresses of elements in the first block of
# A and B can be computed using numpy-style broadcasting
A = A + (rm[:, None] * stride_am + rk[None, :] * stride_ak)
B = B + (rk [:, None] * stride_bk + rn[None, :] * stride_bn)
# initialize and iteratively update accumulator
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(K, 0, -BLOCK_K):
a = tl.load(A)
b = tl.load(B)
# block level matrix multiplication
acc += tl.dot(a, b)
# increment pointers so that the next blocks of A and B
# are loaded during the next iteration
A += BLOCK_K * stride_ak
B += BLOCK_K * stride_bk
# fuse leaky ReLU if desired
# acc = tl.where(acc >= 0, acc, alpha * acc)
# write back result
C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
mask = (rm[:, None] < M) & (rn[None, :] < N)
tl.store(C, acc, mask=mask)
而另一方面,在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
— 完 —
本文系网易新闻•网易号特色内容激励计划签约账号【量子位】原创内容,未经账号授权,禁止随意转载。
「智能汽车」交流群招募中!
欢迎关注智能汽车、自动驾驶的小伙伴们加入社群,与行业大咖交流、切磋,不错过智能汽车行业发展&技术进展。加好友请务必备注您的姓名-公司-职位哦~
点这里👇关注我,记得标星哦~
一键三连「分享」、「点赞」和「在看」
科技前沿进展日日相见~
关键词
代码
性能
内核
矩阵乘法
量子位
最新评论
推荐文章
作者最新文章
你可能感兴趣的文章
Copyright Disclaimer: The copyright of contents (including texts, images, videos and audios) posted above belong to the User who shared or the third-party website which the User shared from. If you found your copyright have been infringed, please send a DMCA takedown notice to [email protected]. For more detail of the source, please click on the button "Read Original Post" below. For other communications, please send to [email protected].
版权声明:以上内容为用户推荐收藏至CareerEngine平台,其内容(含文字、图片、视频、音频等)及知识版权均属用户或用户转发自的第三方网站,如涉嫌侵权,请通知[email protected]进行信息删除。如需查看信息来源,请点击“查看原文”。如需洽谈其它事宜,请联系[email protected]。
版权声明:以上内容为用户推荐收藏至CareerEngine平台,其内容(含文字、图片、视频、音频等)及知识版权均属用户或用户转发自的第三方网站,如涉嫌侵权,请通知[email protected]进行信息删除。如需查看信息来源,请点击“查看原文”。如需洽谈其它事宜,请联系[email protected]。