check if HOST has avx512

Signed-off-by: jiqing-feng <[email protected]>

Author: jiqing-feng

CMakeLists.txt

@@ -271,6 +271,30 @@ target_compile_features(bitsandbytes PUBLIC cxx_std_17)
 target_include_directories(bitsandbytes PUBLIC csrc include)
 
 if (BUILD_CPU)
+    include(CheckCXXSourceRuns)
+    set(AVX512F_TEST_CODE "
+        #include <immintrin.h>
+        int main() {
+            __m512 a = _mm512_setzero_ps();
+            __m512 b = _mm512_add_ps(a, a);
+            return 0;
+        }
+    ")
+    set(AVX512BF16_TEST_CODE "
+        #include <immintrin.h>
+        int main() {
+            __m512 a = _mm512_setzero_ps();
+            __m256bh b = _mm512_cvtneps_pbh(a);
+            return 0;
+        }
+    ")
+    set(CMAKE_REQUIRED_FLAGS "-mavx512f")
+    check_cxx_source_runs("${AVX512F_TEST_CODE}" HOST_HAS_AVX512F)
+    unset(CMAKE_REQUIRED_FLAGS)
+    set(CMAKE_REQUIRED_FLAGS "-mavx512bf16")
+    check_cxx_source_runs("${AVX512BF16_TEST_CODE}" HOST_HAS_AVX512BF16)
+    unset(CMAKE_REQUIRED_FLAGS)
+
     if (OpenMP_CXX_FOUND)
         target_link_libraries(bitsandbytes PRIVATE OpenMP::OpenMP_CXX)
         add_definitions(-DHAS_OPENMP)
@@ -280,13 +304,13 @@ if (BUILD_CPU)
         include(CheckCXXCompilerFlag)
         check_cxx_compiler_flag(-mavx512f HAS_AVX512F_FLAG)
         check_cxx_compiler_flag(-mavx512bf16 HAS_AVX512BF16_FLAG)
-        if (HAS_AVX512F_FLAG)
+        if (HAS_AVX512F_FLAG AND HOST_HAS_AVX512F)
             target_compile_options(bitsandbytes PRIVATE -mavx512f)
             target_compile_options(bitsandbytes PRIVATE -mavx512dq)
             target_compile_options(bitsandbytes PRIVATE -mavx512bw)
             target_compile_options(bitsandbytes PRIVATE -mavx512vl)
         endif()
-        if (HAS_AVX512BF16_FLAG)
+        if (HAS_AVX512BF16_FLAG AND HOST_HAS_AVX512BF16)
             target_compile_options(bitsandbytes PRIVATE -mavx512bf16)
         endif()
         target_compile_options(