api: threadpool: add wait() method for proper async support

Author: mgouicem

include/oneapi/dnnl/dnnl_threadpool_iface.hpp

@@ -57,6 +57,9 @@ struct threadpool_iface {
     /// Returns threadpool behavior flags bit mask (see below).
     virtual uint64_t get_flags() const = 0;
 
+    // Does nothing if SYNCHRONOUS, waits for all jobs for ASYNCHRONOUS
+    virtual void wait() = 0;
+
     /// If set, parallel_for() returns immediately and oneDNN needs implement
     /// waiting for the submitted closures to finish execution on its own.
     static constexpr uint64_t ASYNCHRONOUS = 1;

src/cpu/cpu_stream.hpp

@@ -38,6 +38,15 @@ struct cpu_stream_t : public stream_t {
 
     dnnl::impl::status_t wait() override {
         // CPU execution is synchronous so return immediately
+#if DNNL_CPU_RUNTIME == DNNL_RUNTIME_THREADPOOL
+        dnnl::threadpool_interop::threadpool_iface *tp;
+        auto rc = this->get_threadpool(&tp);
+        if (rc == status::success) {
+            if (tp->get_flags()
+                    & threadpool_interop::threadpool_iface::ASYNCHRONOUS)
+                tp->wait();
+        }
+#endif
         return dnnl::impl::status::success;
     }
 

tests/benchdnn/utils/parallel.cpp

@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2022-2023 Intel Corporation
+* Copyright 2022-2025 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -25,30 +25,39 @@
 #define ACTIVATE_THREADPOOL
 #endif
 
+void synchronize() {
+#if DNNL_CPU_THREADING_RUNTIME == DNNL_RUNTIME_THREADPOOL
+    dnnl::testing::get_threadpool()->wait();
+#endif
+}
+
 // Note: no need in deactivation as `scoped_activation` object will deactivate
 // it automatically at destruction.
-
 void benchdnn_parallel_nd(int64_t D0, const std::function<void(int64_t)> &f) {
     ACTIVATE_THREADPOOL;
     dnnl::impl::parallel_nd(D0, f);
+    synchronize();
 }
 
 void benchdnn_parallel_nd(int64_t D0, int64_t D1,
         const std::function<void(int64_t, int64_t)> &f) {
     ACTIVATE_THREADPOOL;
     dnnl::impl::parallel_nd(D0, D1, f);
+    synchronize();
 }
 
 void benchdnn_parallel_nd(int64_t D0, int64_t D1, int64_t D2,
         const std::function<void(int64_t, int64_t, int64_t)> &f) {
     ACTIVATE_THREADPOOL;
     dnnl::impl::parallel_nd(D0, D1, D2, f);
+    synchronize();
 }
 
 void benchdnn_parallel_nd(int64_t D0, int64_t D1, int64_t D2, int64_t D3,
         const std::function<void(int64_t, int64_t, int64_t, int64_t)> &f) {
     ACTIVATE_THREADPOOL;
     dnnl::impl::parallel_nd(D0, D1, D2, D3, f);
+    synchronize();
 }
 
 void benchdnn_parallel_nd(int64_t D0, int64_t D1, int64_t D2, int64_t D3,
@@ -57,6 +66,7 @@ void benchdnn_parallel_nd(int64_t D0, int64_t D1, int64_t D2, int64_t D3,
                 &f) {
     ACTIVATE_THREADPOOL;
     dnnl::impl::parallel_nd(D0, D1, D2, D3, D4, f);
+    synchronize();
 }
 
 void benchdnn_parallel_nd(int64_t D0, int64_t D1, int64_t D2, int64_t D3,
@@ -65,6 +75,7 @@ void benchdnn_parallel_nd(int64_t D0, int64_t D1, int64_t D2, int64_t D3,
                 int64_t, int64_t, int64_t, int64_t, int64_t, int64_t)> &f) {
     ACTIVATE_THREADPOOL;
     dnnl::impl::parallel_nd(D0, D1, D2, D3, D4, D5, f);
+    synchronize();
 }
 
 int benchdnn_get_max_threads() {

tests/test_thread.cpp

@@ -186,6 +186,10 @@ class threadpool_t : public dnnl::threadpool_interop::threadpool_iface {
 
         counter.Wait();
     };
+
+    void wait() override {
+        // Nothing to do, runtime is synchronous
+    }
 };
 
 } // namespace testing
@@ -210,6 +214,7 @@ class threadpool_t : public dnnl::threadpool_interop::threadpool_iface {
         tbb::parallel_for(
                 0, n, [&](int i) { fn(i, n); }, tbb::static_partitioner());
     }
+    void wait() override {}
 };
 
 } // namespace testing
@@ -282,6 +287,8 @@ class threadpool_t : public dnnl::threadpool_interop::threadpool_iface {
         }
     }
 
+    void wait() override {}
+
 private:
     int num_threads_;
     std::mutex master_mutex_;