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

Use Eigen threadpool for ReduceSum and ReduceMean. (#3441)

Browse source 
* Use Eigen threadpool for ReduceSum and ReduceMean.

* Fix mac build
This commit is contained in:
Pranav Sharma 2020-04-08 11:50:22 -07:00 committed by GitHub
parent f8fa1dde55
commit cdac74b3c3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
onnxruntime/core/providers/cpu/reduction/reduction_ops.cc
View file

@ -5,6 +5,8 @@
#include "core/providers/common.h"
#include "core/util/math_cpuonly.h"
#include "core/providers/cpu/containers.h"
#include "core/platform/threadpool.h"
using namespace std;
namespace onnxruntime {
@ -490,13 +492,10 @@ Status ReduceMean<T>::Compute(OpKernelContext* ctx) const {
if (no_transpose) {
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
#ifdef _OPENMP
#pragma omp parallel for
#endif
for (int64_t i = 0; i < block_size; ++i) {
auto lambda = [input_data, blocks, output_data](ptrdiff_t i) {
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).mean();
}
};
concurrency::ThreadPool::TryBatchParallelFor(ctx->GetOperatorThreadPool(), block_size, lambda, 0);
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().mean();
@ -566,13 +565,10 @@ Status ReduceSum<T>::Compute(OpKernelContext* ctx) const {
if (no_transpose) {
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
#ifdef _OPENMP
#pragma omp parallel for
#endif
for (int64_t i = 0; i < block_size; ++i) {
auto lambda = [input_data, blocks, output_data](ptrdiff_t i) {
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).sum();
}
};
concurrency::ThreadPool::TryBatchParallelFor(ctx->GetOperatorThreadPool(), block_size, lambda, 0);
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().sum();