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


在CUDA的天下,OpenAI开源GPU编程语言Triton,同时支持N卡和A卡
文章图片
高级系统架构
Triton的良好性能得益于以Triton-IR为中心的模块化系统架构 。 Triton-IR是一种基于LLVM的中间表示 , 多维值块(blocksofvalues)是其中最重要的东西 。
在CUDA的天下,OpenAI开源GPU编程语言Triton,同时支持N卡和A卡
文章图片
Triton的高级架构 。
@triton.jit装饰器的工作原理是遍历由Python函数提供的抽象语法树(AST) , 这样一来就能使用通用的SSA构造算法实时生成Triton-IR 。 生成的IR代码随后由编译器后端进行简化、优化和自动并行化 , 然后转换为高质量的LLVM-IR , 最终转换为PTX , 以便在最新的NVIDIAGPU上执行 。 目前Triton还不支持CPU和AMDGPU , 但团队表示对二者的支持正在开发中 。
编译器后端
研究人员发现通过Triton-IR来使用块状程序表示 , 这种方法允许编译器自动执行各种重要的程序优化 。 例如 , 通过查看计算密集型块级操作(例如tl.dot)的操作数 , 数据可以自动存储到共享内存中 , 并使用标准的活跃性分析技术进行数据的分配与同步 。
在CUDA的天下,OpenAI开源GPU编程语言Triton,同时支持N卡和A卡
文章图片
在CUDA的天下,OpenAI开源GPU编程语言Triton,同时支持N卡和A卡】Triton编译器通过分析计算密集型操作中使用的块变量的活动范围来分配共享内存 。
此外 , Triton还可以在SM之间以及SM之内高效、自动地并行化 , 前者通过并发执行不同的内核实例来实现 , 后者通过分析每个块级操作的迭代空间 , 并将其充分划分到不同的SIMD单元来实现 。 如下所示:
在CUDA的天下,OpenAI开源GPU编程语言Triton,同时支持N卡和A卡
文章图片
Triton自动并行化 。 每个块级操作都定义了一个块级迭代空间 , 该空间可以自动并行化以利用SM(StreamingMultiprocessor)上的可用资源 。
参考链接:https://www.openai.com/blog/triton/