Export Kokkos_CUDA_ARCHITECTURES variable with CMake (#5919)

* Export Kokkos_CUDA_ARCHITECTURES variable with CMake * Fixup set KOKKOS_CUDA_ARCHITECTURES in the parent scope as well
sandialabs · Mar 1, 2023 · 569a609 · 569a609
1 parent 416d7b7
commit 569a609
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/cmake/KokkosConfigCommon.cmake.in b/cmake/KokkosConfigCommon.cmake.in
@@ -11,6 +11,10 @@ FOREACH(DEV ${Kokkos_DEVICES})
   SET(Kokkos_ENABLE_${DEV} ON)
 ENDFOREACH()
 
+IF(Kokkos_ENABLE_CUDA)
+  SET(Kokkos_CUDA_ARCHITECTURES @KOKKOS_CUDA_ARCHITECTURES@)
+ENDIF()
+
 IF(NOT Kokkos_FIND_QUIETLY)
   MESSAGE(STATUS "Enabled Kokkos devices: ${Kokkos_DEVICES}")
 ENDIF()

diff --git a/cmake/kokkos_arch.cmake b/cmake/kokkos_arch.cmake
@@ -571,10 +571,14 @@ FUNCTION(CHECK_CUDA_ARCH ARCH FLAG)
       MESSAGE(WARNING "Given CUDA arch ${ARCH}, but Kokkos_ENABLE_CUDA, Kokkos_ENABLE_OPENACC, and Kokkos_ENABLE_OPENMPTARGET are OFF. Option will be ignored.")
       UNSET(KOKKOS_ARCH_${ARCH} PARENT_SCOPE)
     ELSE()
+      IF(KOKKOS_ENABLE_CUDA)
+        STRING(REPLACE "sm_" "" CMAKE_ARCH ${FLAG})
+        SET(KOKKOS_CUDA_ARCHITECTURES ${CMAKE_ARCH})
+        SET(KOKKOS_CUDA_ARCHITECTURES ${CMAKE_ARCH} PARENT_SCOPE)
+      ENDIF()
       SET(KOKKOS_CUDA_ARCH_FLAG ${FLAG} PARENT_SCOPE)
       IF(KOKKOS_ENABLE_COMPILE_AS_CMAKE_LANGUAGE)
-        string(REPLACE "sm_" "" CMAKE_ARCH ${FLAG})
-        SET(CMAKE_CUDA_ARCHITECTURES ${CMAKE_ARCH} PARENT_SCOPE)
+        SET(CMAKE_CUDA_ARCHITECTURES ${KOKKOS_CUDA_ARCHITECTURES} PARENT_SCOPE)
       ELSE()
         IF(KOKKOS_CXX_COMPILER_ID STREQUAL NVHPC)
           STRING(REPLACE "sm_" "cc" NVHPC_CUDA_ARCH ${FLAG})