Skip to content

Commit d5c4073

Browse files
sayantnAmanieu
sayantn
authored and
Amanieu
committed
Added runtime detection
Cannot do a `cupid` test because they don't support `amx`.
1 parent 09215ac commit d5c4073

File tree

4 files changed

+61
-20
lines changed

4 files changed

+61
-20
lines changed

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

+15
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,11 @@ features! {
8484
/// * `"avxneconvert"`
8585
/// * `"avxvnniint8"`
8686
/// * `"avxvnniint16"`
87+
/// * `"amx-tile"`
88+
/// * `"amx-int8"`
89+
/// * `"amx-bf16"`
90+
/// * `"amx-fp16"`
91+
/// * `"amx-complex"`
8792
/// * `"f16c"`
8893
/// * `"fma"`
8994
/// * `"bmi1"`
@@ -196,6 +201,16 @@ features! {
196201
/// AVX-VNNI_INT8 (VNNI with 16-bit Integers)
197202
@FEATURE: #[unstable(feature = "avx512_target_feature", issue = "44839")] avxvnniint8: "avxvnniint8";
198203
/// AVX-VNNI_INT16 (VNNI with 8-bit integers)
204+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_tile: "amx-tile";
205+
/// AMX (Advanced Matrix Extensions) - Tile load/store
206+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_int8: "amx-int8";
207+
/// AMX-INT8 (Operations on 8-bit integers)
208+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_bf16: "amx-bf16";
209+
/// AMX-BF16 (BFloat16 Operations)
210+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_fp16: "amx-fp16";
211+
/// AMX-FP16 (Float16 Operations)
212+
@FEATURE: #[unstable(feature = "x86_amx_intrinsics", issue = "126622")] amx_complex: "amx-complex";
213+
/// AMX-COMPLEX (Complex number Operations)
199214
@FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] f16c: "f16c";
200215
/// F16C (Conversions between IEEE-754 `binary16` and `binary32` formats)
201216
@FEATURE: #[stable(feature = "simd_x86", since = "1.27.0")] fma: "fma";

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

+11
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,7 @@ pub(crate) fn detect_features() -> cache::Initializer {
164164
// * SSE -> `XCR0.SSE[1]`
165165
// * AVX -> `XCR0.AVX[2]`
166166
// * AVX-512 -> `XCR0.AVX-512[7:5]`.
167+
// * AMX -> `XCR0.AMX[18:17]`
167168
//
168169
// by setting the corresponding bits of `XCR0` to `1`.
169170
//
@@ -174,6 +175,8 @@ pub(crate) fn detect_features() -> cache::Initializer {
174175
let os_avx_support = xcr0 & 6 == 6;
175176
// Test `XCR0.AVX-512[7:5]` with the mask `0b1110_0000 == 0xe0`:
176177
let os_avx512_support = xcr0 & 0xe0 == 0xe0;
178+
// Test `XCR0.AMX[18:17]` with the mask `0b110_0000_0000_0000_0000 == 0x60000`
179+
let os_amx_support = xcr0 & 0x60000 == 0x60000;
177180

178181
// Only if the OS and the CPU support saving/restoring the AVX
179182
// registers we enable `xsave` support:
@@ -240,6 +243,14 @@ pub(crate) fn detect_features() -> cache::Initializer {
240243
enable(extended_features_edx, 8, Feature::avx512vp2intersect);
241244
enable(extended_features_edx, 23, Feature::avx512fp16);
242245
enable(extended_features_eax_leaf_1, 5, Feature::avx512bf16);
246+
247+
if os_amx_support {
248+
enable(extended_features_edx, 24, Feature::amx_tile);
249+
enable(extended_features_edx, 25, Feature::amx_int8);
250+
enable(extended_features_edx, 22, Feature::amx_bf16);
251+
enable(extended_features_eax_leaf_1, 21, Feature::amx_fp16);
252+
enable(extended_features_edx_leaf_1, 8, Feature::amx_complex);
253+
}
243254
}
244255
}
245256
}

‎crates/std_detect/tests/cpu-detection.rs

+6-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
#![cfg_attr(target_arch = "powerpc64", feature(stdarch_powerpc_feature_detection))]
66
#![cfg_attr(
77
any(target_arch = "x86", target_arch = "x86_64"),
8-
feature(sha512_sm_x86)
8+
feature(sha512_sm_x86, x86_amx_intrinsics)
99
)]
1010
#![allow(clippy::unwrap_used, clippy::use_debug, clippy::print_stdout)]
1111

@@ -259,6 +259,11 @@ fn x86_all() {
259259
println!("xsaveopt: {:?}", is_x86_feature_detected!("xsaveopt"));
260260
println!("xsaves: {:?}", is_x86_feature_detected!("xsaves"));
261261
println!("xsavec: {:?}", is_x86_feature_detected!("xsavec"));
262+
println!("amx-bf16: {:?}", is_x86_feature_detected!("amx-bf16"));
263+
println!("amx-tile: {:?}", is_x86_feature_detected!("amx-tile"));
264+
println!("amx-int8: {:?}", is_x86_feature_detected!("amx-int8"));
265+
println!("amx-fp16: {:?}", is_x86_feature_detected!("amx-fp16"));
266+
println!("amx-complex: {:?}", is_x86_feature_detected!("amx-complex"));
262267
}
263268

264269
#[test]

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

+29-19
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
#![cfg(any(target_arch = "x86", target_arch = "x86_64"))]
22
#![allow(internal_features)]
3-
#![feature(stdarch_internal, avx512_target_feature, sha512_sm_x86)]
3+
#![feature(
4+
stdarch_internal,
5+
avx512_target_feature,
6+
sha512_sm_x86,
7+
x86_amx_intrinsics
8+
)]
49

510
extern crate cupid;
611
#[macro_use]
@@ -27,34 +32,34 @@ fn dump() {
2732
println!("sha512: {:?}", is_x86_feature_detected!("sha512"));
2833
println!("sm3: {:?}", is_x86_feature_detected!("sm3"));
2934
println!("sm4: {:?}", is_x86_feature_detected!("sm4"));
30-
println!("avx512f {:?}", is_x86_feature_detected!("avx512f"));
31-
println!("avx512cd {:?}", is_x86_feature_detected!("avx512cd"));
32-
println!("avx512er {:?}", is_x86_feature_detected!("avx512er"));
33-
println!("avx512pf {:?}", is_x86_feature_detected!("avx512pf"));
34-
println!("avx512bw {:?}", is_x86_feature_detected!("avx512bw"));
35-
println!("avx512dq {:?}", is_x86_feature_detected!("avx512dq"));
36-
println!("avx512vl {:?}", is_x86_feature_detected!("avx512vl"));
37-
println!("avx512_ifma {:?}", is_x86_feature_detected!("avx512ifma"));
35+
println!("avx512f: {:?}", is_x86_feature_detected!("avx512f"));
36+
println!("avx512cd: {:?}", is_x86_feature_detected!("avx512cd"));
37+
println!("avx512er: {:?}", is_x86_feature_detected!("avx512er"));
38+
println!("avx512pf: {:?}", is_x86_feature_detected!("avx512pf"));
39+
println!("avx512bw: {:?}", is_x86_feature_detected!("avx512bw"));
40+
println!("avx512dq: {:?}", is_x86_feature_detected!("avx512dq"));
41+
println!("avx512vl: {:?}", is_x86_feature_detected!("avx512vl"));
42+
println!("avx512_ifma: {:?}", is_x86_feature_detected!("avx512ifma"));
3843
println!("avx512vbmi {:?}", is_x86_feature_detected!("avx512vbmi"));
3944
println!(
40-
"avx512_vpopcntdq {:?}",
45+
"avx512_vpopcntdq: {:?}",
4146
is_x86_feature_detected!("avx512vpopcntdq")
4247
);
43-
println!("avx512vbmi2 {:?}", is_x86_feature_detected!("avx512vbmi2"));
44-
println!("gfni {:?}", is_x86_feature_detected!("gfni"));
45-
println!("vaes {:?}", is_x86_feature_detected!("vaes"));
46-
println!("vpclmulqdq {:?}", is_x86_feature_detected!("vpclmulqdq"));
47-
println!("avx512vnni {:?}", is_x86_feature_detected!("avx512vnni"));
48+
println!("avx512vbmi2: {:?}", is_x86_feature_detected!("avx512vbmi2"));
49+
println!("gfni: {:?}", is_x86_feature_detected!("gfni"));
50+
println!("vaes: {:?}", is_x86_feature_detected!("vaes"));
51+
println!("vpclmulqdq: {:?}", is_x86_feature_detected!("vpclmulqdq"));
52+
println!("avx512vnni: {:?}", is_x86_feature_detected!("avx512vnni"));
4853
println!(
49-
"avx512bitalg {:?}",
54+
"avx512bitalg: {:?}",
5055
is_x86_feature_detected!("avx512bitalg")
5156
);
52-
println!("avx512bf16 {:?}", is_x86_feature_detected!("avx512bf16"));
57+
println!("avx512bf16: {:?}", is_x86_feature_detected!("avx512bf16"));
5358
println!(
54-
"avx512vp2intersect {:?}",
59+
"avx512vp2intersect: {:?}",
5560
is_x86_feature_detected!("avx512vp2intersect")
5661
);
57-
println!("avx512fp16 {:?}", is_x86_feature_detected!("avx512fp16"));
62+
println!("avx512fp16: {:?}", is_x86_feature_detected!("avx512fp16"));
5863
println!("fma: {:?}", is_x86_feature_detected!("fma"));
5964
println!("abm: {:?}", is_x86_feature_detected!("abm"));
6065
println!("bmi: {:?}", is_x86_feature_detected!("bmi1"));
@@ -82,6 +87,11 @@ fn dump() {
8287
"avxvnniint16: {:?}",
8388
is_x86_feature_detected!("avxvnniint16")
8489
);
90+
println!("amx-bf16: {:?}", is_x86_feature_detected!("amx-bf16"));
91+
println!("amx-tile: {:?}", is_x86_feature_detected!("amx-tile"));
92+
println!("amx-int8: {:?}", is_x86_feature_detected!("amx-int8"));
93+
println!("amx-fp16: {:?}", is_x86_feature_detected!("amx-fp16"));
94+
println!("amx-complex: {:?}", is_x86_feature_detected!("amx-complex"));
8595
}
8696

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

0 commit comments

Comments
 (0)