tests: add asynchronous Eigen threadpool runtime implementation

Author: dzarukin

cmake/Threadpool.cmake

@@ -32,18 +32,28 @@ if("${DNNL_CPU_THREADING_RUNTIME}" STREQUAL "THREADPOOL")
         endif()
     endif()
 
-    if("${_DNNL_TEST_THREADPOOL_IMPL}" STREQUAL "EIGEN")
+    if(_DNNL_TEST_THREADPOOL_IMPL MATCHES "^(EIGEN|EIGEN_ASYNC)$")
         find_package(Eigen3 5.0 REQUIRED NO_MODULE)
         if(Eigen3_FOUND)
             list(APPEND EXTRA_STATIC_LIBS Eigen3::Eigen)
             message(STATUS "Found Eigen: ${PACKAGE_PREFIX_DIR}")
         endif()
 
-        find_package(absl REQUIRED CONFIG)
+        # Note: not every UNIX OS search in lib64 by default, but this path is
+        # generated by absl packaging procedure.
+        find_package(absl REQUIRED CONFIG PATH_SUFFIXES "lib64/cmake/absl")
         if(absl_FOUND)
             list(APPEND EXTRA_STATIC_LIBS absl::synchronization)
             message(STATUS "Found abseil-cpp: ${PACKAGE_PREFIX_DIR}")
         endif()
+
+        if("${_DNNL_TEST_THREADPOOL_IMPL}" STREQUAL "EIGEN_ASYNC")
+            find_package(xla REQUIRED CONFIG PATH_SUFFIXES "lib64/cmake/xla")
+            if(xla_FOUND)
+                list(APPEND EXTRA_STATIC_LIBS xla)
+                message(STATUS "Found XLA: ${PACKAGE_PREFIX_DIR}")
+            endif()
+        endif()
     endif()
 
     if("${_DNNL_TEST_THREADPOOL_IMPL}" STREQUAL "STANDALONE")

cmake/options.cmake

@@ -266,8 +266,8 @@ endif()
 set(_DNNL_TEST_THREADPOOL_IMPL "STANDALONE" CACHE STRING
     "specifies which threadpool implementation to use when
     DNNL_CPU_RUNTIME=THREADPOOL is selected. Valid values: STANDALONE, EIGEN,
-    TBB")
-if(NOT "${_DNNL_TEST_THREADPOOL_IMPL}" MATCHES "^(STANDALONE|TBB|EIGEN)$")
+    EIGEN_ASYNC, TBB")
+if(NOT "${_DNNL_TEST_THREADPOOL_IMPL}" MATCHES "^(STANDALONE|TBB|EIGEN|EIGEN_ASYNC)$")
     message(FATAL_ERROR
         "Unsupported threadpool implementation: ${_DNNL_TEST_THREADPOOL_IMPL}")
 endif()

doc/build/build_options.md

@@ -236,29 +236,38 @@ oneDNN has functional limitations if built with TBB:
 
 #### Threadpool
 To build oneDNN with support for threadpool threading, set `ONEDNN_CPU_RUNTIME`
-to `THREADPOOL`
+to `THREADPOOL`:
 
 ~~~sh
 $ cmake -DONEDNN_CPU_RUNTIME=THREADPOOL ..
 ~~~
 
+Threadpool threading support has the same limitations as TBB plus more:
+* As threadpools are attached to streams which are only passed during primitive
+  execution, work decomposition is performed statically at primitive creation
+  time. At the primitive execution time, the threadpool is responsible for
+  balancing this decomposition across available worker threads.
+
+##### Threadpool validation
 The `_ONEDNN_TEST_THREADPOOL_IMPL` CMake variable controls which of the three
-threadpool implementations would be used for testing: `STANDALONE`, `TBB`, or
-`EIGEN`. The latter two require also passing `TBBROOT` or `Eigen3_DIR` paths
-to CMake. For example:
+threadpool implementations would be used for testing: `STANDALONE`, `TBB`,
+`EIGEN`, `EIGEN_ASYNC`.
+
+The `TBB` requires passing `TBBROOT` for CMake to find a package.
+
+The `EIGEN` requires Eigen 5.0 or higher and Abseil-CPP packages to be
+discoverable by CMake.
 
+The `EIGEN_ASYNC` has same requirements as `EIGEN` and additionally requires
+OpenXLA threadpool package, however, additional actions might be required to
+compile tests since this threadpool implementation relies on internal OpenXLA
+headers.
+
+For example:
 ~~~sh
-$ cmake -DONEDNN_CPU_RUNTIME=THREADPOOL -D_ONEDNN_TEST_THREADPOOL_IMPL=EIGEN -DEigen3_DIR=/path/to/eigen/share/eigen3/cmake ..
+$ cmake -DONEDNN_CPU_RUNTIME=THREADPOOL -D_ONEDNN_TEST_THREADPOOL_IMPL=EIGEN -DCMAKE_PREFIX_PATH="/path/to/eigen/share/eigen3/cmake;/path/to/absl/lib64/cmake" ..
 ~~~
 
-Threadpool threading support is experimental and has the same limitations as
-TBB plus more:
-* As threadpools are attached to streams which are only passed during
-  primitive execution, work decomposition is performed statically at the
-  primitive creation time. At the primitive execution time, the threadpool is
-  responsible for balancing the static decomposition from the previous item
-  across available worker threads.
-
 ### AArch64 Options
 
 oneDNN includes experimental support for Arm 64-bit Architecture (AArch64).

tests/test_thread.cpp

@@ -195,6 +195,84 @@ class threadpool_t : public dnnl::threadpool_interop::threadpool_iface {
 } // namespace testing
 } // namespace dnnl
 
+#elif defined(DNNL_TEST_THREADPOOL_USE_EIGEN_ASYNC)
+
+// absl sources define its own version of `CHECK` macro. oneDNN's version is not
+// needed further the file, thus, disable it for compilation reason.
+#undef CHECK
+
+#define EIGEN_USE_THREADS
+#include "Eigen/ThreadPool"
+
+#include "xla/backends/cpu/runtime/work_queue.h"
+#include "xla/tsl/concurrency/async_value_ref.h"
+#include "xla/tsl/concurrency/chain.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+
+namespace dnnl {
+namespace testing {
+
+static tsl::AsyncValueRef<tsl::Chain> OkDoneEventSingleton() {
+    static std::unique_ptr<tsl::AsyncValueOwningRef<tsl::Chain>> singleton =
+            [] {
+                static auto storage = std::make_unique<
+                        tsl::internal::AsyncValueStorage<tsl::Chain>>();
+                return std::make_unique<tsl::AsyncValueOwningRef<tsl::Chain>>(
+                        tsl::MakeAvailableAsyncValueRef<tsl::Chain>(*storage));
+            }();
+    return singleton->AsRef();
+}
+
+class threadpool_t : public dnnl::threadpool_interop::threadpool_iface {
+private:
+    // Original `OneDnnThreadPool` at
+    // `xla/backends/cpu/runtime/onednn/onednn_threadpool.h` takes
+    // `Eigen::ThreadPoolInterface` instead. Since `Eigen::ThreadPool` is
+    // a parent class, which is an alias to `NonBlockingThreadPool`, it fits
+    // the need.
+    std::unique_ptr<Eigen::ThreadPool> thread_pool_;
+
+    // Async value that signals completion of the last scheduled parallel loop.
+    // This is used only when is_async_ is true.
+    tsl::AsyncValueRef<tsl::Chain> done_event_;
+
+public:
+    explicit threadpool_t(int num_threads = 0) {
+        if (num_threads <= 0) num_threads = read_num_threads_from_env();
+        thread_pool_.reset(new Eigen::ThreadPool(num_threads));
+        done_event_ = OkDoneEventSingleton();
+    }
+    int get_num_threads() const override { return thread_pool_->NumThreads(); }
+    bool get_in_parallel() const override { return false; }
+    uint64_t get_flags() const override { return ASYNCHRONOUS; }
+    void parallel_for(int n, const std::function<void(int, int)> &fn) override {
+        // If we are using oneDNN with async support, we need to schedule the
+        // parallel loop using the done_event_. This allows us to return
+        // immediately and not block the caller thread.
+        auto parallelize = [this, n, fn](tsl::Chain) {
+            return xla::cpu::Worker::Parallelize(thread_pool_.get(),
+                    thread_pool_->NumThreads(), n,
+                    [fn, n](size_t i) { fn(static_cast<int>(i), n); });
+        };
+
+        done_event_ = done_event_.FlatMap(parallelize);
+    }
+    void wait() override {
+        // While performing asynchronous execution, wait() method is needed to
+        // notify the user that the output is ready. oneDNN will not call wait()
+        // inside the library to avoid deadlock.
+        tsl::BlockUntilReady(done_event_);
+    }
+
+    tsl::AsyncValueRef<tsl::Chain> done_event() const { return done_event_; }
+};
+
+} // namespace testing
+} // namespace dnnl
+
 #elif defined(DNNL_TEST_THREADPOOL_USE_TBB)
 #include "tbb/parallel_for.h"
 #include "tbb/task_arena.h"