add Rint operation (tensorflow#4113)

* add rint op

* wip on custom rounding op

* fix Makefile build for the rint op

* rebase, add test, fix examples

Author: Mistobaan

tensorflow/contrib/makefile/proto_text_cc_files.txt

@@ -11,6 +11,7 @@ tensorflow/core/platform/posix/env.cc
 tensorflow/core/platform/posix/load_library.cc
 tensorflow/core/platform/file_system.cc
 tensorflow/core/platform/env.cc
+tensorflow/core/platform/setround.cc
 tensorflow/core/platform/denormal.cc
 tensorflow/core/platform/default/tracing.cc
 tensorflow/core/platform/default/logging.cc

tensorflow/core/kernels/cwise_op_gpu_rint.cu.cc

@@ -0,0 +1,26 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA
+
+#include "tensorflow/core/kernels/cwise_ops_gpu_common.cu.h"
+
+namespace tensorflow {
+namespace functor {
+DEFINE_UNARY2(rint, float, double);
+}  // namespace functor
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA

tensorflow/core/kernels/cwise_op_rint.cc

@@ -0,0 +1,23 @@
+/* Copyright 2016 Google Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/kernels/cwise_ops_common.h"
+
+namespace tensorflow {
+REGISTER2(UnaryOp, CPU, "Rint", functor::rint, float, double);
+#if GOOGLE_CUDA
+REGISTER2(UnaryOp, GPU, "Rint", functor::rint, float, double);
+#endif
+}  // namespace tensorflow

tensorflow/core/kernels/cwise_ops.h

@@ -521,6 +521,26 @@ struct round : base<T, Eigen::internal::scalar_round_op_google<T>> {};
 template <typename T>
 struct ceil : base<T, Eigen::internal::scalar_ceil_op<T>> {};
 
+/** this should go in Eigen
+  * \brief Template functor to compute the round to int value of a scalar
+  */
+template<typename Scalar> struct scalar_rint_op {
+  EIGEN_EMPTY_STRUCT_CTOR(scalar_rint_op)
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
+    const Scalar operator() (const Scalar& a) const {
+    #if defined(__CUDACC__)
+        return ::rint(a);
+    #elif defined(PLATFORM_POSIX_ANDROID)
+        return rint(a);
+    #else
+        return std::rint(a);
+    #endif
+  }
+};
+
+template <typename T>
+struct rint : base<T, scalar_rint_op<T>> {};
+
 ////////////////////////////////////////////////////////////////////////////////
 // Binary functors
 ////////////////////////////////////////////////////////////////////////////////

tensorflow/core/kernels/cwise_ops_test.cc

@@ -59,6 +59,11 @@ BM_UNARY(gpu, Conj, std::complex<float>, DT_COMPLEX64);
 BM_UNARY(cpu, Conj, std::complex<double>, DT_COMPLEX128);
 BM_UNARY(gpu, Conj, std::complex<double>, DT_COMPLEX128);
 
+BM_UNARY(cpu, Rint, double, DT_DOUBLE);
+BM_UNARY(gpu, Rint, double, DT_DOUBLE);
+BM_UNARY(cpu, Rint, float, DT_FLOAT);
+BM_UNARY(gpu, Rint, float, DT_FLOAT);
+
 // data func scalar.
 static Graph* BinaryScalar(int num, const string& func) {
   Graph* g = new Graph(OpRegistry::Global());

tensorflow/core/lib/core/threadpool.cc

@@ -21,9 +21,11 @@ limitations under the License.
 #include "tensorflow/core/platform/denormal.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/setround.h"
 #include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/platform/types.h"
 
+
 namespace tensorflow {
 namespace thread {
 
@@ -50,6 +52,8 @@ struct EigenEnvironment {
     return env_->StartThread(thread_options_, name_, [=]() {
       // Set the processor flag to flush denormals to zero
       port::ScopedFlushDenormal flush;
+      // Set the C++ rounding mode to ROUND TO NEAREST
+      port::ScopedSetRound round;
       f();
     });
   }

tensorflow/core/ops/math_ops.cc

@@ -454,6 +454,25 @@ REGISTER_OP("Ceil")
 Returns element-wise smallest integer in not less than x.
 )doc");
 
+REGISTER_OP("Rint")
+    .Input("x: T")
+    .Output("y: T")
+    .Attr("T: {float, double}")
+    .SetShapeFn(shape_inference::UnchangedShape)
+    .Doc(R"doc(
+Returns element-wise integer closest to x.
+
+If the result is midway between two representable values,
+the even representable is chosen.
+For example:
+
+```
+rint(-1.5) ==> -2.0
+rint(0.5000001) ==> 1.0
+rint([-1.7, -1.5, -0.2, 0.2, 1.5, 1.7, 2.0]) ==> [-2., -2., -0., 0., 2., 2., 2.]
+```
+)doc");
+
 // Declares cwise binary operations signature: 't, 't -> 't.
 
 #define BINARY_MORE()                              \

tensorflow/core/platform/setround.cc

@@ -0,0 +1,35 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/platform/setround.h"
+
+#ifdef __STDC_IEC_559__
+#include <fenv.h> // fesetround, FE_*
+#endif
+
+namespace tensorflow {
+namespace port {
+
+ScopedSetRound::ScopedSetRound() {
+#ifdef __STDC_IEC_559__
+   std::fesetround(FE_TONEAREST);
+#endif
+}
+
+ScopedSetRound::~ScopedSetRound() {
+}
+
+}  // namespace port
+}  // namespace tensorflow

tensorflow/core/platform/setround.h

@@ -0,0 +1,38 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_PLATFORM_SETROUND_H_
+#define TENSORFLOW_PLATFORM_SETROUND_H_
+
+#include "tensorflow/core/platform/macros.h"
+
+namespace tensorflow {
+namespace port {
+
+// While this class is active, floating point numbers are rounded to NEAREST
+// to zero.  The destructor restores the original flags.
+class ScopedSetRound {
+ public:
+  ScopedSetRound();
+  ~ScopedSetRound();
+
+ private:
+  TF_DISALLOW_COPY_AND_ASSIGN(ScopedSetRound);
+};
+
+}  // namespace port
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PLATFORM_SETROUN_H_

tensorflow/python/kernel_tests/cwise_ops_test.py

@@ -1778,9 +1778,17 @@ def testSqrt(self):
 
 class RoundingTest(tf.test.TestCase):
 
-  def _compare(self, x, use_gpu):
+  def _compare_values(self, x, y=None):
+    y = np.rint(x) if y is None else np.asarray(y)
+    with self.test_session() as sess:
+      tf_rint = tf.rint(x)
+      np_rint = sess.run(tf_rint)
+    self.assertAllEqual(y, np_rint)
+    self.assertShapeEqual(y, tf_rint)
+
+  def _compare(self, x):
     np_floor, np_ceil = np.floor(x), np.ceil(x)
-    with self.test_session(use_gpu=use_gpu) as sess:
+    with self.test_session() as sess:
       inx = tf.convert_to_tensor(x)
       ofloor, oceil = tf.floor(inx), tf.ceil(inx)
       tf_floor, tf_ceil = sess.run([ofloor, oceil])
@@ -1790,9 +1798,20 @@ def _compare(self, x, use_gpu):
     self.assertShapeEqual(np_ceil, oceil)
 
   def _testDtype(self, dtype):
-    data = (np.arange(-3, 3) / 4.).reshape([1, 3, 2]).astype(dtype)
-    self._compare(data, use_gpu=True)
-    self._compare(data, use_gpu=True)
+    data = (np.arange(-3, 3) / 4.).reshape(1, 3, 2).astype(dtype)
+    self._compare(data)
+    # TODO: rint op is not supported for float16
+    if dtype is np.float16:
+      return
+    self._compare_values(data)
+    x = [0.5, 0.5000001]
+    y = [0.0, 1.0]
+    self._compare_values(x, y=y)
+
+    # numpy example
+    x = [-1.7, -1.5, -0.2, 0.2, 1.5, 1.7, 2.0]
+    y = [-2., -2., -0.,  0.,  2.,  2.,  2.]
+    self._compare_values(x, y=y)
 
   def testTypes(self):
     for dtype in [np.float16, np.float32, np.float64]: