Refactor CPU code

kimwalisch · Jun 20, 2024 · 4a44eda · 4a44eda
1 parent 6ffc191
commit 4a44eda
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 17 deletions.
diff --git a/include/cpu_supports_avx512_bmi2.hpp b/include/cpu_supports_avx512_bmi2.hpp
@@ -12,11 +12,15 @@
 #define CPU_SUPPORTS_AVX512_BMI2_HPP
 
 #include <cpuid.hpp>
+#include <stdint.h>
 
 #if defined(_MSC_VER)
   #include <immintrin.h>
 #endif
 
+// CPUID bits documentation:
+// https://en.wikipedia.org/wiki/CPUID
+
 // %ebx bit flags
 #define bit_BMI2    (1 << 8)
 #define bit_AVX512F (1 << 16)
@@ -32,17 +36,17 @@
 namespace {
 
 // Get Value of Extended Control Register
-inline int get_xcr0()
+inline uint64_t get_xcr0()
 {
-  int xcr0;
-
 #if defined(_MSC_VER)
-  xcr0 = (int) _xgetbv(0);
+  return _xgetbv(0);
 #else
-  __asm__ ("xgetbv" : "=a" (xcr0) : "c" (0) : "%edx" );
-#endif
+  uint32_t eax;
+  uint32_t edx;
 
-  return xcr0;
+  __asm__ ("xgetbv" : "=a"(eax), "=d"(edx) : "c"(0));
+  return eax | (uint64_t(edx) << 32);
+#endif
 }
 
 inline bool run_cpuid_avx512_bmi2()
@@ -57,10 +61,9 @@ inline bool run_cpuid_avx512_bmi2()
   if ((abcd[2] & osxsave_mask) != osxsave_mask)
     return false;
 
-  int ymm_mask = XSTATE_SSE | XSTATE_YMM;
-  int zmm_mask = XSTATE_SSE | XSTATE_YMM | XSTATE_ZMM;
-
-  int xcr0 = get_xcr0();
+  uint64_t ymm_mask = XSTATE_SSE | XSTATE_YMM;
+  uint64_t zmm_mask = XSTATE_SSE | XSTATE_YMM | XSTATE_ZMM;
+  uint64_t xcr0 = get_xcr0();
 
   // Check AVX OS support
   if ((xcr0 & ymm_mask) != ymm_mask)

diff --git a/include/cpu_supports_popcnt.hpp b/include/cpu_supports_popcnt.hpp
@@ -43,6 +43,7 @@ inline bool run_cpuid_supports_popcnt()
   run_cpuid(1, 0, abcd);
 
   // %ecx POPCNT bit flag
+  // https://en.wikipedia.org/wiki/CPUID
   int bit_POPCNT = 1 << 23;
   return (abcd[2] & bit_POPCNT) == bit_POPCNT;
 }

diff --git a/include/cpuid.hpp b/include/cpuid.hpp
@@ -27,18 +27,18 @@ inline void run_cpuid(int eax, int ecx, int* abcd)
 
   #if defined(__i386__) && \
       defined(__PIC__)
-    /* in case of PIC under 32-bit EBX cannot be clobbered */
+    // In case of PIC under 32-bit EBX cannot be clobbered
     __asm__ ("movl %%ebx, %%edi;"
              "cpuid;"
              "xchgl %%ebx, %%edi;"
-             : "=D" (ebx),
-               "+a" (eax),
+             : "+a" (eax),
+               "=D" (ebx),
                "+c" (ecx),
                "=d" (edx));
   #else
-    __asm__ ("cpuid;"
-             : "+b" (ebx),
-               "+a" (eax),
+    __asm__ ("cpuid"
+             : "+a" (eax),
+               "+b" (ebx),
                "+c" (ecx),
                "=d" (edx));
   #endif