Triton Linalg编译

简介

triton-linalg是一个将Triton Dialect转换为Linalg Dialect的项目。

当前作为寒武纪Triton编译器的前端，几乎支持了所有Triton语言的特性。

在Triton转换为Linalg的过程中，遵从了几个原则：

• 尽可能使用结构化（Structured）的Operator，避免使用linalg.generic（维持Operator的Coarse-Grained语义）。例如：

  • tt.reduce转换为linalg.reduce。

  • tt.dot转换为linalg.matmul。

  • arith和math中逐元素的Operator转换为对应的linalg.map版本

• 尽早识别Operator的语义。例如：tt.reduce变换为linalg.pool。

源代码下载

clone源代码和submodule（triton）：

git clone --recurse-submodules https://github.com/Cambricon/triton-linalg.git

查看triton-linalg的git log：

查看submodule triton的git log：

LLVM编译

查看配套的llvm的commit id：

clone llvm代码（可以选择任意目录，作者选择放到和triton-linalg平行的目录）：

git clone https://github.com/llvm/llvm-project.git

再checkout到配套的commit id：

git checkout 10dc3a8e916d73291269e5e2b82dd22681489aa1

LLVM编译：

mkdir build && cd build
cmake -G Ninja -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=ON ../llvm -DLLVM_ENABLE_PROJECTS="mlir;llvm;lld" -DLLVM_TARGETS_TO_BUILD="host;NVPTX;AMDGPU"
ninja -j$(nproc)

export LLVM路径，让后续Triton编译使用本地编译的这个LLVM：

export LLVM_BUILD_DIR=$HOME/llvm-project/build
export LLVM_INCLUDE_DIRS=$LLVM_BUILD_DIR/include
export LLVM_LIBRARY_DIR=$LLVM_BUILD_DIR/lib
export LLVM_SYSPATH=$LLVM_BUILD_DIR

Triton&Triton-Linalg编译

创建conda环境并激活：

conda create -n triton-linalg python=3.12
conda activate triton-linalg

安装依赖的包：

python3 -m pip install --upgrade pip
python3 -m pip install cmake==3.24 ninja pytest-xdist
sudo apt-get update -y
sudo apt-get install -y ccache clang lld

将triton-linalg加入triton的插件：

export TRITON_PLUGIN_DIRS=$(pwd)

启动编译：

pip install -e python --no-build-isolation -vvv

在编译过程中会下载编译所需的二进制，路径在~/.triton（老版本triton不支持配置）：