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

use hydroplane::{BackendAll, FloatScalar, Backend, Gang, Kernel, SimdDispatch, dispatch, run_scalar};

fn oracle_dot<T: FloatScalar>(a: &[T], b: &[T]) -> f64 {
    a.iter()
        .zip(b)
        .map(|(&x, &y)| x.into_f64() * y.into_f64())
        .sum()
}

fn dot_zip_reduce<T: FloatScalar, S: Backend<T>>(g: Gang<S>, a: &[T], b: &[T]) -> f64 {
    g.zip_reduce(
        a,
        b,
        T::ZERO,
        T::ZERO,
        g.splat(T::ZERO),
        |acc, x, y| x.fma(y, acc),
        |p, q| p + q,
    )
    .reduce_sum()
    .into_f64()
}

fn sum_reduce<T: FloatScalar, S: Backend<T>>(g: Gang<S>, a: &[T]) -> f64 {
    g.reduce(a, T::ZERO, g.splat(T::ZERO), |acc, x| acc + x, |p, q| p + q)
        .reduce_sum()
        .into_f64()
}

// The unroll factor K is chosen by dispatch (`Unroll<S, K>`), so a kernel can't pin it; the size
// sweep crosses the K*lanes() window and tail boundaries for whichever K this core resolved.
struct Variants<'a, T: FloatScalar> {
    a: &'a [T],
    b: &'a [T],
}
impl<T: FloatScalar> Kernel<T> for Variants<'_, T> {
    type Output = [f64; 2];
    fn run<S: BackendAll + Backend<T>>(self, g: Gang<S>) -> [f64; 2] {
        [dot_zip_reduce(g, self.a, self.b), sum_reduce(g, self.a)]
    }
}

const SIZES: &[usize] = &[0, 1, 3, 5, 8, 15, 16, 31, 64, 257];

fn check_all<T: FloatScalar + SimdDispatch>() {
    for &n in SIZES {
        let a: Vec<T> = (0..n)
            .map(|i| T::from_f64((i % 13) as f64 * 0.25 - 1.5))
            .collect();
        let b: Vec<T> = (0..n)
            .map(|i| T::from_f64((i % 7) as f64 * 0.5 - 1.0))
            .collect();

        let want = oracle_dot(&a, &b);
        let sum_want: f64 = a.iter().map(|&x| x.into_f64()).sum();
        let tol = 1e-3 * (1.0 + want.abs());
        let sum_tol = 1e-3 * (1.0 + sum_want.abs());

        for (label, outs) in [
            ("scalar", run_scalar(Variants { a: &a, b: &b })),
            ("dispatch", dispatch(Variants { a: &a, b: &b })),
        ] {
            assert!(
                (outs[0] - want).abs() <= tol,
                "{label} dot mismatch n={n}: got {}, want {want}",
                outs[0]
            );
            assert!(
                (outs[1] - sum_want).abs() <= sum_tol,
                "{label} sum mismatch n={n}: got {}, want {sum_want}",
                outs[1]
            );
        }
    }
}

#[test]
fn zip_reduce_matches_oracle_f32() {
    check_all::<f32>();
}

#[test]
fn zip_reduce_matches_oracle_f64() {
    check_all::<f64>();
}

struct DetectProbe;
impl Kernel<f32> for DetectProbe {
    type Output = f64;
    fn run<S: BackendAll + Backend<f32>>(self, g: Gang<S>) -> f64 {
        let a: Vec<f32> = (0..257).map(|i| (i % 13) as f32 * 0.25 - 1.5).collect();
        dot_zip_reduce(g, &a, &a)
    }
}

#[test]
fn detection_resolves_and_is_stable() {
    let first = dispatch(DetectProbe);
    let second = dispatch(DetectProbe);
    assert_eq!(first.to_bits(), second.to_bits());
    let k = hydroplane::ilp_detected_for_test();
    assert!(
        matches!(k, 1 | 2 | 4 | 8 | 12 | 16),
        "cached K not a candidate factor: {k}"
    );
}