Remove eigen device from thread pool

include/onnxruntime/core/platform/threadpool.h

@@ -32,7 +32,6 @@ limitations under the License.
 namespace Eigen {
 class Allocator;
 class ThreadPoolInterface;
-struct ThreadPoolDevice;
 }  // namespace Eigen
 
 namespace onnxruntime {
@@ -125,12 +124,12 @@ class ThreadPool {
   // REQUIRES: num_threads > 0
   // The allocator parameter is only used for creating a Eigen::ThreadPoolDevice to be used with Eigen Tensor classes.
   ThreadPool(Env* env, const ThreadOptions& thread_options, const NAME_CHAR_TYPE* name, int num_threads,
-             bool low_latency_hint, Eigen::Allocator* allocator = nullptr);
+             bool low_latency_hint);
   // Constructs a pool that wraps around the thread::ThreadPoolInterface
   // instance provided by the caller. Caller retains ownership of
   // `user_threadpool` and must ensure its lifetime is longer than the
   // ThreadPool instance.
-  ThreadPool(Eigen::ThreadPoolInterface* user_threadpool, Eigen::Allocator* allocator);
+  ThreadPool(Eigen::ThreadPoolInterface* user_threadpool);
 
   // Waits until all scheduled work has finished and then destroy the
   // set of threads.
@@ -315,12 +314,6 @@ class ThreadPool {
 #endif
   }
 
-#ifndef _OPENMP
-  //Deprecated. Please avoid using Eigen Tensor because it will blow up binary size quickly.
-  Eigen::ThreadPoolDevice& Device() {
-    return *threadpool_device_;
-  }
-#endif
   ORT_DISALLOW_COPY_AND_ASSIGNMENT(ThreadPool);
 
  private:
@@ -340,9 +333,7 @@ class ThreadPool {
   // eigen_threadpool_ is instantiated and owned by thread::ThreadPool if
   // user_threadpool is not in the constructor.
   std::unique_ptr<ThreadPoolTempl<Env> > eigen_threadpool_;
-#ifndef _OPENMP
-  std::unique_ptr<Eigen::ThreadPoolDevice> threadpool_device_;
-#endif
+
   // Copied from MlasPartitionWork
   static void PartitionWork(std::ptrdiff_t ThreadId, std::ptrdiff_t ThreadCount, std::ptrdiff_t TotalWork,
                             std::ptrdiff_t* WorkIndex, std::ptrdiff_t* WorkRemaining) {

onnxruntime/contrib_ops/cpu/bert/attention.cc

@@ -4,7 +4,6 @@
 #include "attention.h"
 #include "core/framework/tensorprotoutils.h"
 #include "core/graph/onnx_protobuf.h"
-#include "core/util/eigen_common_wrapper.h"
 #include "core/util/math.h"
 #include "core/util/math_cpuonly.h"
 #include "core/providers/cpu/math/gemm_helper.h"

onnxruntime/core/common/threadpool.cc

@@ -18,7 +18,7 @@ limitations under the License.
 
 #include "core/platform/threadpool.h"
 #include "core/common/common.h"
-#include "core/util/eigen_common_wrapper.h"
+#include "core/common/eigen_common_wrapper.h"
 #include "core/platform/EigenNonBlockingThreadPool.h"
 #include "core/platform/ort_mutex.h"
 
@@ -83,29 +83,17 @@ class BlockingCounter {
 namespace concurrency {
 
 ThreadPool::ThreadPool(Env* env, const ThreadOptions& thread_options, const NAME_CHAR_TYPE* name, int num_threads,
-                       bool low_latency_hint, Eigen::Allocator* allocator)
+                       bool low_latency_hint)
     : thread_options_(thread_options) {
   ORT_ENFORCE(num_threads >= 1);
   eigen_threadpool_ =
       onnxruntime::make_unique<ThreadPoolTempl<Env>>(name, num_threads, low_latency_hint, *env, thread_options_);
   underlying_threadpool_ = eigen_threadpool_.get();
-#ifdef _OPENMP
-  ORT_UNUSED_PARAMETER(allocator);
-#else
-  threadpool_device_ =
-      onnxruntime::make_unique<Eigen::ThreadPoolDevice>(underlying_threadpool_, num_threads, allocator);
-#endif
 }
 
-ThreadPool::ThreadPool(Eigen::ThreadPoolInterface* user_threadpool, Eigen::Allocator* allocator)
+ThreadPool::ThreadPool(Eigen::ThreadPoolInterface* user_threadpool)
     : thread_options_(ThreadOptions()) {
   underlying_threadpool_ = user_threadpool;
-#ifdef _OPENMP
-  ORT_UNUSED_PARAMETER(allocator);
-#else
-  threadpool_device_ = onnxruntime::make_unique<Eigen::ThreadPoolDevice>(
-      underlying_threadpool_, underlying_threadpool_->NumThreads(), allocator);
-#endif
 }
 
 ThreadPool::~ThreadPool() = default;

onnxruntime/core/providers/cpu/generator/random.cc

@@ -27,7 +27,7 @@ limitations under the License.
 
 #include "core/common/safeint.h"
 #include "core/util/math_cpuonly.h"
-#include "core/util/eigen_common_wrapper.h"
+#include "core/common/eigen_common_wrapper.h"
 #include "gsl/gsl"
 using namespace ONNX_NAMESPACE;
 using namespace ::onnxruntime::common;

onnxruntime/core/providers/cpu/nn/pool.cc

@@ -4,7 +4,6 @@
 #include "core/providers/cpu/nn/pool.h"
 #include "core/framework/data_types_internal.h"
 #include "core/platform/threadpool.h"
-#include "core/util/eigen_common_wrapper.h"
 #include "pool_functors.h"
 
 using namespace ::onnxruntime::common;
@@ -12,7 +11,7 @@ using namespace ::onnxruntime::common;
 namespace onnxruntime {
 
 template <typename T>
-inline static void RunLoop(concurrency::ThreadPool* tp, Eigen::Index total_channels, T&& task) {
+inline static void RunLoop(concurrency::ThreadPool* tp, std::ptrdiff_t total_channels, T&& task) {
   concurrency::ThreadPool::TryParallelFor(tp, total_channels, task.Cost(), task);
 }
 

onnxruntime/core/providers/cpu/nn/pool_functors.h

@@ -3,7 +3,6 @@
 #pragma once
 
 #include "core/platform/threadpool.h"
-#include "core/util/eigen_common_wrapper.h"
 #include "core/providers/cpu/nn/pool_base.h"
 namespace onnxruntime {
 
@@ -25,7 +24,7 @@ struct Pool1DTask final {
     return TensorOpCost{loop_count, loop_count, loop_count};
   }
 
-  void operator()(Eigen::Index begin, Eigen::Index end) const {
+  void operator()(std::ptrdiff_t begin, std::ptrdiff_t end) const {
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
@@ -33,7 +32,7 @@ struct Pool1DTask final {
       operator()(c);
     }
   }
-  void operator()(Eigen::Index c) const {
+  void operator()(std::ptrdiff_t c) const {
     const T* x_d = X_data + c * x_step;
     T* y_d = Y_data + c * y_step;
 
@@ -77,7 +76,7 @@ struct Pool2DTask final {
     return TensorOpCost{loop_count, loop_count, loop_count};
   }
 
-  void operator()(Eigen::Index begin, Eigen::Index end) const {
+  void operator()(std::ptrdiff_t begin, std::ptrdiff_t end) const {
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
@@ -86,7 +85,7 @@ struct Pool2DTask final {
     }
   }
 
-  void operator()(Eigen::Index c) const {
+  void operator()(std::ptrdiff_t c) const {
     const T* x_d = X_data + c * x_step;
     T* y_d = Y_data + c * y_step;
 
@@ -143,7 +142,7 @@ struct Pool3DTask final {
     return TensorOpCost{loop_count, loop_count, loop_count};
   }
 
-  void operator()(Eigen::Index begin, Eigen::Index end) const {
+  void operator()(std::ptrdiff_t begin, std::ptrdiff_t end) const {
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
@@ -152,7 +151,7 @@ struct Pool3DTask final {
     }
   }
 
-  void operator()(Eigen::Index c) const {
+  void operator()(std::ptrdiff_t c) const {
     const T* x_d = X_data + c * x_step;
     T* y_d = Y_data + c * y_step;
 
@@ -208,7 +207,7 @@ struct MaxPool1DTask final {
     return TensorOpCost{loop_count, loop_count, loop_count};
   }
 
-  void operator()(Eigen::Index begin, Eigen::Index end) const {
+  void operator()(std::ptrdiff_t begin, std::ptrdiff_t end) const {
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
@@ -216,7 +215,7 @@ struct MaxPool1DTask final {
       operator()(c);
     }
   }
-  void operator()(Eigen::Index c) const {
+  void operator()(std::ptrdiff_t c) const {
     const T* x_d = X_data + c * x_step;
     T* y_d = Y_data + c * y_step;
     int64_t* i_d = I_data ? I_data + c * y_step : nullptr;
@@ -264,7 +263,7 @@ struct MaxPool2DTask final {
     return TensorOpCost{loop_count, loop_count, loop_count};
   }
 
-  void operator()(Eigen::Index begin, Eigen::Index end) const {
+  void operator()(std::ptrdiff_t begin, std::ptrdiff_t end) const {
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
@@ -273,7 +272,7 @@ struct MaxPool2DTask final {
     }
   }
 
-  void operator()(Eigen::Index c) const {
+  void operator()(std::ptrdiff_t c) const {
     const T* x_d = X_data + c * x_step;
     T* y_d = Y_data + c * y_step;
     int64_t* i_d = I_data ? I_data + c * y_step : nullptr;
@@ -333,7 +332,7 @@ struct MaxPool3DTask {
   const std::vector<int64_t>& pads;
   int64_t storage_order;
 
-  void operator()(Eigen::Index begin, Eigen::Index end) const {
+  void operator()(std::ptrdiff_t begin, std::ptrdiff_t end) const {
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
@@ -348,7 +347,7 @@ struct MaxPool3DTask {
     return TensorOpCost{loop_count, loop_count, loop_count};
   }
 
-  void operator()(Eigen::Index c) const {
+  void operator()(std::ptrdiff_t c) const {
     const T* x_d = X_data + c * x_step;
     T* y_d = Y_data + c * y_step;
     int64_t* i_d = I_data ? I_data + c * y_step : nullptr;

onnxruntime/core/providers/cpu/tensor/onehot.cc

@@ -15,7 +15,7 @@ limitations under the License.
 /* Modifications Copyright (c) Microsoft. */
 
 #include "core/providers/cpu/tensor/onehot.h"
-#include "core/util/eigen_common_wrapper.h"
+#include "core/common/eigen_common_wrapper.h"
 #include "core/platform/env.h"
 #include "core/providers/common.h"
 

onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc

@@ -7,7 +7,7 @@
 #endif
 
 #include "core/providers/cpu/tensor/space_depth_ops.h"
-#include "core/util/eigen_common_wrapper.h"
+#include "core/common/eigen_common_wrapper.h"
 #include <array>
 
 namespace onnxruntime {

onnxruntime/core/session/environment.cc

@@ -62,12 +62,12 @@ Status Environment::Initialize(std::unique_ptr<logging::LoggingManager> logging_
     if (to.name == nullptr) {
       to.name = ORT_TSTR("intra-op");
     }
-    intra_op_thread_pool_ = concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP, nullptr);
+    intra_op_thread_pool_ = concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP);
     to = tp_options->inter_op_thread_pool_params;
     if (to.name == nullptr) {
       to.name = ORT_TSTR("inter-op");
     }
-    inter_op_thread_pool_ = concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTER_OP, nullptr);
+    inter_op_thread_pool_ = concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTER_OP);
   }
 
   try {

onnxruntime/core/session/inference_session.cc

@@ -183,7 +183,7 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options,
                              session_options_.execution_mode == ExecutionMode::ORT_SEQUENTIAL &&
                              to.affinity_vec_len == 0;
       thread_pool_ =
-          concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP, nullptr);
+          concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP);
     }
     if (session_options_.execution_mode == ExecutionMode::ORT_PARALLEL) {
       OrtThreadPoolParams to = session_options_.inter_op_param;
@@ -194,7 +194,7 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options,
       if (to.name == nullptr)
         to.name = ORT_TSTR("intra-op");
       inter_op_thread_pool_ =
-          concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTER_OP, nullptr);
+          concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTER_OP);
       if (inter_op_thread_pool_ == nullptr) {
         LOGS(*session_logger_, INFO) << "Failed to create the inter-op thread pool for the parallel executor, setting ExecutionMode to SEQUENTIAL";
         session_options_.execution_mode = ExecutionMode::ORT_SEQUENTIAL;

onnxruntime/core/util/thread_utils.cc

@@ -10,7 +10,7 @@
 namespace onnxruntime {
 namespace concurrency {
 static std::unique_ptr<ThreadPool>
-CreateThreadPoolHelper(Env* env, OrtThreadPoolParams options, Eigen::Allocator* allocator) {
+CreateThreadPoolHelper(Env* env, OrtThreadPoolParams options) {
   if (options.thread_pool_size == 1)
     return nullptr;
   std::vector<size_t> cpu_list;
@@ -28,26 +28,25 @@ CreateThreadPoolHelper(Env* env, OrtThreadPoolParams options, Eigen::Allocator*
   }
 
   return onnxruntime::make_unique<ThreadPool>(env, to, options.name, options.thread_pool_size,
-                                              options.allow_spinning, allocator);
+                                              options.allow_spinning);
 }
 
 std::unique_ptr<ThreadPool>
-CreateThreadPool(Env* env, OrtThreadPoolParams options, ThreadPoolType tpool_type, Eigen::Allocator* allocator) {
+CreateThreadPool(Env* env, OrtThreadPoolParams options, ThreadPoolType tpool_type) {
 // If openmp is enabled we don't want to create any additional threadpools for sequential execution.
 // However, parallel execution relies on the existence of a separate threadpool. Hence we allow eigen threadpools
 // to be created for parallel execution.
 #ifdef _OPENMP
   ORT_UNUSED_PARAMETER(env);
   ORT_UNUSED_PARAMETER(options);
-  ORT_UNUSED_PARAMETER(allocator);
   if (tpool_type != ThreadPoolType::INTER_OP) {
     return nullptr;
   } else {
-    return CreateThreadPoolHelper(env, options, allocator);
+    return CreateThreadPoolHelper(env, options);
   }
 #else
   ORT_UNUSED_PARAMETER(tpool_type);
-  return CreateThreadPoolHelper(env, options, allocator);
+  return CreateThreadPoolHelper(env, options);
 #endif
 }
 

onnxruntime/core/util/thread_utils.h

@@ -43,7 +43,6 @@ enum class ThreadPoolType : uint8_t {
   INTER_OP
 };
 std::unique_ptr<ThreadPool> CreateThreadPool(Env* env, OrtThreadPoolParams options,
-                                             ThreadPoolType tpool_type,
-                                             Eigen::Allocator* allocator = nullptr);
+                                             ThreadPoolType tpool_type);
 }  // namespace concurrency
 }  // namespace onnxruntime

onnxruntime/test/mlas/unittest.cpp

@@ -2111,7 +2111,7 @@ main(
 #if !defined(MLAS_NO_ONNXRUNTIME_THREADPOOL)
         if (threadpool != nullptr)
           threadpool = new onnxruntime::concurrency::ThreadPool(
-              &onnxruntime::Env::Default(), onnxruntime::ThreadOptions(), nullptr, 2, true, nullptr);
+              &onnxruntime::Env::Default(), onnxruntime::ThreadOptions(), nullptr, 2, true);
 #else
         break;
 #endif