mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Use Eigen in Round implementation (#23571)
### Description Attempt to make it more consistent. ### Motivation and Context Customer reports big difference in perf of Round between Windows and Linux.
This commit is contained in:
parent
e8b0bdb127
commit
cddc271b0b
1 changed files with 15 additions and 10 deletions
|
|
@ -9,6 +9,7 @@
|
|||
#include <cmath>
|
||||
#include "core/providers/cpu/math/element_wise_ops.h"
|
||||
#include "core/util/math.h"
|
||||
#include "core/util/math_cpuonly.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
|
|
@ -24,24 +25,28 @@ template <typename T>
|
|||
Status Round<T>::Compute(OpKernelContext* ctx) const {
|
||||
const auto& X = *ctx->Input<Tensor>(0);
|
||||
auto& Y = *ctx->Output(0, X.Shape());
|
||||
auto* input = X.Data<T>();
|
||||
const auto* input = X.Data<T>();
|
||||
auto* output = Y.MutableData<T>();
|
||||
const auto size = X.Shape().Size();
|
||||
for (int64_t i = 0; i < size; ++i, ++output, ++input) {
|
||||
*output = ::rint(*input);
|
||||
}
|
||||
const auto size = narrow<Eigen::Index>(X.Shape().Size());
|
||||
|
||||
EigenArrayMap<T> Y_arr(output, 1, size);
|
||||
ConstEigenArrayMap<T> X_arr(input, 1, size);
|
||||
Y_arr = X_arr.rint();
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
template <>
|
||||
Status Round<MLFloat16>::Compute(OpKernelContext* ctx) const {
|
||||
const auto& X = *ctx->Input<Tensor>(0);
|
||||
auto& Y = *ctx->Output(0, X.Shape());
|
||||
auto* input = X.Data<MLFloat16>();
|
||||
const auto* input = X.Data<MLFloat16>();
|
||||
auto* output = Y.MutableData<MLFloat16>();
|
||||
const auto size = X.Shape().Size();
|
||||
for (int64_t i = 0; i < size; ++i, ++output, ++input) {
|
||||
*output = MLFloat16(static_cast<float>(::rint(input->ToFloat())));
|
||||
}
|
||||
const auto size = narrow<Eigen::Index>(X.Shape().Size());
|
||||
|
||||
ConstEigenArrayMap<Eigen::half> X_arr(reinterpret_cast<const Eigen::half*>(input), 1, size);
|
||||
EigenArrayMap<Eigen::half> Y_arr(reinterpret_cast<Eigen::half*>(output), 1, size);
|
||||
Y_arr = X_arr.rint();
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue