himada-dispatch 0.1.1

Adaptive SIMD dispatch for Himada — auto-selects fastest kernel at runtime
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

//! Auto-precision — automatically select f16, bf16, or f32 for operations
//! based on accuracy requirements and hardware support.

/// Numeric precision level.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Precision {
    F64,
    F32,
    F16,
    Bf16,
}

/// Autoprecision configuration.
pub struct AutoPrecision {
    /// Preferred precision.
    pub preferred: Precision,
    /// Allow downcasting to lower precision for speed.
    pub allow_downcast: bool,
    /// Minimum acceptable precision.
    pub minimum: Precision,
}

impl AutoPrecision {
    pub fn new() -> Self {
        Self {
            preferred: Precision::F32,
            allow_downcast: true,
            minimum: Precision::F32,
        }
    }

    /// Select the actual precision to use based on hardware support.
    pub fn select(&self, supports_f16: bool, supports_bf16: bool) -> Precision {
        if self.allow_downcast {
            match self.preferred {
                Precision::F64 => Precision::F64,
                Precision::F32 => Precision::F32,
                Precision::F16 => {
                    if supports_f16 { Precision::F16 }
                    else if supports_bf16 { Precision::Bf16 }
                    else { Precision::F32 }
                }
                Precision::Bf16 => {
                    if supports_bf16 { Precision::Bf16 }
                    else if supports_f16 { Precision::F16 }
                    else { Precision::F32 }
                }
            }
        } else {
            self.preferred
        }
    }
}

/// Check if the CPU supports f16 (FP16C) or bf16 (AVX512_BF16 / BF16).
pub fn check_f16_support() -> (bool, bool) {
    #[cfg(target_arch = "x86_64")]
    {
        let mut f16c = false;
        let mut bf16 = false;
        #[cfg(target_feature = "f16c")]
        { f16c = true; }
        #[cfg(target_feature = "avx512bf16")]
        { bf16 = true; }
        (f16c, bf16)
    }
    #[cfg(target_arch = "aarch64")]
    {
        // ARM has native half-precision in NEON
        (true, false)
    }
    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
    {
        (false, false)
    }
}

/// Runtime conversion of f32 slices to f16 storage (half-precision).
/// Uses the `half` crate if available, otherwise truncates to f32.
pub fn f32_to_f16(src: &[f32]) -> Vec<u16> {
    #[cfg(feature = "half")]
    {
        src.iter().map(|&x| half::f16::from_f32(x).to_bits()).collect()
    }
    #[cfg(not(feature = "half"))]
    {
        // Lossy truncation as placeholder
        src.iter().map(|&x| {
            let bits = x.to_bits();
            (bits >> 16) as u16
        }).collect()
    }
}

/// Runtime conversion of f16 storage back to f32.
pub fn f16_to_f32(src: &[u16]) -> Vec<f32> {
    #[cfg(feature = "half")]
    {
        src.iter().map(|&x| half::f16::from_bits(x).to_f32()).collect()
    }
    #[cfg(not(feature = "half"))]
    {
        src.iter().map(|&x| f32::from_bits((x as u32) << 16)).collect()
    }
}