diff --git a/onnxruntime/core/providers/cpu/math/round.cc b/onnxruntime/core/providers/cpu/math/round.cc index 86be7cce43..4480829779 100644 --- a/onnxruntime/core/providers/cpu/math/round.cc +++ b/onnxruntime/core/providers/cpu/math/round.cc @@ -9,6 +9,7 @@ #include #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 Status Round::Compute(OpKernelContext* ctx) const { const auto& X = *ctx->Input(0); auto& Y = *ctx->Output(0, X.Shape()); - auto* input = X.Data(); + const auto* input = X.Data(); auto* output = Y.MutableData(); - const auto size = X.Shape().Size(); - for (int64_t i = 0; i < size; ++i, ++output, ++input) { - *output = ::rint(*input); - } + const auto size = narrow(X.Shape().Size()); + + EigenArrayMap Y_arr(output, 1, size); + ConstEigenArrayMap X_arr(input, 1, size); + Y_arr = X_arr.rint(); + return Status::OK(); } template <> Status Round::Compute(OpKernelContext* ctx) const { const auto& X = *ctx->Input(0); auto& Y = *ctx->Output(0, X.Shape()); - auto* input = X.Data(); + const auto* input = X.Data(); auto* output = Y.MutableData(); - const auto size = X.Shape().Size(); - for (int64_t i = 0; i < size; ++i, ++output, ++input) { - *output = MLFloat16(static_cast(::rint(input->ToFloat()))); - } + const auto size = narrow(X.Shape().Size()); + + ConstEigenArrayMap X_arr(reinterpret_cast(input), 1, size); + EigenArrayMap Y_arr(reinterpret_cast(output), 1, size); + Y_arr = X_arr.rint(); + return Status::OK(); }