如何設(shè)計MLIR的Dialect來在GPU上生成高性能的代碼？

前言

為什么又要開一個新坑？原因是，最近在做的項目都是和MLIR有關(guān)，并且發(fā)現(xiàn)自己已經(jīng)在MLIR的研發(fā)道路上越走越遠了。剛剛好前段時間大家都在跟風(fēng)各種GPT，就去看了看openai目前放出來的產(chǎn)品，無意間發(fā)現(xiàn)了triton這把瑞士軍刀。其實早在一些年前就聽過triton，那會的triton代碼還沒有被MLIR進行重構(gòu)，代碼內(nèi)部的某些邏輯寫的也沒有看的很明白，結(jié)合"Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations"這篇論文其實也沒有看出太多新的東西。這次在重新?lián)炱饋砜吹臅r候，發(fā)現(xiàn)其中很多不錯的優(yōu)化，其實還是抱著學(xué)習(xí)如何設(shè)計MLIR的Dialect來在GPU上生成高性能的代碼為初衷，來對triton進行一個深入的分析。

什么是Triton?

Triton是openai針對gpu上的算子優(yōu)化提出的一個programming language & compiler。以NVIDIA GPU為例，使用triton可以越過cuda的閉源生態(tài)，直接將自己的后端接入llvm IR，通過走NVPTX來生成在GPU上的代碼。這樣做相比傳統(tǒng)的手寫cuda代碼的好處是可以不需要借助NVIDIA的nvcc compiler就可以得到在GPU上能跑的machine code。同時，triton的一些設(shè)計理念對于 不管是深度學(xué)習(xí)領(lǐng)域還是其他數(shù)據(jù)科學(xué)領(lǐng)域來做高性能計算來說都可以提供豐富的指導(dǎo)意義。同時，triton不僅僅只支持nv的gpu，它同時對amd的gpu，intel的gpu都會推出后續(xù)的支持方案，這其所就彰顯出了mlir的優(yōu)勢，可以通過設(shè)計Dialect來支持更多的后端。

源碼編譯Triton

接下來帶大家一起從源碼來編譯下triton的代碼，后續(xù)我準備分為幾章，對triton的設(shè)計以及具體的優(yōu)化細節(jié)展開分析，能夠給大家一個較為全面的理解。畢竟triton作為mlir中為數(shù)不多成功的end-to-end的例子，對于編譯技術(shù)和系統(tǒng)優(yōu)化的研究者或者工程師來說，都是不可或缺的好資料了。

0x0 先去官網(wǎng)clone triton的官方repo

$gitclonehttps://github.com/openai/triton.git
$cdtriton
$gitcheckout132fe1bb01e0a734d39c60835c76da257dbe7151

0x1 安裝第三方依賴

Triton 整個源碼編譯的過程中，需要使用到兩個最為重要的依賴，一個是llvm，一個是pybind11，我在編譯和構(gòu)建triton的過程中，都是通過手動將llvm和pybind11編譯安裝好后，在編譯triton的過程中通過CMakLists.txt來指定對應(yīng)的路徑。

0x10 LLVM的下載與配置

為什么要使用llvm？其實大家都知道，這就是triton最吸引人的地方，通過將高層的python代碼一步一步lower到llvm IR，然后通過llvm生態(tài)得到最終可以跑在具體設(shè)備上的machine code，將llvm作為最重要的后端，并且triton內(nèi)部的實現(xiàn)也被MLIR進行重構(gòu)，MLIR剛剛好也是llvm中非常重要的一個子項目。那么，llvm的安裝對于想要基于triton來進行二次開發(fā)的工程師或者研究人員來說就顯得非常重要了。

$gitclonehttps://github.com/llvm/llvm-project
$cdllvm-project
$gitcheckoutf733b4fb9b8b
$mkdirbuild
$cdbuild
$cmake-GNinja../llvm
-DLLVM_ENABLE_PROJECTS=mlir
-DLLVM_BUILD_EXAMPLES=ON
-DLLVM_TARGETS_TO_BUILD="X86;NVPTX;RISCV;AMDGPU"
-DMLIR_ENABLE_CUDA_RUNNER=ON
-DCMAKE_BUILD_TYPE=Release
-DLLVM_ENABLE_ASSERTIONS=ON
-DCMAKE_C_COMPILER=clang
-DCMAKE_CXX_COMPILER=clang++
-DLLVM_ENABLE_RTTI=ON
-DLLVM_INSTALL_UTILS=ON
-DMLIR_INCLUDE_INTEGRATION_TESTS=ON

ninja-j8
sudoninjainstall

經(jīng)過一定時間的等待，就可以將llvm裝好了

0x11 pybind11的下載與配置

為什么要使用pybind11？pybind11已經(jīng)是目前主流的ai開發(fā)工具中必不可少的組件了。大部分的框架都以python的DSL暴露給用戶，然后用戶通過寫對應(yīng)的python語法，調(diào)用已經(jīng)用C++/CUDA或者assemble寫好的高性能組件。那么，裝配pybind11的目的就是為了能夠讓我們通過import triton，然后絲滑調(diào)用對應(yīng)的python api來完成高性能算子生成的任務(wù)。

$pipinstallpytest
$gitclonehttps://github.com/pybind/pybind11.git
$cdpybind11
$mkdirbuild
$cdbuild
$cmake..
$makecheck-j8
$sudomakeinstall

0x2 編譯Triton

$cdtriton
$vimCMakeLists.txt(option(TRITON_BUILD_PYTHON_MODULE"BuildPythonTritonbindings"ON))
$mkdirbuild
$cdbuild
$cmake..
$make-j8

編輯切換為居中

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

可以看到最終生成了一個.so文件，libtriton.so

那么接下來只要將libtriton.so文件移動到triton/python/triton/_C目錄下，將triton的python路徑下入bashrc

exportTRITON_HOME=/home/Documents/compiler/triton
exportPYTHONPATH=$TRITON_HOME/python:${PYTHONPATH}

然后通過簡單的import triton，沒有任何錯誤則可以通過triton進行開發(fā)了

img

編輯切換為居中

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

接下來進入triton/python/tutorials,隨便找一個例子進行驗證，這里我們選擇最常見和實用的03-matrix-multiplication.py，直接python 03-matrix-multiplication.py，稍等片刻則可以得到最終結(jié)果。

img

編輯

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

可以看到，triton最終生成的代碼，在3090上，對應(yīng)單batch的gemm在部分size上已經(jīng)超過了cuBLAS。

同時，可以在build目錄下去檢查對應(yīng)的三個bin tool: triton-opt, triton-reduce, triton-translate

然后將本機下的ptxas復(fù)制到該build目錄下，我的ptxas在(/usr/local/cuda-11.6/bin)下。關(guān)于這些工具的使用將會在后續(xù)的解讀中根據(jù)不同層的dialect之間的conversion來進行詳細介紹。

0x3 為什么采用這樣的編譯方式？

其實有的同學(xué)說，直接按照triton教程里的pip install -e . 不就行了么？這樣做的原因是因為后續(xù)我們需要對triton以及對應(yīng)的llvm進行改進，每次改進后，都需要對triton和llvm分別進行編譯。這種分離的方式，可以使得我們在改進完對應(yīng)的llvm代碼或者triton的源碼后，分步編譯，然后再整合成一個新的shared library (libtriton.so)

在后續(xù)的教程中，我將會從triton的frontend, optimizer，backend來作為切入點，分別講講triton是如何通過將用戶手寫的python dsl編譯成能跑在gpu上的machine code的。

Triton目前的設(shè)計

從triton的源碼來看，triton目前在NV的GPU上已經(jīng)有了一套自己比較成熟的mapping路線，通過先對python語言層，也就是triton DSL進行抽象，得到AST，然后將AST中的每個節(jié)點lower到Triton Dialect上，Triton Dialect則是一個比較貼近上層語言表達的IR，他的主要作用則是為了保持用戶在書寫對應(yīng)算法時的準確性。接下來會進一步被map到TritonGPU Dialect上，那么TritonGPU Dialect則是一個更加貼近GPU層面的IR，它則是為了具體的性能優(yōu)化而設(shè)計。圖中其他的藍色模塊，比如SCF，Arith，Tensor等都是MLIR生態(tài)中已經(jīng)被實現(xiàn)好并且廣為使用的Dialect。

這些Dialect會一起和TritonGPU Dialect共存，然后被lower到對應(yīng)的LLVM Dialect，LLVM Dialect則是最貼近LLVM IR的一層設(shè)計，從LLVM Dialect到LLVM IR的轉(zhuǎn)換是非常容易的，最終代碼就會被接入到LLVM的NVPTX的后端，從而生成后續(xù)能跑在GPU上的高性能machine code.

img

編輯切換為居中

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

Triton 未來的支持

通過下圖可以看到，triton的未來計劃和大多數(shù)的compiler有著一樣的發(fā)展藍圖，向上去支持各種各樣具有不同表達能力的前端。向下對接各種不同廠商的hardware，最終將一個application高效的map到一個硬件上。

從哦的

img

編輯切換為居中

審核編輯：劉清