Introduce tsan pipeline

Author: alekstheod

.github/workflows/ci-ut.yml

@@ -31,6 +31,7 @@ jobs:
           - {name: "py3.13", python_version: "3.13", config: ""}
           - {name: "py3.14", python_version: "3.14", config: ""}
           - {name: "asan", python_version: "3.11", config: "--config=asan"}
+          - {name: "tsan", python_version: "3.11", config: "--config=tsan"}
     container:
       # note this image shall match the one defined in platform/linux:tf_linux_gpu
       image: rocm/tensorflow-build@sha256:7fcfbd36b7ac8f6b0805b37c4248e929e31cf5ee3af766c8409dd70d5ab65faa

jax_rocm_plugin/.bazelrc

@@ -101,6 +101,15 @@ build:rocm --copt=-Wno-gnu-offsetof-extensions
 build:rocm --copt=-Qunused-arguments
 build:rocm --action_env=TF_HIPCC_CLANG="1"
 
+#############################################################################
+# Sanitizer configs
+#############################################################################
+build:sanitizer --linkopt="-L/usr/lib/llvm-18/lib/clang/18/lib/linux"
+build:sanitizer --linkopt="-Wl,-rpath,/usr/lib/llvm-18/lib/clang/18/lib/linux"
+build:sanitizer --run_under=//build/rocm:sanitizer_wrapper
+build:sanitizer --action_env TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a,gfx942"
+
+build:asan --config=sanitizer
 build:asan --strip=never
 build:asan --copt -fsanitize=address
 build:asan --copt -DADDRESS_SANITIZER
@@ -111,14 +120,20 @@ build:asan --copt -gsplit-dwarf
 build:asan --copt -fno-omit-frame-pointer
 build:asan --linkopt -fsanitize=address
 build:asan --linkopt -g1
-
-build:asan --linkopt="-L/usr/lib/llvm-18/lib/clang/18/lib/linux"
-build:asan --linkopt="-Wl,-rpath,/usr/lib/llvm-18/lib/clang/18/lib/linux"
 build:asan --linkopt="-lclang_rt.asan-x86_64"
 build:asan --linkopt="-lclang_rt.asan_cxx-x86_64"
 build:asan --//build/rocm:sanitizer=asan
-build:asan --run_under=//build/rocm:sanitizer_wrapper
-build:asan --action_env TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a,gfx942"
+
+build:tsan --config=sanitizer
+build:tsan --strip=never
+build:tsan --copt -fsanitize=thread
+build:tsan --copt -g
+build:tsan --copt -fno-omit-frame-pointer
+build:tsan --linkopt -fsanitize=thread
+build:tsan --linkopt -g
+build:tsan --linkopt="-lclang_rt.tsan-x86_64"
+build:tsan --linkopt="-lclang_rt.tsan_cxx-x86_64"
+build:tsan --//build/rocm:sanitizer=tsan
 
 #############################################################################
 # Configuration for running RBE builds and tests

jax_rocm_plugin/build/rocm/sanitizer_wrapper.sh

@@ -1,25 +1,32 @@
 #!/usr/bin/env bash
 
 ASAN_RT="/usr/lib/llvm-18/lib/clang/18/lib/linux/libclang_rt.asan-x86_64.so"
-
-# Only set LD_PRELOAD for the test binary, not for Bazel's wrappers
-export LD_PRELOAD="${ASAN_RT}"
+TSAN_RT="/usr/lib/llvm-18/lib/clang/18/lib/linux/libclang_rt.tsan-x86_64.so"
 
 # Resolve suppression files from runfiles
 ASAN_SUPP="$TEST_SRCDIR/jax_rocm_plugin/build/rocm/asan_ignore_list.txt"
 LSAN_SUPP="$TEST_SRCDIR/jax_rocm_plugin/build/rocm/lsan_ignore_list.txt"
+TSAN_SUPP="$TEST_SRCDIR/jax_rocm_plugin/build/rocm/tsan_ignore_list.txt"
 
 ASAN_OPTS="use_sigaltstack=0:detect_leaks=0"
 LSAN_OPTS="use_sigaltstack=0"
+TSAN_OPTS="history_size=7:ignore_noninstrumented_modules=1"
 
 if [[ -f "$ASAN_SUPP" ]]; then
-  ASAN_OPTS="suppressions=${ASAN_SUPP}:${ASAN_OPTS}"
+    export LD_PRELOAD="${ASAN_RT}"
+    ASAN_OPTS="suppressions=${ASAN_SUPP}:${ASAN_OPTS}"
 fi
 if [[ -f "$LSAN_SUPP" ]]; then
-  LSAN_OPTS="suppressions=${LSAN_SUPP}:${LSAN_OPTS}"
+    export LD_PRELOAD="${ASAN_RT}"
+    LSAN_OPTS="suppressions=${LSAN_SUPP}:${LSAN_OPTS}"
+fi
+if [[ -f "$TSAN_SUPP" ]]; then
+    export LD_PRELOAD="${TSAN_RT}"
+    TSAN_OPTS="suppressions=${TSAN_SUPP}:${TSAN_OPTS}"
 fi
 
 export ASAN_OPTIONS="${ASAN_OPTS}"
 export LSAN_OPTIONS="${LSAN_OPTS}"
+export TSAN_OPTIONS="${TSAN_OPTS}"
 
 exec "$@"