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