More guards for AVX2

4 files changed, 24 insertions, 9 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index a065cb4..67f18ef 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -20,26 +20,32 @@ endif()
 try_compile(INTGEMM_COMPILER_SUPPORTS_AVX2
   ${CMAKE_CURRENT_BINARY_DIR}/compile_tests
   ${CMAKE_CURRENT_SOURCE_DIR}/compile_test/avx2.cc)
-if(NOT INTGEMM_COMPILER_SUPPORTS_AVX2)
-  message(WARNING "${Orange}Not building AVX2-based multiplication because your compiler is too old.\nFor details rerun cmake with --debug-trycompile then try to build in compile_tests/CMakeFiles/CMakeTmp.${ColourReset}")
-endif()
 
 # Check if compiler supports AVX512BW
 try_compile(INTGEMM_COMPILER_SUPPORTS_AVX512BW
   ${CMAKE_CURRENT_BINARY_DIR}/compile_tests
   ${CMAKE_CURRENT_SOURCE_DIR}/compile_test/avx512bw.cc)
-if(NOT INTGEMM_COMPILER_SUPPORTS_AVX512BW)
-  message(WARNING "${Orange}Not building AVX512BW-based multiplication because your compiler is too old.\nFor details rerun cmake with --debug-trycompile then try to build in compile_tests/CMakeFiles/CMakeTmp.${ColourReset}")
-endif()
 
 # Check if the compiler supports AVX512VNNI
 try_compile(INTGEMM_COMPILER_SUPPORTS_AVX512VNNI
   ${CMAKE_CURRENT_BINARY_DIR}/compile_tests
   ${CMAKE_CURRENT_SOURCE_DIR}/compile_test/avx512vnni.cc)
-if(NOT INTGEMM_COMPILER_SUPPORTS_AVX512VNNI)
-  message(WARNING "${Orange}Not building AVX512VNNI-based multiplication because your compiler is too old.\nFor details rerun cmake with --debug-trycompile then try to build in compile_tests/CMakeFiles/CMakeTmp.${ColourReset}")
+
+if (NOT INTGEMM_COMPILER_SUPPORTS_AVX2 OR NOT INTGEMM_COMPILER_SUPPORTS_AVX512BW OR NOT INTGEMM_COMPILER_SUPPORTS_AVX512VNNI)
+  set(UNSUPPORTED "Your compiler is too old to support")
+  if (NOT INTGEMM_COMPILER_SUPPORTS_AVX2)
+    set(UNSUPPORTED "${UNSUPPORTED} AVX2")
+  endif()
+  if (NOT INTGEMM_COMPILER_SUPPORTS_AVX512BW)
+    set(UNSUPPORTED "${UNSUPPORTED} AVX512BW")
+  endif()
+  if (NOT INTGEMM_COMPILER_SUPPORTS_AVX512VNNI)
+    set(UNSUPPORTED "${UNSUPPORTED} AVX512VNNI")
+  endif()
+  message(WARNING "${Orange}${UNSUPPORTED}.  Multiplication will be slower on CPUs that support these instructions. For details rerun cmake with --debug-trycompile then try to build in compile_tests/CMakeFiles/CMakeTmp.${ColourReset}")
 endif()
 
+
 add_library(intgemm STATIC intgemm/intgemm.cc)
 
 # Generate configure file
diff --git a/intgemm/intrinsics.h b/intgemm/intrinsics.h
index 31957bc..67b36fc 100644
--- a/intgemm/intrinsics.h
+++ b/intgemm/intrinsics.h
@@ -5,8 +5,10 @@
 
 #include <tmmintrin.h>
 #include <emmintrin.h>
-#include <immintrin.h>
 #include <xmmintrin.h>
+#ifdef INTGEMM_COMPILER_SUPPORTS_AVX2
+#include <immintrin.h>
+#endif
 
 #include <cstdint>
 
diff --git a/intgemm/types.h b/intgemm/types.h
index da0429f..8578c75 100644
--- a/intgemm/types.h
+++ b/intgemm/types.h
@@ -1,6 +1,9 @@
 #pragma once
 #include <exception>
+#ifdef INTGEMM_COMPILER_SUPPORTS_AVX2
 #include <immintrin.h>
+#endif
+#include <xmmintrin.h>
 
 #if defined(_MSC_VER)
 /* MSVC does not appear to have target attributes but is also fine with just
@@ -78,10 +81,12 @@ typedef __m512i Register;
 typedef __m512 FRegister;
 } // namespace avx512bw
 #endif
+#ifdef INTGEMM_COMPILER_SUPPORTS_AVX2
 namespace avx2 {
 typedef __m256i Register;
 typedef __m256 FRegister;
 } // namespace avx2
+#endif
 namespace ssse3 {
 typedef __m128i Register;
 typedef __m128 FRegister;
diff --git a/intgemm/vec_traits.h b/intgemm/vec_traits.h
index 86265b2..948dae1 100644
--- a/intgemm/vec_traits.h
+++ b/intgemm/vec_traits.h
@@ -18,11 +18,13 @@ template <> struct vector_s<CPUType::SSSE3, int16_t> { using type = __m128i; };
 template <> struct vector_s<CPUType::SSSE3, int> { using type = __m128i; };
 template <> struct vector_s<CPUType::SSSE3, float> { using type = __m128; };
 template <> struct vector_s<CPUType::SSSE3, double> { using type = __m128d; };
+#ifdef INTGEMM_COMPILER_SUPPORTS_AVX2
 template <> struct vector_s<CPUType::AVX2, int8_t> { using type = __m256i; };
 template <> struct vector_s<CPUType::AVX2, int16_t> { using type = __m256i; };
 template <> struct vector_s<CPUType::AVX2, int> { using type = __m256i; };
 template <> struct vector_s<CPUType::AVX2, float> { using type = __m256; };
 template <> struct vector_s<CPUType::AVX2, double> { using type = __m256d; };
+#endif
 #ifdef INTGEMM_COMPILER_SUPPORTS_AVX512BW
 template <> struct vector_s<CPUType::AVX512BW, int8_t> { using type = __m512i; };
 template <> struct vector_s<CPUType::AVX512BW, int16_t> { using type = __m512i; };