Skip to content

Commit cfbc131

Browse files
sayantnsayantn
committed
Added runtime detection
Cannot do a `cupid` test because they don't support `amx`.
1 parent 86098df commit cfbc131

File tree

3 files changed

+58
-23
lines changed

3 files changed

+58
-23
lines changed

‎crates/std_detect/src/detect/arch/x86.rs

+15
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,11 @@ features! {
7676
/// * `"avx512bf16"`
7777
/// * `"avx512vp2intersect"`
7878
/// * `"avx512fp16"`
79+
/// * `"amx-tile"`
80+
/// * `"amx-int8"`
81+
/// * `"amx-bf16"`
82+
/// * `"amx-fp16"`
83+
/// * `"amx-complex"`
7984
/// * `"f16c"`
8085
/// * `"fma"`
8186
/// * `"bmi1"`
@@ -172,6 +177,16 @@ features! {
172177
/// AVX-512 P2INTERSECT
173178
@FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] avx512fp16: "avx512fp16";
174179
/// AVX-512 FP16 (FLOAT16 instructions)
180+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_tile: "amx-tile";
181+
/// AMX (Advanced Matrix Extensions) - Tile load/store
182+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_int8: "amx-int8";
183+
/// AMX-INT8 (Operations on 8-bit integers)
184+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_bf16: "amx-bf16";
185+
/// AMX-BF16 (BFloat16 Operations)
186+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_fp16: "amx-fp16";
187+
/// AMX-FP16 (Float16 Operations)
188+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_complex: "amx-complex";
189+
/// AMX-COMPLEX (Complex number Operations)
175190
@FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] f16c: "f16c";
176191
/// F16C (Conversions between IEEE-754 `binary16` and `binary32` formats)
177192
@FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] fma: "fma";

‎crates/std_detect/src/detect/os/x86.rs

+19-4
Original file line numberDiff line numberDiff line change
@@ -74,13 +74,17 @@ pub(crate) fn detect_features() -> cache::Initializer {
7474
extended_features_ecx,
7575
extended_features_edx,
7676
extended_features_eax_leaf_1,
77+
extended_features_edx_leaf_1,
7778
) = if max_basic_leaf >= 7 {
7879
let CpuidResult { ebx, ecx, edx, .. } = unsafe { __cpuid(0x0000_0007_u32) };
79-
let CpuidResult { eax: eax_1, .. } =
80-
unsafe { __cpuid_count(0x0000_0007_u32, 0x0000_0001_u32) };
81-
(ebx, ecx, edx, eax_1)
80+
let CpuidResult {
81+
eax: eax_1,
82+
edx: edx_1,
83+
..
84+
} = unsafe { __cpuid_count(0x0000_0007_u32, 0x0000_0001_u32) };
85+
(ebx, ecx, edx, eax_1, edx_1)
8286
} else {
83-
(0, 0, 0, 0) // CPUID does not support "Extended Features"
87+
(0, 0, 0, 0, 0) // CPUID does not support "Extended Features"
8488
};
8589

8690
// EAX = 0x8000_0000, ECX = 0: Get Highest Extended Function Supported
@@ -157,6 +161,7 @@ pub(crate) fn detect_features() -> cache::Initializer {
157161
// * SSE -> `XCR0.SSE[1]`
158162
// * AVX -> `XCR0.AVX[2]`
159163
// * AVX-512 -> `XCR0.AVX-512[7:5]`.
164+
// * AMX -> `XCR0.AMX[18:17]`
160165
//
161166
// by setting the corresponding bits of `XCR0` to `1`.
162167
//
@@ -167,6 +172,8 @@ pub(crate) fn detect_features() -> cache::Initializer {
167172
let os_avx_support = xcr0 & 6 == 6;
168173
// Test `XCR0.AVX-512[7:5]` with the mask `0b1110_0000 == 224`:
169174
let os_avx512_support = xcr0 & 224 == 224;
175+
// Test `XCR0.AMX[18:17]` with the mask `0b110_0000_0000_0000_0000 == 0x60000`
176+
let os_amx_support = xcr0 & 0x60000 == 0x60000;
170177

171178
// Only if the OS and the CPU support saving/restoring the AVX
172179
// registers we enable `xsave` support:
@@ -225,6 +232,14 @@ pub(crate) fn detect_features() -> cache::Initializer {
225232
enable(extended_features_edx, 8, Feature::avx512vp2intersect);
226233
enable(extended_features_edx, 23, Feature::avx512fp16);
227234
enable(extended_features_eax_leaf_1, 5, Feature::avx512bf16);
235+
236+
if os_amx_support {
237+
enable(extended_features_edx, 24, Feature::amx_tile);
238+
enable(extended_features_edx, 25, Feature::amx_int8);
239+
enable(extended_features_edx, 22, Feature::amx_bf16);
240+
enable(extended_features_eax_leaf_1, 21, Feature::amx_fp16);
241+
enable(extended_features_edx_leaf_1, 8, Feature::amx_complex);
242+
}
228243
}
229244
}
230245
}

‎crates/std_detect/tests/x86-specific.rs

+24-19
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#![cfg(any(target_arch = "x86", target_arch = "x86_64"))]
22
#![allow(internal_features)]
3-
#![feature(stdarch_internal)]
3+
#![feature(stdarch_internal, x86_amx_intrinsics)]
44

55
extern crate cupid;
66
#[macro_use]
@@ -24,34 +24,34 @@ fn dump() {
2424
println!("f16c: {:?}", is_x86_feature_detected!("f16c"));
2525
println!("avx: {:?}", is_x86_feature_detected!("avx"));
2626
println!("avx2: {:?}", is_x86_feature_detected!("avx2"));
27-
println!("avx512f {:?}", is_x86_feature_detected!("avx512f"));
28-
println!("avx512cd {:?}", is_x86_feature_detected!("avx512cd"));
29-
println!("avx512er {:?}", is_x86_feature_detected!("avx512er"));
30-
println!("avx512pf {:?}", is_x86_feature_detected!("avx512pf"));
31-
println!("avx512bw {:?}", is_x86_feature_detected!("avx512bw"));
32-
println!("avx512dq {:?}", is_x86_feature_detected!("avx512dq"));
33-
println!("avx512vl {:?}", is_x86_feature_detected!("avx512vl"));
34-
println!("avx512_ifma {:?}", is_x86_feature_detected!("avx512ifma"));
27+
println!("avx512f: {:?}", is_x86_feature_detected!("avx512f"));
28+
println!("avx512cd: {:?}", is_x86_feature_detected!("avx512cd"));
29+
println!("avx512er: {:?}", is_x86_feature_detected!("avx512er"));
30+
println!("avx512pf: {:?}", is_x86_feature_detected!("avx512pf"));
31+
println!("avx512bw: {:?}", is_x86_feature_detected!("avx512bw"));
32+
println!("avx512dq: {:?}", is_x86_feature_detected!("avx512dq"));
33+
println!("avx512vl: {:?}", is_x86_feature_detected!("avx512vl"));
34+
println!("avx512_ifma: {:?}", is_x86_feature_detected!("avx512ifma"));
3535
println!("avx512vbmi {:?}", is_x86_feature_detected!("avx512vbmi"));
3636
println!(
37-
"avx512_vpopcntdq {:?}",
37+
"avx512_vpopcntdq: {:?}",
3838
is_x86_feature_detected!("avx512vpopcntdq")
3939
);
40-
println!("avx512vbmi2 {:?}", is_x86_feature_detected!("avx512vbmi2"));
41-
println!("gfni {:?}", is_x86_feature_detected!("gfni"));
42-
println!("vaes {:?}", is_x86_feature_detected!("vaes"));
43-
println!("vpclmulqdq {:?}", is_x86_feature_detected!("vpclmulqdq"));
44-
println!("avx512vnni {:?}", is_x86_feature_detected!("avx512vnni"));
40+
println!("avx512vbmi2: {:?}", is_x86_feature_detected!("avx512vbmi2"));
41+
println!("gfni: {:?}", is_x86_feature_detected!("gfni"));
42+
println!("vaes: {:?}", is_x86_feature_detected!("vaes"));
43+
println!("vpclmulqdq: {:?}", is_x86_feature_detected!("vpclmulqdq"));
44+
println!("avx512vnni: {:?}", is_x86_feature_detected!("avx512vnni"));
4545
println!(
46-
"avx512bitalg {:?}",
46+
"avx512bitalg: {:?}",
4747
is_x86_feature_detected!("avx512bitalg")
4848
);
49-
println!("avx512bf16 {:?}", is_x86_feature_detected!("avx512bf16"));
49+
println!("avx512bf16: {:?}", is_x86_feature_detected!("avx512bf16"));
5050
println!(
51-
"avx512vp2intersect {:?}",
51+
"avx512vp2intersect: {:?}",
5252
is_x86_feature_detected!("avx512vp2intersect")
5353
);
54-
println!("avx512fp16 {:?}", is_x86_feature_detected!("avx512fp16"));
54+
println!("avx512fp16: {:?}", is_x86_feature_detected!("avx512fp16"));
5555
println!("fma: {:?}", is_x86_feature_detected!("fma"));
5656
println!("abm: {:?}", is_x86_feature_detected!("abm"));
5757
println!("bmi: {:?}", is_x86_feature_detected!("bmi1"));
@@ -68,6 +68,11 @@ fn dump() {
6868
println!("adx: {:?}", is_x86_feature_detected!("adx"));
6969
println!("rtm: {:?}", is_x86_feature_detected!("rtm"));
7070
println!("movbe: {:?}", is_x86_feature_detected!("movbe"));
71+
println!("amx-bf16: {:?}", is_x86_feature_detected!("amx-bf16"));
72+
println!("amx-tile: {:?}", is_x86_feature_detected!("amx-tile"));
73+
println!("amx-int8: {:?}", is_x86_feature_detected!("amx-int8"));
74+
println!("amx-fp16: {:?}", is_x86_feature_detected!("amx-fp16"));
75+
println!("amx-complex: {:?}", is_x86_feature_detected!("amx-complex"));
7176
}
7277

7378
#[cfg(feature = "std_detect_env_override")]

0 commit comments

Comments
 (0)