tenflowers_core/simd/
capabilities.rs1#[derive(Debug, Clone, Copy)]
8pub struct SimdCapabilities {
9 pub has_auto_vectorization: bool,
10 pub has_target_features: bool,
11 pub recommended_unroll_factor: usize,
12 pub cache_line_size: usize,
13}
14
15pub struct Capabilities;
17
18impl Capabilities {
19 pub fn detect_capabilities() -> SimdCapabilities {
21 SimdCapabilities {
22 has_auto_vectorization: true, has_target_features: Self::has_target_feature_support(),
24 recommended_unroll_factor: 8,
25 cache_line_size: Self::estimate_cache_line_size(),
26 }
27 }
28
29 fn has_target_feature_support() -> bool {
31 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
33 {
34 Self::detect_x86_features()
36 }
37
38 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
39 {
40 true
42 }
43 }
44
45 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
47 fn detect_x86_features() -> bool {
48 #[cfg(target_feature = "sse2")]
50 {
51 true }
53
54 #[cfg(not(target_feature = "sse2"))]
55 {
56 std::arch::is_x86_feature_detected!("sse2")
58 }
59 }
60
61 pub fn estimate_cache_line_size() -> usize {
63 64
65 }
66
67 pub fn optimal_block_size(data_size: usize, element_size: usize) -> usize {
69 let cache_line_size = Self::estimate_cache_line_size();
70 let l1_cache_size = 32 * 1024; let max_block_elements = l1_cache_size / element_size / 4; let cache_line_elements = cache_line_size / element_size;
75
76 let optimal_size = (max_block_elements / cache_line_elements) * cache_line_elements;
78
79 optimal_size.min(data_size).max(cache_line_elements * 4)
81 }
82
83 pub fn optimal_unroll_factor(data_size: usize) -> usize {
85 match data_size {
86 0..=32 => 1, 33..=128 => 2, 129..=512 => 4, 513..=2048 => 8, _ => 16, }
92 }
93
94 pub fn should_use_simd(data_size: usize) -> bool {
96 data_size >= 8
98 }
99
100 pub fn memory_alignment_requirement() -> usize {
102 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
104 {
105 32 }
107
108 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
109 {
110 16 }
112 }
113
114 pub fn is_aligned(ptr: *const f32) -> bool {
116 let alignment = Self::memory_alignment_requirement();
117 (ptr as usize) % alignment == 0
118 }
119
120 pub fn simd_lanes_f32() -> usize {
122 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
123 {
124 if Self::has_avx2_support() {
125 8 } else if Self::has_sse2_support() {
127 4 } else {
129 1 }
131 }
132
133 #[cfg(target_arch = "aarch64")]
134 {
135 4 }
137
138 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
139 {
140 1 }
142 }
143
144 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
146 fn has_avx2_support() -> bool {
147 std::arch::is_x86_feature_detected!("avx2")
148 }
149
150 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
152 fn has_sse2_support() -> bool {
153 std::arch::is_x86_feature_detected!("sse2")
154 }
155
156 pub fn get_performance_hints() -> PerformanceHints {
158 PerformanceHints {
159 prefer_unrolling: Self::should_prefer_unrolling(),
160 prefer_blocking: Self::should_prefer_blocking(),
161 optimal_chunk_size: Self::get_optimal_chunk_size(),
162 memory_prefetch_distance: Self::get_prefetch_distance(),
163 }
164 }
165
166 fn should_prefer_unrolling() -> bool {
168 #[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))]
170 {
171 true
172 }
173
174 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
175 {
176 false }
178 }
179
180 fn should_prefer_blocking() -> bool {
182 true
184 }
185
186 fn get_optimal_chunk_size() -> usize {
188 let simd_lanes = Self::simd_lanes_f32();
189 let cache_line_elements = Self::estimate_cache_line_size() / std::mem::size_of::<f32>();
190
191 let base_chunk = simd_lanes.max(cache_line_elements);
193
194 base_chunk.next_power_of_two().min(64) }
197
198 fn get_prefetch_distance() -> usize {
200 Self::estimate_cache_line_size() * 4
202 }
203}
204
205#[derive(Debug, Clone)]
207pub struct PerformanceHints {
208 pub prefer_unrolling: bool,
209 pub prefer_blocking: bool,
210 pub optimal_chunk_size: usize,
211 pub memory_prefetch_distance: usize,
212}
213
214impl Default for PerformanceHints {
215 fn default() -> Self {
216 Capabilities::get_performance_hints()
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 #[test]
225 fn test_detect_capabilities() {
226 let capabilities = Capabilities::detect_capabilities();
227
228 assert!(capabilities.recommended_unroll_factor > 0);
230 assert!(capabilities.cache_line_size > 0);
231 assert!(capabilities.cache_line_size.is_power_of_two());
232 }
233
234 #[test]
235 fn test_optimal_block_size() {
236 let data_size = 1024;
237 let element_size = std::mem::size_of::<f32>();
238
239 let block_size = Capabilities::optimal_block_size(data_size, element_size);
240
241 assert!(block_size > 0);
243 assert!(block_size <= data_size);
244
245 let cache_line_elements = Capabilities::estimate_cache_line_size() / element_size;
247 assert_eq!(block_size % cache_line_elements, 0);
248 }
249
250 #[test]
251 fn test_optimal_unroll_factor() {
252 assert_eq!(Capabilities::optimal_unroll_factor(16), 1);
254 assert_eq!(Capabilities::optimal_unroll_factor(64), 2);
255 assert_eq!(Capabilities::optimal_unroll_factor(256), 4);
256 assert_eq!(Capabilities::optimal_unroll_factor(1024), 8);
257 assert_eq!(Capabilities::optimal_unroll_factor(4096), 16);
258 }
259
260 #[test]
261 fn test_should_use_simd() {
262 assert!(!Capabilities::should_use_simd(4));
264 assert!(!Capabilities::should_use_simd(7));
265
266 assert!(Capabilities::should_use_simd(8));
268 assert!(Capabilities::should_use_simd(32));
269 assert!(Capabilities::should_use_simd(1024));
270 }
271
272 #[test]
273 fn test_memory_alignment_requirement() {
274 let alignment = Capabilities::memory_alignment_requirement();
275
276 assert!(alignment.is_power_of_two());
278 assert!(alignment >= 8);
279 }
280
281 #[test]
282 fn test_simd_lanes_f32() {
283 let lanes = Capabilities::simd_lanes_f32();
284
285 assert!(lanes >= 1);
287 assert!(lanes.is_power_of_two());
288 }
289
290 #[test]
291 fn test_performance_hints() {
292 let hints = Capabilities::get_performance_hints();
293
294 assert!(hints.optimal_chunk_size > 0);
296 assert!(hints.optimal_chunk_size.is_power_of_two());
297 assert!(hints.memory_prefetch_distance > 0);
298 }
299
300 #[test]
301 fn test_cache_line_size() {
302 let cache_line_size = Capabilities::estimate_cache_line_size();
303
304 assert!(cache_line_size.is_power_of_two());
306 assert!(cache_line_size >= 32);
307 assert!(cache_line_size <= 256);
308 }
309}