From 3c26ae5b6d5ab6024dffae0794ae0774e7e1109b Mon Sep 17 00:00:00 2001
From: Changming Sun <chasun@microsoft.com>
Date: Sun, 6 Oct 2019 22:43:59 -0700
Subject: [PATCH] ThreadPool fix for roialign and CropAndResize (#2020)

---
 onnxruntime/contrib_ops/cpu/crop_and_resize.cc        | 11 ++++++++++-
 .../core/providers/cpu/object_detection/roialign.cc   | 10 +++++++++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/onnxruntime/contrib_ops/cpu/crop_and_resize.cc b/onnxruntime/contrib_ops/cpu/crop_and_resize.cc
index 238e7c3361..c7e8232477 100644
--- a/onnxruntime/contrib_ops/cpu/crop_and_resize.cc
+++ b/onnxruntime/contrib_ops/cpu/crop_and_resize.cc
@@ -43,6 +43,15 @@ namespace contrib {
 
 ADD_TYPED_CROPANDRESIZE_OP(float);
 
+template <typename T>
+static void TryParallelFor(concurrency::ThreadPool* tp, int32_t total, T&& fn) {
+  if (tp != nullptr)
+    return tp->ParallelFor(total, fn);
+  for (int32_t i = 0; i != total; ++i) {
+    fn(i);
+  }
+}
+
 template <typename T>
 void CropAndResizeForward(
     int64_t nthreads,
@@ -177,7 +186,7 @@ void CropAndResizeForward(
       }  // for pw
     }    // for ph
   };     // for n
-  const_cast<ThreadPool*>(ttp)->ParallelFor(static_cast<int32_t>(n_rois), work_object);
+  TryParallelFor(const_cast<ThreadPool*>(ttp), static_cast<int32_t>(n_rois), work_object);
 }
 
 template <typename T>
diff --git a/onnxruntime/core/providers/cpu/object_detection/roialign.cc b/onnxruntime/core/providers/cpu/object_detection/roialign.cc
index 9975eb98d5..64389d08d2 100644
--- a/onnxruntime/core/providers/cpu/object_detection/roialign.cc
+++ b/onnxruntime/core/providers/cpu/object_detection/roialign.cc
@@ -43,6 +43,14 @@ ADD_TYPED_ROIALIGN_OP(double);
 
 namespace {
 template <typename T>
+void TryParallelFor(concurrency::ThreadPool* tp, int32_t total, T&& fn) {
+  if (tp != nullptr)
+    return tp->ParallelFor(total, fn);
+  for (int32_t i = 0; i != total; ++i) {
+    fn(i);
+  }
+}
+template <typename T>
 struct PreCalc {
   int64_t pos1;
   int64_t pos2;
@@ -270,7 +278,7 @@ void RoiAlignForward(
       }    // for ph
     }      // for c
   };       // for n
-  if (ttp != nullptr) const_cast<ThreadPool*>(ttp)->ParallelFor(static_cast<int32_t>(n_rois), work_object);
+  TryParallelFor(const_cast<ThreadPool*>(ttp), static_cast<int32_t>(n_rois), work_object);
 }
 }  // namespace