Allow detecting SIMD types based on compiler macros (kokkos#6188)

Allow detecting SIMD types based on compiler macros
nliber · Oct 4, 2023 · 4ce289b · 4ce289b
1 parent 0bf937c
commit 4ce289b
Showing 1 changed file with 46 additions and 16 deletions.
diff --git a/cmake/kokkos_arch.cmake b/cmake/kokkos_arch.cmake
@@ -418,10 +418,7 @@ IF (KOKKOS_ARCH_SKL)
 ENDIF()
 
 IF (KOKKOS_ARCH_SKX)
-  # FIXME_NVHPC nvc++ doesn't seem to support AVX512.
-  IF (NOT KOKKOS_CXX_HOST_COMPILER_ID STREQUAL NVHPC)
-    SET(KOKKOS_ARCH_AVX512XEON ON)
-  ENDIF()
+  SET(KOKKOS_ARCH_AVX512XEON ON)
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
     Cray    NO-VALUE-SPECIFIED
@@ -433,10 +430,7 @@ IF (KOKKOS_ARCH_SKX)
 ENDIF()
 
 IF (KOKKOS_ARCH_ICL)
-  # FIXME_NVHPC nvc++ doesn't seem to support AVX512.
-  IF (NOT KOKKOS_CXX_HOST_COMPILER_ID STREQUAL NVHPC)
-    SET(KOKKOS_ARCH_AVX512XEON ON)
-  ENDIF()
+  SET(KOKKOS_ARCH_AVX512XEON ON)
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
     MSVC    /arch:AVX512
@@ -445,10 +439,7 @@ IF (KOKKOS_ARCH_ICL)
 ENDIF()
 
 IF (KOKKOS_ARCH_ICX)
-  # FIXME_NVHPC nvc++ doesn't seem to support AVX512.
-  IF (NOT KOKKOS_CXX_HOST_COMPILER_ID STREQUAL NVHPC)
-    SET(KOKKOS_ARCH_AVX512XEON ON)
-  ENDIF()
+  SET(KOKKOS_ARCH_AVX512XEON ON)
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
     MSVC    /arch:AVX512
@@ -457,10 +448,7 @@ IF (KOKKOS_ARCH_ICX)
 ENDIF()
 
 IF (KOKKOS_ARCH_SPR)
-  # FIXME_NVHPC nvc++ doesn't seem to support AVX512.
-  IF (NOT KOKKOS_CXX_HOST_COMPILER_ID STREQUAL NVHPC)
-    SET(KOKKOS_ARCH_AVX512XEON ON)
-  ENDIF()
+  SET(KOKKOS_ARCH_AVX512XEON ON)
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
     MSVC    /arch:AVX512
@@ -495,6 +483,48 @@ IF (KOKKOS_ARCH_POWER9)
   )
 ENDIF()
 
+# If Kokkos_ARCH_NATIVE is enabled, we are trying to autodetect
+# the SIMD capabilities based on compiler macros.
+IF (KOKKOS_ARCH_NATIVE)
+  # Make sure to rerun the checks if compile options have changed
+  IF(NOT "${KOKKOS_COMPILE_OPTIONS}" STREQUAL "${KOKKOS_COMPILE_OPTIONS_SAVED}")
+    SET(KOKKOS_COMPILE_OPTIONS_SAVED "${KOKKOS_COMPILE_OPTIONS}" CACHE INTERNAL "")
+
+    SET(CMAKE_REQUIRED_QUIET ON)
+    SET(CMAKE_REQUIRED_FLAGS "${KOKKOS_COMPILE_OPTIONS}")
+    INCLUDE(CheckCXXSymbolExists)
+
+    UNSET(KOKKOS_COMPILER_HAS_AVX512 CACHE)
+    CHECK_CXX_SYMBOL_EXISTS(__AVX512F__ "" KOKKOS_COMPILER_HAS_AVX512)
+    UNSET(KOKKOS_COMPILER_HAS_AVX2 CACHE)
+    CHECK_CXX_SYMBOL_EXISTS(__AVX2__ "" KOKKOS_COMPILER_HAS_AVX2)
+    UNSET(KOKKOS_COMPILER_HAS_AVX CACHE)
+    CHECK_CXX_SYMBOL_EXISTS(__AVX__ "" KOKKOS_COMPILER_HAS_AVX)
+    SET(CMAKE_REQUIRED_FLAGS "${KOKKOS_COMPILE_OPTIONS}")
+
+    UNSET(CMAKE_REQUIRED_QUIET)
+    UNSET(CMAKE_REQUIRED_FLAGS)
+  ENDIF()
+
+  # Only define one of these macros for now
+  # to be uniform with what we are doing for other architectures.
+  IF(KOKKOS_COMPILER_HAS_AVX512)
+    MESSAGE(STATUS "SIMD: AVX512 detected")
+    SET(KOKKOS_ARCH_AVX512XEON ON)
+  ELSEIF(KOKKOS_COMPILER_HAS_AVX2)
+    MESSAGE(STATUS "SIMD: AVX2 detected")
+    SET(KOKKOS_ARCH_AVX2 ON)
+  ELSEIF(KOKKOS_COMPILER_HAS_AVX)
+    MESSAGE(STATUS "SIMD: AVX detected")
+    SET(KOKKOS_ARCH_AVX ON)
+  ENDIF()
+ENDIF()
+
+# FIXME_NVHPC nvc++ doesn't seem to support AVX512.
+IF (KOKKOS_CXX_HOST_COMPILER_ID STREQUAL NVHPC)
+  SET(KOKKOS_ARCH_AVX512XEON OFF)
+ENDIF()
+
 IF (NOT KOKKOS_COMPILE_LANGUAGE STREQUAL CUDA)
   IF (KOKKOS_ENABLE_CUDA_RELOCATABLE_DEVICE_CODE)
       COMPILER_SPECIFIC_FLAGS(