Add KOKKOS_ARCH_ARM_NEON

tcclevenger · Oct 4, 2023 · 29d4ffd · 29d4ffd
1 parent 4ce289b
commit 29d4ffd
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 4 deletions.
diff --git a/Makefile.kokkos b/Makefile.kokkos
@@ -708,6 +708,7 @@ endif
 
 ifeq ($(KOKKOS_INTERNAL_USE_ARCH_ARMV80), 1)
   tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARMV80")
+  tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARM_NEON")
 
   ifeq ($(KOKKOS_INTERNAL_COMPILER_CRAY), 1)
     KOKKOS_CXXFLAGS +=
@@ -720,6 +721,7 @@ endif
 
 ifeq ($(KOKKOS_INTERNAL_USE_ARCH_ARMV81), 1)
   tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARMV81")
+  tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARM_NEON")
 
   ifeq ($(KOKKOS_INTERNAL_COMPILER_CRAY), 1)
     KOKKOS_CXXFLAGS +=
@@ -732,6 +734,7 @@ endif
 
 ifeq ($(KOKKOS_INTERNAL_USE_ARCH_A64FX), 1)
   tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_A64FX")
+  tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARM_NEON")
 
   KOKKOS_CXXFLAGS += -march=armv8.2-a+sve
   KOKKOS_LDFLAGS += -march=armv8.2-a+sve
@@ -787,6 +790,7 @@ endif
 ifeq ($(KOKKOS_INTERNAL_USE_ARCH_ARMV8_THUNDERX), 1)
   tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARMV80")
   tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARMV8_THUNDERX")
+  tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARM_NEON")
 
   ifeq ($(KOKKOS_INTERNAL_COMPILER_CRAY), 1)
     KOKKOS_CXXFLAGS +=
@@ -800,6 +804,7 @@ endif
 ifeq ($(KOKKOS_INTERNAL_USE_ARCH_ARMV8_THUNDERX2), 1)
   tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARMV81")
   tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARMV8_THUNDERX2")
+  tmp := $(call kokkos_append_header,"$H""define KOKKOS_ARCH_ARM_NEON")
 
   ifeq ($(KOKKOS_INTERNAL_COMPILER_CRAY), 1)
     KOKKOS_CXXFLAGS +=

diff --git a/cmake/KokkosCore_config.h.in b/cmake/KokkosCore_config.h.in
@@ -72,6 +72,7 @@
 #cmakedefine KOKKOS_ARCH_AVX
 #cmakedefine KOKKOS_ARCH_AVX2
 #cmakedefine KOKKOS_ARCH_AVX512XEON
+#cmakedefine KOKKOS_ARCH_ARM_NEON
 #cmakedefine KOKKOS_ARCH_KNC
 #cmakedefine KOKKOS_ARCH_AVX512MIC
 #cmakedefine KOKKOS_ARCH_POWER7

diff --git a/cmake/kokkos_arch.cmake b/cmake/kokkos_arch.cmake
@@ -249,6 +249,7 @@ IF(KOKKOS_ARCH_NATIVE)
 ENDIF()
 
 IF (KOKKOS_ARCH_ARMV80)
+  SET(KOKKOS_ARCH_ARM_NEON ON)
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
     Cray    NO-VALUE-SPECIFIED
@@ -259,6 +260,7 @@ IF (KOKKOS_ARCH_ARMV80)
 ENDIF()
 
 IF (KOKKOS_ARCH_ARMV81)
+  SET(KOKKOS_ARCH_ARM_NEON ON)
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
     Cray    NO-VALUE-SPECIFIED
@@ -269,6 +271,7 @@ IF (KOKKOS_ARCH_ARMV81)
 ENDIF()
 
 IF (KOKKOS_ARCH_ARMV8_THUNDERX)
+  SET(KOKKOS_ARCH_ARM_NEON ON)
   SET(KOKKOS_ARCH_ARMV80 ON) #Not a cache variable
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
@@ -280,6 +283,7 @@ IF (KOKKOS_ARCH_ARMV8_THUNDERX)
 ENDIF()
 
 IF (KOKKOS_ARCH_ARMV8_THUNDERX2)
+  SET(KOKKOS_ARCH_ARM_NEON ON)
   SET(KOKKOS_ARCH_ARMV81 ON) #Not a cache variable
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
@@ -291,6 +295,7 @@ IF (KOKKOS_ARCH_ARMV8_THUNDERX2)
 ENDIF()
 
 IF (KOKKOS_ARCH_A64FX)
+  SET(KOKKOS_ARCH_ARM_NEON ON)
   COMPILER_SPECIFIC_FLAGS(
     COMPILER_ID KOKKOS_CXX_HOST_COMPILER_ID
     Clang   -march=armv8.2-a+sve -msve-vector-bits=512
@@ -498,6 +503,8 @@ IF (KOKKOS_ARCH_NATIVE)
     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_ARM_NEON CACHE)
+    CHECK_CXX_SYMBOL_EXISTS(__ARM_NEON "" KOKKOS_COMPILER_HAS_ARM_NEON)
     UNSET(KOKKOS_COMPILER_HAS_AVX CACHE)
     CHECK_CXX_SYMBOL_EXISTS(__AVX__ "" KOKKOS_COMPILER_HAS_AVX)
     SET(CMAKE_REQUIRED_FLAGS "${KOKKOS_COMPILE_OPTIONS}")
@@ -514,6 +521,9 @@ IF (KOKKOS_ARCH_NATIVE)
   ELSEIF(KOKKOS_COMPILER_HAS_AVX2)
     MESSAGE(STATUS "SIMD: AVX2 detected")
     SET(KOKKOS_ARCH_AVX2 ON)
+  ELSEIF(KOKKOS_COMPILER_HAS_ARM_NEON)
+    MESSAGE(STATUS "SIMD: ARM_NEON detected")
+    SET(KOKKOS_ARCH_ARM_NEON ON)
   ELSEIF(KOKKOS_COMPILER_HAS_AVX)
     MESSAGE(STATUS "SIMD: AVX detected")
     SET(KOKKOS_ARCH_AVX ON)

diff --git a/simd/src/Kokkos_SIMD.hpp b/simd/src/Kokkos_SIMD.hpp
@@ -40,7 +40,7 @@
 #include <Kokkos_SIMD_AVX512.hpp>
 #endif
 
-#ifdef __ARM_NEON
+#if defined(KOKKOS_ARCH_ARM_NEON)
 #include <Kokkos_SIMD_NEON.hpp>
 #endif
 #else  // KOKKOS_ENABLE_OPENMPTARGET
@@ -62,7 +62,10 @@
 #include <Kokkos_SIMD_AVX512.hpp>
 #endif
 
-#ifdef __ARM_NEON
+#if defined(KOKKOS_ARCH_ARM_NEON)
+#if !defined(__ARM_NEON)
+#error "__ARM_NEON must be definded for KOKKOS_ARCH_ARM_NEON"
+#endif
 #include <Kokkos_SIMD_NEON.hpp>
 #endif
 #endif
@@ -82,7 +85,7 @@ namespace Impl {
 using host_native = avx512_fixed_size<8>;
 #elif defined(KOKKOS_ARCH_AVX2)
 using host_native  = avx2_fixed_size<4>;
-#elif defined(__ARM_NEON)
+#elif defined(KOKKOS_ARCH_ARM_NEON)
 using host_native  = neon_fixed_size<2>;
 #else
 using host_native   = scalar;
@@ -185,7 +188,7 @@ using data_type_set = data_types<std::int32_t, std::uint32_t, std::int64_t,
 using host_abi_set = abi_set<simd_abi::scalar, simd_abi::avx2_fixed_size<4>>;
 using data_type_set =
     data_types<std::int32_t, std::int64_t, std::uint64_t, double, float>;
-#elif defined(__ARM_NEON)
+#elif defined(KOKKOS_ARCH_ARM_NEON)
 using host_abi_set = abi_set<simd_abi::scalar, simd_abi::neon_fixed_size<2>>;
 using data_type_set =
     data_types<std::int32_t, std::int64_t, std::uint64_t, double, float>;