OpenAI推新語(yǔ)言和編譯器 性能和CUDA想當(dāng)難度更低

用 CUDA 為 GPU 編程實(shí)在太難了。

為了讓沒(méi)有 CUDA 編程經(jīng)驗(yàn)的人寫出和專家效率相當(dāng)?shù)?GPU 代碼，現(xiàn)在 OpenAI 推出了一種新的語(yǔ)言和編譯器 ——Triton。

它的難度比 CUDA 低，但是性能卻可與之相媲美。

OpenAI 聲稱：

Triton 只要 25 行代碼，就能在 FP16 矩陣乘法上達(dá)到與 cuBLAS 相當(dāng)?shù)男阅堋?/p>

OpenAI 的研究人員已經(jīng)使用 Triton，來(lái)生成比同等 Torch 效率高出 1 倍的內(nèi)核。

Triton 項(xiàng)目的負(fù)責(zé)人 Philippe Tillet 說(shuō)：“我們的目標(biāo)是使 Triton 成為深度學(xué)習(xí) CUDA 的可行替代方案。”

25 行代碼實(shí)現(xiàn)最佳性能

Triton 起源于 Tillet 在 2019 年學(xué)術(shù)會(huì)議 MLPF 上的一篇論文，當(dāng)時(shí)他還是哈佛大學(xué)的一名研究生。

Tillet 解決的問(wèn)題是如何開發(fā)一種 cuDNN 更具表現(xiàn)力的語(yǔ)言，既能夠處理神經(jīng)網(wǎng)絡(luò)中涉及的矩陣的各種操作，同時(shí)兼具可移植性且以及和 cuDNN 相媲美的性能。

現(xiàn)代 GPU 大致分為三個(gè)主要組件 ——DRAM、SRAM、ALU，對(duì)這些資源進(jìn)行調(diào)度管理十分復(fù)雜，即便是熟悉 CUDA 的程序員。

Triton 可以將這些優(yōu)化過(guò)程完全自動(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)類似的過(guò)程需要花費(fèi)更多的精力，甚至可能會(huì)降低性能。

手寫矩陣乘法內(nèi)核的一個(gè)重要優(yōu)點(diǎn)是它們可以根據(jù)需要進(jìn)行定制，以適應(yīng)其輸入和輸出的融合變換。

如果沒(méi)有 Triton，對(duì)于沒(méi)有特殊 GPU 編程經(jīng)驗(yàn)的開發(fā)者來(lái)說(shuō)，矩陣乘法內(nèi)核的修改是非常困難的。

Triton 背后的原理

Triton 的良好性能，來(lái)自于以 Triton-IR 為中心的模塊化系統(tǒng)架構(gòu)，這是一種基于 LLVM 的中間表示。

@triton.jit decorator 通過(guò)遍歷提供 Python 函數(shù)的抽象語(yǔ)法樹(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ù)可以通過(guò)查看計(jì)算密集型塊級(jí)操作(例如 tl.dot)的操作數(shù)自動(dòng)存儲(chǔ)到共享內(nèi)存中，并使用標(biāo)準(zhǔn)活性分析技術(shù)進(jìn)行分配/同步。

另一方面，Triton 程序可以通過(guò)同時(shí)執(zhí)行不同的內(nèi)核實(shí)例跨 SM 進(jìn)行高效和自動(dòng)并行化，以及通過(guò)分析每個(gè)塊級(jí)操作的迭代空間，并在不同的 SIMD 中進(jìn)行充分分區(qū)將 SM 內(nèi)單元并行化。

目前 Triton 僅適用于英偉達(dá) GPU，但官方表示 AMD GPU 以及 CPU 的版本正在開發(fā)中。

關(guān)鍵詞： OpenAI 語(yǔ)言 編譯器 性能 難度更低