check sub-leaf to detect avx512bf16

Signed-off-by: Mingzhuo Yin <[email protected]>
rust-lang · Mar 5, 2024 · e50b2f6 · e50b2f6
1 parent 55af6a2
commit e50b2f6
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/crates/std_detect/src/detect/os/x86.rs b/crates/std_detect/src/detect/os/x86.rs
@@ -67,18 +67,18 @@ pub(crate) fn detect_features() -> cache::Initializer {
         ..
     } = unsafe { __cpuid(0x0000_0001_u32) };
 
-    // EAX = 7, ECX = 0: Queries "Extended Features";
+    // EAX = 7: Queries "Extended Features";
     // Contains information about bmi,bmi2, and avx2 support.
     let (
-        extended_features_eax,
         extended_features_ebx,
         extended_features_ecx,
         extended_features_edx,
+        extended_features_eax_leaf_1,
     ) = if max_basic_leaf >= 7 {
-        let CpuidResult {
-            eax, ebx, ecx, edx, ..
-        } = unsafe { __cpuid(0x0000_0007_u32) };
-        (eax, ebx, ecx, edx)
+        let CpuidResult { ebx, ecx, edx, .. } = unsafe { __cpuid(0x0000_0007_u32) };
+        let CpuidResult { eax: eax_1, .. } =
+            unsafe { __cpuid_count(0x0000_0007_u32, 0x0000_0001_u32) };
+        (ebx, ecx, edx, eax_1)
     } else {
         (0, 0, 0, 0) // CPUID does not support "Extended Features"
     };
@@ -206,7 +206,6 @@ pub(crate) fn detect_features() -> cache::Initializer {
                     // For AVX-512 the OS also needs to support saving/restoring
                     // the extended state, only then we enable AVX-512 support:
                     if os_avx512_support {
-                        enable(extended_features_eax, 5, Feature::avx512bf16);
                         enable(extended_features_ebx, 16, Feature::avx512f);
                         enable(extended_features_ebx, 17, Feature::avx512dq);
                         enable(extended_features_ebx, 21, Feature::avx512ifma);
@@ -225,6 +224,7 @@ pub(crate) fn detect_features() -> cache::Initializer {
                         enable(extended_features_ecx, 14, Feature::avx512vpopcntdq);
                         enable(extended_features_edx, 8, Feature::avx512vp2intersect);
                         enable(extended_features_edx, 23, Feature::avx512fp16);
+                        enable(extended_features_eax_leaf_1, 5, Feature::avx512bf16);
                     }
                 }
             }