From 2a68a3201267217bb7bbda70ce767df3cacfdde8 Mon Sep 17 00:00:00 2001
From: Daniel Lemire <daniel@lemire.me>
Date: Fri, 12 May 2023 19:25:02 -0400
Subject: [PATCH] Checking OS support for AVX-512 (#1998)

* Adding support for AVX-512 on macOS.

* Fix.

* Fix.

* Minor fix

* Setting the variable to zero.

* Fixing include

* Checking if the OS supports AVX-512

* Tweaking.

---------

Co-authored-by: Daniel Lemire <dlemire@lemire.me>
---
 include/simdjson/internal/isadetection.h   | 60 +++++++++++++++++-----
 tests/dom/readme_examples_noexceptions.cpp |  4 +-
 2 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/include/simdjson/internal/isadetection.h b/include/simdjson/internal/isadetection.h
index e916419a5..a11322752 100644
--- a/include/simdjson/internal/isadetection.h
+++ b/include/simdjson/internal/isadetection.h
@@ -57,7 +57,6 @@ POSSIBILITY OF SUCH DAMAGE.
 namespace simdjson {
 namespace internal {
 
-
 enum instruction_set {
   DEFAULT = 0x0,
   NEON = 0x1,
@@ -107,7 +106,10 @@ constexpr uint32_t cpuid_avx512cd_bit = 1 << 28;    ///< @private bit 28 of EBX
 constexpr uint32_t cpuid_avx512bw_bit = 1 << 30;    ///< @private bit 30 of EBX for EAX=0x7
 constexpr uint32_t cpuid_avx512vl_bit = 1U << 31;    ///< @private bit 31 of EBX for EAX=0x7
 constexpr uint32_t cpuid_avx512vbmi2_bit = 1 << 6;  ///< @private bit 6 of ECX for EAX=0x7
+constexpr uint64_t cpuid_avx256_saved = uint64_t(1) << 2; ///< @private bit 2 = AVX
+constexpr uint64_t cpuid_avx512_saved = uint64_t(7) << 5; ///< @private bits 5,6,7 = opmask, ZMM_hi256, hi16_ZMM
 constexpr uint32_t cpuid_sse42_bit = 1 << 20;       ///< @private bit 20 of ECX for EAX=0x1
+constexpr uint32_t cpuid_osxsave = (uint32_t(1) << 26) | (uint32_t(1) << 27); ///< @private bits 26+27 of ECX for EAX=0x1
 constexpr uint32_t cpuid_pclmulqdq_bit = 1 << 1;    ///< @private bit  1 of ECX for EAX=0x1
 }
 
@@ -117,7 +119,7 @@ static inline void cpuid(uint32_t *eax, uint32_t *ebx, uint32_t *ecx,
                          uint32_t *edx) {
 #if defined(_MSC_VER)
   int cpu_info[4];
-  __cpuid(cpu_info, *eax);
+  __cpuidex(cpu_info, *eax, *ecx);
   *eax = cpu_info[0];
   *ebx = cpu_info[1];
   *ecx = cpu_info[2];
@@ -135,10 +137,48 @@ static inline void cpuid(uint32_t *eax, uint32_t *ebx, uint32_t *ecx,
 #endif
 }
 
+
+static inline uint64_t xgetbv() {
+#if defined(_MSC_VER)
+  return _xgetbv(0);
+#else
+  uint32_t xcr0_lo, xcr0_hi;
+  asm volatile("xgetbv\n\t" : "=a" (xcr0_lo), "=d" (xcr0_hi) : "c" (0));
+  return xcr0_lo | (uint64_t(xcr0_hi) << 32);
+#endif
+}
+
 static inline uint32_t detect_supported_architectures() {
   uint32_t eax, ebx, ecx, edx;
   uint32_t host_isa = 0x0;
 
+  // EBX for EAX=0x1
+  eax = 0x1;
+  ecx = 0x0;
+  cpuid(&eax, &ebx, &ecx, &edx);
+
+  if (ecx & cpuid_sse42_bit) {
+    host_isa |= instruction_set::SSE42;
+  } else {
+    return host_isa; // everything after is redundant
+  }
+
+  if (ecx & cpuid_pclmulqdq_bit) {
+    host_isa |= instruction_set::PCLMULQDQ;
+  }
+
+
+  if ((ecx & cpuid_osxsave) != cpuid_osxsave) {
+    return host_isa;
+  }
+
+  // xgetbv for checking if the OS saves registers
+  uint64_t xcr0 = xgetbv();
+
+  if ((xcr0 & cpuid_avx256_saved) == 0) {
+    return host_isa;
+  }
+
   // ECX for EAX=0x7
   eax = 0x7;
   ecx = 0x0;
@@ -154,6 +194,10 @@ static inline uint32_t detect_supported_architectures() {
     host_isa |= instruction_set::BMI2;
   }
 
+  if (!((xcr0 & cpuid_avx512_saved) == cpuid_avx512_saved)) {
+     return host_isa;
+  }
+
   if (ebx & cpuid_avx512f_bit) {
     host_isa |= instruction_set::AVX512F;
   }
@@ -190,18 +234,6 @@ static inline uint32_t detect_supported_architectures() {
     host_isa |= instruction_set::AVX512VBMI2;
   }
 
-  // EBX for EAX=0x1
-  eax = 0x1;
-  cpuid(&eax, &ebx, &ecx, &edx);
-
-  if (ecx & cpuid_sse42_bit) {
-    host_isa |= instruction_set::SSE42;
-  }
-
-  if (ecx & cpuid_pclmulqdq_bit) {
-    host_isa |= instruction_set::PCLMULQDQ;
-  }
-
   return host_isa;
 }
 #else // fallback
diff --git a/tests/dom/readme_examples_noexceptions.cpp b/tests/dom/readme_examples_noexceptions.cpp
index 7bb5b1033..ba070cc1f 100644
--- a/tests/dom/readme_examples_noexceptions.cpp
+++ b/tests/dom/readme_examples_noexceptions.cpp
@@ -81,13 +81,13 @@ void basics_error_3() {
           dom::object innerobj;
           if ((error = key_value.value.get(innerobj))) { cerr << error << endl; exit(1); }
 
-          double va, vb;
+          double va{}, vb{};
           if ((error = innerobj["a"].get(va))) { cerr << error << endl; exit(1); }
           cout << "a: " << va << ", ";
           if ((error = innerobj["b"].get(vb))) { cerr << error << endl; exit(1); }
           cout << "b: " << vb << ", ";
 
-          int64_t vc;
+          int64_t vc{};
           if ((error = innerobj["c"].get(vc))) { cerr << error << endl; exit(1); }
           cout << "c: " << vc << endl;
       }