Fix for floating point predicates (nv-legate#466)

* Fix the templates to instantiate correct variants for complex types * Add complex64 and float16 cases to tests for floating point predicates
jjwilke · Jul 29, 2022 · a203e0c · a203e0c
1 parent 858028e
commit a203e0c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 25 deletions.
diff --git a/src/cunumeric/unary/unary_op_util.h b/src/cunumeric/unary/unary_op_util.h
@@ -658,23 +658,25 @@ struct UnaryOp<UnaryOpCode::ISFINITE, CODE> {
 
   UnaryOp(const std::vector<legate::Store>& args) {}
 
-  template <typename _T = T, std::enable_if_t<!std::is_floating_point<_T>::value>* = nullptr>
+  template <typename _T = T, std::enable_if_t<std::is_integral<_T>::value>* = nullptr>
   constexpr bool operator()(const T& x) const
   {
     return true;
   }
 
   template <typename _T = T, std::enable_if_t<std::is_floating_point<_T>::value>* = nullptr>
-  constexpr bool operator()(const T& x) const
+  __CUDA_HD__ bool operator()(const T& x) const
   {
     return std::isfinite(x);
   }
 
   template <typename _T>
-  constexpr bool operator()(const complex<_T>& x) const
+  __CUDA_HD__ bool operator()(const complex<_T>& x) const
   {
     return std::isfinite(x.imag()) && std::isfinite(x.real());
   }
+
+  __CUDA_HD__ bool operator()(const __half& x) const { return isfinite(static_cast<float>(x)); }
 };
 
 template <legate::LegateTypeCode CODE>
@@ -684,31 +686,23 @@ struct UnaryOp<UnaryOpCode::ISINF, CODE> {
 
   UnaryOp(const std::vector<legate::Store>& args) {}
 
-  template <typename _T = T, std::enable_if_t<!std::is_floating_point<_T>::value>* = nullptr>
+  template <typename _T = T, std::enable_if_t<std::is_integral<_T>::value>* = nullptr>
   constexpr bool operator()(const T& x) const
   {
     return false;
   }
 
   template <typename _T = T, std::enable_if_t<std::is_floating_point<_T>::value>* = nullptr>
-  constexpr bool operator()(const T& x) const
+  __CUDA_HD__ bool operator()(const T& x) const
   {
     return std::isinf(x);
   }
 
   template <typename _T>
-  constexpr bool operator()(const complex<_T>& x) const
+  __CUDA_HD__ bool operator()(const complex<_T>& x) const
   {
     return std::isinf(x.imag()) || std::isinf(x.real());
   }
-};
-
-template <>
-struct UnaryOp<UnaryOpCode::ISINF, legate::LegateTypeCode::HALF_LT> {
-  static constexpr bool valid = true;
-  using T                     = __half;
-
-  UnaryOp(const std::vector<legate::Store>& args) {}
 
   __CUDA_HD__ bool operator()(const __half& x) const { return isinf(x); }
 };
@@ -720,32 +714,24 @@ struct UnaryOp<UnaryOpCode::ISNAN, CODE> {
 
   UnaryOp(const std::vector<legate::Store>& args) {}
 
-  template <typename _T = T, std::enable_if_t<!std::is_floating_point<_T>::value>* = nullptr>
+  template <typename _T = T, std::enable_if_t<std::is_integral<_T>::value>* = nullptr>
   constexpr bool operator()(const T& x) const
   {
     return false;
   }
 
   template <typename _T = T, std::enable_if_t<std::is_floating_point<_T>::value>* = nullptr>
-  constexpr bool operator()(const T& x) const
+  __CUDA_HD__ bool operator()(const T& x) const
   {
     using std::isnan;
     return isnan(x);
   }
 
   template <typename _T>
-  constexpr bool operator()(const complex<_T>& x) const
+  __CUDA_HD__ bool operator()(const complex<_T>& x) const
   {
     return std::isnan(x.imag()) || std::isnan(x.real());
   }
-};
-
-template <>
-struct UnaryOp<UnaryOpCode::ISNAN, legate::LegateTypeCode::HALF_LT> {
-  static constexpr bool valid = true;
-  using T                     = __half;
-
-  UnaryOp(const std::vector<legate::Store>& args) {}
 
   __CUDA_HD__ bool operator()(const __half& x) const { return isnan(x); }
 };

diff --git a/tests/integration/test_unary_ufunc.py b/tests/integration/test_unary_ufunc.py
@@ -196,6 +196,8 @@ def test_all():
         # "isnat",
     ]
     check_ops(ops, (np.array([-np.inf, 0.0, 1.0, np.inf, np.nan]),))
+    check_ops(ops, (np.array([-np.inf, 0.0, 1.0, np.inf, np.nan], dtype="F"),))
+    check_ops(ops, (np.array([-np.inf, 0.0, 1.0, np.inf, np.nan], dtype="e"),))
     check_ops(ops, (np.array(np.inf),))