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:
Dmitri Smirnov 2025-02-03 17:06:12 -08:00 committed by GitHub
parent e8b0bdb127
commit cddc271b0b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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();
}