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

Adjust the grouping logic in ThreadPool::TryBatchParallelFor (#3207)

Browse source 
1. No more plus 1.
2. Use MlasPartitionWork function to calculate the work index range.
This commit is contained in:
Changming Sun 2020-03-13 12:49:17 -07:00 committed by GitHub
parent 5bc0d8be5c
commit 0a1257e467
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 30 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
include/onnxruntime/core/platform/threadpool.h
View file

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

2
onnxruntime/contrib_ops/cpu/activations.h
View file

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

27
onnxruntime/core/common/threadpool.cc
View file

@ -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);
}