Skip to content

Fix TritonSan build error #359

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lechenyu merged 3 commits into from
Oct 27, 2025
Merged

Fix TritonSan build error #359

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
10ad047
Fix triton-san build error
lechenyu Oct 25, 2025
2494bdc
Update triton-san readme
lechenyu Oct 25, 2025
88b5f5d
Add description and TODO for the llvm build error
lechenyu Oct 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion triton-san/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,11 +66,15 @@ Usage: triton-san <sanitizer type> <original command used to launch the triton p
Example: triton-san asan python ./my_triton_program.py
```

**Note: before running TritonSan, please add the following import to the Triton program to specify the use of the CPU backend, which ensures all Triton kernels run on the CPU. The sanitizers require CPU backend in order to work.**
**Note: before running TritonSan, please add the following import to the Triton program to specify the use of the CPU backend, which ensures all Triton kernels run on the CPU. The sanitizers require CPU backend in order to work. In addition, all desired GPU tensors in the Triton program need to be set to CPU**

```python
from triton.backends.triton_shared.driver import CPUDriver
triton.runtime.driver.set_active(CPUDriver())
...

#output = torch.empty((size, )).to("gpu")
output = torch.empty((size, )).to("cpu")
```

## Example
Expand Down
49 changes: 39 additions & 10 deletions triton-san/build.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ function check_required_files() {
}

PARENT_FOLDER="$(realpath "$(dirname "$0")")"
TRITON_SHARED_PATH="$(realpath "${PARENT_FOLDER}/..")"
ROOT="$(realpath "${PARENT_FOLDER}/../..")"
SCRIPT_FOLDER="$(realpath "${PARENT_FOLDER}/script")"
VENV_PATH="${ROOT}/venv"
Expand All @@ -23,9 +24,42 @@ TRITON_SAN_INSTALL_DIR="${ROOT}"
# include utility functions
source "${SCRIPT_FOLDER}/utility.inc"

echo "================= Check out triton ================"
TRITON_HASH_FILE="${TRITON_SHARED_PATH}/triton-hash.txt"
if [ ! -e "${TRITON_HASH_FILE}" ]; then
print_error_and_exit "${TRITON_HASH_FILE} does not exist."
fi

EXPECT_HEAD_COMMIT=$(cat "${TRITON_HASH_FILE}")
if [ -e "${TRITON_PATH}" ]; then
warning_msg=("The path ${TRITON_PATH} already exists.")
print_warning "${warning_msg[@]}"
pushd "${TRITON_PATH}" > /dev/null
HEAD_COMMIT=$(git rev-parse HEAD)
if [ "${HEAD_COMMIT}" != "${EXPECT_HEAD_COMMIT}" ]; then
error_msg=("The head commit of the existing Triton repository does not match the expected commit hash."
"HEAD: ${HEAD_COMMIT}, EXPECTED: ${EXPECT_HEAD_COMMIT}")
print_error_and_exit "${error_msg[@]}"
fi
EXIST_TRITON_BUILD="${TRITON_PATH}/build"
if [ -e "${EXIST_TRITON_BUILD}" ]; then
rm -rf "${TRITON_PATH}/build"
echo "Remove ${TRITON_PATH}/build"
fi
popd > /dev/null
else
mkdir -p "${TRITON_PATH}"
git clone https://github.com/triton-lang/triton.git "${TRITON_PATH}"
pushd "${TRITON_PATH}" > /dev/null
cd "${TRITON_PATH}" && git checkout ${EXPECT_HEAD_COMMIT}
popd > /dev/null
fi

echo -e "\n\n\n"

echo "================= Setup Python virtual environment ================"
if [ ! -e "${VENV_PATH}" ]; then
"${SCRIPT_FOLDER}/setup_prerequisite.sh" "${VENV_PATH}"
"${SCRIPT_FOLDER}/setup_prerequisite.sh" "${TRITON_PATH}" "${VENV_PATH}"
else
required_file=(
"${VENV_PATH}/bin/activate"
Expand All @@ -47,7 +81,7 @@ echo -e "\n\n\n"
echo "=========================== Build LLVM ============================"
if [ ! -e "${LLVM_PATH}" ]; then
mkdir -p "${LLVM_PATH}"
"${SCRIPT_FOLDER}/build_llvm.sh" "${LLVM_PATH}"
"${SCRIPT_FOLDER}/build_llvm.sh" "${TRITON_PATH}" "${LLVM_PATH}"
else
required_file=(
"${LLVM_PATH}/llvm-install/bin"
Expand All @@ -64,14 +98,9 @@ fi
echo -e "\n\n\n"

echo "================== Build trion-shared and triton =================="
if [ -e "${TRITON_PATH}" ]; then
warning_msg=("The path ${TRITON_PATH} already exists and will be overwritten.")
print_warning "${warning_msg[@]}"
rm -rf "${TRITON_PATH}"
fi
mkdir -p "${TRITON_PATH}"
"${SCRIPT_FOLDER}/build_triton_shared_with_triton_san.sh" "${LLVM_BUILD_PATH}" "${VENV_PATH}" "${TRITON_PATH}"
export PATH="${LLVM_PATH}/llvm-install/bin:${PATH}"
"${SCRIPT_FOLDER}/build_triton_shared_with_triton_san.sh" "${LLVM_BUILD_PATH}" "${VENV_PATH}" "${TRITON_PATH}" "${TRITON_SHARED_PATH}"
echo -e "\n\n\n"

echo "======================== Install triton-san ======================="
"${SCRIPT_FOLDER}/install_triton_san.sh" "${TRITON_SAN_INSTALL_DIR}" "${LLVM_INSTALL_DIR}" "${VENV_PATH}" "${TRITON_PATH}"
"${SCRIPT_FOLDER}/install_triton_san.sh" "${LLVM_INSTALL_DIR}" "${VENV_PATH}" "${TRITON_PATH}" "${TRITON_SAN_INSTALL_DIR}"
23 changes: 17 additions & 6 deletions triton-san/script/build_llvm.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,21 +6,21 @@
set -e

if [ "$#" -lt 1 ]; then
echo "Usage: $0 <desired path for LLVM installation>"
echo "Usage: $0 <path to Triton source directory> <desired path for LLVM installation>"
exit 1
fi

PARENT_FOLDER="$(realpath "$(dirname "$0")")"
TRITON_SHARED_PATH="$(realpath "${PARENT_FOLDER}/../..")"
LLVM_PATH="$(realpath "$1")"
TRITON_PATH="$(realpath "$1")"
LLVM_PATH="$(realpath "$2")"

# include utility functions
source "${PARENT_FOLDER}/utility.inc"

print_info "Installing LLVM to path: ${LLVM_PATH}."
cd $LLVM_PATH

LLVM_HASH_FILE="${TRITON_SHARED_PATH}/triton/cmake/llvm-hash.txt"
LLVM_HASH_FILE="${TRITON_PATH}/cmake/llvm-hash.txt"
if [ ! -e "${LLVM_HASH_FILE}" ]; then
print_error_and_exit "${LLVM_HASH_FILE} does not exist."
fi
Expand All @@ -32,7 +32,7 @@ LLVM_SOURCE="${LLVM_SOURCE_DIR}/llvm"

# compiler-rt and clang are the sanitizer-specific LLVM projects
# openmp is used for parallelizing the triton grid for ThreadSanitizer (TSan)
LLVM_PROJECTS="clang;compiler-rt;openmp;mlir"
LLVM_PROJECTS="clang;compiler-rt;openmp;mlir;lld"

# these are the targets supported by the Triton language
# Triton's build script for LLVM uses these exact targets
Expand All @@ -49,6 +49,17 @@ LLVM_HASH=$(cat "${LLVM_HASH_FILE}")
cd "${LLVM_SOURCE_DIR}"
git checkout ${LLVM_HASH}

# Cherry-pick the patch to resolve the OpenMP build error. An incorrect setting in openmp/runtime/src/CMakeLists.txt
# causes the generated omp-tools.h to be placed in Clang's include directory instead of OpenMP's, leading to a
# compilation failure when locating omp-tools.h.
# More details are available at this link: https://github.com/llvm/llvm-project/commit/62ff9ac4c68f48c089528105259c68943ab176de

# TODO: remove this once the LLVM hash for the next LLVM->Triton merge beyond this commit is confirmed.
if ! git merge-base --is-ancestor 62ff9ac HEAD; then
Comment thread
bmyerz0 marked this conversation as resolved.
echo "cherry pick commit 62ff9ac to avoid OpenMP build failure"
git cherry-pick 62ff9ac
fi

export CXXFLAGS="-Wno-unused-command-line-argument $CXXFLAGS"
export CFLAGS="-Wno-unused-command-line-argument $CFLAGS"

Expand All @@ -64,4 +75,4 @@ cmake -GNinja -DCMAKE_BUILD_TYPE=Release \
ninja -j$(nproc --all) -l$(nproc --all) || exit -1
ninja -j$(nproc --all) -l$(nproc --all) install

print_info "LLVM installation succeeded."
print_info "LLVM installation succeeded."
15 changes: 5 additions & 10 deletions triton-san/script/build_triton_shared_with_triton_san.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,15 @@
set -e

if [ "$#" -lt 3 ]; then
echo "Usage: $0 <path to LLVM build directory> <path to Python venv> <path to Triton source directory>"
echo "Usage: $0 <path to LLVM build directory> <path to Python venv> <path to Triton source directory> <path to triton-shared source directory>"
exit 1
fi

PARENT_FOLDER="$(realpath "$(dirname "$0")")"
TRITON_SHARED_PATH="$(realpath "${PARENT_FOLDER}/../..")"
LLVM_BUILD_PATH="$(realpath "$1")"
VENV_PATH="$(realpath "$2")"
TRITON_PATH="$(realpath "$3")"
TRITON_SHARED_PATH="$(realpath "$4")"

# include utility functions
source "${PARENT_FOLDER}/utility.inc"
Expand Down Expand Up @@ -46,15 +46,10 @@ fi
mkdir -p "${TRITON_HOME}"

# build triton-shared and triton with the custom LLVM
TRITON_HASH_FILE="${TRITON_SHARED_PATH}/triton-hash.txt"
if [ ! -e "${TRITON_HASH_FILE}" ]; then
print_error_and_exit "${TRITON_HASH_FILE} does not exist."
fi

git clone https://github.com/triton-lang/triton.git "${TRITON_PATH}"
cd "${TRITON_PATH}" && git checkout $(cat "${TRITON_HASH_FILE}")
pushd "${TRITON_PATH}" > /dev/null
export TRITON_PLUGIN_DIRS="${TRITON_SHARED_PATH}"
export LLVM_BUILD_DIR="${LLVM_BUILD_PATH}"
LLVM_INCLUDE_DIRS="${LLVM_BUILD_DIR}/include" LLVM_LIBRARY_DIR="${LLVM_BUILD_DIR}/lib" LLVM_SYSPATH="${LLVM_BUILD_DIR}" TRITON_BUILD_WITH_CLANG_LLD=true TRITON_BUILD_WITH_CCACHE=false python3 -m pip install --no-build-isolation -vvv '.[tests]'
popd > /dev/null

print_info "triton-shared and triton build completed."
print_info "triton-shared and triton build completed."
12 changes: 6 additions & 6 deletions triton-san/script/install_triton_san.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,16 @@
set -e

if [ "$#" -lt 4 ]; then
echo "Usage: $0 <desired path for triton-san installation> <path to LLVM install directory> <path to Python venv> <path to Triton source directory>"
echo "Usage: $0 <path to LLVM install directory> <path to Python venv> <path to Triton source directory> <desired path for triton-san installation>"
exit 1
fi

PARENT_FOLDER="$(realpath "$(dirname "$0")")"
TRITON_SAN_PATH="$(realpath "${PARENT_FOLDER}/..")"
TRITON_SAN_INSTALL_DIR="$(realpath "$1")"
LLVM_INSTALL_DIR="$(realpath "$2")"
VENV_PATH="$(realpath "$3")"
TRITON_PATH="$(realpath "$4")"
LLVM_INSTALL_DIR="$(realpath "$1")"
VENV_PATH="$(realpath "$2")"
TRITON_PATH="$(realpath "$3")"
TRITON_SAN_INSTALL_DIR="$(realpath "$4")"

# include utility functions
source "${PARENT_FOLDER}/utility.inc"
Expand Down Expand Up @@ -57,4 +57,4 @@ cp -r "${TRITON_SAN_PATH}/examples" "${TRITON_SAN_INSTALL_ROOT}/examples"
info_message=("triton-san has been installed into ${TRITON_SAN_INSTALL_ROOT} successfully."
"To optionally enable global access, add triton-san to your PATH environment variable using:"
'export PATH="${PATH}:'"${TRITON_SAN_INSTALL_ROOT}\"")
print_info "${info_message[@]}"
print_info "${info_message[@]}"
8 changes: 4 additions & 4 deletions triton-san/script/setup_prerequisite.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,14 @@
set -e

if [ "$#" -lt 1 ]; then
echo "Usage: $0 <desired path to venv>"
echo "Usage: $0 <path to Triton source directory> <desired path to venv>"
exit 1
fi

PARENT_FOLDER="$(realpath "$(dirname "$0")")"
TRITON_SHARED_PATH="$(realpath "${PARENT_FOLDER}/../..")"
TRITON_PATH="${TRITON_SHARED_PATH}/triton"
VENV_PATH="$(realpath "$1")"
TRITON_PATH="$(realpath "$1")"
VENV_PATH="$(realpath "$2")"

# include utility functions
source "${PARENT_FOLDER}/utility.inc"
Expand All @@ -36,4 +36,4 @@ python3 -m pip install pytest-xdist torch
# echo to user to start virtual environment
info_message=("Please start the virtual environment: source \"${VENV_PATH}/bin/activate\""
"This export is recommended for subsequent triton-shared builds: export TRITON_PLUGIN_DIRS=\"${TRITON_SHARED_PATH}/triton_shared\"")
print_info "${info_message[@]}"
print_info "${info_message[@]}"
⚡