diff --git a/include/onnxruntime/core/platform/threadpool.h b/include/onnxruntime/core/platform/threadpool.h index b656e92882..4a9d762043 100644 --- a/include/onnxruntime/core/platform/threadpool.h +++ b/include/onnxruntime/core/platform/threadpool.h @@ -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 > 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 diff --git a/onnxruntime/core/providers/cpu/math/top_k.cc b/onnxruntime/core/providers/cpu/math/top_k.cc index 8a3f92f53c..2d9c1bb82b 100644 --- a/onnxruntime/core/providers/cpu/math/top_k.cc +++ b/onnxruntime/core/providers/cpu/math/top_k.cc @@ -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(batch * rows / num_threads); - int64_t end_row = static_cast((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(batch * rows / num_threads); - int64_t end_row = static_cast((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 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(batch * rows / num_threads); - int64_t end_row = static_cast((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 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(comparer, row_offset, num_blocks, block_slice, j, k, sorted, data_holder); diff --git a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h index f951d7ccbb..aea90b9fdf 100644 --- a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h +++ b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h @@ -332,9 +332,8 @@ void TreeEnsembleCommon::ComputeAgg(concurrency::ThreadPool* ttp, num_threads, [this, &agg, &scores, &merge_mutex, num_threads, x_data](ptrdiff_t batch_num) { std::vector> 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::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> 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({0, 0})); for (j = 0; j < roots_.size(); ++j) { agg.ProcessTreeNodePrediction(scores, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride));