hydroplane 0.1.0

Floating but fast: float-agnostic, ISPC-style SPMD/SIMD infrastructure — write one kernel generic over the scalar element (f32/f64/f16/bf16) and run it with runtime ISA dispatch.
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

//! AVX-512-BF16 (x86_64): hardware `bf16`↔`f32` conversion and the `vdpbf16ps` dot product.
//! `bf16` has no native ALU (convert + dot only), so element-wise compute still happens in `f32`.
#![allow(unsafe_op_in_unsafe_fn, clippy::missing_safety_doc)]

#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
use half::bf16;

/// Widen 16 contiguous `bf16` (unaligned) to `f32x16`.
#[target_feature(enable = "avx512bf16,avx512f,avx512bw")]
#[inline]
pub unsafe fn widen(p: *const bf16) -> __m512 {
    let raw: __m256i = core::ptr::read_unaligned(p as *const __m256i);
    _mm512_cvtpbh_ps(core::mem::transmute::<__m256i, __m256bh>(raw))
}

/// Narrow `f32x16` to 16 contiguous `bf16` (unaligned), round-to-nearest-even in hardware.
#[target_feature(enable = "avx512bf16,avx512f,avx512bw")]
#[inline]
pub unsafe fn narrow(v: __m512, p: *mut bf16) {
    let packed: __m256bh = _mm512_cvtneps_pbh(v);
    core::ptr::write_unaligned(p as *mut __m256i, core::mem::transmute::<__m256bh, __m256i>(packed));
}

/// VNNI-pack two K-rows of a 16-wide column block: result lane `n` holds `lo[n]` in its low 16
/// bits and `hi[n]` in its high 16 bits. Built from `vpmovzxwd`, whose element-to-lane mapping is
/// defined, so the layout is correct by construction.
#[target_feature(enable = "avx512bf16,avx512f,avx512bw")]
#[inline]
pub unsafe fn pack_pair(lo: *const bf16, hi: *const bf16) -> __m512i {
    let l: __m256i = core::ptr::read_unaligned(lo as *const __m256i);
    let h: __m256i = core::ptr::read_unaligned(hi as *const __m256i);
    _mm512_or_si512(
        _mm512_cvtepu16_epi32(l),
        _mm512_slli_epi32::<16>(_mm512_cvtepu16_epi32(h)),
    )
}

/// Broadcast one A-pair to every lane: each 32-bit lane holds `a0` (low 16) and `a1` (high 16).
#[target_feature(enable = "avx512bf16,avx512f,avx512bw")]
#[inline]
pub unsafe fn bcast_pair(a0: bf16, a1: bf16) -> __m512i {
    _mm512_set1_epi32((((a1.to_bits() as u32) << 16) | a0.to_bits() as u32) as i32)
}

/// `acc[n] += a[n].lo·b[n].lo + a[n].hi·b[n].hi` lane-wise (`vdpbf16ps`): the bf16 products are
/// exact in f32, then accumulated. Pair `a` from [`bcast_pair`] with `b` from [`pack_pair`].
#[target_feature(enable = "avx512bf16,avx512f,avx512bw")]
#[inline]
pub unsafe fn dp(acc: __m512, a: __m512i, b: __m512i) -> __m512 {
    _mm512_dpbf16_ps(
        acc,
        core::mem::transmute::<__m512i, __m512bh>(a),
        core::mem::transmute::<__m512i, __m512bh>(b),
    )
}