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

Make ThreadPool::PartitionWork a bit more user friendly. Update a few places to use PartitionWork. (#3795)

Browse source
This commit is contained in:
Scott McKay 2020-05-02 17:09:55 +10:00 committed by GitHub
parent 2b8d9ef0fd
commit 15eca74d15
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 44 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -173,10 +173,8 @@ class ThreadPool {
}
#pragma omp parallel for
for (std::ptrdiff_t i = 0; i < num_threads; i++) {
std::ptrdiff_t start, work_remaining;
PartitionWork(i, num_threads, total, &start, &work_remaining);
std::ptrdiff_t end = start + work_remaining;
fn(start, end);
auto work = PartitionWork(i, num_threads, total);
fn(work.start, work.end);
}
#else
if (tp == nullptr) {
@ -203,10 +201,8 @@ class ThreadPool {
}
#pragma omp parallel for
for (std::ptrdiff_t i = 0; i < num_threads; i++) {
std::ptrdiff_t start, work_remaining;
PartitionWork(i, num_threads, total, &start, &work_remaining);
std::ptrdiff_t end = start + work_remaining;
fn(start, end);
auto work = PartitionWork(i, num_threads, total);
fn(work.start, work.end);
}
#else
if (tp == nullptr) {
@ -304,16 +300,38 @@ class ThreadPool {
}
tp->SimpleParallelFor(num_batches, [&](std::ptrdiff_t batch_index) {
std::ptrdiff_t start, work_remaining;
PartitionWork(batch_index, num_batches, total, &start, &work_remaining);
std::ptrdiff_t end = start + work_remaining;
for (std::ptrdiff_t i = start; i < end; i++) {
auto work = PartitionWork(batch_index, num_batches, total);
for (std::ptrdiff_t i = work.start; i < work.end; i++) {
fn(i);
}
});
#endif
}
struct WorkInfo {
std::ptrdiff_t start;
std::ptrdiff_t end;
};
/** Calculate the start and end offsets for a batch.
@remarks Based on MlasPartitionWork
*/
static WorkInfo PartitionWork(std::ptrdiff_t batch_idx, std::ptrdiff_t num_batches, std::ptrdiff_t total_work) {
const std::ptrdiff_t work_per_batch = total_work / num_batches;
const std::ptrdiff_t work_per_batch_extra = total_work % num_batches;
WorkInfo info;
if (batch_idx < work_per_batch_extra) {
info.start = (work_per_batch + 1) * batch_idx;
info.end = info.start + work_per_batch + 1;
} else {
info.start = work_per_batch * batch_idx + work_per_batch_extra;
info.end = info.start + work_per_batch;
}
return info;
}
ORT_DISALLOW_COPY_AND_ASSIGNMENT(ThreadPool);
private:
@ -333,22 +351,7 @@ class ThreadPool {
// eigen_threadpool_ is instantiated and owned by thread::ThreadPool if
// user_threadpool is not in the constructor.
std::unique_ptr<ThreadPoolTempl<Env> > eigen_threadpool_;
// Copied from MlasPartitionWork
static void PartitionWork(std::ptrdiff_t ThreadId, std::ptrdiff_t ThreadCount, std::ptrdiff_t TotalWork,
std::ptrdiff_t* WorkIndex, std::ptrdiff_t* WorkRemaining) {
const std::ptrdiff_t WorkPerThread = TotalWork / ThreadCount;
const std::ptrdiff_t WorkPerThreadExtra = TotalWork % ThreadCount;
if (ThreadId < WorkPerThreadExtra) {
*WorkIndex = (WorkPerThread + 1) * ThreadId;
*WorkRemaining = WorkPerThread + 1;
} else {
*WorkIndex = WorkPerThread * ThreadId + WorkPerThreadExtra;
*WorkRemaining = WorkPerThread;
}
}
}; // namespace concurrency
};
} // namespace concurrency
} // namespace onnxruntime

21
onnxruntime/core/providers/cpu/math/top_k.cc
View file

@ -184,13 +184,12 @@ static void FindTopKElements(const Tensor* input, const TensorShape& input_shape
if (k == 1) {
// just need to compare values and not indexes as the first instance of the best value is always selected
find_top_k =
[num_threads, rows, block_slice, num_blocks, input_data, cols, &values_map, &indices_map](std::ptrdiff_t batch) {
int64_t start_row = static_cast<int64_t>(batch * rows / num_threads);
int64_t end_row = static_cast<int64_t>((batch + 1) * rows / num_threads);
[num_threads, rows, block_slice, num_blocks, input_data, cols,
&values_map, &indices_map](std::ptrdiff_t batch) {
auto work = concurrency::ThreadPool::PartitionWork(batch, num_threads, rows);
Comparator comparer(input_data);
for (int64_t i = start_row; i < end_row; ++i) {
for (auto i = work.start; i < work.end; ++i) {
auto row_offset = i * cols;
for (int64_t j = 0; j < block_slice; ++j) {
int64_t cur_idx = row_offset + j;
@ -219,9 +218,7 @@ static void FindTopKElements(const Tensor* input, const TensorShape& input_shape
find_top_k =
[num_threads, rows, block_slice, num_blocks, k, sorted,
input_data, cols, &values_map, &indices_map](std::ptrdiff_t batch) {
int64_t start_row = static_cast<int64_t>(batch * rows / num_threads);
int64_t end_row = static_cast<int64_t>((batch + 1) * rows / num_threads);
auto work = concurrency::ThreadPool::PartitionWork(batch, num_threads, rows);
Comparator comparer(input_data);
// the heap is stored in indices_data. each iteration overwrites the old data when it adds the
@ -229,7 +226,7 @@ static void FindTopKElements(const Tensor* input, const TensorShape& input_shape
std::vector<int64_t> indices_data(k);
int64_t* indices = indices_data.data(); // raw pointer is slightly faster for HeapifyIthPosition
for (int64_t i = start_row; i < end_row; ++i) {
for (auto i = work.start; i < work.end; ++i) {
const auto row_offset = i * cols;
for (int64_t j = 0; j < block_slice; ++j) {
@ -292,16 +289,14 @@ static void FindTopKElements(const Tensor* input, const TensorShape& input_shape
[num_threads, rows, block_slice, num_blocks, k, sorted,
input_data, cols,
&values_map, &indices_map](std::ptrdiff_t batch) {
int64_t start_row = static_cast<int64_t>(batch * rows / num_threads);
int64_t end_row = static_cast<int64_t>((batch + 1) * rows / num_threads);
auto work = concurrency::ThreadPool::PartitionWork(batch, num_threads, rows);
Comparator comparer(input_data);
// we re-use a single data_holder for performance. avoids allocating memory on each iteration.
// the call to SelectTopK overwrites any existing data so we don't need to clear on each iteration.
std::vector<int64_t> data_holder(num_blocks);
for (int64_t i = start_row; i < end_row; ++i) {
for (auto i = work.start; i < work.end; ++i) {
auto row_offset = i * cols;
for (int64_t j = 0; j < block_slice; ++j) {
SelectTopK<Comparator>(comparer, row_offset, num_blocks, block_slice, j, k, sorted, data_holder);

11
onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h
View file

@ -332,9 +332,8 @@ void TreeEnsembleCommon<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp,
num_threads,
[this, &agg, &scores, &merge_mutex, num_threads, x_data](ptrdiff_t batch_num) {
std::vector<ScoreValue<OTYPE>> private_scores(n_targets_or_classes_, {0, 0});
int64_t j = batch_num * n_trees_ / num_threads;
int64_t end = (batch_num + 1) * n_trees_ / num_threads;
for (; j < end; ++j) {
auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, n_trees_);
for (auto j = work.start; j < work.end; ++j) {
agg.ProcessTreeNodePrediction(private_scores, *ProcessTreeNodeLeave(roots_[j], x_data));
}
@ -368,9 +367,9 @@ void TreeEnsembleCommon<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp,
[this, &agg, num_threads, x_data, z_data, label_data, N, stride](ptrdiff_t batch_num) {
size_t j;
std::vector<ScoreValue<OTYPE>> scores(n_targets_or_classes_);
int64_t i = batch_num * N / num_threads;
int64_t end = (batch_num + 1) * N / num_threads;
for (; i < end; ++i) {
auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, N);
for (auto i = work.start; i < work.end; ++i) {
std::fill(scores.begin(), scores.end(), ScoreValue<OTYPE>({0, 0}));
for (j = 0; j < roots_.size(); ++j) {
agg.ProcessTreeNodePrediction(scores, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride));