Kokkos::nan() -> KokkosKernels::Impl::quiet_NaN_v

kokkos · May 13, 2024 · 0d872ef · 0d872ef
1 parent 41e3e9b
commit 0d872ef
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 6 deletions.
diff --git a/common/impl/KokkosKernels_NaN.hpp b/common/impl/KokkosKernels_NaN.hpp
@@ -0,0 +1,43 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#ifndef KOKKOSKERNELS_NAN_HPP
+#define KOKKOSKERNELS_NAN_HPP
+
+#include <Kokkos_ArithTraits.hpp>
+#include <Kokkos_NumericTraits.hpp>
+
+namespace KokkosKernels::Impl {
+
+// This could be constexpr if Kokkos::complex ctor was
+template <typename T>
+KOKKOS_INLINE_FUNCTION T quiet_NaN() {
+  if constexpr (std::is_same_v<double, T>) {
+    return double(Kokkos::Experimental::quiet_NaN_v<float>);  // Kokkos::Experimetnal::quiet_NaN_v<double>
+                                                              // is undefined in
+                                                              // device code
+  } else if constexpr (Kokkos::ArithTraits<T>::is_complex) {
+    using value_type = typename T::value_type;
+    return T(quiet_NaN<value_type>(),
+             quiet_NaN<value_type>());  // Kokkos::complex ctor is not constexpr
+  } else {
+    return Kokkos::Experimental::quiet_NaN_v<T>;
+  }
+}
+
+}  // namespace KokkosKernels::Impl
+
+#endif  // KOKKOSKERNELS_NAN_HPP
diff --git a/sparse/unit_test/Test_Sparse_spmv.hpp b/sparse/unit_test/Test_Sparse_spmv.hpp
@@ -23,6 +23,7 @@
 #include <KokkosKernels_IOUtils.hpp>
 #include <KokkosSparse_IOUtils.hpp>
 #include <KokkosKernels_Utils.hpp>
+#include <KokkosKernels_NaN.hpp>
 
 #include "KokkosKernels_default_types.hpp"
 
@@ -451,13 +452,13 @@ void test_spmv(KokkosSparse::SPMVAlgorithm algo, lno_t numRows, size_type nnz,
   Kokkos::parallel_for(
       y_policy(0, input_y_nans.extent(0)), KOKKOS_LAMBDA(const size_t i) {
         if (0 == (i % 19)) {
-          input_y_nans(i) = NAN;
+          input_y_nans(i) = KokkosKernels::Impl::quiet_NaN<scalar_t>();
         }
       });
   Kokkos::parallel_for(
       y_policy(0, input_yt_nans.extent(0)), KOKKOS_LAMBDA(const size_t i) {
         if (0 == (i % 23)) {
-          input_yt_nans(i) = Kokkos::nan("");
+          input_yt_nans(i) = KokkosKernels::Impl::quiet_NaN<scalar_t>();
         }
       });
 

diff --git a/sparse/unit_test/Test_Sparse_spmv_bsr.hpp b/sparse/unit_test/Test_Sparse_spmv_bsr.hpp
@@ -353,13 +353,13 @@ auto random_vecs_for_spmv(const char *mode, const Bsr &a,
     Kokkos::parallel_for(
         policy_type(0, x.extent(0)), KOKKOS_LAMBDA(size_t i) {
           if (0 == (i % 17)) {
-            x(i) = Kokkos::nan("");
+            x(i) = KokkosKernels::Impl::quiet_NaN<scalar_type>();
           }
         });
     Kokkos::parallel_for(
         policy_type(0, y.extent(0)), KOKKOS_LAMBDA(size_t i) {
           if (0 == (i % 17)) {
-            y(i) = Kokkos::nan("");
+            y(i) = KokkosKernels::Impl::quiet_NaN<scalar_type>();
           }
         });
   }
@@ -615,15 +615,15 @@ auto random_multivecs_for_spm_mv(const char *mode, const Bsr &a,
         policy_type(0, x.extent(0)), KOKKOS_LAMBDA(size_t i) {
           for (size_t j = 0; j < x.extent(1); ++j) {
             if (0 == ((i * x.extent(1) + j) % 13)) {
-              x(i, j) = Kokkos::nan("");
+              x(i, j) = KokkosKernels::Impl::quiet_NaN<scalar_type>();
             }
           }
         });
     Kokkos::parallel_for(
         policy_type(0, y.extent(0)), KOKKOS_LAMBDA(size_t i) {
           for (size_t j = 0; j < y.extent(1); ++j) {
             if (0 == ((i * y.extent(1) + j) % 17)) {
-              y(i, j) = Kokkos::nan("");
+              y(i, j) = KokkosKernels::Impl::quiet_NaN<scalar_type>();
             }
           }
         });