前言

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

什么是Triton?

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

源碼編譯Triton

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

0x0 先去官網clone triton的官方repo

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

0x1 安裝第三方依賴

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

0x10 LLVM的下載與配置

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

$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

經過一定時間的等待，就可以將llvm裝好了

0x11 pybind11的下載與配置

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

$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進行開發了

img

編輯切換為居中

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

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

img

編輯

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

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

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

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

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

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

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

Triton目前的設計

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

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

img

編輯切換為居中

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

Triton 未來的支持

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

從哦的

img

編輯切換為居中

審核編輯：劉清