過去十年中，深度神經(jīng)網(wǎng)絡(luò) (DNN) 已成為最重要的機(jī)器學(xué)習(xí)模型之一，創(chuàng)造了從自然語言處理到計算機(jī)視覺、計算神經(jīng)科學(xué)等許多領(lǐng)域的 SOTA 實現(xiàn)。DNN 模型的優(yōu)勢來自于它的層次結(jié)構(gòu)，這一特征導(dǎo)致其計算量巨大，但也會產(chǎn)生大量高度并行化的工作，特別適合多核和眾核處理器。

深度學(xué)習(xí)領(lǐng)域的新研究思路往往是結(jié)合原生框架 operator 來實現(xiàn)的，這種方法雖然方便，但需要創(chuàng)建或移動許多臨時張量，因此可能會造成神經(jīng)網(wǎng)絡(luò)的性能損失。編寫專門的 GPU 內(nèi)核或許可以解決這個問題，但 GPU 編程的確是一件相當(dāng)復(fù)雜的事。

DNN 計算潛力與 GPU 編程困難之間的矛盾由來已久。英偉達(dá)在 2007 年發(fā)布了 CUDA 的初始版本，CUDA 平臺是一個軟件層，使用者可以直接訪問 GPU 的虛擬指令集和并行計算單元，用于執(zhí)行計算內(nèi)核。近年來，主流深度學(xué)習(xí)框架幾乎都是基于 CUDA 進(jìn)行加速，英偉達(dá)也一直在完善 CUDA 工具包，但對于一般的開發(fā)者來說，CUDA 還是「不那么容易上手」。

今天，OpenAI 正式推出 Triton 1.0，這是一種類 Python 的開源編程語言。即使沒有 CUDA 經(jīng)驗的研究人員，也能夠高效編寫 GPU 代碼。例如，它可以用不到 25 行代碼寫出與 cuBLAS 性能相匹配的 FP16 矩陣乘法內(nèi)核，后者是許多專業(yè)的 GPU 編程者尚且無法做到的。此外，OpenAI 的研究者已經(jīng)使用 Triton 成功生成了比 PyTorch 同類實現(xiàn)效率高 2 倍的內(nèi)核。

代碼地址：https://github.com/openai/triton

Triton 的最初想法來源于現(xiàn)任 OpenAI 科學(xué)家的 Philippe Tillet 2019 年在哈佛大學(xué)攻讀研究生學(xué)位時發(fā)表的一篇論文，當(dāng)時他的導(dǎo)師是 H. T. Kung 和 David Cox。

論文鏈接：http://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf

Tillet 希望解決的問題是打造一種比英偉達(dá)的 CUDA 等特定供應(yīng)商庫更好用的庫，能夠處理神經(jīng)網(wǎng)絡(luò)中涉及矩陣的各種操作，具備可移植性，且性能可與 cuDNN 或類似的供應(yīng)商庫相媲美。團(tuán)隊表示：「直接用 CUDA 進(jìn)行 GPU 編程太難了，比如為 GPU 編寫原生內(nèi)核或函數(shù)這件事，會因為 GPU 編程的復(fù)雜性而出奇困難?！?/p>

Facebook AI 研究中心科學(xué)家 Soumith Chintala 也在推特上表達(dá)了自己對 Triton 的期待：

新發(fā)布的 Triton 可以為一些核心的神經(jīng)網(wǎng)絡(luò)任務(wù)（例如矩陣乘法）提供顯著的易用性優(yōu)勢?！肝覀兊哪繕?biāo)是使其成為深度學(xué)習(xí) CUDA 的可行替代方案，」Philippe Tillet 作為 Triton 項目負(fù)責(zé)人如此表示。

GPU 編程面臨的挑戰(zhàn)

現(xiàn)代 GPU 的架構(gòu)大致可以分為三個主要組件：DRAM、SRAM 和 ALU。優(yōu)化 CUDA 代碼時，必須考慮到每一個組件：

• 來自 DRAM 的內(nèi)存?zhèn)鬏敱仨毢喜⑦M(jìn)大型事務(wù)，以利用現(xiàn)代內(nèi)存接口的總線位寬；
• 必須在數(shù)據(jù)重新使用之前手動存儲到 SRAM 中，并進(jìn)行管理以最大限度地減少檢索時共享內(nèi)存庫沖突；
• 計算必須在流處理器（SM）內(nèi)部或之間細(xì)致分區(qū)和調(diào)度，以促進(jìn)指令 / 線程級的并行以及專用算術(shù)邏輯單元（ALU）的利用。

種種因素導(dǎo)致 GPU 編程難度驟增，即使對于具有多年經(jīng)驗的 CUDA 程序員也是如此。Triton 的目的是將這些優(yōu)化過程自動化，以此讓開發(fā)人員更專注于并行代碼的高級邏輯。出于對泛用能力的考量，Triton 不會自動調(diào)度跨流處理器的工作，而是將一些重要的算法考慮因素（例如 tiling、SM 間同步）留給開發(fā)者自行決定。

編程模型

在所有可用的領(lǐng)域?qū)Ｓ谜Z言和 JIT 編譯器中，Triton 或許與 Numba 最相似：內(nèi)核被定義為修飾過的 Python 函數(shù)，并與實例網(wǎng)格上不同的 program_id 的同時啟動。但不同之處值得注意：如下圖代碼片段所示，Triton 通過對 block 的操作來展示 intra-instance 并行，此處 block 是維數(shù)為 2 的冪的數(shù)組，而不是單指令多線程（SIMT）執(zhí)行模型。如此一來，Triton 高效地抽象出了與 CUDA 線程 block 內(nèi)的并發(fā)相關(guān)的所有問題（比如內(nèi)存合并、共享內(nèi)存同步 / 沖突、張量核心調(diào)度）。

雖然這對 embarrassingly 并行（即 element-wise）計算可能沒什么幫助，但是可以簡化更復(fù)雜的 GPU 程序的開發(fā)。例如，在融合 softmax 核的情況下，對于每個輸入張量 X∈R^M×N 來說，每個實例對給定輸入張量的不同行進(jìn)行歸一化。這種并行化策略的標(biāo)準(zhǔn) CUDA 實現(xiàn)可能難以編寫，需要線程之間的顯式同步，因為這種策略并發(fā)地減少 X 的同一行。而 Triton 很大程度上消除了這種復(fù)雜性，每個內(nèi)核實例加載感興趣的行，并使用類似 NumPy 的原語順序?qū)ζ溥M(jìn)行規(guī)范化。

import triton 
import triton.language as tl 
@triton.jit 
def softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N): 
    # row index 
    m = tl.program_id(0) 
    # col indices 
    # this specific kernel only works for matrices that  
    # have less than BLOCK_SIZE columns 
    BLOCK_SIZE = 1024 
    n = tl.arange(0, BLOCK_SIZE) 
    # the memory address of all the elements 
    # that we want to load can be computed as follows 
    X = X + m * stride_xm + n * stride_xn 
    # load input data; pad out-of-bounds elements with 0  
    x = tl.load(X, mask=n < N, other=-float('inf')) 
    # compute numerically-stable softmax 
    z = x - tl.max(x, axis=0) 
    num = tl.exp(z) 
    denom = tl.sum(num, axis=0) 
    y = num / denom 
    # write back to Y 
    Y = Y + m * stride_ym + n * stride_yn 
    tl.store(Y, y, mask=n < N) 
import torch 
# Allocate input/output tensors 
X = torch.normal(0, 1, size=(583, 931), device='cuda') 
Y = torch.empty_like(X) 
# SPMD launch grid 
grid = (X.shape[0], ) 
# enqueue GPU kernel 
softmax[grid](Y, Y.stride(0), Y.stride(1),  
              X, X.stride(0), X.stride(1), 
              X.shape[0]    , X.shape[1]) 

在 Triton 中融合 softmax

Triton JIT 把 X、Y 當(dāng)作指針而不是張量。最重要的是，softmax 這種特殊實現(xiàn)方式在整個規(guī)范化過程中保持 SRAM 中 X 的行不變，從而在適用時最大限度地實現(xiàn)數(shù)據(jù)重用（約 32K 列）。這與 PyTorch 的內(nèi)部 CUDA 代碼不同，后者使用臨時內(nèi)存使其更通用，但速度明顯變慢（見下圖）。

Torch (v1.9) JIT 較低的性能突出了從高級張量操作序列自動生成 CUDA 代碼的難度。

@torch.jit.script 
def softmax(x): 
    x_max = x.max(dim=1)[0] 
    z = x - x_max[:, None] 
    numerator = torch.exp(x) 
    denominator = numerator.sum(dim=1) 
    return numerator / denominator[:, None] 

融合 softmax 與 Torch JIT

矩陣乘法

能夠為元素操作（element-wise operation）和規(guī)約操作（reduction operation）編寫融合內(nèi)核是很重要的，但考慮到神經(jīng)網(wǎng)絡(luò)中矩陣乘法的重要性，這還不夠。事實證明，Triton 在這些方面表現(xiàn)很好，僅用大約 25 行 Python 代碼就能達(dá)到最佳性能。相比之下，CUDA 效率就沒有那么高了。

手寫矩陣乘法內(nèi)核的一個重要優(yōu)點是它們可以根據(jù)需要進(jìn)行定制，以適應(yīng)其輸入（例如切片）和輸出（例如 Leaky ReLU）的融合變換。假如不存在 Triton 這樣的系統(tǒng)，那么對于沒有出色的 GPU 編程專業(yè)知識的開發(fā)人員來說，矩陣乘法內(nèi)核將很難大改。

高級系統(tǒng)架構(gòu)

Triton 的良好性能得益于以 Triton-IR 為中心的模塊化系統(tǒng)架構(gòu)。Triton-IR 是一種基于 LLVM 的中間表示，多維值塊（blocks of values）是其中最重要的東西。

@triton.jit 裝飾器的工作原理是遍歷由 Python 函數(shù)提供的抽象語法樹（AST），這樣一來就能使用通用的 SSA 構(gòu)造算法實時生成 Triton-IR。生成的 IR 代碼隨后由編譯器后端進(jìn)行簡化、優(yōu)化和自動并行化，然后轉(zhuǎn)換為高質(zhì)量的 LLVM-IR，最終轉(zhuǎn)換為 PTX，以便在最新的 NVIDIA GPU 上執(zhí)行。目前 Triton 還不支持 CPU 和 AMD GPU，但團(tuán)隊表示對二者的支持正在開發(fā)中。

編譯器后端

研究人員發(fā)現(xiàn)通過 Triton-IR 來使用塊狀程序表示，這種方法允許編譯器自動執(zhí)行各種重要的程序優(yōu)化。例如，通過查看計算密集型塊級操作（例如 tl.dot）的操作數(shù)，數(shù)據(jù)可以自動存儲到共享內(nèi)存中，并使用標(biāo)準(zhǔn)的活躍性分析技術(shù)進(jìn)行數(shù)據(jù)的分配與同步。

此外，Triton 還可以在 SM 之間以及 SM 之內(nèi)高效、自動地并行化，前者通過并發(fā)執(zhí)行不同的內(nèi)核實例來實現(xiàn)，后者通過分析每個塊級操作的迭代空間，并將其充分劃分到不同的 SIMD 單元來實現(xiàn)。如下所示：

Triton 自動并行化。每個塊級操作都定義了一個塊級迭代空間，該空間可以自動并行化以利用 SM（Streaming Multiprocessor） 上的可用資源。