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
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

//! AVX-512-FP16 (x86_64): 32-wide half-precision arithmetic via raw `asm!`, since the intrinsics
//! and the `f16` primitive are still unstable. Data rides in [`__m512i`], a stable 512-bit
//! carrier holding the 32x`half::f16` bit pattern.
#![allow(unsafe_op_in_unsafe_fn, clippy::missing_safety_doc)]

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

/// Extended compare predicates (ordered, quiet) for `vcmpph`, matching the `_CMP_*_OQ` imms.
pub const CMP_LT_OQ: i32 = 0x11;
pub const CMP_LE_OQ: i32 = 0x12;
pub const CMP_GE_OQ: i32 = 0x1d;
pub const CMP_GT_OQ: i32 = 0x1e;

/// Load 32 contiguous `f16` (unaligned) into a register.
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn load(p: *const f16) -> __m512i {
    core::ptr::read_unaligned(p as *const __m512i)
}

/// Store a register to 32 contiguous `f16` (unaligned).
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn store(p: *mut f16, v: __m512i) {
    core::ptr::write_unaligned(p as *mut __m512i, v)
}

/// Broadcast one `f16` to all 32 lanes.
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn splat(v: f16) -> __m512i {
    _mm512_set1_epi16(v.to_bits() as i16)
}

/// Negate every lane by flipping its sign bit (exact for `±0`/NaN, unlike `0 - x`).
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn neg(a: __m512i) -> __m512i {
    _mm512_xor_si512(a, _mm512_set1_epi16(0x8000u16 as i16))
}

macro_rules! binop {
    ($name:ident, $mnem:literal) => {
        #[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
        #[inline]
        pub unsafe fn $name(a: __m512i, b: __m512i) -> __m512i {
            let r;
            asm!(
                concat!($mnem, " {r}, {a}, {b}"),
                r = lateout(zmm_reg) r, a = in(zmm_reg) a, b = in(zmm_reg) b,
                options(pure, nomem, nostack, preserves_flags),
            );
            r
        }
    };
}

binop!(add, "vaddph");
binop!(sub, "vsubph");
binop!(mul, "vmulph");
binop!(div, "vdivph");

/// IEEE minimumNumber: `vminph` already yields `b` when `a` is NaN; the masked move patches the
/// b-is-NaN case (bare `vminph` would return the NaN). The `3` imm is `_CMP_UNORD_Q`.
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn min(a: __m512i, b: __m512i) -> __m512i {
    let r;
    asm!(
        "vminph {r}, {a}, {b}",
        "vcmpph {k}, {b}, {b}, 3",
        "vmovdqu16 {r}{{{k}}}, {a}",
        r = out(zmm_reg) r, a = in(zmm_reg) a, b = in(zmm_reg) b, k = out(kreg) _,
        options(pure, nomem, nostack, preserves_flags),
    );
    r
}

/// IEEE maximumNumber; see [`min`].
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn max(a: __m512i, b: __m512i) -> __m512i {
    let r;
    asm!(
        "vmaxph {r}, {a}, {b}",
        "vcmpph {k}, {b}, {b}, 3",
        "vmovdqu16 {r}{{{k}}}, {a}",
        r = out(zmm_reg) r, a = in(zmm_reg) a, b = in(zmm_reg) b, k = out(kreg) _,
        options(pure, nomem, nostack, preserves_flags),
    );
    r
}

#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn sqrt(a: __m512i) -> __m512i {
    let r;
    asm!(
        "vsqrtph {r}, {a}",
        r = lateout(zmm_reg) r, a = in(zmm_reg) a,
        options(pure, nomem, nostack, preserves_flags),
    );
    r
}

/// Fused multiply-add: `a * b + c`, lane-wise, single rounding.
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn fma(a: __m512i, b: __m512i, c: __m512i) -> __m512i {
    // vfmadd213ph dst, s2, s3 => dst = s2*dst + s3; with dst=a, s2=b, s3=c => a*b + c.
    let mut dst = a;
    asm!(
        "vfmadd213ph {dst}, {b}, {c}",
        dst = inout(zmm_reg) dst, b = in(zmm_reg) b, c = in(zmm_reg) c,
        options(pure, nomem, nostack, preserves_flags),
    );
    dst
}

/// Lane-wise compare to a 32-bit mask (one bit per lane), with an ordered-quiet predicate.
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn cmp<const IMM: i32>(a: __m512i, b: __m512i) -> u32 {
    let m: u32;
    asm!(
        "vcmpph k1, {a}, {b}, {imm}",
        "kmovd {m:e}, k1",
        a = in(zmm_reg) a, b = in(zmm_reg) b, m = lateout(reg) m, imm = const IMM,
        out("k1") _,
        options(pure, nomem, nostack, preserves_flags),
    );
    m
}

/// Blend by mask, lane-wise: `mask_bit ? a : b` (16-bit granularity via `vpblendmw`).
#[target_feature(enable = "avx512fp16,avx512f,avx512bw")]
#[inline]
pub unsafe fn select(mask: u32, a: __m512i, b: __m512i) -> __m512i {
    // vpblendmw dst{k}, s1, s2 => dst = k ? s2 : s1; with s1=b, s2=a => k ? a : b.
    let r;
    asm!(
        "kmovd k1, {m:e}",
        "vpblendmw {r}{{k1}}, {b}, {a}",
        m = in(reg) mask, a = in(zmm_reg) a, b = in(zmm_reg) b, r = lateout(zmm_reg) r,
        out("k1") _,
        options(pure, nomem, nostack, preserves_flags),
    );
    r
}