From cddc271b0b9c5fdfa6d05e72751ab2516d6ed642 Mon Sep 17 00:00:00 2001
From: Dmitri Smirnov <yuslepukhin@users.noreply.github.com>
Date: Mon, 3 Feb 2025 17:06:12 -0800
Subject: [PATCH] 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.
---
 onnxruntime/core/providers/cpu/math/round.cc | 25 ++++++++++++--------
 1 file changed, 15 insertions(+), 10 deletions(-)

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