ENH: Create boolean and integer ufuncs for isnan, isinf, and isfinite.

Previously, boolean values would be routed through the half implementations of
these functions, which added considerable overhead. Creating specialized
ufuncs improves performance by ~250x

Additionally, enable autovectorization of new isnan, isinf, and isfinite ufuncs.

5 files changed, 88 insertions, 3 deletions

diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
index 687a8467b..de0bb81fe 100644
--- a/numpy/core/code_generators/generate_umath.py
+++ b/numpy/core/code_generators/generate_umath.py
@@ -827,7 +827,7 @@ defdict = {
     Ufunc(1, 1, None,
           docstrings.get('numpy.core.umath.isnan'),
           None,
-          TD(inexact, out='?'),
+          TD(nodatetime_or_obj, out='?'),
           ),
 'isnat':
     Ufunc(1, 1, None,
@@ -839,13 +839,13 @@ defdict = {
     Ufunc(1, 1, None,
           docstrings.get('numpy.core.umath.isinf'),
           None,
-          TD(inexact, out='?'),
+          TD(nodatetime_or_obj, out='?'),
           ),
 'isfinite':
     Ufunc(1, 1, None,
           docstrings.get('numpy.core.umath.isfinite'),
           None,
-          TD(inexact, out='?'),
+          TD(nodatetime_or_obj, out='?'),
           ),
 'signbit':
     Ufunc(1, 1, None,
diff --git a/numpy/core/src/umath/fast_loop_macros.h b/numpy/core/src/umath/fast_loop_macros.h
index 37656dcf5..e3cfa1f72 100644
--- a/numpy/core/src/umath/fast_loop_macros.h
+++ b/numpy/core/src/umath/fast_loop_macros.h
@@ -64,6 +64,8 @@
 #define IS_UNARY_CONT(tin, tout) (steps[0] == sizeof(tin) && \
                                   steps[1] == sizeof(tout))
 
+#define IS_OUTPUT_CONT(tout) (steps[1] == sizeof(tout))
+
 #define IS_BINARY_REDUCE ((args[0] == args[2])\
         && (steps[0] == steps[2])\
         && (steps[0] == 0))
@@ -84,6 +86,29 @@
 
 /*
  * loop with contiguous specialization
+ * op should be the code storing the result in `tout * out`
+ * combine with NPY_GCC_OPT_3 to allow autovectorization
+ * should only be used where its worthwhile to avoid code bloat
+ */
+#define BASE_OUTPUT_LOOP(tout, op) \
+    OUTPUT_LOOP { \
+        tout * out = (tout *)op1; \
+        op; \
+    }
+#define OUTPUT_LOOP_FAST(tout, op) \
+    do { \
+    /* condition allows compiler to optimize the generic macro */ \
+    if (IS_OUTPUT_CONT(tout)) { \
+        BASE_OUTPUT_LOOP(tout, op) \
+    } \
+    else { \
+        BASE_OUTPUT_LOOP(tout, op) \
+    } \
+    } \
+    while (0)
+
+/*
+ * loop with contiguous specialization
  * op should be the code working on `tin in` and
  * storing the result in `tout * out`
  * combine with NPY_GCC_OPT_3 to allow autovectorization
diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src
index 04e6cbdee..1e4ab350b 100644
--- a/numpy/core/src/umath/loops.c.src
+++ b/numpy/core/src/umath/loops.c.src
@@ -644,6 +644,19 @@ BOOL__ones_like(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UN
 }
 
 
+/**begin repeat
+ * #kind = isnan, isinf, isfinite#
+ * #func = npy_isnan, npy_isinf, npy_isfinite#
+ * #val = NPY_FALSE, NPY_FALSE, NPY_TRUE#
+ **/
+NPY_NO_EXPORT void
+BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    OUTPUT_LOOP_FAST(npy_bool, *out = @val@);
+}
+
+/**end repeat**/
+
 /*
  *****************************************************************************
  **                           INTEGER LOOPS
@@ -875,6 +888,18 @@ NPY_NO_EXPORT void
     }
 }
 
+/**begin repeat1
+ * #kind = isnan, isinf, isfinite#
+ * #func = npy_isnan, npy_isinf, npy_isfinite#
+ * #val = NPY_FALSE, NPY_FALSE, NPY_TRUE#
+ **/
+NPY_NO_EXPORT void
+@TYPE@_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    OUTPUT_LOOP_FAST(npy_bool, *out = @val@);
+}
+/**end repeat1**/
+
 /**end repeat**/
 
 /**begin repeat
diff --git a/numpy/core/src/umath/loops.h.src b/numpy/core/src/umath/loops.h.src
index 5264a6533..f48319056 100644
--- a/numpy/core/src/umath/loops.h.src
+++ b/numpy/core/src/umath/loops.h.src
@@ -38,6 +38,13 @@ BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED
 NPY_NO_EXPORT void
 BOOL__ones_like(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data));
 
+/**begin repeat
+ * #kind = isnan, isinf, isfinite#
+ **/
+NPY_NO_EXPORT void
+BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
+/**end repeat**/
+
 /*
  *****************************************************************************
  **                           INTEGER LOOPS
@@ -146,6 +153,13 @@ NPY_NO_EXPORT void
 NPY_NO_EXPORT void
 @S@@TYPE@_lcm(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
 
+/**begin repeat2
+ * #kind = isnan, isinf, isfinite#
+ **/
+NPY_NO_EXPORT void
+@S@@TYPE@_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
+/**end repeat2**/
+
 /**end repeat1**/
 
 /**end repeat**/
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index 478a08397..b6b68d922 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -1915,3 +1915,24 @@ class TestUfunc(object):
         exc = pytest.raises(TypeError, np.sqrt, None)
         # minimally check the exception text
         assert 'loop of ufunc does not support' in str(exc)
+
+    @pytest.mark.parametrize('nat', [np.datetime64('nat'), np.timedelta64('nat')])
+    def test_nat_is_not_finite(self, nat):
+        try:
+            assert not np.isfinite(nat)
+        except TypeError:
+            pass  # ok, just not implemented
+
+    @pytest.mark.parametrize('nat', [np.datetime64('nat'), np.timedelta64('nat')])
+    def test_nat_is_nan(self, nat):
+        try:
+            assert np.isnan(nat)
+        except TypeError:
+            pass  # ok, just not implemented
+
+    @pytest.mark.parametrize('nat', [np.datetime64('nat'), np.timedelta64('nat')])
+    def test_nat_is_not_inf(self, nat):
+        try:
+            assert not np.isinf(nat)
+        except TypeError:
+            pass  # ok, just not implemented