rlx_runtime/
nan_check.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! NaN/inf check epilogue (plan #18).
6//!
7//! Borrowed from MAX's `Mogg/MOGGKernelAPI/nan_check.mojo` pattern.
8//! When the `nan-check` Cargo feature is on, [`scan`] reports the
9//! first NaN or inf in a slice — useful as a debug epilogue on
10//! every output buffer to localize precision blow-ups to the op
11//! that introduced them.
12//!
13//! Always present in the API surface so callers can compile against
14//! it; the feature flag controls whether it's a real scan or a
15//! no-op (returns `Ok(())` immediately).
16
17/// What was found in a buffer that fails the check.
18#[derive(Debug, Clone, Copy)]
19pub enum BadValue {
20    Nan,
21    PosInf,
22    NegInf,
23}
24
25#[derive(Debug)]
26pub struct NanCheckError {
27    pub kind: BadValue,
28    pub index: usize,
29    pub label: String,
30}
31
32impl std::fmt::Display for NanCheckError {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        let what = match self.kind {
35            BadValue::Nan => "NaN",
36            BadValue::PosInf => "+inf",
37            BadValue::NegInf => "-inf",
38        };
39        write!(f, "{} at index {} of `{}`", what, self.index, self.label)
40    }
41}
42
43impl std::error::Error for NanCheckError {}
44
45/// Scan `data` for the first NaN or infinity. With the `nan-check`
46/// feature OFF, returns `Ok(())` immediately (the optimizer
47/// eliminates the call). With it ON, walks the slice — the cost is
48/// O(n) but only paid when a caller opts in.
49#[cfg(feature = "nan-check")]
50#[inline(always)]
51pub fn scan(label: &str, data: &[f32]) -> Result<(), NanCheckError> {
52    for (i, &v) in data.iter().enumerate() {
53        if v.is_nan() {
54            return Err(NanCheckError {
55                kind: BadValue::Nan,
56                index: i,
57                label: label.to_string(),
58            });
59        }
60        if v.is_infinite() {
61            let kind = if v > 0.0 {
62                BadValue::PosInf
63            } else {
64                BadValue::NegInf
65            };
66            return Err(NanCheckError {
67                kind,
68                index: i,
69                label: label.to_string(),
70            });
71        }
72    }
73    Ok(())
74}
75
76#[cfg(not(feature = "nan-check"))]
77#[inline(always)]
78pub fn scan(_label: &str, _data: &[f32]) -> Result<(), NanCheckError> {
79    Ok(())
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn clean_data_passes() {
88        let data = [1.0, 2.0, -3.5, 0.0];
89        assert!(scan("clean", &data).is_ok());
90    }
91
92    #[cfg(feature = "nan-check")]
93    #[test]
94    fn detects_nan() {
95        let data = [1.0, f32::NAN, 3.0];
96        let err = scan("nan", &data).unwrap_err();
97        assert!(matches!(err.kind, BadValue::Nan));
98        assert_eq!(err.index, 1);
99    }
100
101    #[cfg(feature = "nan-check")]
102    #[test]
103    fn detects_pos_inf() {
104        let data = [f32::INFINITY, 0.0];
105        let err = scan("inf", &data).unwrap_err();
106        assert!(matches!(err.kind, BadValue::PosInf));
107        assert_eq!(err.index, 0);
108    }
109}
rlx_runtime/nan_check.rs

rlx_runtime/
nan_check.rs
