hydroplane/
lib.rs

1//! `hydroplane`: float-agnostic, ISPC-style SPMD/SIMD infrastructure.
2//!
3//! Write one kernel generic over the scalar element ([`Scalar`]: `f32`, `f64`, `f16`, `bf16`);
4//! [`dispatch()`] runs it on the widest backend runtime CPU detection finds, else the portable
5//! [`ScalarBackend`] (also the rust-gpu/SPIR-V lowering target). Kernels never name a backend.
6//! Stable Rust, `no_std`-compatible, no SIMD-crate dependency; `f16`/`bf16` come from the
7//! `half` crate, re-exported as [`struct@f16`]/[`bf16`].
8#![cfg_attr(not(feature = "std"), no_std)]
9// The subgroup backend's SubgroupSize reader uses inline SPIR-V assembly, gated behind
10// `asm_experimental_arch`. This is the crate's only nightly requirement and applies solely to
11// the rust-gpu/SPIR-V build; every CPU target compiles on stable.
12#![cfg_attr(target_arch = "spirv", feature(asm_experimental_arch))]
13
14#[cfg(feature = "alloc")]
15extern crate alloc;
16
17/// Padding granularity: the widest lane count we target (AVX-512-FP16 is 32-wide f16). Every
18/// backend's lane count (1, 2, 4, 8, 16, 32) divides this, so a full-register loop never has
19/// a remainder, and a single inline `[T; MAX_LANES]` buffer holds any one register.
20pub const MAX_LANES: usize = 32;
21
22/// Upper bound on a backend's [`UNROLL`](Backend::UNROLL): the most independent accumulator chains
23/// any reduction unrolls to. Sizes the fixed `[init; MAX_UNROLL]` chain array, so a backend's
24/// `UNROLL` must not exceed it.
25pub const MAX_UNROLL: usize = 16;
26
27pub mod backend;
28pub mod arch;
29pub mod dispatch;
30// The runtime unroll cache; gone when `build.rs` resolved the factor at compile time.
31#[cfg(not(hp_resolved_unroll))]
32pub(crate) mod ilp;
33pub mod dense;
34pub mod matrix;
35pub mod scalar;
36pub mod varying;
37
38#[cfg(feature = "alloc")]
39pub mod cols;
40#[cfg(feature = "alloc")]
41pub mod soa;
42
43/// Opt-in `glam`-aware wide-vector helpers ([`Vec3Wide`](glam_ext::Vec3Wide) etc.). Behind the
44/// `glam` feature so the core stays geometry-free.
45#[cfg(feature = "glam")]
46pub mod glam_ext;
47
48pub use backend::{Backend, BackendAll, ScalarBackend};
49pub use dispatch::{Kernel, SimdDispatch, dispatch, run_scalar};
50
51/// The on-device entry point (rust-gpu / SPIR-V target): mirrors [`dispatch`], but branches on
52/// work size (subgroup-distributed vs. a single sequential invocation) instead of CPU ISA.
53#[cfg(target_arch = "spirv")]
54pub use backend::subgroup::dispatch_subgroup;
55pub use dense::{Diag, Mat, Side, Trans, Uplo, col_sums, fro_norm, gemv, row_sums};
56#[cfg(feature = "alloc")]
57pub use dense::{MatMut, gemm, potrf, syrk, trsm};
58pub use matrix::{
59    Accumulator, Layout, MatrixA, MatrixB, MatrixBackend, MatrixDispatch, MatrixKernel, Role, Tile,
60    Tiles, dispatch_matrix, run_matrix_scalar,
61};
62pub use scalar::{FloatScalar, IntScalar, Scalar};
63pub use varying::{ChunksExact, Varying, VaryingI32, VaryingU32, Mask, Gang};
64
65/// The `#[kernel]` attribute: write a [`Kernel`]/[`MatrixKernel`] as a plain generic function.
66/// Contexts, tuning flags (`tiny`, `noalias`, `unroll = N`), the generated `<name>_on` companion
67/// for calling one kernel from another without re-dispatching, and the `matrix` form are all
68/// documented on [the attribute itself](macro@kernel).
69pub use hydroplane_macros::kernel;
70
71/// `f16`/`bf16` element types (from the `half` crate), usable anywhere a [`Scalar`] is expected.
72pub use half::{bf16, f16};
73
74/// Re-export of the crate supplying [`Scalar`]'s numeric supertrait
75/// ([`FloatCore`](num_traits::float::FloatCore)), so generic kernels can name its traits without
76/// depending on `num-traits` themselves.
77pub use num_traits;
78
79#[cfg(feature = "alloc")]
80pub use cols::Cols;
81#[cfg(feature = "alloc")]
82pub use soa::Soa;
83
84#[cfg(feature = "glam")]
85pub use glam_ext::{GangGlamExt, Mat3Wide, MatWide, Vec3Wide};
86
87/// The combo-dispatch tier tokens and codes, reachable by `#[kernel]`-generated wrappers.
88/// Implementation detail: application code never names a backend, the generated match does.
89#[doc(hidden)]
90pub mod towers {
91    pub use crate::dispatch::tier::*;
92    pub use crate::ScalarBackend;
93    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
94    pub use crate::backend::{
95        avx1::Avx1, avx2::Avx2, avx512::Avx512, avx512bf16::Avx512Bf16, avx512fp16::Avx512Fp16,
96        sse4::Sse4,
97    };
98    #[cfg(target_arch = "aarch64")]
99    pub use crate::backend::neon::Neon;
100    #[cfg(target_arch = "aarch64")]
101    pub use crate::backend::sve::Sve;
102    #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
103    pub use crate::backend::wasm::Simd128;
104    #[cfg(all(target_arch = "wasm32", target_feature = "simd128", target_feature = "relaxed-simd"))]
105    pub use crate::backend::wasm::RelaxedSimd;
106}
107
108/// Test hook: the unroll factor in effect — `0` until the first dispatch resolves it via the runtime
109/// sweep, or the `build.rs`-baked constant when it was resolved at compile time. Not part of the
110/// stable surface.
111#[doc(hidden)]
112pub fn ilp_detected_for_test() -> u8 {
113    #[cfg(not(hp_resolved_unroll))]
114    {
115        ilp::cached()
116    }
117    #[cfg(hp_resolved_unroll)]
118    {
119        varying::STATIC_UNROLL as u8
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    /// A tiny generic kernel: are any of `xs` within `r` of the origin along one axis?
128    /// Written once, runs on any `Backend<T>`.
129    fn any_within<T: Scalar, S: Backend<T>>(s: S, xs: &[T], r: T) -> bool {
130        let lanes = s.lanes();
131        let rv = s.splat(r);
132        let mut i = 0;
133        while i + lanes <= xs.len() {
134            let x = s.load(&xs[i..i + lanes]);
135            let ax = s.max(x, s.neg(x)); // |x|
136            if s.any(s.le(ax, rv)) {
137                return true;
138            }
139            i += lanes;
140        }
141        // scalar tail
142        while i < xs.len() {
143            if xs[i] <= r && xs[i].neg() <= r {
144                return true;
145            }
146            i += 1;
147        }
148        false
149    }
150
151    #[test]
152    fn scalar_backend_smoke() {
153        let xs = [3.0f32, -2.0, 5.0, 0.5];
154        assert!(any_within(ScalarBackend, &xs, 1.0));
155        assert!(!any_within(ScalarBackend, &xs, 0.4));
156
157        let xd = [3.0f64, -2.0, 5.0, 0.5];
158        assert!(any_within(ScalarBackend, &xd, 1.0));
159        assert!(!any_within(ScalarBackend, &xd, 0.4));
160    }
161
162    /// A `Kernel` wrapping `any_within`, run through `dispatch` (whatever backend is best).
163    struct AnyWithin<'a, T: Scalar> {
164        xs: &'a [T],
165        r: T,
166    }
167    impl<T: Scalar> Kernel<T> for AnyWithin<'_, T> {
168        type Output = bool;
169        fn run<S: backend::BackendAll + Backend<T>>(self, simd: Gang<S>) -> bool {
170            any_within(simd.backend(), self.xs, self.r)
171        }
172    }
173
174    #[test]
175    fn dispatch_matches_scalar_oracle() {
176        // The dispatched backend (AVX2 when present) must agree with the scalar path.
177        let xs: Vec<f32> = (0..1000).map(|i| (i as f32 % 13.0) - 6.0).collect();
178        for &r in &[0.1f32, 0.5, 1.0, 3.0] {
179            let dispatched = dispatch(AnyWithin { xs: &xs, r });
180            let oracle = any_within(ScalarBackend, &xs, r);
181            assert_eq!(dispatched, oracle, "r={r}");
182        }
183    }
184}
hydroplane/lib.rs

hydroplane/
lib.rs
