numr 0.5.2

High-performance numerical computing with multi-backend GPU acceleration (CPU/CUDA/WebGPU)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

//! Shared SIMD mathematical functions
//!
//! This module provides optimized SIMD implementations of transcendental functions
//! (exp, tanh, log, sin, cos, tan) that are used across multiple SIMD kernel modules.
//! By centralizing these implementations, we ensure consistency and eliminate code
//! duplication.
//!
//! # Supported Functions
//!
//! | Function | f32 | f64 | Algorithm |
//! |----------|-----|-----|-----------|
//! | exp      | ✓   | ✓   | Range reduction + Taylor series |
//! | tanh     | ✓   | ✓   | Based on exp: (e^2x - 1)/(e^2x + 1) |
//! | log      | ✓   | ✓   | Exponent extraction + minimax polynomial |
//! | sin      | ✓   | ✓   | Range reduction + Taylor series |
//! | cos      | ✓   | ✓   | sin(x + π/2) |
//! | tan      | ✓   | ✓   | Range reduction + minimax polynomial |
//!
//! # Architecture Support
//!
//! | ISA | File | Vector Width | Notes |
//! |-----|------|--------------|-------|
//! | AVX2+FMA | avx2.rs | 256-bit | 8 f32 / 4 f64 |
//! | AVX-512F | avx512.rs | 512-bit | 16 f32 / 8 f64, native 64-bit ops |
//!
//! # Accuracy
//!
//! These approximations prioritize speed over full IEEE precision:
//! - Relative error: < 1e-6 for f32, < 1e-12 for f64
//! - Valid input range: [-88, 88] for f32 exp, [-709, 709] for f64 exp
//! - Trigonometric functions: Accurate for |x| < 2^20; larger inputs may lose precision
//!
//! # Design Notes
//!
//! All polynomial coefficients are centralized in `common.rs` to ensure algorithm
//! consistency between AVX2 and AVX-512 implementations. The implementations differ
//! only in:
//! - Intrinsic prefixes (`_mm256` vs `_mm512`)
//! - Vector types (`__m256` vs `__m512`)
//! - Mask handling (AVX2 uses blendv, AVX-512 uses native masks)
//! - 64-bit operations (AVX-512 has native support, AVX2 requires workarounds)

pub mod common;

#[cfg(target_arch = "x86_64")]
pub mod avx2;
#[cfg(target_arch = "x86_64")]
pub mod avx512;

#[cfg(target_arch = "aarch64")]
pub mod aarch64;

#[cfg(test)]
mod tests {
    #[cfg(target_arch = "x86_64")]
    use super::*;

    #[test]
    #[cfg(target_arch = "x86_64")]
    fn test_exp_f32_accuracy() {
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            return;
        }

        let inputs: Vec<f32> = (-40..=40).map(|x| x as f32 * 0.1).collect();
        let mut outputs = vec![0.0f32; inputs.len()];

        unsafe {
            for (i, &x) in inputs.iter().enumerate() {
                let v = std::arch::x86_64::_mm256_set1_ps(x);
                let result = avx2::exp_f32(v);
                let mut arr = [0.0f32; 8];
                std::arch::x86_64::_mm256_storeu_ps(arr.as_mut_ptr(), result);
                outputs[i] = arr[0];
            }
        }

        for (i, (&input, &output)) in inputs.iter().zip(outputs.iter()).enumerate() {
            let expected = input.exp();
            let rel_err = (output - expected).abs() / expected.abs().max(1e-10);
            assert!(
                rel_err < 1e-5,
                "exp({}) = {} (expected {}), rel_err = {} at index {}",
                input,
                output,
                expected,
                rel_err,
                i
            );
        }
    }

    #[test]
    #[cfg(target_arch = "x86_64")]
    fn test_tanh_f32_accuracy() {
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            return;
        }

        let inputs: Vec<f32> = (-30..=30).map(|x| x as f32 * 0.1).collect();
        let mut outputs = vec![0.0f32; inputs.len()];

        unsafe {
            for (i, &x) in inputs.iter().enumerate() {
                let v = std::arch::x86_64::_mm256_set1_ps(x);
                let result = avx2::tanh_f32(v);
                let mut arr = [0.0f32; 8];
                std::arch::x86_64::_mm256_storeu_ps(arr.as_mut_ptr(), result);
                outputs[i] = arr[0];
            }
        }

        for (&input, &output) in inputs.iter().zip(outputs.iter()) {
            let expected = input.tanh();
            let abs_err = (output - expected).abs();
            assert!(
                abs_err < 1e-5,
                "tanh({}) = {} (expected {}), abs_err = {}",
                input,
                output,
                expected,
                abs_err
            );
        }
    }
}