Added python binding. Added test.

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

docs/api/python/intrin.rst

@@ -35,6 +35,7 @@ tvm.intrin
    tvm.ceil
    tvm.trunc
    tvm.round
+   tvm.nearbyint
    tvm.abs
    tvm.isnan
 
@@ -52,5 +53,6 @@ tvm.intrin
 .. autofunction:: tvm.ceil
 .. autofunction:: tvm.trunc
 .. autofunction:: tvm.round
+.. autofunction:: tvm.nearbyint
 .. autofunction:: tvm.abs
 .. autofunction:: tvm.isnan

python/tvm/intrin.py

@@ -434,6 +434,29 @@ def round(x):
     return _make.round(x)
 
 
+def nearbyint(x):
+    """Round elements of the array to the nearest integer.
+    This intrinsic uses llvm.nearbyint instead of llvm.round
+    which is faster but will results different from tvm.round.
+    Notably nearbyint rounds according to the rounding mode,
+    whereas tvm.round (llvm.round) ignores that.
+    For differences between the two see:
+    https://en.cppreference.com/w/cpp/numeric/math/round
+    https://en.cppreference.com/w/cpp/numeric/math/nearbyint
+
+    Parameters
+    ----------
+    x : Expr
+        Input argument.
+
+    Returns
+    -------
+    y : Expr
+        The result.
+    """
+    return _make.nearbyint(x)
+
+
 def isnan(x):
     """Check if input value is Nan.
 

src/api/api_ir.cc

@@ -50,6 +50,9 @@ TVM_REGISTER_API("make.ceil")
 TVM_REGISTER_API("make.round")
 .set_body_typed(tvm::round);
 
+TVM_REGISTER_API("make.nearbyint")
+.set_body_typed(tvm::nearbyint);
+
 TVM_REGISTER_API("make.trunc")
 .set_body_typed(tvm::trunc);
 

tests/python/unittest/test_tvm_intrin.py

@@ -0,0 +1,48 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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.
+import tvm
+import topi
+from tvm.contrib import util, clang
+import numpy as np
+import ctypes
+import math
+
+
+def test_nearbyint():
+    m = tvm.var("m",)
+    A = tvm.placeholder((m,), name='A')
+    A_rounded = tvm.compute((m,), lambda *i: tvm.nearbyint(A(*i)), name='A')
+    s = tvm.create_schedule(A_rounded.op)
+    f = tvm.build(s, [A, A_rounded], "llvm")
+    ctx = tvm.cpu(0)
+    n = 10
+    a = tvm.nd.array(np.random.uniform(high=100, size=n).astype(A.dtype), ctx)
+    a_rounded = tvm.nd.array( \
+            np.random.uniform(size=n).astype(A_rounded.dtype), ctx)
+    f(a, a_rounded)
+    # Note that numpys rint rounds to nearest integer with
+    # ties to halfway is broken by rounding to even.
+    # So that 1.5 and 2.5 will round 2.
+    # This is the default rounding mode with libc as well.
+    # However one can set a different rounding mode and in that
+    # case numpy result might differ.
+    tvm.testing.assert_allclose(
+        a_rounded.asnumpy(), np.rint(a.asnumpy()))
+
+
+if __name__ == "__main__":
+    test_nearbyint()