前言

上一章的反響還不錯，很多人都私信催更想看Triton的具體優(yōu)化有哪些，為什么它能夠得到比cuBLAS更好的性能。大家不用急，這也是我為什么要寫這一系列文章的初衷，來帶著大家從Triton的DSL前端一步一步到最終的machine code生成有一個清晰的理解，從而為大家展示編譯在高性能計算中所起到的作用。先來看看openai對Triton所打的廣告:

"An open-source python-like programming language which enables researchers with no CUDA experience to write highly efficient GPU code -- most of the time on par with what an expert would be able to produce"

確實(shí)宣傳的很強(qiáng)，Triton作為一個完全開源的編譯流，提供了python-based的前端，我們這里把它稱為DSL，也就是這篇文章要介紹的主要內(nèi)容。DSL的全稱是Domain Specific Language，為什么要設(shè)計DSL?其實(shí)設(shè)計DSL的目的就是為了能夠讓使用該套工具的人能夠以一個低成本的入門代價來體驗(yàn)到該套工具或者軟件棧能夠帶來的性能提升。類比PyTorch，TensorFlow，MXNet，Taichi，TVM等，其都是給使用這些工具的人提供了一套比較清晰的python api，然后用戶只需要花一定的時間學(xué)習(xí)下這些python api的使用規(guī)范以及常用的python開發(fā)流程，就能夠在不摸清軟件或者框架的底層細(xì)節(jié)的同時帶來極致的開發(fā)體驗(yàn)。當(dāng)然，這里多講講有關(guān)DSL的設(shè)計理念，以我這些年對于軟件系統(tǒng)的開發(fā)和實(shí)戰(zhàn)經(jīng)驗(yàn)來看，DSL的設(shè)計最注重的是靈活性，也即一些關(guān)于編程語言設(shè)計的論文中經(jīng)常提到的flexibility，靈活的編程方式能夠給用戶帶來不一樣的使用體驗(yàn)。對于深度學(xué)習(xí)算法從業(yè)者來說，最常見的例子其實(shí)就是pytorch的出現(xiàn)，對于caffe和tensorflow 1.x 所帶來的編程方式的顛覆，前者則是以imperative的編程范式為用戶提供一個簡單易控制的開發(fā)體驗(yàn)，方便用戶的調(diào)試和已有代碼的交互，可以在model構(gòu)建的過程中，隨時可視化和運(yùn)行已有的部分計算圖，將更多的優(yōu)化細(xì)節(jié)進(jìn)行隱藏。而后者則是以declarative的編程范式，讓用戶首先通過特定的placeholder api去構(gòu)建一個完整的計算圖，然后對該計算圖進(jìn)行全局范圍的優(yōu)化，這樣做的好處自然而然帶來了更多的優(yōu)化空間。但是，它所帶來的問題也顯而易見，那就是對于一些經(jīng)驗(yàn)不是很足的用戶，一旦出現(xiàn)編程中的bug，很難去快速的定位到具體的問題。那么，再次回到Triton，Triton給我們提供了一個什么樣的編程范式呢？

Triton給我們提供的更像是一種imperative的編程范式，但是Triton每次可以操作的粒度卻是一個Block級別。有些同學(xué)可能會問，什么是Block呢？這里Block的概念和CUDA編程中的thread-Block是具有相同的概念，也即是說，當(dāng)我們在編寫CUDA代碼的過程中，需要對thread-Block中的每個thread進(jìn)行精確的編程。說的深入一些，其實(shí)目前像tvm這樣的代碼生成工具，或者cutlass這樣的模板庫在做for-loop tiling的過程，對于inter-level級別，也就是thread-Block的concurrently運(yùn)行這個層面的優(yōu)化來說，做的已經(jīng)很不錯了。但是對于每個thread-Block內(nèi)部的intra-level級別的并行來說，還是有很多優(yōu)化的空間包括memory coalescing, 共享內(nèi)存的sync以及bank conflict的處理，在包括更加細(xì)粒度的register-level的tensor core的調(diào)度上。上面這些優(yōu)化，如果你不是一個經(jīng)驗(yàn)十分老道的高性能工程師，對于GPU的架構(gòu)和CUDA的設(shè)計沒有比較深入的研究和經(jīng)驗(yàn)，是很難在段時間內(nèi)寫出媲美cuBLAS的高性能算子庫的。同時，我們可以像寫pytorch一樣，通過Triton給予的DSL就可以完美的定義出自己想要的算子，然后通過將其嵌入到其他已有的框架的后端中，作為一個codegen來使用。所以，關(guān)于Triton的定位，我更更傾向于將其定位成一個以python DSL來生成高性能GPU算子的瑞士軍刀。當(dāng)然，Triton的使用是具有一定門檻的，假如你先前從來沒有寫過類似CUDA或者OpenCL這類編程GPU加速器的語言，直接讓你通過學(xué)習(xí)Triton就寫出能夠一定程度媲美cuBLAS的代碼，我覺得還是有一定的難度的。這里以NV的顯卡為例，感覺Triton的定位更像是先前對于CUDA的優(yōu)化有一定的基礎(chǔ)，并且想通過隱藏更多的高級優(yōu)化細(xì)節(jié)，僅僅通過python層面的描述將整個算法的流程定義清楚，同時將編譯與高級的優(yōu)化細(xì)節(jié)交給codegen幫你去做的用戶。

說了這么多， 總結(jié)下來，Triton DSL能夠幫助用戶做什么樣的事情？

Embedded In Python: 使用python的裝飾器來定義需要優(yōu)化的kernel

Pointer Arithmetics: 使用pointer arithmetic的方式去操作DRAM上的多維數(shù)據(jù)

Optimizing Compiler：以Block為粒度的編程方式，為用戶隱藏更多的優(yōu)化細(xì)節(jié)，將這些優(yōu)化工作交給編譯器

Triton DSL基礎(chǔ)

Triton官方對于DSL并沒有像pytorch，tf或者tvm那些工具一樣有一個比較詳細(xì)的說明和介紹，對于新手入門還是稍微有一些門檻的，官方關(guān)于DSL的文檔如下地址所示：

triton.language - Triton documentationtriton-lang.org/main/python-api/triton.language.html

由于我在使用Triton進(jìn)行二次開發(fā)的過程中，發(fā)現(xiàn)有些東西可能已經(jīng)過時，文檔還沒來得及更新，我們就以目前Triton的main分支的代碼進(jìn)行介紹。有關(guān)Triton這門編程語言的大部分東西都位于/python/triton的目錄下，該目錄下的compiler，language，runtime是定義有關(guān)Triton DSL描述對應(yīng)具體workload，到中間代碼的生成，以及最終通過自動調(diào)優(yōu)將最優(yōu)的實(shí)現(xiàn)找到的過程。要使用Triton的DSL，在最開始的時候，我們需要通過如下代碼將Triton引入我們的開發(fā)環(huán)境中，這就類似以前寫pytorch時候使用的import torch

importtriton
importtriton.languageastl

那么接下來，一旦tl被我們import進(jìn)來了，就可以開始使用Triton DSL來構(gòu)建各種各樣的workload了。關(guān)于tl的所有操作，可以在python/triton/language/init.py中的__all__下查到，總共定義了95個常用的操作。

__all__=[
"abs",
"advance",
"arange",
"argmin",
"argmax",
"atomic_add",
"atomic_and",
"atomic_cas",
"atomic_max",
"atomic_min",
"atomic_or",
"atomic_xchg",
"atomic_xor",
"bfloat16",
"block_type",
"broadcast",
"broadcast_to",
"builtin",
"cat",
"cdiv",
"constexpr",
"cos",
"debug_barrier",
"device_assert",
"device_print",
"dot",
"dtype",
"exp",
"expand_dims",
"extra",
"fdiv",
"float16",
"float32",
"float64",
"float8e4",
"float8e5",
"full",
"function_type",
"int1",
"int16",
"int32",
"int64",
"int8",
"ir",
"math",
"load",
"log",
"make_block_ptr",
"max",
"max_contiguous",
"maximum",
"min",
"minimum",
"multiple_of",
"num_programs",
"pair_uniform_to_normal",
"philox",
"philox_impl",
"pi32_t",
"pointer_type",
"program_id",
"rand",
"rand4x",
"randint",
"randint4x",
"randn",
"randn4x",
"ravel",
"reduce",
"reshape",
"sigmoid",
"sin",
"softmax",
"sqrt",
"static_range",
"static_assert",
"static_print",
"store",
"sum",
"swizzle2d",
"tensor",
"trans",
"triton",
"uint16",
"uint32",
"uint32_to_uniform_float",
"uint64",
"uint8",
"umulhi",
"view",
"void",
"where",
"xor_sum",
"zeros",
"zeros_like",
]

那么，關(guān)于triton的所有操作，可以在/python/triton/init.py下進(jìn)行查看，總共定義了19個常用的操作。

__all__=[
"autotune",
"cdiv",
"CompilationError",
"compile",
"Config",
"heuristics",
"impl",
"jit",
"JITFunction",
"KernelInterface",
"language",
"MockTensor",
"next_power_of_2",
"ops",
"OutOfResources",
"reinterpret",
"runtime",
"TensorWrapper",
"testing",
"program_ids_from_grid",
]

接下來，我們就來講講如何通過這95+19中的常用操作來定義一個完整的關(guān)于“矩陣乘法”的優(yōu)化流程

Triton DSL做矩陣乘法

首先，和編寫CUDA的kernel的流程類似，首先定義需要進(jìn)行運(yùn)算的輸入tensor和輸出tensor，然后launch kernel進(jìn)行計算，最終對計算結(jié)果和golden data進(jìn)行比較進(jìn)行單元測試。

0x0 定義kernel準(zhǔn)備工作

defmatmul(a,b):
#Checkconstraints.
asserta.shape[1]==b.shape[0],"Incompatibledimensions"
asserta.is_contiguous(),"MatrixAmustbecontiguous"
assertb.is_contiguous(),"MatrixBmustbecontiguous"
M,K=a.shape
K,N=b.shape
#Allocatesoutput.
c=torch.empty((M,N),device=a.device,dtype=a.dtype)
#1Dlaunchkernelwhereeachblockgetsitsownprogram.
grid=lambdaMETA:(
triton.cdiv(M,META['BLOCK_SIZE_M'])*triton.cdiv(N,META['BLOCK_SIZE_N']),
)
matmul_kernel[grid](
a,b,c,
M,N,K,
a.stride(0),a.stride(1),
b.stride(0),b.stride(1),
c.stride(0),c.stride(1),
ACTIVATION=activation
)
returnc

上述代碼片段中，我們可以看到，唯一比較陌生的應(yīng)該就是如下關(guān)于grid和matmul_kernel的定義

grid=lambdaMETA:(
triton.cdiv(M,META['BLOCK_SIZE_M'])*triton.cdiv(N,META['BLOCK_SIZE_N']),
)
matmul_kernel[grid](
a,b,c,
M,N,K,
a.stride(0),a.stride(1),
b.stride(0),b.stride(1),
c.stride(0),c.stride(1),
ACTIVATION=activation
)

這里完全可以類比到CUDA編程中在main函數(shù)中所寫到的關(guān)于怎么去launch一個kernel的環(huán)節(jié)，類比如下代碼

dim3block(BLOCK_SIZE_M,BLOCK_SIZE_N);
dim3grid((M+BLOCK_SIZE_M-1)/BLOCK_SIZE_M,(N+BLOCK_SIZE_N-1)/BLOCK_SIZE_N);
matmul_kernel<<>>(Ad,Bd,Cd,M,N,K);

其中，grid表示的是每個grid中所含有的thread-Blocks的個數(shù)，block表示的則是每個thread-Blocks所啟動的threads的個數(shù)。在上述Triton的程序中，在matmul_kernel<<< >>>的后面，我們本質(zhì)是將"BLOCK_SIZE_M"和"BLOCK_SIZE_N"這兩個維度進(jìn)行了合并，也即后面準(zhǔn)備通過一組id進(jìn)行訪問， triton.cdiv表示來做除法操作。接下來，我們就來看看最為關(guān)鍵的matmul_kernel是如何定義的。

0x1 Triton Kernel的編寫

@triton.jit
defmatmul_kernel(
#Pointerstomatrices
a_ptr,b_ptr,c_ptr,
#Matrixdimensions
M,N,K,
#Thestridevariablesrepresenthowmuchtoincreasetheptrbywhenmovingby1
#elementinaparticulardimension.E.g.`stride_am`ishowmuchtoincrease`a_ptr`
#bytogettheelementonerowdown(AhasMrows).
stride_am,stride_ak,
stride_bk,stride_bn,
stride_cm,stride_cn,
#Meta-parameters
BLOCK_SIZE_M:tl.constexpr,BLOCK_SIZE_N:tl.constexpr,BLOCK_SIZE_K:tl.constexpr,
GROUP_SIZE_M:tl.constexpr,
ACTIVATION:tl.constexpr,
):
"""KernelforcomputingthematmulC=AxB.
Ahasshape(M,K),Bhasshape(K,N)andChasshape(M,N)
"""
#-----------------------------------------------------------
#Mapprogramids`pid`totheblockofCitshouldcompute.
#ThisisdoneinagroupedorderingtopromoteL2datareuse.
#Seeabove`L2CacheOptimizations`sectionfordetails.
pid=tl.program_id(axis=0)
num_pid_m=tl.cdiv(M,BLOCK_SIZE_M)
num_pid_n=tl.cdiv(N,BLOCK_SIZE_N)
num_pid_in_group=GROUP_SIZE_M*num_pid_n
group_id=pid//num_pid_in_group
first_pid_m=group_id*GROUP_SIZE_M
group_size_m=min(num_pid_m-first_pid_m,GROUP_SIZE_M)
pid_m=first_pid_m+(pid%group_size_m)
pid_n=(pid%num_pid_in_group)//group_size_m

#----------------------------------------------------------
#CreatepointersforthefirstblocksofAandB.
#WewilladvancethispointeraswemoveintheKdirection
#andaccumulate
#`a_ptrs`isablockof[BLOCK_SIZE_M,BLOCK_SIZE_K]pointers
#`b_ptrs`isablockof[BLOCK_SIZE_K,BLOCK_SIZE_N]pointers
#Seeabove`PointerArithmetics`sectionfordetails
offs_am=(pid_m*BLOCK_SIZE_M+tl.arange(0,BLOCK_SIZE_M))%M
offs_bn=(pid_n*BLOCK_SIZE_N+tl.arange(0,BLOCK_SIZE_N))%N
offs_k=tl.arange(0,BLOCK_SIZE_K)
a_ptrs=a_ptr+(offs_am[:,None]*stride_am+offs_k[None,:]*stride_ak)
b_ptrs=b_ptr+(offs_k[:,None]*stride_bk+offs_bn[None,:]*stride_bn)

#-----------------------------------------------------------
#IteratetocomputeablockoftheCmatrix.
#Weaccumulateintoa`[BLOCK_SIZE_M,BLOCK_SIZE_N]`block
#offp32valuesforhigheraccuracy.
#`accumulator`willbeconvertedbacktofp16aftertheloop.
accumulator=tl.zeros((BLOCK_SIZE_M,BLOCK_SIZE_N),dtype=tl.float32)
forkinrange(0,tl.cdiv(K,BLOCK_SIZE_K)):
#LoadthenextblockofAandB,generateamaskbycheckingtheKdimension.
#Ifitisoutofbounds,setitto0.
a=tl.load(a_ptrs,mask=offs_k[None,:]

	

	上述代碼則對應(yīng)了matmul_kernel的具體實(shí)現(xiàn)細(xì)節(jié)，我們可以將上述代碼分成三個部分來進(jìn)行學(xué)習(xí)。

	第一個部分，先來看看matmul_kernel的輸入參數(shù)有哪些？首先在Triton中定義一個kernel的時候，需要使用@triton.jit對其進(jìn)行裝飾。a_ptr, b_ptr, c_ptr指的是輸入tensor和輸出tensor所對應(yīng)的首地址，M,N,K則表示需要計算的tensor的維度分別為[M, K] x [K, N]。stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn則表示的是分別針對a,b,c這三個tensor來說，訪問一個元素所需要移動的步長。而后面的BLOCK_SIZE_M, BLOCK_SIZE_N等被定義為tl.constexpr的變量都屬于是自動調(diào)優(yōu)系統(tǒng)中可以被枚舉的knob，如果你用過autotvm的話，應(yīng)該不會很陌生。

	第二部分，則是將id對應(yīng)到輸出tensor的每個block上，這塊的內(nèi)容在tutorial中講到是為了提高L2 Cache的命中率。在文中，openai使用了一個叫做"super-grouping"的名字來表示一個block中所含有的block的個數(shù)。其實(shí)super-grouping的原理很簡單，看下圖所示

	img

	編輯切換為居中

	添加圖片注釋，不超過 140 字（可選）

	當(dāng)我們在進(jìn)行AxB=C的時候，如果在load A的數(shù)據(jù)的時候，以行優(yōu)先的方式，一次性讀取9個block，那么如果要得到C矩陣的第一行結(jié)果，并且C的存儲方式也是以行優(yōu)先的方式進(jìn)行，總共需要進(jìn)行9+81=90次對block的load的操作，9次對block的write的操作才能得到所要的結(jié)果。但是，如果我們采用了“super-grouping”的方式，也就是說同樣為了得到得到C矩陣中的9次block的write操作，那么對于A矩陣來說，進(jìn)行93次load操作，B矩陣也同樣進(jìn)行93次的load操作，對block總的load操作則為27+27=54次。前后對比下，第一種方式則總共進(jìn)行了90次load+9次write，而第二種采用了super-grouping技術(shù)則進(jìn)行了54次load和9次write。并且openai還在備注中說明了可以在A100上由220TFLOPS提升到245TFLOPS。等后面可以對該技術(shù)專門寫一個章節(jié)進(jìn)行介紹和測試。

	第三部分，則比較常規(guī)，對應(yīng)到CUDA編程中，其實(shí)就是在探索如何通過Triton DSL去訪問每個block，然后通過一個accumulator變量來記錄tl.dot(a, b)的結(jié)果，mask的作用是來判斷迭代的過程中，是否越界，如果超過了界限的范圍，就將對應(yīng)的block置為0。最終再將結(jié)果按位寫會到對應(yīng)的c矩陣則完成了對應(yīng)的操作。

	0x2 單元測試

	單元測試的編寫就顯而易見了，為的是比較通過Triton生成的代碼和通過pytorch的torch.mm算出的結(jié)果是否對齊

	
torch.manual_seed(0)
a=torch.randn((512,512),device='cuda',dtype=torch.float16)
b=torch.randn((512,512),device='cuda',dtype=torch.float16)
triton_output=matmul(a,b)
torch_output=torch.matmul(a,b)
print(f"triton_output={triton_output}")
print(f"torch_output={torch_output}")
iftorch.allclose(triton_output,torch_output,atol=1e-2,rtol=0):
print("TritonandTorchmatch")
else:
print("TritonandTorchdiffer")


	

	Triton的自動調(diào)優(yōu)

	這里不會對Triton的自動調(diào)優(yōu)技術(shù)進(jìn)行過多的解讀，僅僅通過一些小的實(shí)驗(yàn)來表明，定義不同的搜索空間還是可以很大程度上提高matmul最終的TFLOPS的。那么，對于Triton的自動調(diào)優(yōu)以及到底應(yīng)該如何定義一個高效的搜索空間將會在后面的內(nèi)容中進(jìn)行詳細(xì)的講解。所有實(shí)驗(yàn)都是在NV 3090 GPU上，batch = 1， datatype = fp16.

	在openai給出的默認(rèn)自動調(diào)優(yōu)空間下

	
@triton.autotune(
configs=[
triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':256,'BLOCK_SIZE_K':64,'GROUP_SIZE_M':8},num_stages=3,num_warps=8),
triton.Config({'BLOCK_SIZE_M':64,'BLOCK_SIZE_N':256,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':128,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':64,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':64,'BLOCK_SIZE_N':128,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':32,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':64,'BLOCK_SIZE_N':32,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=5,num_warps=2),
triton.Config({'BLOCK_SIZE_M':32,'BLOCK_SIZE_N':64,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=5,num_warps=2),
],
key=['M','N','K'],
)

img

	

	當(dāng)我們?nèi)フ{(diào)整對應(yīng)的調(diào)優(yōu)空間

	
@triton.autotune(
configs=[
triton.Config({'BLOCK_SIZE_M':32,'BLOCK_SIZE_N':64,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=5,num_warps=2),
],
key=['M','N','K'],
)

img

	

	編輯切換為居中

	添加圖片注釋，不超過 140 字（可選）

	當(dāng)我們繼續(xù)對搜索空間進(jìn)行調(diào)整

	
@triton.autotune(
configs=[
triton.Config({'BLOCK_SIZE_M':64,'BLOCK_SIZE_N':32,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=5,num_warps=2),
],
key=['M','N','K'],
)

img

	

	進(jìn)一步的進(jìn)行修改

	
@triton.autotune(
configs=[
triton.Config({'BLOCK_SIZE_M':64,'BLOCK_SIZE_N':256,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
],
key=['M','N','K'],
)

img

	

	通過上面簡單的實(shí)驗(yàn)，可以看出，要想得到一個比較好的TFLOPS的數(shù)值，對于"BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K", "num_stages", "num_warps"都需要一個很好的調(diào)整，才能夠得到一個媲美甚至超過cuBLAS的性能上界。

	總結(jié)

	通過上述對于Triton DSL的解讀，以及通過Triton DSL來完成矩陣乘法的操作，我們可以看到，用戶只需要懂一些基礎(chǔ)的python語法和寫pytorch，然后將先前使用CUDA的經(jīng)驗(yàn)?zāi)眠^來，使用一些和pytorch很像的api，就可以在NV的顯卡上，使用Triton就可以很輕松的生成性能媲美cuBLAS的高性能算子。如果你能夠通過triton熟練的寫出matmul和flashAttention的話，那么像深度學(xué)習(xí)中的大部分算子你都可以輕松的通過Triton來幫你cover，后面的教程中，我會將重點(diǎn)放在Triton在使用MLIR進(jìn)行重構(gòu)的過程中，所采取的一些工程上的組織以及Triton本身的內(nèi)部設(shè)計使得其能夠生成媲美NV的cublas的高性能算法庫。

	

	   審核編輯：彭靜