diff --git a/shortfin/CMakeLists.txt b/shortfin/CMakeLists.txt index b05dd353a..0c69a08a0 100644 --- 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 index 7147a0888..aedcbc331 100644 --- 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(); }