mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-12 17:57:38 +00:00
a02638eb46
1. Do not reuse the main thread. 2. Do not plus one when mlas calculate the number of tasks to schedule. (It was me put the plus one there) This is the second try of #1839 It's known that this change has negative performance impact on some of the models.
120 lines
3 KiB
C++
120 lines
3 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#include "core/platform/threadpool.h"
|
|
#include "core/common/common.h"
|
|
|
|
#include <cassert>
|
|
|
|
#if defined(__GNUC__)
|
|
#pragma GCC diagnostic push
|
|
#pragma GCC diagnostic ignored "-Wunused-parameter"
|
|
#else
|
|
#pragma warning(push)
|
|
#pragma warning(disable : 4267)
|
|
#endif
|
|
#include <unsupported/Eigen/CXX11/src/ThreadPool/Barrier.h>
|
|
#if defined(__GNUC__)
|
|
#pragma GCC diagnostic pop
|
|
#else
|
|
#pragma warning(pop)
|
|
#endif
|
|
|
|
using Eigen::Barrier;
|
|
|
|
namespace onnxruntime {
|
|
|
|
namespace concurrency {
|
|
//
|
|
// ThreadPool
|
|
//
|
|
ThreadPool::ThreadPool(const std::string&, int num_threads) : impl_(num_threads) {}
|
|
|
|
void ThreadPool::Schedule(std::function<void()> fn) { impl_.Schedule(fn); }
|
|
|
|
void ThreadPool::ParallelFor(int32_t total, std::function<void(int32_t)> fn) {
|
|
if (total <= 0)
|
|
return;
|
|
|
|
if (total == 1) {
|
|
fn(0);
|
|
return;
|
|
}
|
|
|
|
// TODO: Eigen supports a more efficient ThreadPoolDevice mechanism
|
|
// We will simply rely on the work queue and stealing in the short term.
|
|
Barrier barrier(static_cast<unsigned int>(total));
|
|
std::function<void(int32_t)> handle_iteration = [&barrier, &fn](int iteration) {
|
|
fn(iteration);
|
|
barrier.Notify();
|
|
};
|
|
|
|
for (int32_t id = 0; id < total; ++id) {
|
|
Schedule([=, &handle_iteration]() { handle_iteration(id); });
|
|
}
|
|
|
|
barrier.Wait();
|
|
}
|
|
|
|
void ThreadPool::BatchParallelFor(int32_t total, std::function<void(int32_t)> fn, int32_t num_batches) {
|
|
if (total <= 0)
|
|
return;
|
|
|
|
if (total == 1) {
|
|
fn(0);
|
|
return;
|
|
}
|
|
|
|
if (num_batches <= 1) {
|
|
for (int i = 0; i < total; i++) {
|
|
fn(i);
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (num_batches >= total) {
|
|
ParallelFor(total, fn);
|
|
return;
|
|
}
|
|
|
|
ParallelFor(num_batches, [&](int batch_index) {
|
|
int start = batch_index * total / num_batches;
|
|
int end = (batch_index + 1) * total / num_batches;
|
|
for (int i = start; i < end; i++) {
|
|
fn(i);
|
|
}
|
|
});
|
|
}
|
|
|
|
void ThreadPool::ParallelForRange(int64_t first, int64_t last, std::function<void(int64_t, int64_t)> fn) {
|
|
if (last <= first) return;
|
|
if (last - first == 1) {
|
|
fn(first, last);
|
|
return;
|
|
}
|
|
|
|
// TODO: Eigen supports a more efficient ThreadPoolDevice mechanism
|
|
// We will simply rely on the work queue and stealing in the short term.
|
|
Barrier barrier(static_cast<unsigned int>(last - first));
|
|
std::function<void(int64_t, int64_t)> handle_range = [&barrier, &fn](int64_t first, int64_t last) {
|
|
fn(first, last);
|
|
barrier.Notify();
|
|
};
|
|
|
|
for (int64_t id = first + 1; id <= last; ++id) {
|
|
Schedule([=, &handle_range]() { handle_range(id, id + 1); });
|
|
}
|
|
|
|
fn(first, first + 1);
|
|
barrier.Wait();
|
|
}
|
|
|
|
// void ThreadPool::SetStealPartitions(const std::vector<std::pair<unsigned, unsigned>>& partitions) {
|
|
// impl_->SetStealPartitions(partitions);
|
|
// }
|
|
|
|
int ThreadPool::NumThreads() const { return impl_.NumThreads(); }
|
|
|
|
int ThreadPool::CurrentThreadId() const { return impl_.CurrentThreadId(); }
|
|
} // namespace concurrency
|
|
} // namespace onnxruntime
|