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

Make the reduction ops more consistent in checking if no transpose is required and skipping the copy of the input data if that is the case. Significantly better performance when this is done (2x faster for model calling ReduceSumSquare with input of {2048,10}). (#3265)

Browse source
This commit is contained in:
Scott McKay 2020-03-20 06:55:38 +10:00 committed by GitHub
parent 8f00147c14
commit 6dc25a60f8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 61 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -312,12 +312,21 @@ Status ReduceL1<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).cwiseAbs().rowwise().sum();
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).cwiseAbs().sum();
}
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).cwiseAbs().rowwise().sum();
}
return Status::OK();
}
@ -328,12 +337,21 @@ Status ReduceL2<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().norm();
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).norm();
}
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().norm();
}
return Status::OK();
}
@ -344,12 +362,22 @@ Status ReduceLogSum<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().sum();
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).sum();
}
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().sum();
}
for (int j = 0; j < block_size; ++j) {
*(output_data) = static_cast<T>(std::log(*(output_data)));
++output_data;
@ -463,12 +491,21 @@ Status ReduceProd<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().prod();
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).prod();
}
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().prod();
}
return Status::OK();
}
@ -506,12 +543,21 @@ Status ReduceSumSquare<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().squaredNorm();
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).squaredNorm();
}
} else {
EigenVectorMap<T> out_vec(output_data, block_size);
out_vec = ConstEigenMatrixMap<T>(&transposedInputData[0], block_size, blocks).rowwise().squaredNorm();
}
return Status::OK();
}