Fix a data race and add a CMake TSan option. (#848)

The C++ bit solves a data race reported by TSan, see #847. The CMake bit mimics what exists for ASAN and removes an incorrect `# Enable more ASAN checks` comment. The `IREE_SANITIZER_` are used in core logic to check against mismatches between the IREE runtime and the libraries it loads. https://github.com/iree-org/iree/blob/db129e500f35ead4debc363eb15ebfb929ee512b/runtime/src/iree/hal/local/loaders/system_library_loader.c#L168-L178 Fixes #847 Signed-off-by: Benoit Jacob <[email protected]>
nod-ai · Jan 20, 2025 · 3d36fe8 · 3d36fe8
1 parent 26bb250
commit 3d36fe8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/shortfin/CMakeLists.txt b/shortfin/CMakeLists.txt
@@ -118,11 +118,16 @@ option(SHORTFIN_ENABLE_ASAN "Enable ASAN" OFF)
 if(SHORTFIN_ENABLE_ASAN)
   add_compile_options(-fsanitize=address)
   add_link_options(-fsanitize=address)
-
-  # Enable more ASAN checks.
   add_compile_definitions(IREE_SANITIZER_ADDRESS)
 endif()
 
+option(SHORTFIN_ENABLE_TSAN "Enable TSAN" OFF)
+if(SHORTFIN_ENABLE_TSAN)
+  add_compile_options(-fsanitize=thread)
+  add_link_options(-fsanitize=thread)
+  add_compile_definitions(IREE_SANITIZER_THREAD)
+endif()
+
 # Thread safety annotations: Enabled if the compiler supports it.
 check_cxx_compiler_flag("-Wthread-safety" SHORTFIN_HAS_THREAD_SAFETY_ANNOTATIONS)
 if(SHORTFIN_HAS_THREAD_SAFETY)

diff --git a/shortfin/src/shortfin/support/blocking_executor.cc b/shortfin/src/shortfin/support/blocking_executor.cc
@@ -157,7 +157,11 @@ int BlockingExecutor::ThreadInstance::RunOnThread() noexcept {
     auto status = iree_wait_source_wait_one(signal_transact.await(),
                                             iree_infinite_timeout());
     IREE_CHECK_OK(status);
-    Task exec_task = std::move(current_task);
+    Task exec_task;
+    {
+      iree::slim_mutex_lock_guard g(executor->control_mu_);
+      exec_task = std::move(current_task);
+    }
     if (exec_task) {
       exec_task();
     }