saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-21 19:18:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

ThreadPool fix for roialign and CropAndResize (#2020)

Browse source
This commit is contained in:
Changming Sun 2019-10-06 22:43:59 -07:00 committed by GitHub
parent 4cdb95e436
commit 3c26ae5b6d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/contrib_ops/cpu/crop_and_resize.cc
View file

@ -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>

10
onnxruntime/core/providers/cpu/object_detection/roialign.cc
View file

@ -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