Improve macro definitions

nliber · Jul 26, 2023 · 5a7039f · 5a7039f
1 parent 00ff0a5
commit 5a7039f
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/algorithms/src/std_algorithms/impl/Kokkos_ForEachForEachN.hpp b/algorithms/src/std_algorithms/impl/Kokkos_ForEachForEachN.hpp
@@ -24,18 +24,16 @@
 #include <string>
 
 #define KOKKOS_DO_PRAGMA(x) _Pragma(#x)
-#if defined __NVCC__
-#ifdef __NVCC_DIAG_PRAGMA_SUPPORT__
+#if defined __NVCC__ && defined __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__
+#elif defined __NVCC__
 #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
 #else
 #define KOKKOS_NVCC_PRAGMA_PUSH
 #define KOKKOS_NVCC_PRAGMA_SUPPRESS(WARNING)