fix bug with isnan and isinf

owensgroup · Sep 27, 2024 · 4808274 · 4808274
1 parent 582f259
commit 4808274
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 5 deletions.
diff --git a/include/rxmesh/diff/scalar.h b/include/rxmesh/diff/scalar.h
@@ -403,14 +403,12 @@ struct Scalar
 
  __host__ __device__ friend bool isnan(const Scalar& a)
  {
-
- return std::isnan(a.val);
+ return is_nan(a.val);
  }
 
  __host__ __device__ friend bool isinf(const Scalar& a)
  {
-
- return std::isinf(a.val);
+ return is_inf(a.val);
  }
 
  __host__ __device__ friend bool isfinite(const Scalar& a)

diff --git a/include/rxmesh/diff/util.h b/include/rxmesh/diff/util.h
@@ -78,6 +78,35 @@ __device__ __host__ __inline__ bool is_same_matrix(
  return true;
 }
 
+/**
+ * @brief check if a variable is +/- infinity. Works both on host and device
+ */
+template <typename T>
+__device__ __host__ __inline__ bool is_inf(const T& A)
+{
+
+#ifdef __CUDA_ARCH__
+ return ::isinf(A);
+#else
+ return std::isinf(A);
+#endif
+}
+
+/**
+ * @brief check if a variable is nan. Works both on host and device
+ */
+template <typename T>
+__device__ __host__ __inline__ bool is_nan(const T& A)
+{
+
+#ifdef __CUDA_ARCH__
+ return ::isnan(A);
+#else
+ return std::isnan(A);
+#endif
+}
+
+
 /**
  * @brief check if a variable is finite. Works both on host and device
  */

diff --git a/tests/RXMesh_test/test_scalar.cu b/tests/RXMesh_test/test_scalar.cu
@@ -633,7 +633,7 @@ __global__ static void test_is_nan_is_inf(int* d_err, T eps = 1e-9)
  RX_ASSERT_TRUE(isfinite(a), d_err);
  RX_ASSERT_TRUE(!isfinite(b), d_err);
  RX_ASSERT_TRUE(!isfinite(c), d_err);
- RX_ASSERT_TRUE(isfinite(d), d_err);
+ RX_ASSERT_TRUE(!isfinite(d), d_err);
 }
 
 template <typename T, bool WithHessian>