From 0a1257e46731d10edb25b59c852b97c07555ad5e Mon Sep 17 00:00:00 2001
From: Changming Sun <chasun@microsoft.com>
Date: Fri, 13 Mar 2020 12:49:17 -0700
Subject: [PATCH] Adjust the grouping logic in ThreadPool::TryBatchParallelFor
 (#3207)

1. No more plus 1.
2. Use MlasPartitionWork function to calculate the work index range.
---
 .../onnxruntime/core/platform/threadpool.h    |  7 ++---
 onnxruntime/contrib_ops/cpu/activations.h     |  2 +-
 onnxruntime/core/common/threadpool.cc         | 27 +++++++++++++++++--
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/include/onnxruntime/core/platform/threadpool.h b/include/onnxruntime/core/platform/threadpool.h
index b3f9d42c85..cacfae22ea 100644
--- a/include/onnxruntime/core/platform/threadpool.h
+++ b/include/onnxruntime/core/platform/threadpool.h
@@ -64,13 +64,14 @@ class ThreadPool {
   // void SetStealPartitions(const std::vector<std::pair<unsigned, unsigned>>& partitions);
 
   /**
-  Tries to call the given function in parallel, with calls split into (num_batches) batches.
-  **/
+   * Tries to call the given function in parallel, with calls split into (num_batches) batches.
+   * If tp is NULL and OpenMP is enabled, no grouping
+   */
   template <typename F>
   inline static void TryBatchParallelFor(concurrency::ThreadPool* tp, int32_t total, F&& fn, int32_t num_batches = 0) {
     if (tp != nullptr) {
       if (num_batches <= 0) {
-        num_batches = tp->NumThreads() + 1;
+        num_batches = std::min(total, tp->NumThreads());
       }
       tp->BatchParallelFor(total, std::forward<F>(fn), num_batches);
     } else {
diff --git a/onnxruntime/contrib_ops/cpu/activations.h b/onnxruntime/contrib_ops/cpu/activations.h
index cc13805f9f..7218026170 100644
--- a/onnxruntime/contrib_ops/cpu/activations.h
+++ b/onnxruntime/contrib_ops/cpu/activations.h
@@ -43,7 +43,7 @@ class Gelu : public OpKernel {
     if (nullptr != tp) {
       const T* input = X->template Data<T>();
       T* output = Y->template MutableData<T>();
-      int task_count = tp->NumThreads() + 1;
+      int task_count = tp->NumThreads();
       int64_t elem_count = X->Shape().Size();
       if (elem_count > task_count) {
         tp->ParallelFor(task_count, [input,
diff --git a/onnxruntime/core/common/threadpool.cc b/onnxruntime/core/common/threadpool.cc
index d4eba8a0ec..5c0199bcec 100644
--- a/onnxruntime/core/common/threadpool.cc
+++ b/onnxruntime/core/common/threadpool.cc
@@ -22,6 +22,28 @@
 
 using Eigen::Barrier;
 
+namespace {
+//Copied from MlasPartitionWork
+inline void PartitionWork(
+    int32_t ThreadId,
+    int32_t ThreadCount,
+    int32_t TotalWork,
+    int32_t* WorkIndex,
+    int32_t* WorkRemaining) {
+  const int32_t WorkPerThread = TotalWork / ThreadCount;
+  const int32_t WorkPerThreadExtra = TotalWork % ThreadCount;
+
+  if (ThreadId < WorkPerThreadExtra) {
+    *WorkIndex = (WorkPerThread + 1) * ThreadId;
+    *WorkRemaining = WorkPerThread + 1;
+  } else {
+    *WorkIndex = WorkPerThread * ThreadId + WorkPerThreadExtra;
+    *WorkRemaining = WorkPerThread;
+  }
+}
+
+}  // namespace
+
 namespace onnxruntime {
 
 namespace concurrency {
@@ -78,8 +100,9 @@ void ThreadPool::BatchParallelFor(int32_t total, std::function<void(int32_t)> fn
   }
 
   ParallelFor(num_batches, [&](int batch_index) {
-    int start = batch_index * total / num_batches;
-    int end = (batch_index + 1) * total / num_batches;
+    int start, work_remaining;
+    PartitionWork(batch_index, num_batches, total, &start, &work_remaining);
+    int end = start + work_remaining;
     for (int i = start; i < end; i++) {
       fn(i);
     }