Skip to main content

sparse_ir_core/
fpu_check.rs

1//! FPU state checking and correction for numerical stability
2//!
3//! This module detects dangerous FPU settings (particularly Flush-to-Zero and
4//! Denormals-Are-Zero flags) that can cause incorrect SVD results when called
5//! from Intel Fortran programs compiled with `-O3` without `-fp-model precise`.
6//!
7//! # Background
8//!
9//! Intel Fortran's `-O3` optimization may set the MXCSR register's FZ and DAZ bits
10//! at program startup for performance. However, this causes problems for SVD
11//! computations that rely on proper handling of denormalized (subnormal) numbers.
12//!
13//! # Usage
14//!
15//! The [`FpuGuard`] RAII guard automatically saves, corrects, and restores FPU state:
16//!
17//! ```
18//! use sparse_ir::fpu_check::FpuGuard;
19//!
20//! let _guard = FpuGuard::new_protect_computation();
21//! // Computation here - FZ/DAZ are disabled
22//! // FPU state is automatically restored when _guard is dropped
23//! ```
24//!
25//! # Performance
26//!
27//! The `stmxcsr`/`ldmxcsr` instructions are very lightweight (a few CPU cycles),
28//! so the overhead of checking and restoring FPU state is negligible compared
29//! to actual matrix computations.
30
31use once_cell::sync::Lazy;
32use std::sync::atomic::{AtomicBool, Ordering};
33
34/// MXCSR bit positions
35const MXCSR_FZ_BIT: u32 = 15; // Flush to Zero
36const MXCSR_DAZ_BIT: u32 = 6; // Denormals Are Zero
37
38/// Flag to track if warning has been shown (only show once per process)
39static WARNING_SHOWN: AtomicBool = AtomicBool::new(false);
40
41/// Lazy initialization to check FPU state at library load time
42static FPU_CHECK_INIT: Lazy<bool> = Lazy::new(|| {
43    let state = get_fpu_state();
44    if state.is_dangerous() {
45        print_fpu_warning(&state);
46        WARNING_SHOWN.store(true, Ordering::SeqCst);
47        true // dangerous state detected
48    } else {
49        false
50    }
51});
52
53/// Result of FPU state check
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub struct FpuState {
56    /// Raw MXCSR register value
57    pub mxcsr: u32,
58    /// Flush to Zero flag
59    pub flush_to_zero: bool,
60    /// Denormals Are Zero flag
61    pub denormals_are_zero: bool,
62}
63
64impl FpuState {
65    /// Check if FPU state is dangerous for numerical computation
66    pub fn is_dangerous(&self) -> bool {
67        self.flush_to_zero || self.denormals_are_zero
68    }
69}
70
71impl std::fmt::Display for FpuState {
72    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73        write!(
74            f,
75            "MXCSR=0x{:08X}, FZ={}, DAZ={}",
76            self.mxcsr, self.flush_to_zero as u8, self.denormals_are_zero as u8
77        )
78    }
79}
80
81/// Print FPU warning message (called only once)
82fn print_fpu_warning(state: &FpuState) {
83    eprintln!();
84    eprintln!("================================================================================");
85    eprintln!("sparse-ir WARNING: Dangerous FPU settings detected!");
86    eprintln!("================================================================================");
87    eprintln!();
88    eprintln!("  Current FPU state: {}", state);
89    eprintln!();
90    eprintln!("  Problem: Flush-to-Zero (FZ) or Denormals-Are-Zero (DAZ) is enabled.");
91    eprintln!("           This causes subnormal numbers to be treated as zero, which");
92    eprintln!("           can produce INCORRECT results in SVD/SVE computations.");
93    eprintln!();
94    eprintln!("  Common cause: Intel Fortran compiler (ifort/ifx) with -O3 optimization");
95    eprintln!("                sets FZ/DAZ flags at program startup for performance.");
96    eprintln!();
97    eprintln!("  Solution: Add '-fp-model precise' flag when compiling your Fortran code:");
98    eprintln!();
99    eprintln!("      ifort -O3 -fp-model precise your_program.f90");
100    eprintln!("      ifx   -O3 -fp-model precise your_program.f90");
101    eprintln!();
102    eprintln!("  For Quantum ESPRESSO/EPW, add to make.inc:");
103    eprintln!();
104    eprintln!("      FFLAGS += -fp-model precise");
105    eprintln!();
106    eprintln!("  Action: sparse-ir will temporarily disable FZ/DAZ during each computation");
107    eprintln!("          and restore the original settings afterward.");
108    eprintln!("          Results will be correct, but please fix the compiler flags");
109    eprintln!("          to avoid this warning.");
110    eprintln!();
111    eprintln!("================================================================================");
112    eprintln!();
113}
114
115/// Get current FPU state (x86/x86_64 only)
116#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
117pub fn get_fpu_state() -> FpuState {
118    let mut mxcsr: u32 = 0;
119    unsafe {
120        std::arch::asm!(
121            "stmxcsr [{}]",
122            in(reg) &mut mxcsr,
123            options(nostack)
124        );
125    }
126
127    FpuState {
128        mxcsr,
129        flush_to_zero: (mxcsr >> MXCSR_FZ_BIT) & 1 != 0,
130        denormals_are_zero: (mxcsr >> MXCSR_DAZ_BIT) & 1 != 0,
131    }
132}
133
134/// Get current FPU state (non-x86 fallback)
135#[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
136pub fn get_fpu_state() -> FpuState {
137    // On non-x86 platforms, assume safe defaults
138    FpuState {
139        mxcsr: 0,
140        flush_to_zero: false,
141        denormals_are_zero: false,
142    }
143}
144
145/// Set MXCSR register value (x86/x86_64 only)
146#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
147fn set_mxcsr(value: u32) {
148    unsafe {
149        std::arch::asm!(
150            "ldmxcsr [{}]",
151            in(reg) &value,
152            options(nostack)
153        );
154    }
155}
156
157/// Set MXCSR register value (non-x86 fallback - no-op)
158#[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
159fn set_mxcsr(_value: u32) {
160    // No-op on non-x86 platforms
161}
162
163/// Initialize FPU check (call this early to trigger warning if needed)
164///
165/// This function triggers the lazy initialization which checks FPU state
166/// and prints a warning if dangerous settings are detected.
167/// The warning is only printed once per process.
168pub fn init_fpu_check() {
169    let _ = *FPU_CHECK_INIT;
170}
171
172/// RAII guard that protects a computation from dangerous FPU settings
173///
174/// On creation:
175/// 1. Triggers one-time FPU check and warning (if not already done)
176/// 2. If FZ or DAZ is enabled, disables them temporarily
177///
178/// On drop:
179/// - Restores the original FPU state
180///
181/// # Example
182///
183/// ```
184/// use sparse_ir::fpu_check::{FpuGuard, get_fpu_state};
185///
186/// let before = get_fpu_state();
187/// {
188///     let _guard = FpuGuard::new_protect_computation();
189///     // Computation here - FZ/DAZ are disabled
190///     assert!(!get_fpu_state().is_dangerous());
191///     let subnormal = std::hint::black_box(f64::MIN_POSITIVE) / 2.0;
192///     assert!(subnormal > 0.0); // not flushed to zero
193/// } // Original FPU state is restored here
194/// let after = get_fpu_state();
195/// assert_eq!(after.flush_to_zero, before.flush_to_zero);
196/// assert_eq!(after.denormals_are_zero, before.denormals_are_zero);
197/// ```
198///
199/// # Performance
200///
201/// The overhead is negligible (a few CPU cycles for stmxcsr/ldmxcsr).
202pub struct FpuGuard {
203    original_mxcsr: u32,
204    needs_restore: bool,
205}
206
207impl FpuGuard {
208    /// Create a new guard that protects computation from FZ/DAZ
209    ///
210    /// - Triggers one-time warning if dangerous FPU settings are detected
211    /// - Temporarily disables FZ/DAZ if they are enabled
212    /// - Restores original state when dropped
213    pub fn new_protect_computation() -> Self {
214        // Trigger one-time FPU check and warning
215        let _ = *FPU_CHECK_INIT;
216
217        let state = get_fpu_state();
218        let original_mxcsr = state.mxcsr;
219
220        if state.is_dangerous() {
221            // Clear FZ and DAZ bits
222            let safe_mxcsr = original_mxcsr & !((1 << MXCSR_FZ_BIT) | (1 << MXCSR_DAZ_BIT));
223            set_mxcsr(safe_mxcsr);
224
225            Self {
226                original_mxcsr,
227                needs_restore: true,
228            }
229        } else {
230            Self {
231                original_mxcsr,
232                needs_restore: false,
233            }
234        }
235    }
236
237    /// Check if the guard needed to modify FPU state
238    pub fn was_modified(&self) -> bool {
239        self.needs_restore
240    }
241}
242
243impl Drop for FpuGuard {
244    fn drop(&mut self) {
245        if self.needs_restore {
246            // Restore original FPU state
247            set_mxcsr(self.original_mxcsr);
248        }
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255
256    #[test]
257    fn test_get_fpu_state() {
258        let state = get_fpu_state();
259        // Just verify we can read the state without panicking
260        println!("Current FPU state: {}", state);
261    }
262
263    #[test]
264    fn test_fpu_guard_creation() {
265        let guard = FpuGuard::new_protect_computation();
266        // Guard should be created successfully
267        drop(guard);
268    }
269
270    #[test]
271    fn test_fpu_state_display() {
272        let state = FpuState {
273            mxcsr: 0x1F80,
274            flush_to_zero: false,
275            denormals_are_zero: false,
276        };
277        let display = format!("{}", state);
278        assert!(display.contains("MXCSR=0x00001F80"));
279        assert!(display.contains("FZ=0"));
280        assert!(display.contains("DAZ=0"));
281    }
282
283    #[test]
284    fn test_fpu_state_dangerous() {
285        let safe_state = FpuState {
286            mxcsr: 0x1F80,
287            flush_to_zero: false,
288            denormals_are_zero: false,
289        };
290        assert!(!safe_state.is_dangerous());
291
292        let dangerous_fz = FpuState {
293            mxcsr: 0x9F80,
294            flush_to_zero: true,
295            denormals_are_zero: false,
296        };
297        assert!(dangerous_fz.is_dangerous());
298
299        let dangerous_daz = FpuState {
300            mxcsr: 0x1FC0,
301            flush_to_zero: false,
302            denormals_are_zero: true,
303        };
304        assert!(dangerous_daz.is_dangerous());
305    }
306
307    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
308    #[test]
309    fn test_fpu_guard_restores_state() {
310        let original_state = get_fpu_state();
311
312        {
313            let _guard = FpuGuard::new_protect_computation();
314            // State might be modified here
315        }
316
317        let restored_state = get_fpu_state();
318        assert_eq!(original_state.mxcsr, restored_state.mxcsr);
319    }
320}