經(jīng)過前面幾章關(guān)于triton在nv gpu上調(diào)優(yōu)的講解，我們這章開始來看看triton的一個(gè)third_party庫，該庫是為了讓triton去支持更多其他的backend。該項(xiàng)目的地址如下所示，并且已經(jīng)在triton的main分支中，作為third_party進(jìn)行了官方支持，在clone triton的時(shí)候，只需要帶上recursive的flag就可以完成對triton-shared的使用。

什么是triton-shared？

關(guān)于triton-shared的官方具體實(shí)現(xiàn)，如下github repo所示：

GitHub - microsoft/triton-shared: Shared Middle-Layer for Triton Compilationgithub.com/microsoft/triton-shared

如下所示為官方對triton-shared的解釋:

Asharedmiddle-layerfortheTritonCompiler.

Currentlythemiddlelayerisnotcompletebuthasenoughfunctionalitytodemonstratehowitcanwork.ThegeneralideaisthatTritonIRisloweredintoanMLIRcoredialecttoallowittobebothsharedacrossTritontargetsaswellasallowback-endstobesharedwithotherlanguages.

Thebasicintendedarchitecturelookslikethis:

[TritonIR]->[MiddleLayer]->[HWspecificIR]

Themiddle-layerusesMLIR'sLinalgandTenorDialectsforoperationsonTritonblockvalues.OperationsonTritonpointersusetheMemrefDialect.

triton-shared其實(shí)就是為了提供一個(gè)膠水一樣的中間層，通過對middle-layer的設(shè)計(jì)來方便我們的編程語言或者編譯器對接到下游不同的硬件生態(tài)，因?yàn)閠riton自身已經(jīng)把nv和amd這兩個(gè)比較常見的GPU后端實(shí)現(xiàn)了，如果第三方的廠商想通過復(fù)用triton的前端來對自己的芯片搞一套編譯flow，那么triton-shared就起到了決定性的作用。下面這個(gè)圖是triton的codebase所希望支持的一個(gè)愿景，可以看出來，中間這條垂直下來的分支就是triton所支持的nv gpu的優(yōu)化路線，當(dāng)用戶寫完的triton dsl會被翻譯成python的AST，然后再從AST到對應(yīng)的triton dialect，從這一步開始，也就正式將用戶手寫的成分轉(zhuǎn)到了MLIR這套生態(tài)，然后再從triton dialect進(jìn)一步優(yōu)化到triton gpu dialect,從trition gpu dialect開始，就走了比較標(biāo)準(zhǔn)的LLVM代碼生成，從LLVM IR一路lower到PTX，再到SASS，最終可以成功運(yùn)行在NV的GPU上，這套codegen的路線相比TVM等其他編譯框架來說更加的激進(jìn)，直接越過了nvcc compiler，從而使得整個(gè)過程都變成了透明的，對于性能優(yōu)化來說帶來了更多的可能。

img

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

triton-shared其實(shí)主要是用來cover最右邊的分支，因?yàn)槭煜LIR的朋友都知道，在右邊的分支中，Linalg dialect是一個(gè)非常重要dialect，該dialect可以去承接很多不同的backend，在主流一些backend的編譯優(yōu)化環(huán)節(jié)，都會將Linalg作為主要的dialect來進(jìn)行上下游不同dialect之間的轉(zhuǎn)換與對接。

Triton-shared的安裝

Triton-shared的安裝其實(shí)也很簡單，只需要一開始通過recursive來clone整個(gè)triton的主分支，然后使用

exportTRITON_CODEGEN_TRITON_SHARED=1

來指明，我們在build triton整個(gè)項(xiàng)目的過程中需要使用到triton-shared這個(gè)第三方的庫。接下來的流程按照triton官方repo的readme一步一步進(jìn)行即可，有關(guān)LLVM我是使用的具體commit id下手動編譯得到的llvm

LLVMcommitid:b1115f8ccefb380824a9d997622cc84fc0d84a89
Tritoncommitid:1c2d2405bf04dca2de140bccd65480c3d02d995e

為什么要選擇如上兩個(gè)固定的commit id，其實(shí)理由很簡單，因?yàn)槲仪懊孀鲞^一些關(guān)于triton和llvm的開發(fā)都是基于上面兩個(gè)id做的，所以后面我的所有教程以及案例展示都是以這兩個(gè)commit id為主進(jìn)行。如果不知道怎么從0開始編譯triton，可以參考我之前的教程：

科研敗犬丶：OpenAI/Triton MLIR 第零章: 源碼編譯70 贊同 · 7 評論文章

Triton-shared的使用

講解完了什么是triton-shared，以及triton-shared怎么安裝，接下來，我們就來談?wù)勅绾问褂靡呀?jīng)被編譯好的triton-shared。當(dāng)你按照我的上述流程編譯好triton后，會在該路徑下：

/triton/build/tools/triton-shared-opt

看到一個(gè)triton-shared-opt的可執(zhí)行文件，熟悉MLIR的同學(xué)可能很快發(fā)現(xiàn)該方法其實(shí)就是MLIR中最基本的opt，該二進(jìn)制文件可以完成從一個(gè)dialect向另外一個(gè)dialect的lowering，那么我們使用--help來看看triton-shared-opt的所有功能。如果能在終端中輸出如下所示的信息，說明你的triton-shared已經(jīng)全部安裝完畢了。

OVERVIEW:Triton-Sharedtestdriver

AvailableDialects:arith,builtin,cf,gpu,math,scf,triton_gpu,tt
USAGE:triton-shared-opt[options]

OPTIONS:

ColorOptions:

--color-Usecolorsinoutput(default=autodetect)

Generaloptions:

--abort-on-max-devirt-iterations-reached-AbortwhenthemaxiterationsfordevirtualizationCGSCCrepeatpassisreached
--allow-unregistered-dialect-Allowoperationwithnoregistereddialects
Compilerpassestorun
Passes:
--affine-data-copy-generate-Generateexplicitcopyingforaffinememoryoperations
--fast-mem-capacity=-SetfastmemoryspacecapacityinKiB(default:unlimited)
--fast-mem-space=-Fastmemoryspaceidentifierforcopygeneration(default:1)
--generate-dma-GenerateDMAinsteadofpoint-wisecopy
--min-dma-transfer=-MinimumDMAtransfersizesupportedbythetargetinbytes
--skip-non-unit-stride-loops-Testingpurposes:avoidnon-unitstrideloopchoicedepthsforcopyplacement
--slow-mem-space=-Slowmemoryspaceidentifierforcopygeneration(default:0)
--tag-mem-space=-Tagmemoryspaceidentifierforcopygeneration(default:0)
--affine-expand-index-ops-Loweraffineoperationsoperatingonindicesintomorefundamentaloperations
--affine-loop-coalescing-Coalescenestedloopswithindependentboundsintoasingleloop
--affine-loop-fusion-Fuseaffineloopnests
...

這里先來展示

--triton-to-linalg-ConvertTritontoLinalgdialect

這個(gè)pass的使用，因?yàn)閠riton-shared主要就是用來做該優(yōu)化的。他表示的就是將triton dialect作為輸入，然后經(jīng)過triton-to-linalg這個(gè)pass，將其lowering到具有相同語義的linalg dialect上，那triton dialect從哪里來得到呢？不要慌，triton-shared的repo為我們提供了很多MLIR格式的文件來方便我們使用該功能，具體路徑如下：

/triton/third_party/triton_shared/test/Conversion/TritonToLinalg/*

在該教程中，我們使用dot.mlir作為案例進(jìn)行分析，具體代碼如下所示：

//RUN:triton-shared-opt--triton-to-linalg%s|FileCheck%s
module{
tt.func@kernel(
%arg0:!tt.ptr,
%arg1:!tt.ptr,
%arg2:!tt.ptr
)
{
%0=tt.make_range{end=128:i32,start=0:i32}:tensor<128xi32>
%c64=arith.constant128:i32
%1=tt.splat%c64:(i32)->tensor<128xi32>
%2=arith.muli%0,%1:tensor<128xi32>
%3=tt.expand_dims%2{axis=1:i32}:(tensor<128xi32>)->tensor<128x1xi32>
%4=tt.broadcast%3:(tensor<128x1xi32>)->tensor<128x64xi32>
%5=tt.make_range{end=64:i32,start=0:i32}:tensor<64xi32>
%6=tt.expand_dims%5{axis=0:i32}:(tensor<64xi32>)->tensor<1x64xi32>
%7=tt.broadcast%6:(tensor<1x64xi32>)->tensor<128x64xi32>
%8=arith.addi%4,%7:tensor<128x64xi32>
%10=tt.make_range{end=256:i32,start=0:i32}:tensor<256xi32>
%11=tt.expand_dims%10{axis=1:i32}:(tensor<256xi32>)->tensor<256x1xi32>
%12=tt.broadcast%11:(tensor<256x1xi32>)->tensor<256x64xi32>
%13=tt.make_range{end=64:i32,start=0:i32}:tensor<64xi32>
%c256=arith.constant256:i32
%14=tt.splat%c256:(i32)->tensor<64xi32>
%15=arith.muli%13,%14:tensor<64xi32>
%16=tt.expand_dims%15{axis=0:i32}:(tensor<64xi32>)->tensor<1x64xi32>
%17=tt.broadcast%16:(tensor<1x64xi32>)->tensor<256x64xi32>
%18=arith.addi%12,%17:tensor<256x64xi32>
%20=tt.splat%c256:(i32)->tensor<128xi32>
%21=arith.muli%0,%20:tensor<128xi32>
%22=tt.expand_dims%21{axis=1:i32}:(tensor<128xi32>)->tensor<128x1xi32>
%23=tt.broadcast%22:(tensor<128x1xi32>)->tensor<128x256xi32>
%24=tt.expand_dims%10{axis=0:i32}:(tensor<256xi32>)->tensor<1x256xi32>
%25=tt.broadcast%24{axis=0:i32}:(tensor<1x256xi32>)->tensor<128x256xi32>
%26=arith.addi%23,%25:tensor<128x256xi32>
%30=tt.splat%arg0:(!tt.ptr)->tensor<128x64x!tt.ptr>
%31=tt.addptr%30,%8:tensor<128x64x!tt.ptr>,tensor<128x64xi32>
%32=tt.load%31{cache=1:i32,evict=1:i32,isVolatile=false}:tensor<128x64xbf16>
%40=tt.splat%arg1:(!tt.ptr)->tensor<256x64x!tt.ptr>
%41=tt.addptr%40,%18:tensor<256x64x!tt.ptr>,tensor<256x64xi32>
%42=tt.load%41{cache=1:i32,evict=1:i32,isVolatile=false}:tensor<256x64xbf16>
%43=tt.trans%42:(tensor<256x64xbf16>)->tensor<64x256xbf16>
%50=tt.splat%arg2:(!tt.ptr)->tensor<128x256x!tt.ptr>
%51=tt.addptr%50,%26:tensor<128x256x!tt.ptr>,tensor<128x256xi32>
%52=tt.load%51{cache=1:i32,evict=1:i32,isVolatile=false}:tensor<128x256xbf16>
%60=tt.dot%32,%43,%52{allowTF32=false,maxNumImpreciseAcc=0:i32}:tensor<128x64xbf16>*tensor<64x256xbf16>->tensor<128x256xbf16>
tt.store%51,%60:tensor<128x256xbf16>
tt.return
}
}

上述MLIR其實(shí)很容易看懂，在%0->%10其實(shí)都是triton dialect的內(nèi)容，該內(nèi)容表示的就是從上層的triton dsl通過lower轉(zhuǎn)換到對應(yīng)的triton dialect的過程。其中tt就是表示的該MLIR所處的dialect是triton dialect，然后tt.xxx則表示了該dialect所支持的所有operation，有關(guān)如何定義一個(gè)MLIR dialect，我準(zhǔn)備拿一個(gè)單獨(dú)的教程來講。

接下來，只需要在終端中輸入

./triton-shared-opt--triton-to-linalg/triton/third_party/triton_shared/test/Conversion/TritonToLinalg/dot.mlir

就可以得到從triton dialect轉(zhuǎn)到linag dialect部分對應(yīng)的內(nèi)容

#map=affine_map<(d0,?d1)?->(d0,d1)>
module{
func.func@kernel(%arg0:memref<*xbf16>,%arg1:memref<*xbf16>,%arg2:memref<*xbf16>,%arg3:i32,%arg4:i32,%arg5:i32,%arg6:i32,%arg7:i32,%arg8:i32){
%c256=arith.constant256:index
%c128=arith.constant128:index
%reinterpret_cast=memref.reinterpret_cast%arg0tooffset:[0],sizes:[128,64],strides:[%c128,1]:memref<*xbf16>tomemref<128x64xbf16,?strided<[?,?1]>>
%alloc=memref.alloc():memref<128x64xbf16>
memref.copy%reinterpret_cast,%alloc:memref<128x64xbf16,?strided<[?,?1]>>tomemref<128x64xbf16>
%0=bufferization.to_tensor%allocrestrictwritable:memref<128x64xbf16>
%reinterpret_cast_0=memref.reinterpret_cast%arg1tooffset:[0],sizes:[256,64],strides:[1,%c256]:memref<*xbf16>tomemref<256x64xbf16,?strided<[1,??]>>
%alloc_1=memref.alloc():memref<256x64xbf16>
memref.copy%reinterpret_cast_0,%alloc_1:memref<256x64xbf16,?strided<[1,??]>>tomemref<256x64xbf16>
%1=bufferization.to_tensor%alloc_1restrictwritable:memref<256x64xbf16>
%2=tensor.empty():tensor<64x256xbf16>
%transposed=linalg.transposeins(%1:tensor<256x64xbf16>)outs(%2:tensor<64x256xbf16>)permutation=[1,0]
%reinterpret_cast_2=memref.reinterpret_cast%arg2tooffset:[0],sizes:[128,256],strides:[%c256,1]:memref<*xbf16>tomemref<128x256xbf16,?strided<[?,?1]>>
%alloc_3=memref.alloc():memref<128x256xbf16>
memref.copy%reinterpret_cast_2,%alloc_3:memref<128x256xbf16,?strided<[?,?1]>>tomemref<128x256xbf16>
%3=bufferization.to_tensor%alloc_3restrictwritable:memref<128x256xbf16>
%4=tensor.empty():tensor<128x256xbf16>
%5=linalg.matmulins(%0,%transposed:tensor<128x64xbf16>,tensor<64x256xbf16>)outs(%4:tensor<128x256xbf16>)->tensor<128x256xbf16>
%6=linalg.generic{indexing_maps=[#map,#map,#map],iterator_types=["parallel","parallel"]}ins(%5,%3:tensor<128x256xbf16>,tensor<128x256xbf16>)outs(%5:tensor<128x256xbf16>){
^bb0(%in:bf16,%in_4:bf16,%out:bf16):
%7=arith.addf%in,%in_4:bf16
linalg.yield%7:bf16
}->tensor<128x256xbf16>
memref.tensor_store%6,%reinterpret_cast_2:memref<128x256xbf16,?strided<[?,?1]>>
return
}
}

關(guān)于其他更加具體的operator，我們可以都按照上述流程來進(jìn)行操作，一旦你的編譯框架是基于MLIR來開發(fā)的，那么如果能很好的轉(zhuǎn)到Linalg，那么就說明了后續(xù)在接入自己的backend以及適配一些ISA的過程就會方便不少，這也從另外一個(gè)角度彰顯了為什么現(xiàn)在的趨勢都是將自己的compiler通過MLIR進(jìn)行重構(gòu)。最重要的原因，其實(shí)就是以最小的開發(fā)成本方便的接入各種軟件或者硬件的生態(tài)。

后記

對triton的研究已經(jīng)有一段時(shí)間了，由于當(dāng)時(shí)學(xué)triton也是基于源碼一步一步硬吃過來的，并且triton也沒有比較好的中文教程，所以后面會利用空閑時(shí)間將我目前對于使用triton來做codegen的各種優(yōu)化方法(不同backend以及不同IR層面的pass)和細(xì)節(jié)(底層layout的設(shè)計(jì))進(jìn)行一個(gè)詳細(xì)的梳理，來幫助更多想要使用triton來做codegen的同學(xué)。