兽性总裁的爱奴,武道至尊帝临小说,神武八荒一颗小说

前言

上一章的反響還不錯，很多人都私信催更想看Triton的具體優化有哪些，為什么它能夠得到比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"

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

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

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

Embedded In Python: 使用python的裝飾器來定義需要優化的kernel

Pointer Arithmetics: 使用pointer arithmetic的方式去操作DRAM上的多維數據

Optimizing Compiler：以Block為粒度的編程方式，為用戶隱藏更多的優化細節，將這些優化工作交給編譯器

Triton DSL基礎

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

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

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

importtriton
importtriton.languageastl

那么接下來，一旦tl被我們import進來了，就可以開始使用Triton DSL來構建各種各樣的workload了。關于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",
]

那么，關于triton的所有操作，可以在/python/triton/init.py下進行查看，總共定義了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中的常用操作來定義一個完整的關于“矩陣乘法”的優化流程

Triton DSL做矩陣乘法

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

0x0 定義kernel準備工作

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

上述代碼片段中，我們可以看到，唯一比較陌生的應該就是如下關于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函數中所寫到的關于怎么去launch一個kernel的環節，類比如下代碼

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的個數，block表示的則是每個thread-Blocks所啟動的threads的個數。在上述Triton的程序中，在matmul_kernel<<< >>>的后面，我們本質是將"BLOCK_SIZE_M"和"BLOCK_SIZE_N"這兩個維度進行了合并，也即后面準備通過一組id進行訪問， triton.cdiv表示來做除法操作。接下來，我們就來看看最為關鍵的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,:]

	

	上述代碼則對應了matmul_kernel的具體實現細節，我們可以將上述代碼分成三個部分來進行學習。

	第一個部分，先來看看matmul_kernel的輸入參數有哪些？首先在Triton中定義一個kernel的時候，需要使用@triton.jit對其進行裝飾。a_ptr, b_ptr, c_ptr指的是輸入tensor和輸出tensor所對應的首地址，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的變量都屬于是自動調優系統中可以被枚舉的knob，如果你用過autotvm的話，應該不會很陌生。

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

	img

	編輯切換為居中

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

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

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

	0x2 單元測試

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

	
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的自動調優

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

	在openai給出的默認自動調優空間下

	
@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

	

	當我們去調整對應的調優空間

	
@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 字（可選）

	當我們繼續對搜索空間進行調整

	
@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

	

	進一步的進行修改

	
@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

	

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

	總結

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

	

	   審核編輯：彭靜