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

前言

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

什么是Triton?

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

源碼編譯Triton

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

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

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

0x1 安裝第三方依賴

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

0x10 LLVM的下載與配置

為什么要使用llvm？其實(shí)大家都知道，這就是triton最吸引人的地方，通過將高層的python代碼一步一步lower到llvm IR，然后通過llvm生態(tài)得到最終可以跑在具體設(shè)備上的machine code，將llvm作為最重要的后端，并且triton內(nèi)部的實(shí)現(xiàn)也被MLIR進(jìn)行重構(gòu)，MLIR剛剛好也是llvm中非常重要的一個(gè)子項(xiàng)目。那么，llvm的安裝對(duì)于想要基于triton來進(jìn)行二次開發(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)過一定時(shí)間的等待，就可以將llvm裝好了

0x11 pybind11的下載與配置

為什么要使用pybind11？pybind11已經(jīng)是目前主流的ai開發(fā)工具中必不可少的組件了。大部分的框架都以python的DSL暴露給用戶，然后用戶通過寫對(duì)應(yīng)的python語法，調(diào)用已經(jīng)用C++/CUDA或者assemble寫好的高性能組件。那么，裝配pybind11的目的就是為了能夠讓我們通過import triton，然后絲滑調(diào)用對(duì)應(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 字（可選）

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

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

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

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

img

編輯切換為居中

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

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

img

編輯

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

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

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

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

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

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

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

Triton目前的設(shè)計(jì)

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

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

img

編輯切換為居中

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

Triton 未來的支持

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

從哦的

img

編輯切換為居中

審核編輯：劉清