From 895da3c96e8ed3ae9d921a951aa74ae69068355f Mon Sep 17 00:00:00 2001
From: usamoi <usamoi@outlook.com>
Date: Tue, 12 Dec 2023 17:04:15 +0800
Subject: [PATCH] feat: std_detect avx512fp16

Signed-off-by: usamoi <usamoi@outlook.com>
---
 crates/std_detect/src/detect/arch/x86.rs |  3 +++
 crates/std_detect/src/detect/os/x86.rs   | 14 ++++++++------
 crates/std_detect/tests/cpu-detection.rs |  1 +
 crates/std_detect/tests/x86-specific.rs  |  1 +
 4 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/crates/std_detect/src/detect/arch/x86.rs b/crates/std_detect/src/detect/arch/x86.rs
index 828ac5c38a..f4f45750ed 100644
--- a/crates/std_detect/src/detect/arch/x86.rs
+++ b/crates/std_detect/src/detect/arch/x86.rs
@@ -75,6 +75,7 @@ features! {
     /// * `"avx512bitalg"`
     /// * `"avx512bf16"`
     /// * `"avx512vp2intersect"`
+    /// * `"avx512fp16"`
     /// * `"f16c"`
     /// * `"fma"`
     /// * `"bmi1"`
@@ -169,6 +170,8 @@ features! {
     /// AVX-512 BF16 (BFLOAT16 instructions)
     @FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] avx512vp2intersect: "avx512vp2intersect";
     /// AVX-512 P2INTERSECT
+    @FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] avx512fp16: "avx512fp16";
+    /// AVX-512 FP16 (FLOAT16 instructions)
     @FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] f16c: "f16c";
     /// F16C (Conversions between IEEE-754 `binary16` and `binary32` formats)
     @FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] fma: "fma";
diff --git a/crates/std_detect/src/detect/os/x86.rs b/crates/std_detect/src/detect/os/x86.rs
index d8dd84db49..4ff9ac5f13 100644
--- a/crates/std_detect/src/detect/os/x86.rs
+++ b/crates/std_detect/src/detect/os/x86.rs
@@ -69,12 +69,13 @@ pub(crate) fn detect_features() -> cache::Initializer {
 
     // EAX = 7, ECX = 0: Queries "Extended Features";
     // Contains information about bmi,bmi2, and avx2 support.
-    let (extended_features_ebx, extended_features_ecx) = if max_basic_leaf >= 7 {
-        let CpuidResult { ebx, ecx, .. } = unsafe { __cpuid(0x0000_0007_u32) };
-        (ebx, ecx)
-    } else {
-        (0, 0) // CPUID does not support "Extended Features"
-    };
+    let (extended_features_ebx, extended_features_ecx, extended_features_edx) =
+        if max_basic_leaf >= 7 {
+            let CpuidResult { ebx, ecx, edx, .. } = unsafe { __cpuid(0x0000_0007_u32) };
+            (ebx, ecx, edx)
+        } else {
+            (0, 0, 0) // CPUID does not support "Extended Features"
+        };
 
     // EAX = 0x8000_0000, ECX = 0: Get Highest Extended Function Supported
     // - EAX returns the max leaf value for extended information, that is,
@@ -217,6 +218,7 @@ pub(crate) fn detect_features() -> cache::Initializer {
                         enable(extended_features_ecx, 11, Feature::avx512vnni);
                         enable(extended_features_ecx, 12, Feature::avx512bitalg);
                         enable(extended_features_ecx, 14, Feature::avx512vpopcntdq);
+                        enable(extended_features_edx, 23, Feature::avx512fp16);
                     }
                 }
             }
diff --git a/crates/std_detect/tests/cpu-detection.rs b/crates/std_detect/tests/cpu-detection.rs
index 1ad897a2e2..1053de3a82 100644
--- a/crates/std_detect/tests/cpu-detection.rs
+++ b/crates/std_detect/tests/cpu-detection.rs
@@ -233,6 +233,7 @@ fn x86_all() {
         "avx512vp2intersect {:?}",
         is_x86_feature_detected!("avx512vp2intersect")
     );
+    println!("avx512fp16 {:?}", is_x86_feature_detected!("avx512fp16"));
     println!("f16c: {:?}", is_x86_feature_detected!("f16c"));
     println!("fma: {:?}", is_x86_feature_detected!("fma"));
     println!("bmi1: {:?}", is_x86_feature_detected!("bmi1"));
diff --git a/crates/std_detect/tests/x86-specific.rs b/crates/std_detect/tests/x86-specific.rs
index 1d8d8ba2ec..ae7f677ed4 100644
--- a/crates/std_detect/tests/x86-specific.rs
+++ b/crates/std_detect/tests/x86-specific.rs
@@ -51,6 +51,7 @@ fn dump() {
         "avx512vp2intersect {:?}",
         is_x86_feature_detected!("avx512vp2intersect")
     );
+    println!("avx512fp16 {:?}", is_x86_feature_detected!("avx512fp16"));
     println!("fma: {:?}", is_x86_feature_detected!("fma"));
     println!("abm: {:?}", is_x86_feature_detected!("abm"));
     println!("bmi: {:?}", is_x86_feature_detected!("bmi1"));