fix cmake check

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

Author: jiqing-feng

CMakeLists.txt

@@ -85,6 +85,7 @@ endif()
 if (BUILD_CPU)
     set(CMAKE_CXX_STANDARD 17)
     set(CMAKE_CXX_STANDARD_REQUIRED ON)
+    string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" HOST_ARCH)
     find_package(OpenMP)
 endif()
 
@@ -270,18 +271,48 @@ target_compile_features(bitsandbytes PUBLIC cxx_std_17)
 target_include_directories(bitsandbytes PUBLIC csrc include)
 
 if (BUILD_CPU)
-    target_link_libraries(bitsandbytes PRIVATE OpenMP::OpenMP_CXX)
-    include(CheckCXXCompilerFlag)
-
-    check_cxx_compiler_flag(-mavx512f HAS_AVX512F)
-    check_cxx_compiler_flag(-mavx512bf16 HAS_AVX512BF16)
+    if (OpenMP_CXX_FOUND)
+        target_link_libraries(bitsandbytes PRIVATE OpenMP::OpenMP_CXX)
+        add_definitions(-DHAS_OPENMP)
+    else()
+        add_definitions(-DNO_OPENMP)
+    endif()
 
-    if(HAS_AVX512F)
-    target_compile_options(bitsandbytes PRIVATE -mavx512f)
+    if (HOST_ARCH MATCHES "x86_64|amd64")
+        include(CheckCXXCompilerFlag)
+        check_cxx_compiler_flag(-mavx512f HAS_AVX512F_FLAG)
+        check_cxx_compiler_flag(-mavx512bf16 HAS_AVX512BF16_FLAG)
+        if (HAS_AVX512F_FLAG)
+            target_compile_options(bitsandbytes PRIVATE -mavx512f)
+            add_definitions(-DHAS_AVX512F)
+        endif()
+        if (HAS_AVX512BF16_FLAG)
+            target_compile_options(bitsandbytes PRIVATE -mavx512bf16)
+            add_definitions(-DHAS_AVX512BF16)
+        else()
+            add_definitions(-DNO_AVX512BF16)
+        endif()
     endif()
+endif()
 
-    if(HAS_AVX512BF16)
-    target_compile_options(bitsandbytes PRIVATE -mavx512bf16)
+# --- Windows MSVC specific AVX512BF16 probe (after add_library) ---
+if (MSVC AND BUILD_CPU)
+    include(CheckCXXSourceCompiles)
+    set(_AVX512BF16_TEST "
+        #include <immintrin.h>
+        int main(){
+            __m512bh a{}, b{};
+            auto c = _mm512_dpbf16_ps(_mm512_setzero_ps(), a, b);
+            (void)c;
+            return 0;
+        }")
+    check_cxx_source_compiles("${_AVX512BF16_TEST}" MSVC_HAS_AVX512BF16)
+    if (MSVC_HAS_AVX512BF16)
+        # /arch:AVX512;
+        target_compile_options(bitsandbytes PRIVATE /arch:AVX512)
+        add_definitions(-DHAS_AVX512BF16)
+    else()
+        add_definitions(-DNO_AVX512BF16)
     endif()
 endif()
 

csrc/cpu_ops.cpp

@@ -3,6 +3,13 @@
 #include <cpu_ops.h>
 #include <thread>
 
+#ifdef HAS_OPENMP
+#include <omp.h>
+#define BNB_OMP_PARALLEL_FOR _Pragma("omp parallel for")
+#else
+#define BNB_OMP_PARALLEL_FOR
+#endif
+
 using namespace BinSearch;
 
 
@@ -99,7 +106,7 @@ void dequantizeBlockwise4bitCpu(unsigned char* A,
     if (dim_1 % VEC_LEN == 0 && blocksize >= VEC_LEN) {
         __m512 lut = DATA_TYPE == 1 ? set_fp4_lut() : set_nf4_lut();
         constexpr auto k_step = VEC_LEN / 2; // 8
-        #pragma omp parallel for
+        BNB_OMP_PARALLEL_FOR
         for (int block_idx = 0; block_idx < dim_0; ++block_idx) {
             for (int k = 0; k < input_dim_1; k += k_step) {
                 // Load 64 bits of nf4 data and a single scale data
@@ -141,7 +148,7 @@ void dequantizeBlockwise4bitCpu(unsigned char* A,
 #endif
     // Scalar fallback branch
     long long total = m * n;
-    #pragma omp parallel for
+    BNB_OMP_PARALLEL_FOR
     for (long long block_idx = 0; block_idx < total; block_idx += blocksize) {
         long long valid_items = (total - block_idx >= blocksize ? blocksize : total - block_idx);
         float scale = absmax[block_idx / blocksize];
@@ -187,7 +194,7 @@ void dequantizeBlockwise8bitCpu(float* code,
                             long long n) {
     if (blocksize <= 0 || n <= 0) return;
     // 8-bit path
-    #pragma omp parallel for
+    BNB_OMP_PARALLEL_FOR
     for (long long block_idx = 0; block_idx < n; block_idx += blocksize) {
         long long valid_items = (n - block_idx >= blocksize ? blocksize : n - block_idx);
         long long block_end   = block_idx + valid_items;