Triton-XDNA

This repository contains a plugin for building AIR as Triton's compiler backend.

Usage

Clone the repository

git clone https://github.com/AARInternal/triton-xdna.git
cd triton-xdna
git submodule update --init

Install XRT

Please follow the instructions in mlir-aie project on how to install the XDNA driver.

Setup build environment

Option 1: Install Pre-built Wheel (Recommended)

The easiest way to get started is to install the pre-built wheel from GitHub Releases:

python3 -m venv sandbox
source sandbox/bin/activate
python3 -m pip install --upgrade pip

# Install triton-xdna from GitHub Releases
pip install triton-xdna \
  --find-links https://github.com/AARInternal/triton-xdna/releases/expanded_assets/latest-wheels \
  --find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-no-rtti \
  --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly \
  --find-links https://github.com/Xilinx/mlir-air/releases/expanded_assets/latest-air-wheels-no-rtti

Note: To install from a local wheel file:

pip install /path/to/triton_xdna-*.whl \
  --find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-no-rtti \
  --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly \
  --find-links https://github.com/Xilinx/mlir-air/releases/expanded_assets/latest-air-wheels-no-rtti

Option 2: Build from Source (Using Pip)

Starting from the root of the repository:

python3 -m venv sandbox
source sandbox/bin/activate
python3 -m pip install --upgrade pip
pip install cmake pybind11 nanobind wheel ninja pytest setuptools Cython

# Install triton-xdna from source and all dependencies automatically
pip install . \
  --find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-no-rtti \
  --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly \
  --find-links https://github.com/Xilinx/mlir-air/releases/expanded_assets/latest-air-wheels-no-rtti

This will automatically install all required dependencies:

• mlir-aie
• llvm-aie
• mlir-air

The versions are managed in utils/mlir-aie-hash.txt, utils/llvm-aie-hash.txt, and utils/mlir-air-hash.txt.

Option 3: Build from Source (Using Cmake)

python3 -m venv sandbox
source sandbox/bin/activate
python3 -m pip install --upgrade pip
pip install cmake pybind11 nanobind wheel ninja pytest setuptools Cython
source utils/env_setup.sh

cmake cmake -GNinja -S . -Bbuild
cd build
ninja

Cmake shall install the C++ binaries under third_party/triton/python/build. A triton python package with a new amd_triton_npu backend is also pip installed to the virtual environment sandbox.

Run examples

Please make sure to run source {path_to_xrt}/setup.sh before running examples. The test also depends on PyTorch as CPU reference.

cd examples/matmul
AIR_TRANSFORM_TILING_SCRIPT=transform_aie2.mlir python matmul.py

Note: The transform_aie2.mlir transform dialect IR is specifically designed for the AIE2 architecture. For AIE2P architecture, use transform_aie2p.mlir instead.