gtests, benchdnn: explicitly finalize TBB scheduler

oneapi-src · Jan 20, 2023 · ab7938f · ab7938f
1 parent 31c9e7b
commit ab7938f
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/tests/benchdnn/dnnl_common.cpp b/tests/benchdnn/dnnl_common.cpp
@@ -358,6 +358,7 @@ void get_gpu_profiling_info(uint64_t &nsec, double &freq, int mode) {
 
 void finalize() {
     reset_gpu_profiling();
+    finalize_tbb();
 }
 
 inline bool should_stop(const timer::timer_t &t) {

diff --git a/tests/gtests/main.cpp b/tests/gtests/main.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2016-2020 Intel Corporation
+* Copyright 2016-2022 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -91,6 +91,7 @@ int main(int argc, char *argv[]) {
     system("where /q umdh && del post_cpu.txt");
 #endif
 
+    finalize_tbb();
     return result;
 }
 

diff --git a/tests/test_thread.hpp b/tests/test_thread.hpp
@@ -80,13 +80,16 @@ struct thr_ctx_t {
 #error "src/common/dnnl_thread.hpp" has an unexpected header guard
 #endif
 
+#ifdef TBB_INTERFACE_VERSION
 // tbb constraints on core type appear in 2021.2
 // tbb constraints on max_concurrency appear in 2020
 // we check only for 2021.2 to enable thread context knobs
-#ifdef TBB_INTERFACE_VERSION
 #define DNNL_TBB_CONSTRAINTS_ENABLED (TBB_INTERFACE_VERSION >= 12020)
+// API to do explicit finalization was introduced in 2021.6.
+#define DNNL_TBB_NEED_EXPLICIT_FINALIZE (TBB_INTERFACE_VERSION >= 12060)
 #else
 #define DNNL_TBB_CONSTRAINTS_ENABLED 0
+#define DNNL_TBB_NEED_EXPLICIT_FINALIZE 0
 #endif
 
 #define DNNL_TBB_THREADING_WITH_CONSTRAINTS \
@@ -318,6 +321,21 @@ auto execute_in_thr_ctx(const thr_ctx_t &ctx, F &&f, Args_t &...args)
 #error __FILE__"(" __LINE__ ")" "unsupported threading runtime!"
 #endif
 
+// TBB runtime may crash when it is used under CTest. This is a known TBB
+// limitation that can be worked around by doing explicit finalization.
+// The API to do that was introduced in 2021.6.0. When using an older TBB
+// runtime the crash may still happen.
+#if DNNL_TBB_NEED_EXPLICIT_FINALIZE
+#include "tbb/global_control.h"
+inline void finalize_tbb() {
+    oneapi::tbb::task_scheduler_handle handle
+            = oneapi::tbb::task_scheduler_handle {oneapi::tbb::attach {}};
+    oneapi::tbb::finalize(handle);
+}
+#else
+inline void finalize_tbb() {};
+#endif
+
 #undef ALIAS_TO_RUN_IN_THR_CTX
 #undef THR_CTX_ASSERT
 #undef DNNL_TBB_THREADING_WITHOUT_CONSTRAINTS