Create N-1 threads in intra-op pool, given main thread now active (#4493)

Create N-1 threads in a thread pool when configured with intra-op parallelism of N. This ensures we have N active threads, given that the main thread also runs work. To avoid ambiguity on the value returned, rename ThreadPool::NumThreads method to ThreadPool::DegreeOfParallelism, and make corresponding updates in MLAS and operators.

docs/NotesOnThreading.md

@@ -12,7 +12,7 @@ Examples of these abstractions are: ([threadpool.h](https://github.com/microsoft
 * TryBatchParallelFor
 * TryParallelFor
 * TrySimpleParallelFor
-* static version of NumThreads
+* DegreeOfParallelism
 
 **Please do not write #ifdef pragma omp in operator code**.
 

include/onnxruntime/core/platform/threadpool.h

@@ -48,7 +48,7 @@ class ThreadPoolTempl;
 namespace concurrency {
 
 class ExtendedThreadPoolInterface;
-class BatchHandle;
+class LoopCounter;
 
 class ThreadPool {
  public:
@@ -118,27 +118,30 @@ class ThreadPool {
 #else
   using NAME_CHAR_TYPE = char;
 #endif
-  // Constructs a pool that contains "num_threads" threads with specified
-  // "name". env->StartThread() is used to create individual threads with the
-  // given ThreadOptions. If "low_latency_hint" is true the thread pool
+  // Constructs a pool for running with with "degree_of_parallelism" threads with
+  // specified "name". env->StartThread() is used to create individual threads
+  // with the given ThreadOptions. If "low_latency_hint" is true the thread pool
   // implementation may use it as a hint that lower latency is preferred at the
   // cost of higher CPU usage, e.g. by letting one or more idle threads spin
   // wait. Conversely, if the threadpool is used to schedule high-latency
   // operations like I/O the hint should be set to false.
   //
-  // REQUIRES: num_threads > 0
+  // REQUIRES: degree_of_parallelism > 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,
+             int degree_of_parallelism,
              bool low_latency_hint);
 
   // Waits until all scheduled work has finished and then destroy the
   // set of threads.
   ~ThreadPool();
 
-  // Schedules fn() for execution in the pool of threads.
+  // Schedules fn() for execution in the pool of threads.  The function may run
+  // synchronously if it cannot be enqueued.  This will occur if the thread pool's
+  // degree-of-parallelism is 1, but it may also occur for implementation-dependent
+  // reasons such as if queues used for buffering work are full.
   void Schedule(std::function<void()> fn);
 
   // Returns the number of shards used by ParallelForFixedBlockSizeScheduling
@@ -171,7 +174,7 @@ class ThreadPool {
                              const std::function<void(std::ptrdiff_t first, std::ptrdiff_t last)>& fn) {
 #ifdef _OPENMP
     ORT_UNUSED_PARAMETER(cost_per_unit);
-    std::ptrdiff_t num_threads = concurrency::ThreadPool::NumThreads(tp);
+    std::ptrdiff_t num_threads = concurrency::ThreadPool::DegreeOfParallelism(tp);
     if (total < num_threads) {
       num_threads = total;
     }
@@ -199,7 +202,7 @@ class ThreadPool {
                              const std::function<void(std::ptrdiff_t first, std::ptrdiff_t last)>& fn) {
 #ifdef _OPENMP
     ORT_UNUSED_PARAMETER(scheduling_params);
-    std::ptrdiff_t num_threads = concurrency::ThreadPool::NumThreads(tp);
+    std::ptrdiff_t num_threads = concurrency::ThreadPool::DegreeOfParallelism(tp);
     if (total < num_threads) {
       num_threads = total;
     }
@@ -217,16 +220,15 @@ class ThreadPool {
 #endif
   }
 
-  // Prefer using this API to get the number of threads unless you know what you're doing.
-  // This API takes into account if openmp is enabled/disabled and if the thread pool ptr is nullptr.
-  static int NumThreads(const concurrency::ThreadPool* tp);
-
-  // Returns the number of threads in the pool. Preferably use the static version of this API instead.
-  int NumThreads() const;
-
-  // Returns current thread id between 0 and NumThreads() - 1, if called from a
-  // thread in the pool. Returns -1 otherwise.
-  int CurrentThreadId() const;
+  // Return the degree of parallelism that code should assume when using the thread pool.
+  // This API takes into account if OpenMP is enabled/disabled, and if the thread pool ptr is
+  // nullptr.  It decouples the degree of parallelism for use with the thread pool from
+  // the implementation choice of whether this matches the number of threads created in
+  // the pool.
+  //
+  // Currently, a loop with degree-of-parallelism N is supported by a pool of N-1 threads
+  // working in combination with the thread initiating the loop.
+  static int DegreeOfParallelism(const concurrency::ThreadPool* tp);
 
   // Directly schedule the 'total' tasks to the underlying threadpool, without
   // cutting them by halves
@@ -254,7 +256,7 @@ class ThreadPool {
 
   /**
    * Tries to call the given function in parallel, with calls split into (num_batches) batches.
-   *\param num_batches If it is zero, it will be replaced to the value of NumThreads().
+   *\param num_batches If it is zero, it will be replaced to the value of DegreeOfParallelism().
    *\param fn A std::function or STL style functor with signature of "void f(int32_t);"
    * Pitfall: Caller should cap `num_batches` to a reasonable value based on the cost of `fn` and the value of `total`.
    *For example, if fn is as simple as: int sum=0; fn = [&](int i){sum +=i;} and `total` is 100, then num_batches should
@@ -288,7 +290,7 @@ class ThreadPool {
     }
 
     if (num_batches <= 0) {
-      num_batches = std::min<ptrdiff_t>(total, tp->NumThreads());
+      num_batches = std::min<ptrdiff_t>(total, DegreeOfParallelism(tp));
     }
 
     if (num_batches <= 1) {
@@ -334,6 +336,16 @@ class ThreadPool {
   ORT_DISALLOW_COPY_AND_ASSIGNMENT(ThreadPool);
 
  private:
+  friend class LoopCounter;
+
+  // Returns the number of threads created in the pool.  This may be different from the
+  // value returned by DegreeOfParallelism to code using the pool.
+  int NumThreads() const;
+
+  // Returns current thread id between 0 and NumThreads() - 1, if called from a
+  // thread in the pool. Returns -1 otherwise.
+  int CurrentThreadId() const;
+
   // Run fn with up to n degree-of-parallelism enlisting the thread pool for
   // help.  The degree-of-parallelism includes the caller, and so if n==1
   // then the function will run directly in the caller.  The fork-join
@@ -359,11 +371,14 @@ class ThreadPool {
                              const std::ptrdiff_t block_size = 1) const;
 
   ThreadOptions thread_options_;
-  // underlying_threadpool_ is the user_threadpool if user_threadpool is
-  // provided in the constructor. Otherwise it is the eigen_threadpool_.
-  ExtendedThreadPoolInterface* underlying_threadpool_;
-  // eigen_threadpool_ is instantiated and owned by thread::ThreadPool if
-  // user_threadpool is not in the constructor.
+
+  // If a thread pool is created with degree_of_parallelism != 1 then an underlying
+  // EigenThreadPool is used to create OS threads and handle work distribution to them.
+  // If degree_of_parallelism == 1 then underlying_threadpool_ is left as nullptr
+  // and parallel work is run directly by the caller.
+  ExtendedThreadPoolInterface* underlying_threadpool_ = nullptr;
+
+  // If used, underlying_threadpool_ is instantiated and owned by the ThreadPool.
   std::unique_ptr<ThreadPoolTempl<Env> > extended_eigen_threadpool_;
 };
 

onnxruntime/core/common/threadpool.cc

@@ -74,19 +74,19 @@ public:
     // does not need to be unique, but we aim for a good distribution, particularly in the case where
     // most/all of the thread pool's threads are active in the loop.  Threads outside the pool may
     // also be claiming work, with CurrentThreadId -1.
-    int num_threads = _tp.NumThreads();
-    int my_thread_idx = (_tp.CurrentThreadId() + 1) % num_threads;
-    assert(my_thread_idx >= 0 && my_thread_idx < num_threads);
+    int d_of_p = ThreadPool::DegreeOfParallelism(&_tp);
+    int my_thread_idx = (_tp.CurrentThreadId() + 1) % d_of_p;
+    assert(my_thread_idx >= 0 && my_thread_idx < d_of_p);
 
     int home_shard;
-    if (num_threads >= NUM_SHARDS) {
+    if (d_of_p >= NUM_SHARDS) {
       // More threads than shards => allocate them home shards round-robin, aiming to sprace the load across
       // the shards
       home_shard = my_thread_idx % NUM_SHARDS;
     } else {
       // Fewer threads than shards => spread the threads evenly across the shards, so each will work
       // on a run of successive shards before contention
-      home_shard = (my_thread_idx * NUM_SHARDS) / num_threads;
+      home_shard = (my_thread_idx * NUM_SHARDS) / d_of_p;
     }
     assert(home_shard >= 0 && home_shard < NUM_SHARDS);
     return home_shard;
@@ -126,13 +126,26 @@ private:
 #pragma warning(pop) /* Padding added in LoopCounterShard, LoopCounter */
 #endif
 
-ThreadPool::ThreadPool(Env* env, const ThreadOptions& thread_options, const NAME_CHAR_TYPE* name, int num_threads,
+ThreadPool::ThreadPool(Env* env,
+                       const ThreadOptions& thread_options,
+                       const NAME_CHAR_TYPE* name,
+                       int degree_of_parallelism,
                        bool low_latency_hint)
     : thread_options_(thread_options) {
-  ORT_ENFORCE(num_threads >= 1);
-  extended_eigen_threadpool_ =
-      onnxruntime::make_unique<ThreadPoolTempl<Env>>(name, num_threads, low_latency_hint, *env, thread_options_);
-  underlying_threadpool_ = extended_eigen_threadpool_.get();
+  // In the current implementation, a thread pool with degree_of_parallelism==1 uses
+  // the caller as one of the threads for executing work.  Hence we only create
+  // additional thread(s) for degree_of_parallelism>=2.
+  ORT_ENFORCE(degree_of_parallelism >= 1);
+  if (degree_of_parallelism >= 2) {
+    int threads_to_create = degree_of_parallelism - 1;
+    extended_eigen_threadpool_ =
+        onnxruntime::make_unique<ThreadPoolTempl<Env>>(name,
+                                                       threads_to_create,
+                                                       low_latency_hint,
+                                                       *env,
+                                                       thread_options_);
+    underlying_threadpool_ = extended_eigen_threadpool_.get();
+  }
 }
 
 ThreadPool::~ThreadPool() = default;
@@ -153,8 +166,8 @@ void ThreadPool::ParallelForFixedBlockSizeScheduling(const std::ptrdiff_t total,
 
   // Split the work across threads in the pool.  Each work item will run a loop claiming iterations,
   // hence we need at most one for each thread, even if the numberof blocks of iterations is larger.
-  int num_threads = NumThreads();
-  int num_work_items = static_cast<int>(std::min(static_cast<std::ptrdiff_t>(num_threads), total));
+  auto d_of_p = DegreeOfParallelism(this);
+  int num_work_items = static_cast<int>(std::min(static_cast<std::ptrdiff_t>(d_of_p), total));
   assert(num_work_items > 0);
 
   LoopCounter lc(*this, total, block_size);
@@ -184,12 +197,20 @@ void ThreadPool::SimpleParallelFor(std::ptrdiff_t total, const std::function<voi
 
 void ThreadPool::Schedule(std::function<void()> fn) {
   ORT_ENFORCE(fn != nullptr);
-  underlying_threadpool_->Schedule(std::move(fn));
+  if (underlying_threadpool_) {
+    underlying_threadpool_->Schedule(std::move(fn));
+  } else {
+    fn();
+  }
 }
 
 void ThreadPool::RunInParallel(std::function<void()> fn, int n) {
   ORT_ENFORCE(fn != nullptr);
-  underlying_threadpool_->RunInParallel(std::move(fn), n);
+  if (underlying_threadpool_) {
+    underlying_threadpool_->RunInParallel(std::move(fn), n);
+  } else {
+    fn();
+  }
 }
 
 bool ThreadPool::ShouldParallelizeLoop(const std::ptrdiff_t num_iterations,
@@ -201,9 +222,10 @@ bool ThreadPool::ShouldParallelizeLoop(const std::ptrdiff_t num_iterations,
 
   // Do not parallelize loops with only a single thread available.  If the
   // caller is outside the current pool (ID == -1) then we parallelize
-  // via the pool's thread(s).  If the caller is inside the current pool
+  // if the pool has any threads.  If the caller is inside the current pool
   // (ID != -1) then we require at least one additional thread in the pool.
-  if (CurrentThreadId() != -1 && NumThreads() == 1) {
+  if ((CurrentThreadId() == -1 && NumThreads() == 0) ||
+      (CurrentThreadId() != -1 && NumThreads() == 1)) {
     return false;
   }
 
@@ -304,14 +326,17 @@ void ThreadPool::ParallelFor(std::ptrdiff_t n, const TensorOpCost& c,
                              const std::function<void(std::ptrdiff_t first, std::ptrdiff_t)>& f) {
   ORT_ENFORCE(n >= 0);
   Eigen::TensorOpCost cost{c.bytes_loaded, c.bytes_stored, c.compute_cycles};
+  auto d_of_p = DegreeOfParallelism(this);
   // Compute small problems directly in the caller thread.
   if ((!ShouldParallelizeLoop(n)) ||
-      Eigen::TensorCostModel<Eigen::ThreadPoolDevice>::numThreads(static_cast<double>(n), cost, static_cast<int>(NumThreads())) == 1) {
+      Eigen::TensorCostModel<Eigen::ThreadPoolDevice>::numThreads(static_cast<double>(n),
+                                                                  cost,
+                                                                  d_of_p) == 1) {
     f(0, n);
     return;
   }
 
-  ptrdiff_t block = CalculateParallelForBlock(n, cost, nullptr, NumThreads());
+  ptrdiff_t block = CalculateParallelForBlock(n, cost, nullptr, d_of_p);
   ParallelForFixedBlockSizeScheduling(n, block, f);
 }
 
@@ -320,23 +345,38 @@ void ThreadPool::ParallelFor(std::ptrdiff_t total, double cost_per_unit,
   ParallelFor(total, TensorOpCost{0, 0, static_cast<double>(cost_per_unit)}, fn);
 }
 
-int ThreadPool::NumThreads(const concurrency::ThreadPool* tp) {
+int ThreadPool::DegreeOfParallelism(const concurrency::ThreadPool* tp) {
 #ifdef _OPENMP
+  // When using OpenMP, omp_get_num_threads() returns the number of threads in the
+  // current parallel region.  Hence if this is 1 then we aim to parallelise
+  // across the number of threads configured.  Otherwise, given that we do not
+  // use nested parallelism, we do not parallelise further.
   ORT_UNUSED_PARAMETER(tp);
   return (omp_get_num_threads() == 1) ? omp_get_max_threads() : 1;
 #else
-  return tp ? tp->NumThreads() : 1;
+  // When not using OpenMP, we parallelise over the N threads created by the pool
+  // tp, plus 1 for the thread entering a loop.
+  return tp ? (tp->NumThreads()+1) : 1;
 #endif
 }
 
+// Return the number of threads created by the pool.
 int ThreadPool::NumThreads() const {
-  return underlying_threadpool_->NumThreads();
+  if (underlying_threadpool_) {
+    return underlying_threadpool_->NumThreads();
+  } else {
+    return 0;
+  }
 }
 
 // Return ID of the current thread within this pool.  Returns -1 for a thread outside the
 // current pool.
 int ThreadPool::CurrentThreadId() const {
-  return underlying_threadpool_->CurrentThreadId();
+  if (underlying_threadpool_) {
+    return underlying_threadpool_->CurrentThreadId();
+  } else {
+    return -1;
+  }
 }
 
 }  // namespace concurrency

onnxruntime/core/mlas/lib/mlasi.h

@@ -774,7 +774,7 @@ MlasGetMaximumThreadCount(
     return 1;
 #endif
 #else
-	return onnxruntime::concurrency::ThreadPool::NumThreads(ThreadPool);
+    return onnxruntime::concurrency::ThreadPool::DegreeOfParallelism(ThreadPool);
 #endif
 }
 

onnxruntime/core/providers/cpu/math/top_k.cc

@@ -164,7 +164,7 @@ static void FindTopKElements(const Tensor* input, const TensorShape& input_shape
   const int64_t num_blocks = input_shape[axis_parsed];
   const int64_t block_slice = reduced_cols / k;
 
-  int64_t tp_threads = concurrency::ThreadPool::NumThreads(threadpool);
+  int64_t tp_threads = concurrency::ThreadPool::DegreeOfParallelism(threadpool);
   int64_t num_threads = std::min(tp_threads, rows);  // split on rows so can't have more threads than rows
 
   // rough attempt to make sure there's enough work for each thread. if there's insufficient work the usage of

onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h

@@ -326,7 +326,7 @@ void TreeEnsembleCommon<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp,
       } else {
         // split the work into one block per thread so we can re-use the 'private_scores' vector as much as possible
         // TODO: Refine the number of threads used
-        auto num_threads = std::min<int32_t>(concurrency::ThreadPool::NumThreads(ttp), SafeInt<int32_t>(n_trees_));
+        auto num_threads = std::min<int32_t>(concurrency::ThreadPool::DegreeOfParallelism(ttp), SafeInt<int32_t>(n_trees_));
         OrtMutex merge_mutex;
         concurrency::ThreadPool::TrySimpleParallelFor(
             ttp,
@@ -361,7 +361,7 @@ void TreeEnsembleCommon<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp,
       } else {
         // split the work into one block per thread so we can re-use the 'scores' vector as much as possible
         // TODO: Refine the number of threads used.
-        auto num_threads = std::min<int32_t>(concurrency::ThreadPool::NumThreads(ttp), SafeInt<int32_t>(N));
+        auto num_threads = std::min<int32_t>(concurrency::ThreadPool::DegreeOfParallelism(ttp), SafeInt<int32_t>(N));
         concurrency::ThreadPool::TrySimpleParallelFor(
             ttp,
             num_threads,

onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc

@@ -1069,7 +1069,7 @@ void UniDirectionalLstm<T>::GateComputations(
 
 template <typename T>
 void UniDirectionalLstm<T>::SetNumThreads() {
-  int threads = concurrency::ThreadPool::NumThreads(thread_pool_);
+  int threads = concurrency::ThreadPool::DegreeOfParallelism(thread_pool_);
 
   if (threads < 1)
     threads = 1;