Use OpenMP-like synchronization patterns in Eigen thread pool (#4236)

Updates the thread pool implementation to make work distribution over the Eigen thread pool more closely resemble techniques used in OpenMP. In particular:

(1) A thread entering a parallel loop works on the iterations itself, rather than requiring a thread switch to/from a thread in the pool, if called from outside the thread pool.

(2) To support this, work items pushed to the thread pool run a loop to claim iterations from a shared counter via atomic-fetch-and-add, as opposed to having work items themselves represent individual batches of iterations. This means that any thread working on the loop can execute any batch of iterations, including having the main thread run through all of the batches itself if the loop turns out to be short-running.

(3) As with OpenMP active scheduling, the worker loop spins waiting for work prior to blocking. This avoids OS blocking / wake-up paths in workloads with series of short-running parallel sections.

include/onnxruntime/core/platform/Barrier.h

@@ -25,13 +25,14 @@ class Barrier {
   }
 #endif
 
-  void Notify() {
-    unsigned int v = state_.fetch_sub(2, std::memory_order_acq_rel) - 2;
+  void Notify(unsigned int c = 1) {
+    unsigned int delta = c << 1;
+    unsigned int v = state_.fetch_sub(delta, std::memory_order_acq_rel) - delta;
     if (v != 1) {
       // Clear the lowest bit (waiter flag) and check that the original state
       // value was not zero. If it was zero, it means that notify was called
       // more times than the original count.
-      assert(((v + 2) & ~1) != 0);
+      assert(((v + delta) & ~1) != 0);
       return;  // either count has not dropped to 0, or waiter is not waiting
     }
     std::unique_lock<OrtMutex> l(mu_);

include/onnxruntime/core/platform/threadpool.h

@@ -44,7 +44,12 @@ struct TensorOpCost {
 
 template <typename Environment>
 class ThreadPoolTempl;
+
 namespace concurrency {
+
+class ExtendedThreadPoolInterface;
+class BatchHandle;
+
 class ThreadPool {
  public:
   // Scheduling strategies for ParallelFor. The strategy governs how the given
@@ -123,13 +128,11 @@ 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,
+  ThreadPool(Env* env,
+             const ThreadOptions& thread_options,
+             const NAME_CHAR_TYPE* name,
+             int num_threads,
              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);
 
   // Waits until all scheduled work has finished and then destroy the
   // set of threads.
@@ -140,7 +143,8 @@ class ThreadPool {
 
   // Returns the number of shards used by ParallelForFixedBlockSizeScheduling
   // with these parameters.
-  int NumShardsUsedByFixedBlockSizeScheduling(std::ptrdiff_t total, std::ptrdiff_t block_size);
+  int NumShardsUsedByFixedBlockSizeScheduling(std::ptrdiff_t total,
+                                              std::ptrdiff_t block_size) const;
 
   // ParallelFor shards the "total" units of work assuming each unit of work
   // having roughly "cost_per_unit" cost, in cycles. Each unit of work is
@@ -224,11 +228,6 @@ class ThreadPool {
   // thread in the pool. Returns -1 otherwise.
   int CurrentThreadId() const;
 
-  // If ThreadPool implementation is compatible with Eigen::ThreadPoolInterface,
-  // returns a non-null pointer. The caller does not own the object the returned
-  // pointer points to, and should not attempt to delete.
-  Eigen::ThreadPoolInterface* AsEigenThreadPool() const;
-
   // Directly schedule the 'total' tasks to the underlying threadpool, without
   // cutting them by halves
   void SimpleParallelFor(std::ptrdiff_t total, const std::function<void(std::ptrdiff_t)>& fn);
@@ -335,6 +334,13 @@ class ThreadPool {
   ORT_DISALLOW_COPY_AND_ASSIGNMENT(ThreadPool);
 
  private:
+  // 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
+  // synchronization is handled in the thread pool, and so any state captured
+  // by fn() is safe from concurrent access once RunWithHelp returns.
+  void RunInParallel(std::function<void()> fn, int n);
+
   // Divides the work represented by the range [0, total) into k shards.
   // Calls fn(i*block_size, (i+1)*block_size) from the ith shard (0 <= i < k).
   // Each shard may be executed on a different thread in parallel, depending on
@@ -344,13 +350,21 @@ class ThreadPool {
   // Requires 0 < block_size <= total.
   void ParallelForFixedBlockSizeScheduling(std::ptrdiff_t total, std::ptrdiff_t block_size,
                                            const std::function<void(std::ptrdiff_t, std::ptrdiff_t)>& fn);
+
+
+  // Return whether or not the calling thread should run a loop of
+  // num_iterations divided in chunks of block_size in parallel.  If not,
+  // the caller should run the loop sequentially.
+  bool ShouldParallelizeLoop(const std::ptrdiff_t num_iterations,
+                             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_.
-  Eigen::ThreadPoolInterface* underlying_threadpool_;
+  ExtendedThreadPoolInterface* underlying_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_;
+  std::unique_ptr<ThreadPoolTempl<Env> > extended_eigen_threadpool_;
 };
 
 }  // namespace concurrency

onnxruntime/core/common/threadpool.cc

@@ -1,4 +1,3 @@
-
 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -23,101 +22,164 @@ limitations under the License.
 #include "core/platform/ort_mutex.h"
 
 namespace onnxruntime {
-namespace {
-class BlockingCounter {
- public:
-  BlockingCounter(int initial_count) : state_(initial_count << 1), notified_(false) {
-    ORT_ENFORCE(initial_count >= 0);
-#ifndef NDEBUG
-    ORT_ENFORCE(((initial_count << 1) >> 1) == initial_count);
-#endif
-  }
 
-  ~BlockingCounter() = default;
-
-  inline void DecrementCount() {
-    unsigned int v = state_.fetch_sub(2, std::memory_order_acq_rel) - 2;
-    if (v != 1) {
-#ifndef NDEBUG
-      ORT_ENFORCE(((v + 2) & ~1) != 0);
-#endif
-      return;  // either count has not dropped to 0, or waiter is not waiting
-    }
-    std::lock_guard<OrtMutex> l(mu_);
-    notified_ = true;
-    cond_var_.notify_all();
-  }
-
-  inline void Wait() {
-    unsigned int v = state_.fetch_or(1, std::memory_order_acq_rel);
-    if ((v >> 1) == 0)
-      return;
-    std::unique_lock<OrtMutex> l(mu_);
-    while (!notified_) {
-      cond_var_.wait(l);
-    }
-  }
-  // Wait for the specified time, return false iff the count has not dropped to
-  // zero before the timeout expired.
-  inline bool WaitFor(std::chrono::milliseconds ms) {
-    unsigned int v = state_.fetch_or(1, std::memory_order_acq_rel);
-    if ((v >> 1) == 0)
-      return true;
-    std::unique_lock<OrtMutex> l(mu_);
-    while (!notified_) {
-      const std::cv_status status = cond_var_.wait_for(l, ms);
-      if (status == std::cv_status::timeout) {
-        return false;
-      }
-    }
-    return true;
-  }
-
- private:
-  OrtMutex mu_;
-  OrtCondVar cond_var_;
-  std::atomic<int> state_;  // low bit is waiter flag
-  bool notified_;
-};
-}  // namespace
 namespace concurrency {
 
+// A sharded loop counter distributes loop iterations between a set of worker threads.  The iteration space of
+// the loop is divided (perhaps unevenly) between the shards.  Each thread has a home shard (perhaps not uniquely
+// to it), and it claims iterations via atomic operations on its home shard.  It then proceeds through the other
+// shards until all of the shards' iterations are complete.  This approach serves to purposes.  First, compared
+// with atomic operations on a single counter, it reduces contention on a single counter in the case of loops with
+// large numbers of short-running iteration.  Second, by having a thread work on its home shard initially, it
+// promotes affinity between the work that a thread performs in one loop and the work that it performs in the next.
+
+#ifdef _MSC_VER
+#pragma warning(push)
+#pragma warning(disable : 4324) /* Padding added to LoopCounterShard, LoopCounter for alignment */
+#endif
+
+static constexpr int CACHE_LINE_BYTES = 64;
+static constexpr int NUM_SHARDS = 8;
+
+struct alignas(CACHE_LINE_BYTES) LoopCounterShard {
+  ::std::atomic<uint64_t> _next;
+  uint64_t _end;
+};
+
+class alignas(CACHE_LINE_BYTES) LoopCounter {
+public:
+ LoopCounter(const ThreadPool& tp,
+             uint64_t num_iterations,
+             uint64_t block_size = 1) : _tp(tp),
+                                        _block_size(block_size) {
+   assert(sizeof(LoopCounterShard) == 64);
+   assert(block_size != 0);
+
+   // Divide the iteration space into NUM_SHARDS pieces.  If the iteration space does not
+   // divide evenly into shards of multiples of block_size then the final shard is left uneven.
+   double iterations_per_shard = static_cast<double>(num_iterations) / NUM_SHARDS;
+   uint64_t split = 0;
+   for (uint64_t shard = 0; shard < NUM_SHARDS; shard++) {
+     _shards[shard]._next = split;
+     split = (static_cast<uint64_t>((shard + 1) * iterations_per_shard) / block_size) * block_size;
+     _shards[shard]._end = split;
+   }
+
+   // Ensure that the final shard finishes precisely at the end of the iteration space
+   _shards[NUM_SHARDS - 1]._end = num_iterations;
+ }
+
+  int GetHomeShard() const {
+    // Allocate each thread to a home shard, from which it starts claiming iterations.  The allocation
+    // 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 home_shard;
+    if (num_threads >= 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;
+    }
+    assert(home_shard >= 0 && home_shard < NUM_SHARDS);
+    return home_shard;
+  }
+
+  // Attempt to claim iterations from the sharded counter.  The function either
+  // returns true, along with a block of exactly block_size iterations, or it returns false
+  // if all of the iterations have been claimed.
+  bool ClaimIterations(int my_home_shard,
+                       int& my_shard,
+                       uint64_t& my_start,
+                       uint64_t& my_end) {
+    do {
+      if (_shards[my_shard]._next < _shards[my_shard]._end) {
+        // Appears to be work in the current shard, try to claim with atomic fetch-and-add
+        uint64_t temp_start = _shards[my_shard]._next.fetch_add(_block_size);
+        if (temp_start < _shards[my_shard]._end) {
+          my_start = temp_start;
+          my_end = std::min(_shards[my_shard]._end, temp_start + _block_size);
+          return true;
+        }
+      }
+      // Work in the current shard is exhausted, move to the next shard, until
+      // we are back at the home shard.
+      my_shard = (my_shard + 1) % NUM_SHARDS;
+    } while (my_shard != my_home_shard);
+    return false;
+  }
+
+private:
+  alignas(CACHE_LINE_BYTES) LoopCounterShard _shards[NUM_SHARDS];
+  const ThreadPool& _tp;
+  const uint64_t _block_size;
+};
+
+#ifdef _MSC_VER
+#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,
                        bool low_latency_hint)
     : thread_options_(thread_options) {
   ORT_ENFORCE(num_threads >= 1);
-  eigen_threadpool_ =
+  extended_eigen_threadpool_ =
       onnxruntime::make_unique<ThreadPoolTempl<Env>>(name, num_threads, low_latency_hint, *env, thread_options_);
-  underlying_threadpool_ = eigen_threadpool_.get();
-}
-
-ThreadPool::ThreadPool(Eigen::ThreadPoolInterface* user_threadpool)
-    : thread_options_(ThreadOptions()) {
-  underlying_threadpool_ = user_threadpool;
+  underlying_threadpool_ = extended_eigen_threadpool_.get();
 }
 
 ThreadPool::~ThreadPool() = default;
 
-void ThreadPool::SimpleParallelFor(std::ptrdiff_t total, const std::function<void(std::ptrdiff_t)>& fn) {
+// Base case for parallel loops, running iterations 0..total, divided into blocks
+// of block_size iterations, and calling into a function that takes a start..end
+// range of indices to run.
+void ThreadPool::ParallelForFixedBlockSizeScheduling(const std::ptrdiff_t total,
+                                                     const std::ptrdiff_t block_size,
+                                                     const std::function<void(std::ptrdiff_t, std::ptrdiff_t)>& fn) {
   if (total <= 0)
     return;
 
-  if (total == 1) {
-    fn(0);
+  if (total <= block_size) {
+    fn(0, total);
     return;
   }
 
-  Barrier barrier(static_cast<unsigned int>(total));
-  std::function<void(std::ptrdiff_t)> handle_iteration = [&barrier, &fn](std::ptrdiff_t iteration) {
-    fn(iteration);
-    barrier.Notify();
+  // 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));
+  assert(num_work_items > 0);
+
+  LoopCounter lc(*this, total, block_size);
+  std::function<void()> run_work = [&]() {
+    int my_home_shard = lc.GetHomeShard();
+    int my_shard = my_home_shard;
+    uint64_t my_iter_start, my_iter_end;
+    while (lc.ClaimIterations(my_home_shard, my_shard, my_iter_start, my_iter_end)) {
+      fn(static_cast<std::ptrdiff_t>(my_iter_start),
+         static_cast<std::ptrdiff_t>(my_iter_end));
+    }
   };
 
-  for (std::ptrdiff_t id = 0; id < total; ++id) {
-    Schedule([=, &handle_iteration]() { handle_iteration(id); });
-  }
+  // Run the work in the thread pool (and in the current thread).  Synchronization with helping
+  // threads is handled within RunInParallel, hence we can deallocate lc and other state captured by
+  // run_work.
+  RunInParallel(run_work, num_work_items);
+}
 
-  barrier.Wait();
+void ThreadPool::SimpleParallelFor(std::ptrdiff_t total, const std::function<void(std::ptrdiff_t)>& fn) {
+  ParallelForFixedBlockSizeScheduling(total, 1, [&](std::ptrdiff_t first, std::ptrdiff_t last) {
+    for (std::ptrdiff_t idx = first; idx < last; idx++) {
+      fn(idx);
+    }
+  });
 }
 
 void ThreadPool::Schedule(std::function<void()> fn) {
@@ -125,12 +187,37 @@ void ThreadPool::Schedule(std::function<void()> fn) {
   underlying_threadpool_->Schedule(std::move(fn));
 }
 
-int ThreadPool::NumShardsUsedByFixedBlockSizeScheduling(const std::ptrdiff_t total, const std::ptrdiff_t block_size) {
-  if (block_size <= 0 || total <= 1 || total <= block_size || NumThreads() == 1) {
-    return 1;
+void ThreadPool::RunInParallel(std::function<void()> fn, int n) {
+  ORT_ENFORCE(fn != nullptr);
+  underlying_threadpool_->RunInParallel(std::move(fn), n);
+}
+
+bool ThreadPool::ShouldParallelizeLoop(const std::ptrdiff_t num_iterations,
+                                       const std::ptrdiff_t block_size) const {
+  // Do not parallelize trivial loops, with only a single block of work
+  if (block_size <= 0 || num_iterations <= block_size) {
+    return false;
+  }
+
+  // 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
+  // (ID != -1) then we require at least one additional thread in the pool.
+  if (CurrentThreadId() != -1 && NumThreads() == 1) {
+    return false;
+  }
+
+  return true;
+}
+
+int ThreadPool::NumShardsUsedByFixedBlockSizeScheduling(const std::ptrdiff_t total,
+                                                        const std::ptrdiff_t block_size) const {
+  if (!ShouldParallelizeLoop(total, block_size)) {
+    return 1;
+  } else {
+    // TODO:check overflow?
+    return static_cast<int>((total + block_size - 1) / block_size);
   }
-  // TODO:check overflow?
-  return static_cast<int>((total + block_size - 1) / block_size);
 }
 
 void ThreadPool::ParallelFor(std::ptrdiff_t total, const SchedulingParams& scheduling_params,
@@ -151,41 +238,6 @@ void ThreadPool::ParallelFor(std::ptrdiff_t total, const SchedulingParams& sched
   }
 }
 
-// This functionality is similar to parallelFor, except that reasoning about
-// the number of shards used is significantly easier.
-void ThreadPool::ParallelForFixedBlockSizeScheduling(const std::ptrdiff_t total, const std::ptrdiff_t block_size,
-                                                     const std::function<void(std::ptrdiff_t, std::ptrdiff_t)>& fn) {
-  const int num_shards_used = NumShardsUsedByFixedBlockSizeScheduling(total, block_size);
-  if (num_shards_used == 1) {
-    fn(0, total);
-    return;
-  }
-
-  // Adapted from Eigen's parallelFor implementation.
-  BlockingCounter counter(num_shards_used);
-  std::function<void(ptrdiff_t, ptrdiff_t)> handle_range = [=, &handle_range, &counter, &fn](std::ptrdiff_t first,
-                                                                                             std::ptrdiff_t last) {
-    while (last - first > block_size) {
-      // Find something near the midpoint which is a multiple of block size.
-      const std::ptrdiff_t mid = first + ((last - first) / 2 + block_size - 1) / block_size * block_size;
-      Schedule([=, &handle_range]() { handle_range(mid, last); });
-      last = mid;
-    }
-    // Single block or less, execute directly.
-    fn(first, last);
-    counter.DecrementCount();  // The shard is done.
-  };
-
-  // Execute the root in the thread pool to avoid running work on more than
-  // numThreads() threads.
-  Schedule([=, &handle_range]() { handle_range(0, total); });
-  counter.Wait();
-}
-
-struct ParallelForBlock {
-  ptrdiff_t size;   // block size
-  ptrdiff_t count;  // number of blocks
-};
 using CostModel = Eigen::TensorCostModel<Eigen::ThreadPoolDevice>;
 
 // Calculates block size based on (1) the iteration cost and (2) parallel
@@ -193,8 +245,8 @@ using CostModel = Eigen::TensorCostModel<Eigen::ThreadPoolDevice>;
 // overheads; not too large to mitigate tail effect and potential load
 // imbalance and we also want number of blocks to be evenly dividable across
 // threads.
-static ParallelForBlock CalculateParallelForBlock(const ptrdiff_t n, const Eigen::TensorOpCost& cost,
-                                                  std::function<ptrdiff_t(ptrdiff_t)> block_align, int num_threads) {
+static ptrdiff_t CalculateParallelForBlock(const ptrdiff_t n, const Eigen::TensorOpCost& cost,
+                                           std::function<ptrdiff_t(ptrdiff_t)> block_align, int num_threads) {
   const double block_size_f = 1.0 / CostModel::taskSize(1, cost);
   const ptrdiff_t max_oversharding_factor = 4;
   ptrdiff_t block_size = Eigen::numext::mini(
@@ -245,7 +297,7 @@ static ParallelForBlock CalculateParallelForBlock(const ptrdiff_t n, const Eigen
     }
   }
 
-  return {block_size, block_count};
+  return block_size;
 }
 
 void ThreadPool::ParallelFor(std::ptrdiff_t n, const TensorOpCost& c,
@@ -253,35 +305,16 @@ void ThreadPool::ParallelFor(std::ptrdiff_t n, const TensorOpCost& c,
   ORT_ENFORCE(n >= 0);
   Eigen::TensorOpCost cost{c.bytes_loaded, c.bytes_stored, c.compute_cycles};
   // Compute small problems directly in the caller thread.
-  if (n <= 1 || NumThreads() == 1 ||
+  if ((!ShouldParallelizeLoop(n)) ||
       Eigen::TensorCostModel<Eigen::ThreadPoolDevice>::numThreads(static_cast<double>(n), cost, static_cast<int>(NumThreads())) == 1) {
     f(0, n);
     return;
   }
 
-  // Compute block size and total count of blocks.
-  ParallelForBlock block = CalculateParallelForBlock(n, cost, nullptr, NumThreads());
-
-  // Recursively divide size into halves until we reach block_size.
-  // Division code rounds mid to block_size, so we are guaranteed to get
-  // block_count leaves that do actual computations.
-  Barrier barrier(static_cast<unsigned int>(block.count));
-  std::function<void(ptrdiff_t, ptrdiff_t)> handleRange;
-  handleRange = [=, &handleRange, &barrier, &f](ptrdiff_t firstIdx, ptrdiff_t lastIdx) {
-    while (lastIdx - firstIdx > block.size) {
-      // Split into halves and schedule the second half on a different thread.
-      const ptrdiff_t midIdx = firstIdx + Eigen::divup((lastIdx - firstIdx) / 2, block.size) * block.size;
-      underlying_threadpool_->Schedule([=, &handleRange]() { handleRange(midIdx, lastIdx); });
-      lastIdx = midIdx;
-    }
-    // Single block or less, execute directly.
-    f(firstIdx, lastIdx);
-    barrier.Notify();
-  };
-
-  underlying_threadpool_->Schedule([=, &handleRange]() { handleRange(0, n); });
-  barrier.Wait();
+  ptrdiff_t block = CalculateParallelForBlock(n, cost, nullptr, NumThreads());
+  ParallelForFixedBlockSizeScheduling(n, block, f);
 }
+
 void ThreadPool::ParallelFor(std::ptrdiff_t total, double cost_per_unit,
                              const std::function<void(std::ptrdiff_t first, std::ptrdiff_t)>& fn) {
   ParallelFor(total, TensorOpCost{0, 0, static_cast<double>(cost_per_unit)}, fn);
@@ -300,13 +333,11 @@ int ThreadPool::NumThreads() const {
   return underlying_threadpool_->NumThreads();
 }
 
+// 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();
 }
 
-Eigen::ThreadPoolInterface* ThreadPool::AsEigenThreadPool() const {
-  ORT_ENFORCE(underlying_threadpool_ != nullptr);
-  return underlying_threadpool_;
-}
 }  // namespace concurrency
 }  // namespace onnxruntime

onnxruntime/test/perftest/performance_runner.cc

@@ -81,13 +81,13 @@ void PerformanceResult::DumpToFile(const std::basic_string<ORTCHAR_T>& path, boo
     std::sort(sorted_time.begin(), sorted_time.end());
 
     auto output_stats = [&](std::ostream& ostream) {
-      ostream << "Min Latency is " << sorted_time[0] << "sec\n";
-      ostream << "Max Latency is " << sorted_time[total - 1] << "sec\n";
-      ostream << "P50 Latency is " << sorted_time[n50] << "sec\n";
-      ostream << "P90 Latency is " << sorted_time[n90] << "sec\n";
-      ostream << "P95 Latency is " << sorted_time[n95] << "sec\n";
-      ostream << "P99 Latency is " << sorted_time[n99] << "sec\n";
-      ostream << "P999 Latency is " << sorted_time[n999] << "sec" << std::endl;
+      ostream << "Min Latency: " << sorted_time[0] << " s\n";
+      ostream << "Max Latency: " << sorted_time[total - 1] << " s\n";
+      ostream << "P50 Latency: " << sorted_time[n50] << " s\n";
+      ostream << "P90 Latency: " << sorted_time[n90] << " s\n";
+      ostream << "P95 Latency: " << sorted_time[n95] << " s\n";
+      ostream << "P99 Latency: " << sorted_time[n99] << " s\n";
+      ostream << "P999 Latency: " << sorted_time[n999] << " s" << std::endl;
     };
 
     if (have_file) {
@@ -132,12 +132,16 @@ Status PerformanceRunner::Run() {
   // if (!performance_test_config_.run_config.profile_file.empty()) session_object->EndProfiling();
   std::chrono::duration<double> inference_duration = performance_result_.end - performance_result_.start;
 
-  std::cout << "Session creation time cost:" << session_create_duration.count() << " s" << std::endl
-            << "Total inference time cost:" << performance_result_.total_time_cost << " s" << std::endl  // sum of time taken by each request
-            << "Total inference requests:" << performance_result_.time_costs.size() << std::endl
-            << "Average inference time cost:" << performance_result_.total_time_cost / performance_result_.time_costs.size() * 1000 << " ms" << std::endl
+  std::cout << "Session creation time cost: " << session_create_duration.count() << " s\n"
+            << "Total inference time cost: " << performance_result_.total_time_cost << " s\n"  // sum of time taken by each request
+            << "Total inference requests: " << performance_result_.time_costs.size() << "\n"
+            << "Average inference time cost: " << performance_result_.total_time_cost / performance_result_.time_costs.size() * 1000 << " ms\n"
             // Time between start and end of run. Less than Total time cost when running requests in parallel.
-            << "Total inference run time:" << inference_duration.count() << " s" << std::endl;
+            << "Total inference run time: " << inference_duration.count() << " s\n"
+            << "Avg CPU usage: " << performance_result_.average_CPU_usage << " %\n"
+            << "Peak working set size: " << performance_result_.peak_workingset_size << " bytes"
+            << std::endl;
+
   return Status::OK();
 }
 

onnxruntime/test/platform/threadpool_test.cc

@@ -67,6 +67,43 @@ void TestBatchParallelFor(const std::string& name, int num_threads, int num_task
   ValidateTestData(*test_data);
 }
 
+void TestMultipleParallelFor(const std::string& name, int num_threads, int num_concurrent, int num_tasks) {
+  // Test running multiple concurrent loops over the same thread pool.  This aims to provoke a
+  // more diverse mix of interleavings than with a single loop running at a time.
+  for (int rep = 0; rep < 5; rep++) {
+    CreateThreadPoolAndTest(name, num_threads, [&](ThreadPool* tp) {
+      std::vector<std::unique_ptr<TestData>> td;
+      onnxruntime::Barrier b(num_concurrent - 1);
+
+      // Each concurrent tests runs with its own set of counters
+      for (int c = 0; c < num_concurrent; c++) {
+        td.push_back(CreateTestData(num_tasks));
+      }
+
+      // For a range of scenarios, run some tests via the thread pool, and one directly
+      for (int c = 0; c < num_concurrent - 1; c++) {
+        tp->Schedule([&, c]() {
+          tp->SimpleParallelFor(num_tasks, [&](std::ptrdiff_t i) {
+            IncrementElement(*td[c], i);
+          });
+          b.Notify();
+        });
+      }
+
+      tp->SimpleParallelFor(num_tasks, [&](std::ptrdiff_t i) {
+        IncrementElement(*td[num_concurrent - 1], i);
+      });
+
+      // Validate all outputs
+      b.Wait();
+      for (int c = 0; c < num_concurrent; c++) {
+        ValidateTestData(*td[c]);
+      }
+      td.clear();
+    });
+  }
+}
+
 }  // namespace
 
 namespace onnxruntime {
@@ -102,6 +139,69 @@ TEST(ThreadPoolTest, TestBatchParallelFor_2_Thread_81_Task_20_Batch) {
   TestBatchParallelFor("TestBatchParallelFor_2_Thread_81_Task_20_Batch", 2, 81, 20);
 }
 
+TEST(ThreadPoolTest, TestMultipleParallelFor_1Thread_1Conc_0Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_1Thread_1Conc_0Tasks", 1, 1, 0);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_1Thread_1Conc_1Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_1Thread_1Conc_1Tasks", 1, 1, 1);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_1Thread_1Conc_8Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_1Thread_1Conc_8Tasks", 1, 1, 8);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_1Thread_1Conc_1MTasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_1Thread_1Conc_1MTasks", 1, 1, 1000000);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_1Thread_4Conc_0Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_1Thread_4Conc_0Tasks", 1, 4, 0);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_1Thread_4Conc_1Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_1Thread_4Conc_1Tasks", 1, 4, 1);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_1Thread_4Conc_8Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_1Thread_4Conc_8Tasks", 1, 4, 8);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_1Thread_4Conc_1MTasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_1Thread_4Conc_1MTasks", 1, 4, 1000000);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_1Conc_0Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_0Tasks", 4, 1, 0);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_1Conc_1Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_1Tasks", 4, 1, 1);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_1Conc_8Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_8Tasks", 4, 1, 8);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_1Conc_1MTasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_1MTasks", 4, 1, 1000000);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_4Conc_0Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_0Tasks", 4, 4, 0);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_4Conc_1Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_1Tasks", 4, 4, 1);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_4Conc_8Tasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_8Tasks", 4, 4, 8);
+}
+
+TEST(ThreadPoolTest, TestMultipleParallelFor_4Thread_4Conc_1MTasks) {
+  TestMultipleParallelFor("TestMultipleParallelFor_4Thread_4Conc_1MTasks", 4, 4, 1000000);
+}
 #ifdef _WIN32
 TEST(ThreadPoolTest, TestStackSize) {
   ThreadOptions to;