用 CUDA 為 GPU 編程實在太難了。
為了讓沒有 CUDA 編程經(jīng)驗的人寫出和專家效率相當?shù)?GPU 代碼，現(xiàn)在 OpenAI 推出了一種新的語言和編譯器 ——Triton。
它的難度比 CUDA 低，但是性能卻可與之相媲美。
OpenAI 聲稱：
Triton 只要 25 行代碼，就能在 FP16 矩陣乘法上達到與 cuBLAS 相當?shù)男阅堋?/p>
OpenAI 的研究人員已經(jīng)使用 Triton，來生成比同等 Torch 效率高出 1 倍的內(nèi)核。
Triton 項目的負責人 Philippe Tillet 說：“我們的目標是使 Triton 成為深度學習 CUDA 的可行替代方案?！?/p>
25 行代碼實現(xiàn)最佳性能
Triton 起源于 Tillet 在 2019 年學術會議 MLPF 上的一篇論文，當時他還是哈佛大學的一名研究生。
Tillet 解決的問題是如何開發(fā)一種 cuDNN 更具表現(xiàn)力的語言，既能夠處理神經(jīng)網(wǎng)絡中涉及的矩陣的各種操作，同時兼具可移植性且以及和 cuDNN 相媲美的性能。
現(xiàn)代 GPU 大致分為三個主要組件 ——DRAM、SRAM、ALU，對這些資源進行調(diào)度管理十分復雜，即便是熟悉 CUDA 的程序員。
Triton 可以將這些優(yōu)化過程完全自動化，讓開發(fā)者可以更好地專注于并行代碼的高級邏輯。
以矩陣乘法為例，能夠為逐元素運算和歸約編寫融合內(nèi)核很重要，但考慮到神經(jīng)網(wǎng)絡中矩陣乘法任務的重要性，這還不夠。
Triton 非常適合這些應用，只需約 25 行 Python 代碼即可實現(xiàn)最佳性能。
而另一方面，在 CUDA 中實現(xiàn)類似的過程需要花費更多的精力，甚至可能會降低性能。
手寫矩陣乘法內(nèi)核的一個重要優(yōu)點是它們可以根據(jù)需要進行定制，以適應其輸入和輸出的融合變換。
如果沒有 Triton，對于沒有特殊 GPU 編程經(jīng)驗的開發(fā)者來說，矩陣乘法內(nèi)核的修改是非常困難的。
Triton 背后的原理
Triton 的良好性能，來自于以 Triton-IR 為中心的模塊化系統(tǒng)架構，這是一種基于 LLVM 的中間表示。
@triton.jit decorator 通過遍歷提供 Python 函數(shù)的抽象語法樹（AST），產(chǎn)生的 Triton-IR 使用通用 SSA 構建算法上的動態(tài)。
生成的 IR 代碼隨后由編譯器后端進行簡化、優(yōu)化和自動并行化，然后轉換為高質量的 LLVM-IR（最終轉換為 PTX）。
研究人員發(fā)現(xiàn)，數(shù)據(jù)可以通過查看計算密集型塊級操作（例如 tl.dot）的操作數(shù)自動存儲到共享內(nèi)存中，并使用標準活性分析技術進行分配/同步。
另一方面，Triton 程序可以通過同時執(zhí)行不同的內(nèi)核實例跨 SM 進行高效和自動并行化，以及通過分析每個塊級操作的迭代空間，并在不同的 SIMD 中進行充分分區(qū)將 SM 內(nèi)單元并行化。
目前 Triton 僅適用于英偉達 GPU，但官方表示 AMD GPU 以及 CPU 的版本正在開發(fā)中。
開源地址：
https://github.com/openai/triton
論文：
https://dl.acm.org/doi/abs/10.1145/3315508.3329973
廣告聲明：文內(nèi)含有的對外跳轉鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結果僅供參考，IT之家所有文章均包含本聲明。