Enable LLVM ThreadSanitizer for triton-shared (#301)

This PR builds on top of
#294, which enabled LLVM
AddressSanitizer for triton-shared. Enabling ThreadSanitizer support
involves the same infrastructure, with some small differences:
- Requires `-fsanitize=thread` during compiling and linking
- Requires parallelizing the Triton grid using OpenMP in order to detect
data races. The LLVM build script has also been changed to reflect this
- Requires some TSan-specific suppressions and libraries (archer), added
to the run script

---------

Co-authored-by: Tim Lu <[email protected]>

Author: timthlu

README.md

@@ -206,6 +206,9 @@ $ ls /tmp/ir_dumps
 ll.ir  ll.mlir  tt.mlir  ttshared.mlir
 ```
 
+## Debugging Triton Programs
+Triton-shared includes a build option that enables LLVM-sanitizers - AddressSanitizer (ASan) and ThreadSanitizer (TSan) - to help detect memory safety and concurrency issues in Triton programs. These sanitizers dynamically analyze the program during execution, identifying bugs such as buffer overflows and data races respectively. For more details and setup instructions, refer [here](https://github.com/microsoft/triton-shared/blob/main/scripts/SANITIZER.md).
+
 ## Contributing
 
 This project welcomes contributions and suggestions.  Most contributions require you to agree to a

backend/compiler.py

@@ -36,10 +36,10 @@ def _dump_ir_if_needed(files):
 
 def _get_sanitizer_type():
     # returns "" if not set
-    # throws error if set to something other than "asan"
+    # throws error if set to something other than "asan" or "tsan"
     sanitizer_type = os.getenv("TRITON_SHARED_SANITIZER_TYPE", "")
 
-    if sanitizer_type != "" and sanitizer_type != "asan":
+    if sanitizer_type != "" and sanitizer_type != "asan" and sanitizer_type != "tsan":
         # throw error
         raise Exception(f"TRITON_SHARED_SANITIZER_TYPE {sanitizer_type} is invalid.")
 
@@ -164,6 +164,8 @@ def _llir_to_bin(llir: str, metadata):
 
             if sanitizer_type == "asan":
                 subprocess_args.extend(["-g", "-fsanitize=address", "-mllvm", "-asan-stack=0"])
+            elif sanitizer_type == "tsan":
+                subprocess_args.extend(["-g", "-fsanitize=thread"])
 
             subprocess.check_call(subprocess_args)
         else:

backend/driver.py

@@ -20,10 +20,10 @@ def _get_llvm_bin_path(bin_name: str) -> str:
 
 def _get_sanitizer_type():
     # returns "" if not set
-    # throws error if set to something other than "asan"
+    # throws error if set to something other than "asan" or "tsan"
     sanitizer_type = os.getenv("TRITON_SHARED_SANITIZER_TYPE", "")
 
-    if sanitizer_type != "" and sanitizer_type != "asan":
+    if sanitizer_type != "" and sanitizer_type != "asan" and sanitizer_type != "tsan":
         # throw error
         raise Exception(f"TRITON_SHARED_SANITIZER_TYPE {sanitizer_type} is invalid.")
 
@@ -114,6 +114,9 @@ def _generate_launcher(constants, signature, kernel_name):
 static void _launch(int gridX, int gridY, int gridZ, {arg_decls}) {{
   if (gridX*gridY*gridZ > 0) {{
     // Cast "function" to the real function type.
+    // apply parallelization to the triton grid when using ThreadSanitizer (TSan) 
+    // to help detect potential data races across program instances during kernel execution
+    {"#pragma omp parallel for collapse(3)" if _get_sanitizer_type() == "tsan" else ""}
     for(int x = 0; x < gridX; x++) {{
       for(int y = 0; y < gridY; y++) {{
         for(int z = 0; z < gridZ; z++) {{
@@ -317,6 +320,16 @@ def launch(
 
                       if sanitizer_type == "asan":
                           subprocess_args.extend(["-g", "-fsanitize=address", "-mllvm", "-asan-stack=0"])
+                      elif sanitizer_type == "tsan":
+                          # ensure that openmp is available
+                          libomp_path = next(Path(Path(_get_llvm_bin_path("")).parent).rglob("libomp.so"), None)
+
+                          if not libomp_path:
+                              raise Exception(f"libomp.so does not exist.")
+
+                          libomp_path = str(libomp_path.parent)
+
+                          subprocess_args.extend(["-g", "-fsanitize=thread", "-fopenmp", f"-Wl,-rpath,{libomp_path}"])
 
                       subprocess.check_call(subprocess_args)
                   else:

scripts/SANITIZER.md

@@ -38,5 +38,5 @@ For building LLVM and triton_shared for usage with sanitizers, run the scripts i
 
 For runtime setup (one-time per shell): `source setup_runtime_for_sanitizers.sh <existing path to venv> <existing path to llvm install dir> <existing path to triton shared>`
 
-For running a python program with sanitizers enabled: `run_triton_with_sanitizers.sh <sanitizer type> python program.py`. Currently, the only supported `<sanitizer type>` is `asan`.
+For running a python program with sanitizers enabled: `run_triton_with_sanitizers.sh <sanitizer type> python program.py`. Currently, the supported `<sanitizer type>`s are `asan` and `tsan`.
 

scripts/build_llvm_for_sanitizers.sh

@@ -40,7 +40,8 @@ LLVM_SOURCE_DIR="${LLVM_PATH}/llvm-project"
 LLVM_SOURCE="${LLVM_SOURCE_DIR}/llvm"
 
 # compiler-rt and clang are the sanitizer-specific LLVM projects
-LLVM_PROJECTS="clang;compiler-rt;mlir"
+# openmp is used for parallelizing the triton grid for ThreadSanitizer (TSan)
+LLVM_PROJECTS="clang;compiler-rt;openmp;mlir"
 
 # these are the targets supported by the Triton language
 # Triton's build script for LLVM uses these exact targets

scripts/build_triton_shared_for_sanitizers.sh

@@ -29,12 +29,19 @@ if [ ! -e "$TRITON_SHARED_PATH" ]; then
   exit 1
 fi
 
+# check whether ~/.triton is empty
+# llvm being used from before may be cached and may be different from the custom llvm
+# will cause linking errors during the triton-shared build
+if [ -e "~/.triton" ]; then
+  echo "Error: Please remove ~/.triton and run this script again."
+  exit 1
+fi
+
 cd "$TRITON_SHARED_PATH"
 
 # prepare for triton_shared build
 export PATH="${LLVM_INSTALL_PATH}/bin:${PATH}"
 which clang
-rm -rf ~/.triton
 
 # build triton-shared with the custom LLVM
 cd "${TRITON_SHARED_PATH}/triton"

scripts/run_triton_with_sanitizers.sh

@@ -3,14 +3,14 @@
 # use LLVM_BINARY_DIR to obtain the locations of the .so files used for LD_PRELOAD
 
 if [ "$#" -eq 0 ]; then
-    echo "Usage: $0 <sanitizer type> python program.py. <sanitizer type> can be \"asan\""
+    echo "Usage: $0 <sanitizer type> python program.py. <sanitizer type> can be \"asan\", \"tsan\""
     exit 1
 fi
 
 sanitizer_type=$1
 
-if [ "$sanitizer_type" != "asan" ]; then
-    echo "Error: Unsupported <sanitizer type> $sanitizer_type. Usage: $0 <sanitizer type> python program.py. <sanitizer type> can be \"asan\""
+if [ "$sanitizer_type" != "asan" ] && [ "$sanitizer_type" != "tsan" ]; then
+    echo "Error: Unsupported <sanitizer type> $sanitizer_type. Usage: $0 <sanitizer type> python program.py. <sanitizer type> can be \"asan\", \"tsan\""
     exit 1
 fi
 
@@ -42,6 +42,57 @@ ASAN_OPTIONS=\"detect_leaks=0\""
 
     # shift command line arguments to the left by 1 to account for "asan"
     shift 1
+elif [ "${sanitizer_type}" = "tsan" ]; then
+    # find path to tsan shared library
+    tsan_dir="$(find "$llvm_install_dir" -type f -name "libclang_rt.tsan.so")"
+
+    if [ -z "$tsan_dir" ]; then
+        echo "Error: unable to find libclang_rt.tsan.so in $llvm_install_dir"
+        exit 1
+    fi
+
+    count=$(echo "$tsan_dir" | wc -l)
+
+    if [ "$count" -gt 1 ]; then
+        echo "Error: multiple libclang_rt.tsan.so found in $llvm_install_dir"
+        echo "$tsan_dir"
+        exit 1
+    fi
+
+    # find path to archer library
+    archer_dir="$(find "$llvm_install_dir" -type f -name "libarcher.so")"
+
+    if [ -z "$archer_dir" ]; then
+        echo "Error: unable to find libarcher.so in $llvm_install_dir"
+        exit 1
+    fi
+
+    count=$(echo "$archer_dir" | wc -l)
+
+    if [ "$count" -gt 1 ]; then
+        echo "Error: multiple libarcher.so found in $llvm_install_dir"
+        echo "$archer_dir"
+        exit 1
+    fi
+
+    # make new suppression.txt file if it doesn't exist already
+    if [ ! -f "suppression.txt" ]; then
+        echo "called_from_lib:libomp.so
+called_from_lib:libtorch_python.so
+called_from_lib:libtorch_cpu.so
+called_from_lib:libtorch_cuda.so" > "./suppression.txt"
+    fi
+
+    env_args="LD_PRELOAD=\"$tsan_dir\" \
+TRITON_ALWAYS_COMPILE=1 \
+TRITON_SHARED_SANITIZER_TYPE=\"tsan\" \
+TSAN_OPTIONS=\"ignore_noninstrumented_modules=0:suppressions=suppression.txt\" \
+OMP_NUM_THREADS=16 \
+OMP_TOOL_LIBRARIES=\"$archer_dir\" \
+ARCHER_OPTIONS=\"verbose=1\""
+
+    # shift command line arguments to the left by 1 to account for "tsan"
+    shift 1
 fi
 
 # invoke python function