1use bitflags::bitflags;
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, thiserror::Error)]
12pub enum FeatureError {
13 #[error("Feature detection not supported on this architecture")]
14 UnsupportedArch,
15 #[error("Failed to detect feature {0}")]
16 DetectionFailed(String),
17}
18
19bitflags! {
20 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
22 pub struct X86Features: u64 {
23 const SSE = 1 << 0;
24 const SSE2 = 1 << 1;
25 const SSE3 = 1 << 2;
26 const SSSE3 = 1 << 3;
27 const SSE4_1 = 1 << 4;
28 const SSE4_2 = 1 << 5;
29 const AVX = 1 << 6;
30 const AVX2 = 1 << 7;
31 const FMA = 1 << 8;
32 const BMI1 = 1 << 9;
33 const BMI2 = 1 << 10;
34 const F16C = 1 << 11;
35 const POPCNT = 1 << 12;
36 const AES = 1 << 13;
37 const AVX512F = 1 << 14;
38 const AVX512BW = 1 << 15;
39 const AVX512CD = 1 << 16;
40 const AVX512DQ = 1 << 17;
41 const AVX512VL = 1 << 18;
42 }
43}
44
45bitflags! {
46 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
48 pub struct ArmFeatures: u64 {
49 const NEON = 1 << 0;
50 const AES = 1 << 1;
51 const PMULL = 1 << 2;
52 const SHA1 = 1 << 3;
53 const SHA2 = 1 << 4;
54 const CRC32 = 1 << 5;
55 const ATOMICS = 1 << 6;
56 const FP = 1 << 7;
57 const ASIMD = 1 << 8;
58 const FPHP = 1 << 9;
59 const ASIMDHP = 1 << 10;
60 const ASIMDDP = 1 << 11;
61 const ASIMDFHM = 1 << 12;
62 }
63}
64
65#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
71pub fn detect_features() -> Result<X86Features, FeatureError> {
72 let mut features = X86Features::empty();
73
74 if cfg!(target_arch = "x86") || cfg!(target_arch = "x86_64") {
76 if std::is_x86_feature_detected!("sse") {
77 features |= X86Features::SSE;
78 }
79 if std::is_x86_feature_detected!("sse2") {
80 features |= X86Features::SSE2;
81 }
82 if std::is_x86_feature_detected!("sse3") {
83 features |= X86Features::SSE3;
84 }
85 if std::is_x86_feature_detected!("ssse3") {
86 features |= X86Features::SSSE3;
87 }
88 if std::is_x86_feature_detected!("sse4.1") {
89 features |= X86Features::SSE4_1;
90 }
91 if std::is_x86_feature_detected!("sse4.2") {
92 features |= X86Features::SSE4_2;
93 }
94 if std::is_x86_feature_detected!("avx") {
95 features |= X86Features::AVX;
96 }
97 if std::is_x86_feature_detected!("avx2") {
98 features |= X86Features::AVX2;
99 }
100 if std::is_x86_feature_detected!("fma") {
101 features |= X86Features::FMA;
102 }
103 if std::is_x86_feature_detected!("bmi1") {
104 features |= X86Features::BMI1;
105 }
106 if std::is_x86_feature_detected!("bmi2") {
107 features |= X86Features::BMI2;
108 }
109 if std::is_x86_feature_detected!("f16c") {
110 features |= X86Features::F16C;
111 }
112 if std::is_x86_feature_detected!("popcnt") {
113 features |= X86Features::POPCNT;
114 }
115 if std::is_x86_feature_detected!("aes") {
116 features |= X86Features::AES;
117 }
118 if std::is_x86_feature_detected!("avx512f") {
119 features |= X86Features::AVX512F;
120 }
121 if std::is_x86_feature_detected!("avx512bw") {
122 features |= X86Features::AVX512BW;
123 }
124 if std::is_x86_feature_detected!("avx512cd") {
125 features |= X86Features::AVX512CD;
126 }
127 if std::is_x86_feature_detected!("avx512dq") {
128 features |= X86Features::AVX512DQ;
129 }
130 if std::is_x86_feature_detected!("avx512vl") {
131 features |= X86Features::AVX512VL;
132 }
133 }
134
135 Ok(features)
136}
137
138#[cfg(target_arch = "aarch64")]
144pub fn detect_features() -> Result<ArmFeatures, FeatureError> {
145 let mut features = ArmFeatures::empty();
146
147 if cfg!(target_arch = "aarch64") {
149 if std::arch::is_aarch64_feature_detected!("neon") {
150 features |= ArmFeatures::NEON;
151 }
152 if std::arch::is_aarch64_feature_detected!("aes") {
153 features |= ArmFeatures::AES;
154 }
155 if std::arch::is_aarch64_feature_detected!("pmull") {
156 features |= ArmFeatures::PMULL;
157 }
158 if std::arch::is_aarch64_feature_detected!("sha2") {
159 features |= ArmFeatures::SHA2;
160 }
161 if std::arch::is_aarch64_feature_detected!("crc") {
162 features |= ArmFeatures::CRC32;
163 }
164 if std::arch::is_aarch64_feature_detected!("lse") {
165 features |= ArmFeatures::ATOMICS;
166 }
167 if std::arch::is_aarch64_feature_detected!("fp") {
168 features |= ArmFeatures::FP;
169 }
170 if std::arch::is_aarch64_feature_detected!("asimd") {
171 features |= ArmFeatures::ASIMD;
172 }
173 }
175
176 Ok(features)
177}
178
179#[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
181pub fn detect_features() -> Result<(), FeatureError> {
182 Err(FeatureError::UnsupportedArch)
183}
184
185#[cfg(test)]
186mod tests {
187 use super::*;
188
189 #[test]
190 fn test_x86_features_flags() {
191 let features = X86Features::SSE | X86Features::SSE2;
192 assert!(features.contains(X86Features::SSE));
193 assert!(features.contains(X86Features::SSE2));
194 assert!(!features.contains(X86Features::AVX));
195 }
196
197 #[test]
198 fn test_arm_features_flags() {
199 let features = ArmFeatures::NEON | ArmFeatures::AES;
200 assert!(features.contains(ArmFeatures::NEON));
201 assert!(features.contains(ArmFeatures::AES));
202 assert!(!features.contains(ArmFeatures::SHA2));
203 }
204
205 }