Suppress warnings

nliber · Jul 25, 2023 · 8ef6cb5 · 8ef6cb5
1 parent 44bf945
commit 8ef6cb5
Showing 1 changed file with 41 additions and 5 deletions.
diff --git a/algorithms/src/std_algorithms/impl/Kokkos_ForEachForEachN.hpp b/algorithms/src/std_algorithms/impl/Kokkos_ForEachForEachN.hpp
@@ -23,6 +23,30 @@
 #include <std_algorithms/Kokkos_Distance.hpp>
 #include <string>
 
+#define KOKKOS_DO_PRAGMA(x) _Pragma(#x)
+#if defined __NVCC__
+#ifdef __NVCC_DIAG_PRAGMA_SUPPORT__
+#define KOKKOS_NVCC_PRAGMA_PUSH KOKKOS_DO_PRAGMA(nv_diagnostic push)
+#define KOKKOS_NVCC_PRAGMA_SUPPRESS(WARNING) \
+  KOKKOS_DO_PRAGMA(nv_diag_suppress = WARNING)
+#define KOKKOS_NVCC_PRAGMA_POP KOKKOS_DO_PRAGMA(nv_diagnostic pop)
+#elif defined __CUDA_ARCH__
+#define KOKKOS_NVCC_PRAGMA_PUSH KOKKOS_DO_PRAGMA(diagnostic push)
+#define KOKKOS_NVCC_PRAGMA_SUPPRESS(WARNING) \
+  KOKKOS_DO_PRAGMA(diag_suppress WARNING)
+#define KOKKOS_NVCC_PRAGMA_POP KOKKOS_DO_PRAGMA(diagnostic pop)
+#endif
+#elif defined __NVCOMPILER
+#define KOKKOS_NVCC_PRAGMA_PUSH KOKKOS_DO_PRAGMA(diagnostic push)
+#define KOKKOS_NVCC_PRAGMA_SUPPRESS(WARNING) \
+  KOKKOS_DO_PRAGMA(diag_suppress = WARNING)
+#define KOKKOS_NVCC_PRAGMA_POP KOKKOS_DO_PRAGMA(diagnostic pop)
+#else
+#define KOKKOS_NVCC_PRAGMA_PUSH
+#define KOKKOS_NVCC_PRAGMA_SUPPRESS(WARNING)
+#define KOKKOS_NVCC_PRAGMA_POP
+#endif
+
 namespace Kokkos {
 namespace Experimental {
 namespace Impl {
@@ -41,6 +65,10 @@ struct StdForEachFunctor {
       : m_first(std::move(_first)), m_functor(std::move(_functor)) {}
 };
 
+KOKKOS_NVCC_PRAGMA_PUSH
+KOKKOS_NVCC_PRAGMA_SUPPRESS(20011)
+KOKKOS_NVCC_PRAGMA_SUPPRESS(20014)
+
 template <class HandleType, class IteratorType, class UnaryFunctorType>
 KOKKOS_FUNCTION UnaryFunctorType
 for_each_impl([[maybe_unused]] const char* label, const HandleType& handle,
@@ -52,11 +80,12 @@ for_each_impl([[maybe_unused]] const char* label, const HandleType& handle,
   // run
   const auto num_elements = Kokkos::Experimental::distance(first, last);
   if constexpr (is_execution_space_v<HandleType>) {
-    KOKKOS_IF_ON_HOST(
-        (::Kokkos::parallel_for(
-             label, RangePolicy<HandleType>(handle, 0, num_elements),
-             StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
-         handle.fence("Kokkos::for_each: fence after operation");))
+    KOKKOS_IF_ON_DEVICE(
+        (Kokkos::abort("Can't use an execution space on the device!");))
+    ::Kokkos::parallel_for(
+        label, RangePolicy<HandleType>(handle, 0, num_elements),
+        StdForEachFunctor<IteratorType, UnaryFunctorType>(first, functor));
+    handle.fence("Kokkos::for_each: fence after operation");
   } else {
     ::Kokkos::parallel_for(
         TeamThreadRange(handle, 0, num_elements),
@@ -67,6 +96,8 @@ for_each_impl([[maybe_unused]] const char* label, const HandleType& handle,
   return functor;
 }
 
+KOKKOS_NVCC_PRAGMA_POP
+
 template <class ExecutionSpace, class IteratorType, class SizeType,
           class UnaryFunctorType>
 IteratorType for_each_n_exespace_impl(const std::string& label,
@@ -129,4 +160,9 @@ for_each_n_team_impl(const TeamHandleType& teamHandle, IteratorType first,
 }  // namespace Experimental
 }  // namespace Kokkos
 
+#undef KOKKOS_DO_PRAGMA
+#undef KOKKOS_NVCC_PRAGMA_PUSH
+#undef KOKKOS_NVCC_PRAGMA_SUPPRESS
+#undef KOKKOS_NVCC_PRAGMA_POP
+
 #endif