Add checks for xsave and _xgetbv in cpuid checks

Author: r-devulap

lib/x86simdsortcpuid.h

@@ -8,31 +8,66 @@
 
 static std::unordered_map<std::string, bool> xss_cpu_features;
 
-inline void xss_cpu_init()
+static bool os_supports_avx()
 {
-    int cpuInfo[4] = {0};
-    // Check AVX2
-    __cpuid(cpuInfo, 0);
-    int nIds = cpuInfo[0];
+    int cpuInfo[4];
     __cpuid(cpuInfo, 1);
-    bool osxsave = (cpuInfo[2] & (1 << 27)) != 0;
-    bool avx = (cpuInfo[2] & (1 << 28)) != 0;
-    __cpuid(cpuInfo, 7);
-    bool avx2 = (cpuInfo[1] & (1 << 5)) != 0;
-    bool avx512f = (cpuInfo[1] & (1 << 16)) != 0;
-    bool avx512dq = (cpuInfo[1] & (1 << 17)) != 0;
-    bool avx512bw = (cpuInfo[1] & (1 << 30)) != 0;
-    bool avx512vl = (cpuInfo[1] & (1 << 31)) != 0;
-    bool avx512vbmi2 = (cpuInfo[2] & (1 << 6)) != 0;
-    bool avx512fp16 = (cpuInfo[3] & (1 << 23)) != 0;
-    // Store results
-    xss_cpu_features["avx2"] = avx2;
-    xss_cpu_features["avx512f"] = avx512f;
-    xss_cpu_features["avx512dq"] = avx512dq;
-    xss_cpu_features["avx512bw"] = avx512bw;
-    xss_cpu_features["avx512vl"] = avx512vl;
-    xss_cpu_features["avx512vbmi2"] = avx512vbmi2;
-    xss_cpu_features["avx512fp16"] = avx512fp16;
+
+    bool osxsaveSupported = (cpuInfo[2] & (1 << 27)) != 0; // OSXSAVE bit
+    bool avxSupported     = (cpuInfo[2] & (1 << 28)) != 0; // AVX bit
+    if (!(avxSupported && osxsaveSupported))
+        return false;
+
+    // Check XCR0[2:1] (XMM and YMM state)
+    unsigned long long xcr0 = _xgetbv(0);
+    return (xcr0 & 0x6) == 0x6;
+}
+
+static bool os_supports_avx512()
+{
+    if (!os_supports_avx())
+        return false;
+
+    // Need XCR0[7:5] = opmask/ZMM/YMM state enabled
+    unsigned long long xcr0 = _xgetbv(0);
+    return (xcr0 & 0xE0) == 0xE0;
+}
+
+void xss_cpu_init()
+{
+    int cpuInfo[4];
+    __cpuid(cpuInfo, 0);
+    int maxLeaf = cpuInfo[0];
+
+    bool hasAVX2 = false;
+    bool hasAVX512F = false, hasAVX512DQ = false, hasAVX512BW = false, hasAVX512VL = false;
+    bool hasAVX512VBMI2 = false, hasAVX512FP16 = false;
+
+    if (maxLeaf >= 7)
+    {
+        __cpuidex(cpuInfo, 7, 0);
+
+        // EBX bits
+        hasAVX2      = os_supports_avx()     && (cpuInfo[1] & (1 << 5));
+        hasAVX512F   = os_supports_avx512() && (cpuInfo[1] & (1 << 16));
+        hasAVX512DQ  = os_supports_avx512() && (cpuInfo[1] & (1 << 17));
+        hasAVX512BW  = os_supports_avx512() && (cpuInfo[1] & (1 << 30));
+        hasAVX512VL  = os_supports_avx512() && (cpuInfo[1] & (1 << 31));
+
+        // ECX bits
+        hasAVX512VBMI2 = os_supports_avx512() && (cpuInfo[2] & (1 << 6));
+
+        // EDX bits
+        hasAVX512FP16  = os_supports_avx512() && (cpuInfo[3] & (1 << 23));
+    }
+
+    xss_cpu_features["avx2"]        = hasAVX2;
+    xss_cpu_features["avx512f"]     = hasAVX512F;
+    xss_cpu_features["avx512dq"]    = hasAVX512DQ;
+    xss_cpu_features["avx512bw"]    = hasAVX512BW;
+    xss_cpu_features["avx512vl"]    = hasAVX512VL;
+    xss_cpu_features["avx512vbmi2"] = hasAVX512VBMI2;
+    xss_cpu_features["avx512fp16"]  = hasAVX512FP16;
 }
 
 inline bool xss_cpu_supports(const char *feature)