NEON backend for VectorF32 on AArch64: one float32x4_t per vector
(Len * Packing <= 4), with a vectorised sin/cos port of the SSE helpers.

VectorF16 gets a native float16x8_t path when the target has ARMv8.2-FP16
arithmetic. Without it VectorF16 aliases VectorF32, which converts
fp16 <-> fp32 on Load/Store.

The scalar fallback gains the batched Dot that MatrixRowMajor relies on.
Test tiers added for armv8-a and armv8.2-a+fp16.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Jorijn van der Graaf 2026-10-06 17:13:54 +02:00
commit 0c55e76873
Signed by: jorijnvdgraaf
GPG key ID: 2937E59CDCC1BCFB
6 changed files with 825 additions and 17 deletions

View file

@ -11,6 +11,9 @@ module;
#ifdef __riscv_vector
#include <riscv_vector.h>
#endif
#if defined(__aarch64__) && defined(__ARM_NEON)
#include <arm_neon.h>
#endif
export module Crafter.Math:VectorF32;
import std;
import :Common;
@ -1931,10 +1934,318 @@ namespace Crafter {
return r;
}
};
#elif defined(__aarch64__) && defined(__ARM_NEON)
// AArch64 NEON implementation. VectorType is always a float32x4_t: NEON
// registers are 128 bits wide, so Common.cppm caps Len*Packing at 4
// floats, as on WASM. Shuffles and blends have compile-time lane maps and
// go through __builtin_shufflevector, which lets the compiler pick the
// cheapest of dup/ext/rev/zip/uzp/ins/tbl per map instead of always
// paying for a tbl and its index load.
export template <std::uint8_t Len, std::uint8_t Packing>
struct VectorF32 : public VectorBase<Len, Packing, float> {
template <std::uint8_t Len2, std::uint8_t Packing2>
friend struct VectorF32;
using Base = VectorBase<Len, Packing, float>;
static constexpr std::uint8_t NElems = Base::AlignmentElement;
static_assert(NElems == 4, "NEON VectorF32 assumes 4-lane vectors");
private:
// Source lane for each of the four output lanes: 0..3 pick from the
// first operand of Permute, 4..7 from the second. Every map names a
// real source lane for all four outputs, padding included, so no lane
// of a result is ever undefined.
using LaneMap = std::array<int, 4>;
template <LaneMap Lanes>
static constexpr float32x4_t Permute(float32x4_t a, float32x4_t b) {
return __builtin_shufflevector(a, b, Lanes[0], Lanes[1], Lanes[2], Lanes[3]);
}
// Per-pack shuffle; padding lanes stay where they are.
template <std::array<std::uint8_t, Len> ShuffleValues>
static consteval LaneMap GetShuffleLanes() {
LaneMap lanes{0, 1, 2, 3};
for (std::uint8_t p = 0; p < Packing; ++p)
for (std::uint8_t i = 0; i < Len; ++i)
lanes[p * Len + i] = p * Len + ShuffleValues[i];
return lanes;
}
// Per-pack blend: the second operand where the mask is set, the first
// elsewhere.
template <std::array<bool, Len> BlendValues>
static consteval LaneMap GetBlendLanes() {
LaneMap lanes{0, 1, 2, 3};
for (std::uint8_t p = 0; p < Packing; ++p)
for (std::uint8_t i = 0; i < Len; ++i)
if (BlendValues[i]) lanes[p * Len + i] += 4;
return lanes;
}
// Moves every sub-vector to its slot in a (ToLen, ToPacking) layout,
// carrying over the components both layouts have. The remaining lanes
// are left as they were.
template <std::uint8_t ToLen, std::uint8_t ToPacking>
static consteval LaneMap GetResizeLanes() {
LaneMap lanes{0, 1, 2, 3};
for (std::uint8_t p = 0; p < std::min(Packing, ToPacking); ++p)
for (std::uint8_t i = 0; i < std::min(Len, ToLen); ++i)
lanes[p * ToLen + i] = p * Len + i;
return lanes;
}
// Normalize keeps one scale per sub-vector in a VectorF32<1, N*Packing>;
// this spreads the scales of input `Index` over the Len lanes of each
// of its sub-vectors. Padding lanes repeat the last scale, so the map
// is a plain dup wherever it can be.
template <std::uint8_t Index>
static consteval LaneMap GetSpreadLanes() {
LaneMap lanes{};
for (std::uint8_t i = 0; i < 4; ++i)
lanes[i] = Index * Packing + std::min(i / Len, Packing - 1);
return lanes;
}
// Collapses every sub-vector of the BatchSize registers into the sum
// of its components: result lane i*Packing+p is the sum of lanes
// [p*Len, p*Len+Len) of m[i]. faddp adds the adjacent lane pairs of
// two registers in one go, so two levels of it reduce four 4-lane
// registers to four sums without a transpose.
static constexpr float32x4_t SumPacks(std::array<float32x4_t, Base::BatchSize> m) {
if constexpr (Packing == 1 && Len >= 3) {
if constexpr (Len == 3) {
// Lane 3 is padding and must not reach the sum.
for (float32x4_t& product : m) product = vsetq_lane_f32(0.0f, product, 3);
}
return vpaddq_f32(vpaddq_f32(m[0], m[1]), vpaddq_f32(m[2], m[3]));
} else if constexpr (Packing == 1 && Len == 2) {
return vpaddq_f32(
vcombine_f32(vget_low_f32(m[0]), vget_low_f32(m[1])),
vcombine_f32(vget_low_f32(m[2]), vget_low_f32(m[3])));
} else if constexpr (Packing == 1) {
return vcombine_f32(
vget_low_f32(vzip1q_f32(m[0], m[1])),
vget_low_f32(vzip1q_f32(m[2], m[3])));
} else if constexpr (Packing == 2 && Len == 2) {
return vpaddq_f32(m[0], m[1]);
} else if constexpr (Packing == 2) {
return vcombine_f32(vget_low_f32(m[0]), vget_low_f32(m[1]));
} else {
// Len == 1 with a lone input: its lanes already are the sums.
return m[0];
}
}
public:
constexpr VectorF32() = default;
constexpr VectorF32(float32x4_t vv) { this->v = vv; }
constexpr VectorF32(const float* vB) { Load(vB); }
constexpr VectorF32(const _Float16* vB) { Load(vB); }
constexpr VectorF32(float val) { this->v = vdupq_n_f32(val); }
constexpr void Load(const float* vB) { this->v = vld1q_f32(vB); }
constexpr void Store(float* vB) const { vst1q_f32(vB, this->v); }
// fp16 <-> fp32 conversion (fcvtl / fcvtn) is part of the AArch64
// baseline. The halves move through memcpy because arm_neon.h types
// its fp16 pointers as __fp16, which _Float16 may not alias.
constexpr void Load(const _Float16* vB) {
float16x4_t half;
std::memcpy(&half, vB, sizeof(half));
this->v = vcvt_f32_f16(half);
}
constexpr void Store(_Float16* vB) const {
float16x4_t half = vcvt_f16_f32(this->v);
std::memcpy(vB, &half, sizeof(half));
}
template<typename T>
constexpr std::array<T, NElems> Store() const {
std::array<T, NElems> r;
Store(r.data());
return r;
}
template <std::uint8_t BLen, std::uint8_t BPacking>
constexpr operator VectorF32<BLen, BPacking>() const {
return VectorF32<BLen, BPacking>(Permute<GetResizeLanes<BLen, BPacking>()>(this->v, this->v));
}
constexpr VectorF32<Len, Packing> operator+(VectorF32<Len, Packing> b) const { return VectorF32<Len, Packing>(vaddq_f32(this->v, b.v)); }
constexpr VectorF32<Len, Packing> operator-(VectorF32<Len, Packing> b) const { return VectorF32<Len, Packing>(vsubq_f32(this->v, b.v)); }
constexpr VectorF32<Len, Packing> operator*(VectorF32<Len, Packing> b) const { return VectorF32<Len, Packing>(vmulq_f32(this->v, b.v)); }
constexpr VectorF32<Len, Packing> operator/(VectorF32<Len, Packing> b) const { return VectorF32<Len, Packing>(vdivq_f32(this->v, b.v)); }
constexpr void operator+=(VectorF32<Len, Packing> b) { this->v = vaddq_f32(this->v, b.v); }
constexpr void operator-=(VectorF32<Len, Packing> b) { this->v = vsubq_f32(this->v, b.v); }
constexpr void operator*=(VectorF32<Len, Packing> b) { this->v = vmulq_f32(this->v, b.v); }
constexpr void operator/=(VectorF32<Len, Packing> b) { this->v = vdivq_f32(this->v, b.v); }
constexpr VectorF32<Len, Packing> operator+(float b) const { return *this + VectorF32<Len, Packing>(b); }
constexpr VectorF32<Len, Packing> operator-(float b) const { return *this - VectorF32<Len, Packing>(b); }
constexpr VectorF32<Len, Packing> operator*(float b) const { return *this * VectorF32<Len, Packing>(b); }
constexpr VectorF32<Len, Packing> operator/(float b) const { return *this / VectorF32<Len, Packing>(b); }
constexpr void operator+=(float b) { *this += VectorF32<Len, Packing>(b); }
constexpr void operator-=(float b) { *this -= VectorF32<Len, Packing>(b); }
constexpr void operator*=(float b) { *this *= VectorF32<Len, Packing>(b); }
constexpr void operator/=(float b) { *this /= VectorF32<Len, Packing>(b); }
constexpr VectorF32<Len, Packing> operator-() const { return VectorF32<Len, Packing>(vnegq_f32(this->v)); }
// Only the Len*Packing lanes in use take part; padding lanes can hold
// anything.
constexpr bool operator==(VectorF32<Len, Packing> b) const {
uint32x4_t equal = vceqq_f32(this->v, b.v);
if constexpr (Len * Packing == 4) {
return vminvq_u32(equal) == ~0u;
} else if constexpr (Len * Packing == 3) {
return vminvq_u32(vcopyq_laneq_u32(equal, 3, equal, 0)) == ~0u;
} else if constexpr (Len * Packing == 2) {
return vgetq_lane_u64(vreinterpretq_u64_u32(equal), 0) == ~0ull;
} else {
return vgetq_lane_u32(equal, 0) == ~0u;
}
}
constexpr bool operator!=(VectorF32<Len, Packing> b) const { return !(*this == b); }
template<std::uint32_t ExtractLen>
constexpr VectorF32<ExtractLen, Packing> ExtractLo() const {
return VectorF32<ExtractLen, Packing>(Permute<GetResizeLanes<ExtractLen, Packing>()>(this->v, this->v));
}
constexpr VectorF32<Len, Packing> Cos() const {
return VectorF32<Len, Packing>(Base::cos_f32x4(this->v));
}
constexpr VectorF32<Len, Packing> Sin() const {
return VectorF32<Len, Packing>(Base::sin_f32x4(this->v));
}
constexpr std::tuple<VectorF32<Len, Packing>, VectorF32<Len, Packing>> SinCos() const {
float32x4_t s, c;
Base::sincos_f32x4(this->v, s, c);
return { VectorF32<Len, Packing>(s), VectorF32<Len, Packing>(c) };
}
template <std::array<bool, Len> values>
constexpr VectorF32<Len, Packing> Negate() const {
constexpr std::array<float, NElems> mask = Base::template GetNegateMask<values>();
return VectorF32<Len, Packing>(vreinterpretq_f32_u32(veorq_u32(vreinterpretq_u32_f32(this->v), vreinterpretq_u32_f32(vld1q_f32(mask.data())))));
}
// vfmaq_f32(a, b, c) is a + b*c.
static constexpr VectorF32<Len, Packing> MulitplyAdd(VectorF32<Len, Packing> a, VectorF32<Len, Packing> b, VectorF32<Len, Packing> add) {
return VectorF32<Len, Packing>(vfmaq_f32(add.v, a.v, b.v));
}
static constexpr VectorF32<Len, Packing> MulitplySub(VectorF32<Len, Packing> a, VectorF32<Len, Packing> b, VectorF32<Len, Packing> sub) {
return VectorF32<Len, Packing>(vfmaq_f32(vnegq_f32(sub.v), a.v, b.v));
}
// a × b = (a * b.yzx - b * a.yzx).yzx, one shuffle less than building
// both yzx and zxy of each input. The rotations let the padding lane
// repeat y, which is two ext instructions; keeping that lane in place,
// as Shuffle does, takes three. vfmsq_f32(a, b, c) is a - b*c.
constexpr static VectorF32<Len, Packing> Cross(VectorF32<Len, Packing> a, VectorF32<Len, Packing> b) requires(Len == 3) {
constexpr LaneMap yzx{1, 2, 0, 1};
float32x4_t aYzx = Permute<yzx>(a.v, a.v);
float32x4_t bYzx = Permute<yzx>(b.v, b.v);
float32x4_t zxy = vfmsq_f32(vmulq_f32(a.v, bYzx), b.v, aYzx);
return VectorF32<Len, Packing>(Permute<yzx>(zxy, zxy));
}
template <const std::array<std::uint8_t, Len> ShuffleValues>
constexpr VectorF32<Len, Packing> Shuffle() const {
return VectorF32<Len, Packing>(Permute<GetShuffleLanes<ShuffleValues>()>(this->v, this->v));
}
template <std::array<bool, Len> ShuffleValues>
constexpr static VectorF32<Len, Packing> Blend(VectorF32<Len, Packing> a, VectorF32<Len, Packing> b) {
return VectorF32<Len, Packing>(Permute<GetBlendLanes<ShuffleValues>()>(a.v, b.v));
}
template<typename... Rest>
requires((std::is_same_v<Rest, VectorF32<Len, Packing>> && ...) &&
(1 + sizeof...(Rest) == VectorBase<Len, Packing, float>::BatchSize))
constexpr static auto LengthSq(VectorF32<Len, Packing> first, Rest... rest) {
constexpr std::uint8_t N = VectorBase<Len, Packing, float>::BatchSize;
std::array<VectorF32<Len, Packing>, N> args{ first, rest... };
std::array<float32x4_t, N> squares;
for (std::uint8_t i = 0; i < N; ++i) squares[i] = vmulq_f32(args[i].v, args[i].v);
return VectorF32<1, static_cast<std::uint8_t>(Packing * N)>(SumPacks(squares));
}
template<typename... Rest>
requires((std::is_same_v<Rest, VectorF32<Len, Packing>> && ...) &&
(1 + sizeof...(Rest) == VectorBase<Len, Packing, float>::BatchSize))
constexpr static auto Length(VectorF32<Len, Packing> first, Rest... rest) {
auto sq = LengthSq(first, rest...);
sq.v = vsqrtq_f32(sq.v);
return sq;
}
// Pairwise dot products across BatchSize pairs: Dot(a0, b0, a1, b1,
// ...) returns one lane per sub-vector pair.
template<typename... Rest>
requires((std::is_same_v<Rest, VectorF32<Len, Packing>> && ...) &&
(1 + sizeof...(Rest) == 2 * VectorBase<Len, Packing, float>::BatchSize))
constexpr static auto Dot(VectorF32<Len, Packing> first, Rest... rest) {
constexpr std::uint8_t N = VectorBase<Len, Packing, float>::BatchSize;
std::array<VectorF32<Len, Packing>, 2 * N> args{ first, rest... };
std::array<float32x4_t, N> products;
for (std::uint8_t i = 0; i < N; ++i) products[i] = vmulq_f32(args[2 * i].v, args[2 * i + 1].v);
return VectorF32<1, static_cast<std::uint8_t>(Packing * N)>(SumPacks(products));
}
// One sqrt and one divide cover the whole batch; each input then
// multiplies by its own lanes of the result.
template<typename... Rest>
requires((std::is_same_v<Rest, VectorF32<Len, Packing>> && ...) &&
(1 + sizeof...(Rest) == VectorBase<Len, Packing, float>::BatchSize))
constexpr static auto Normalize(VectorF32<Len, Packing> first, Rest... rest) {
constexpr std::uint8_t N = VectorBase<Len, Packing, float>::BatchSize;
std::array<VectorF32<Len, Packing>, N> args{ first, rest... };
float32x4_t scale = vdivq_f32(vdupq_n_f32(1.0f), Length(first, rest...).v);
return [&]<std::size_t... Is>(std::index_sequence<Is...>) {
return std::array<VectorF32<Len, Packing>, N>{
VectorF32<Len, Packing>(vmulq_f32(args[Is].v, Permute<GetSpreadLanes<Is>()>(scale, scale)))...
};
}(std::make_index_sequence<N>{});
}
constexpr static VectorF32<Len, Packing> Rotate(VectorF32<3, Packing> v, VectorF32<4, Packing> q) requires(Len == 3) {
VectorF32<3, Packing> qv = q;
VectorF32<3, Packing> qw = q.template Shuffle<{{3,3,3,3}}>();
VectorF32<3, Packing> t = Cross(qv, v) * 2.0f;
return MulitplyAdd(t, qw, v) + Cross(qv, t);
}
constexpr static VectorF32<3, Packing> RotatePivot(VectorF32<3, Packing> v, VectorF32<4, Packing> q, VectorF32<3, Packing> pivot) requires(Len == 3) {
VectorF32<3, Packing> translated = v - pivot;
return Rotate(translated, q) + pivot;
}
// With (roll, pitch, yaw) half angles in x, y, z:
// x = sx*cy*cz - cx*sy*sz y = cx*sy*cz + sx*cy*sz
// z = cx*cy*sz - sx*sy*cz w = cx*cy*cz + sx*sy*sz
// Each factor row below is one two-source shuffle of the sin (lanes
// 0..3) and cos (lanes 4..7) registers.
constexpr static VectorF32<4, Packing> QuanternionFromEuler(VectorF32<3, Packing> EulerHalf) requires(Len == 4) {
float32x4_t s, c;
Base::sincos_f32x4(EulerHalf.v, s, c);
float32x4_t left = vmulq_f32(
vmulq_f32(
Permute<LaneMap{0, 4, 4, 4}>(s, c), // sx cx cx cx
Permute<LaneMap{5, 1, 5, 5}>(s, c)), // cy sy cy cy
Permute<LaneMap{6, 6, 2, 6}>(s, c)); // cz cz sz cz
float32x4_t right = vmulq_f32(
Permute<LaneMap{4, 0, 0, 0}>(s, c), // cx sx sx sx
Permute<LaneMap{1, 5, 1, 1}>(s, c)); // sy cy sy sy
VectorF32<4, Packing> rightZ(Permute<LaneMap{2, 2, 6, 2}>(s, c)); // sz sz cz sz
rightZ = rightZ.template Negate<{{true, false, true, false}}>();
return VectorF32<4, Packing>(vfmaq_f32(left, right, rightZ.v));
}
};
#else
// Scalar software fallback for non-x86_64 targets. Future arches can swap
// in their own intrinsic implementation by adding an arch-specific branch
// above and gating this one out.
// Scalar software fallback for targets without an intrinsic backend above.
// Future arches can swap in their own intrinsic implementation by adding
// an arch-specific branch above and gating this one out.
export template <std::uint8_t Len, std::uint8_t Packing>
struct VectorF32 : public VectorBase<Len, Packing, float> {
template <std::uint8_t Len2, std::uint8_t Packing2>
@ -2136,6 +2447,28 @@ namespace Crafter {
return sq;
}
// Batched dot product over BatchSize (a, b) pairs: Dot(a0, b0, a1, b1,
// ...) returns one lane per pair. MatrixRowMajor relies on it, so
// without it this fallback (every target that has no intrinsic
// backend above) does not compile. Only the first Len lanes of each
// input are summed, as in the other backends.
template<typename... Rest>
requires((std::is_same_v<Rest, VectorF32<Len, Packing>> && ...) &&
(1 + sizeof...(Rest) == 2 * VectorBase<Len, Packing, float>::BatchSize))
constexpr static auto Dot(VectorF32<Len, Packing> first, Rest... rest) {
constexpr std::uint8_t N = VectorBase<Len, Packing, float>::BatchSize;
VectorF32<1, static_cast<std::uint8_t>(Packing * N)> r;
std::array<VectorF32<Len, Packing>, 2 * N> args{ first, rest... };
for (std::uint8_t i = 0; i < N; ++i)
for (std::uint8_t p = 0; p < Packing; ++p) {
float acc = 0.0f;
for (std::uint8_t k = 0; k < Len; ++k)
acc += args[2 * i].v[p * Len + k] * args[2 * i + 1].v[p * Len + k];
r.v[i * Packing + p] = acc;
}
return r;
}
template<typename... Rest>
requires((std::is_same_v<Rest, VectorF32<Len, Packing>> && ...) &&
(1 + sizeof...(Rest) == VectorBase<Len, Packing, float>::BatchSize))