Fix TritonSan build error with triton version e44bd1c (#359)

This PR resolves the TritonSan build error. The build script was broken
due to several issues:

- Removal of the Triton submodule: We now use triton-hash.txt to record
the Triton version. The previous logic for locating Triton submodule was
invalidated by this change.

- LLVM-related issue: In the corresponding LLVM version, there is a
known bug that causes [a compilation error when building
OpenMP](llvm/llvm-project@62ff9ac).
The build script now cherry-picks this patch to address the problem.

This PR fixes all these issues. The build script will now automatically
check out Triton and place it alongside triton-shared, LLVM, and the
Python virtual environment. There are no changes to the usage of the
build script or the TritonSan driver script.

Author: lechenyu

triton-san/README.md

@@ -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

triton-san/build.sh

@@ -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"
@@ -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"
@@ -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"
@@ -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}"

triton-san/script/build_llvm.sh

@@ -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
@@ -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
@@ -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
+  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" 
 
@@ -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."

triton-san/script/build_triton_shared_with_triton_san.sh

@@ -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"
@@ -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."

triton-san/script/install_triton_san.sh

@@ -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"
@@ -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[@]}"

triton-san/script/setup_prerequisite.sh

@@ -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"
@@ -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[@]}"