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

ReduceMax/Min performance improvements on CPU (#1925)

Browse source 
On behalf of Tracy
This commit is contained in:
KeDengMS 2019-09-26 00:31:28 -07:00 committed by GitHub
parent c9b9e79e38
commit 9ed85987e3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 24 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -102,8 +102,9 @@ bool PrepareForReduce(OpKernelContext* ctx,
if (axes.empty()) {
// This is the default case for non-arg kind reductions. Reduce on all dimensions.
for (size_t i = 0; i < ndim; i++)
for (size_t i = 0; i < ndim; i++) {
axes.push_back(i);
}
}
std::sort(axes.begin(), axes.end());
@ -320,12 +321,20 @@ Status ReduceMax<T>::Compute(OpKernelContext* ctx) const {
int64_t block_size;
int64_t blocks;
Tensor* reduced;
PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_);
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
T* output_data = reduced->template MutableData<T>();
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().maxCoeff();
if (no_transpose) {
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
for (int64_t i = 0; i < block_size; ++i) {
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).maxCoeff();
}
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().maxCoeff();
}
return Status::OK();
}
@ -363,12 +372,20 @@ Status ReduceMin<T>::Compute(OpKernelContext* ctx) const {
int64_t block_size;
int64_t blocks;
Tensor* reduced;
PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_);
bool no_transpose = PrepareForReduce<T>(ctx, transposedInputData, &reduced, block_size, blocks, axes_, keepdims_, true);
T* output_data = reduced->template MutableData<T>();
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().minCoeff();
if (no_transpose) {
const T* input_data = ctx->Input<Tensor>(0)->template Data<T>();
for (int64_t i = 0; i < block_size; ++i) {
output_data[i] = ConstEigenVectorMap<T>(input_data + (i * blocks), blocks).minCoeff();
}
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().minCoeff();
}
return Status::OK();
}