autozig_build/
simd.rs

1//! SIMD Optimization Detection
2//!
3//! This module detects CPU features at compile-time and configures Zig
4//! compilation to use appropriate SIMD instructions.
5
6use std::env;
7
8/// CPU Architecture detected from Rust target
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum CpuArch {
11    X86_64,
12    Aarch64,
13    Arm,
14    Other,
15}
16
17/// SIMD feature sets available on x86_64
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum X86SimdLevel {
20    /// Baseline x86_64 (SSE2)
21    Baseline,
22    /// SSE4.2 support
23    SSE4_2,
24    /// AVX support
25    AVX,
26    /// AVX2 support
27    AVX2,
28    /// AVX-512 support
29    AVX512,
30}
31
32/// SIMD feature sets available on ARM
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ArmSimdLevel {
35    /// No SIMD
36    None,
37    /// NEON support
38    NEON,
39    /// SVE support (Scalable Vector Extension)
40    SVE,
41}
42
43/// SIMD configuration for Zig compilation
44#[derive(Debug, Clone)]
45pub struct SimdConfig {
46    /// Target architecture
47    pub arch: CpuArch,
48    /// Zig CPU flag to pass to compiler
49    pub zig_cpu_flag: String,
50    /// Human-readable description
51    pub description: String,
52    /// Whether native CPU features should be used
53    pub use_native: bool,
54}
55
56impl SimdConfig {
57    /// Detect SIMD configuration from Rust environment
58    pub fn detect() -> Self {
59        let target = env::var("TARGET").unwrap_or_else(|_| "unknown".to_string());
60        let rust_flags = env::var("RUSTFLAGS").unwrap_or_default();
61        let target_features = env::var("CARGO_CFG_TARGET_FEATURE").unwrap_or_default();
62
63        // Check if native CPU optimization is requested
64        let use_native = rust_flags.contains("target-cpu=native")
65            || rust_flags.contains("-C native")
66            || rust_flags.contains("--target-cpu=native");
67
68        let arch = Self::detect_arch(&target);
69
70        let (zig_cpu_flag, description) = if use_native {
71            // Use native CPU features
72            ("-mcpu=native".to_string(), format!("{:?} with native CPU features", arch))
73        } else {
74            // Detect explicit features
75            Self::detect_features(&target, &target_features, arch)
76        };
77
78        SimdConfig {
79            arch,
80            zig_cpu_flag,
81            description,
82            use_native,
83        }
84    }
85
86    /// Detect CPU architecture from target triple
87    fn detect_arch(target: &str) -> CpuArch {
88        if target.contains("x86_64") {
89            CpuArch::X86_64
90        } else if target.contains("aarch64") {
91            CpuArch::Aarch64
92        } else if target.contains("arm") {
93            CpuArch::Arm
94        } else {
95            CpuArch::Other
96        }
97    }
98
99    /// Detect SIMD features from target and cargo configuration
100    fn detect_features(target: &str, features: &str, arch: CpuArch) -> (String, String) {
101        match arch {
102            CpuArch::X86_64 => Self::detect_x86_features(target, features),
103            CpuArch::Aarch64 | CpuArch::Arm => Self::detect_arm_features(features),
104            CpuArch::Other => {
105                ("-mcpu=baseline".to_string(), "baseline (unknown architecture)".to_string())
106            },
107        }
108    }
109
110    /// Detect x86_64 SIMD features
111    fn detect_x86_features(target: &str, features: &str) -> (String, String) {
112        // Check for AVX-512 first (most advanced)
113        if features.contains("avx512") {
114            return ("-mcpu=x86_64_v4".to_string(), "x86_64 v4 (AVX-512)".to_string());
115        }
116
117        // Check for AVX2
118        if features.contains("avx2") {
119            return ("-mcpu=x86_64_v3".to_string(), "x86_64 v3 (AVX2, FMA)".to_string());
120        }
121
122        // Check for AVX
123        if features.contains("avx") {
124            return ("-mcpu=x86_64+avx".to_string(), "x86_64 with AVX".to_string());
125        }
126
127        // Check for SSE4.2
128        if features.contains("sse4.2") || features.contains("sse4_2") {
129            return ("-mcpu=x86_64+sse4.2".to_string(), "x86_64 with SSE4.2".to_string());
130        }
131
132        // Default x86_64 baseline (SSE2)
133        // x86_64 v1 is the baseline with SSE2
134        if target.contains("x86_64") {
135            ("-mcpu=x86_64".to_string(), "x86_64 baseline (SSE2)".to_string())
136        } else {
137            ("-mcpu=baseline".to_string(), "baseline".to_string())
138        }
139    }
140
141    /// Detect ARM SIMD features
142    fn detect_arm_features(features: &str) -> (String, String) {
143        // Check for SVE
144        if features.contains("sve") {
145            return ("-mcpu=generic+sve".to_string(), "ARM with SVE".to_string());
146        }
147
148        // Check for NEON
149        if features.contains("neon") {
150            return ("-mcpu=generic+neon".to_string(), "ARM with NEON".to_string());
151        }
152
153        // ARM baseline
154        ("-mcpu=generic".to_string(), "ARM generic".to_string())
155    }
156
157    /// Generate a compiler report for display during build
158    pub fn report(&self) -> String {
159        format!(
160            r#"
161╔════════════════════════════════════════════╗
162║     AutoZig SIMD Optimization Report      ║
163╚════════════════════════════════════════════╝
164
165Architecture:    {:?}
166Configuration:   {}
167Zig CPU Flag:    {}
168Native Features: {}
169
170TIP: Set RUSTFLAGS="-C target-cpu=native" for maximum performance
171"#,
172            self.arch,
173            self.description,
174            self.zig_cpu_flag,
175            if self.use_native {
176                "ENABLED"
177            } else {
178                "disabled"
179            }
180        )
181    }
182
183    /// Get the Zig compiler flag as a string
184    pub fn as_zig_flag(&self) -> &str {
185        &self.zig_cpu_flag
186    }
187}
188
189/// Detect and print SIMD configuration
190///
191/// This should be called from build.rs to configure Zig compilation
192pub fn detect_and_report() -> SimdConfig {
193    let config = SimdConfig::detect();
194
195    // Print report to cargo output
196    println!("cargo:warning={}", config.report());
197
198    // Emit configuration for use in source code
199    println!("cargo:rustc-env=AUTOZIG_SIMD_ARCH={:?}", config.arch);
200    println!("cargo:rustc-env=AUTOZIG_SIMD_CONFIG={}", config.description);
201    println!("cargo:rustc-env=AUTOZIG_ZIG_CPU_FLAG={}", config.zig_cpu_flag);
202
203    config
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    #[test]
211    fn test_detect_arch() {
212        assert_eq!(SimdConfig::detect_arch("x86_64-unknown-linux-gnu"), CpuArch::X86_64);
213        assert_eq!(SimdConfig::detect_arch("aarch64-apple-darwin"), CpuArch::Aarch64);
214        assert_eq!(SimdConfig::detect_arch("armv7-unknown-linux"), CpuArch::Arm);
215    }
216
217    #[test]
218    fn test_x86_feature_detection() {
219        let (flag, desc) = SimdConfig::detect_x86_features("x86_64-pc-windows-msvc", "avx2,fma");
220        assert!(flag.contains("x86_64"));
221        assert!(desc.contains("AVX2") || desc.contains("v3"));
222
223        let (flag, desc) = SimdConfig::detect_x86_features("x86_64-pc-windows-msvc", "sse4.2");
224        assert!(flag.contains("sse4.2") || flag.contains("x86_64"));
225        assert!(desc.contains("SSE") || desc.contains("baseline"));
226    }
227
228    #[test]
229    fn test_arm_feature_detection() {
230        let (flag, desc) = SimdConfig::detect_arm_features("neon");
231        assert!(flag.contains("neon") || flag.contains("generic"));
232        assert!(desc.contains("NEON") || desc.contains("ARM"));
233    }
234
235    #[test]
236    fn test_config_detection() {
237        // Just ensure it doesn't panic
238        let config = SimdConfig::detect();
239        assert!(!config.zig_cpu_flag.is_empty());
240        assert!(!config.description.is_empty());
241    }
242
243    #[test]
244    fn test_report_generation() {
245        let config = SimdConfig::detect();
246        let report = config.report();
247        assert!(report.contains("AutoZig SIMD"));
248        assert!(report.contains("Architecture"));
249    }
250}
autozig_build/simd.rs

autozig_build/
simd.rs
