f8 0.2.0

A no_std, one-byte UNORM with exact rounding, saturating arithmetic, and SIMD conversion
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

//! AVX2 encoding. Compiled only for x86-64; all other targets stay in Rust.

use crate::f8;
use core::arch::{asm, x86_64::_mm256_set1_epi32};

#[cfg(target_feature = "avx2")]
#[inline]
pub(crate) fn available() -> bool {
    true
}

#[cfg(not(target_feature = "avx2"))]
pub(crate) fn available() -> bool {
    use core::{
        arch::x86_64::{__cpuid, __cpuid_count, _xgetbv},
        sync::atomic::{AtomicU8, Ordering},
    };

    // 0 = unprobed, 1 = unavailable, 2 = available. Races only repeat the probe;
    // no other data is published, so relaxed ordering is sufficient.
    static SUPPORT: AtomicU8 = AtomicU8::new(0);
    let cached = SUPPORT.load(Ordering::Relaxed);
    if cached != 0 {
        return cached == 2;
    }
    // SAFETY: CPUID exists on x86-64. XGETBV is executed only with XSAVE and
    // OSXSAVE enabled. XCR0 must enable both XMM and YMM state before AVX2 use.
    let supported = unsafe {
        __cpuid(0).eax >= 7
            && __cpuid(1).ecx & 0x1c00_0000 == 0x1c00_0000
            && _xgetbv(0) & 6 == 6
            && __cpuid_count(7, 0).ebx & (1 << 5) != 0
    };
    SUPPORT.store(if supported { 2 } else { 1 }, Ordering::Relaxed);
    supported
}

/// The caller must establish AVX2 support and equal slice lengths.
#[target_feature(enable = "avx2")]
pub(crate) unsafe fn encode(src: &[f32], dst: &mut [f8]) {
    // SAFETY: AVX2 is required by this function. Only complete eight-element
    // chunks are accessed; unaligned loads/stores are intentional. The slices
    // are valid and disjoint by their shared/exclusive Rust borrows. Every
    // register modified by the assembly is declared as an output.
    unsafe {
        let one = _mm256_set1_epi32(1);
        let min = _mm256_set1_epi32(0x3b00_0000);
        let max = _mm256_set1_epi32(0x3f80_0000);
        let infinity = _mm256_set1_epi32(0x7f80_0000);
        let mantissa = _mm256_set1_epi32(0x007f_ffff);
        let hidden = _mm256_set1_epi32(0x0080_0000);
        let bias = _mm256_set1_epi32(149);
        let end = src.len() / 8 * 8;
        let mut offset = 0;
        while offset < end {
            // This is the scalar guard/sticky-bit algorithm in eight lanes.
            // Integer classification also handles signaling NaNs without
            // touching MXCSR, and variable shifts preserve exact bin edges.
            asm!(
                "vmovdqu {value}, [{src}]",
                "vpsrad {invalid}, {value}, 31",
                "vpcmpgtd {temp}, {value}, {infinity}",
                "vpor {invalid}, {invalid}, {temp}",
                "vpandn {value}, {invalid}, {value}",
                "vpminud {value}, {value}, {max}",
                "vpmaxud {value}, {value}, {min}",
                "vpsrld {shift}, {value}, 23",
                "vpsubd {shift}, {bias}, {shift}",
                "vpand {value}, {value}, {mantissa}",
                "vpor {value}, {value}, {hidden}",
                "vpslld {temp}, {value}, 8",
                "vpsubd {value}, {temp}, {value}",
                "vpsrlvd {guard}, {value}, {shift}",
                "vpsllvd {temp}, {one}, {shift}",
                "vpsubd {temp}, {temp}, {one}",
                "vpand {value}, {value}, {temp}",
                "vpxor {temp}, {temp}, {temp}",
                "vpcmpeqd {value}, {value}, {temp}",
                "vpandn {value}, {value}, {one}",
                "vpsrld {invalid}, {guard}, 1",
                "vpor {value}, {value}, {invalid}",
                "vpand {value}, {value}, {guard}",
                "vpand {value}, {value}, {one}",
                "vpaddd {value}, {invalid}, {value}",
                "vextracti128 {temp:x}, {value}, 1",
                "vpackusdw {value:x}, {value:x}, {temp:x}",
                "vpackuswb {value:x}, {value:x}, {value:x}",
                "vmovq [{dst}], {value:x}",
                src = in(reg) src.as_ptr().add(offset),
                dst = in(reg) dst.as_mut_ptr().add(offset),
                one = in(ymm_reg) one,
                min = in(ymm_reg) min,
                max = in(ymm_reg) max,
                infinity = in(ymm_reg) infinity,
                mantissa = in(ymm_reg) mantissa,
                hidden = in(ymm_reg) hidden,
                bias = in(ymm_reg) bias,
                value = out(ymm_reg) _,
                invalid = out(ymm_reg) _,
                temp = out(ymm_reg) _,
                shift = out(ymm_reg) _,
                guard = out(ymm_reg) _,
                options(nostack, preserves_flags),
            );
            offset += 8;
        }
        for (&value, out) in src[end..].iter().zip(&mut dst[end..]) {
            *out = f8::from_f32(value);
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate std;

    #[test]
    fn dispatch_agrees_with_standard_library_detection() {
        assert_eq!(super::available(), std::is_x86_feature_detected!("avx2"));
        assert_eq!(super::available(), std::is_x86_feature_detected!("avx2"));
    }
}