numkong 7.0.0

package/include/numkong/curved/genoa.h

@@ -0,0 +1,182 @@
+/**
+ *  @brief SIMD-accelerated Curved Space Similarity for Genoa.
+ *  @file include/numkong/curved/genoa.h
+ *  @author Ash Vardanian
+ *  @date January 14, 2026
+ *
+ *  @sa include/numkong/curved.h
+ *
+ *  Implements bf16 bilinear forms using AVX-512 with BF16 extensions.
+ */
+#ifndef NK_CURVED_GENOA_H
+#define NK_CURVED_GENOA_H
+
+#if NK_TARGET_X86_
+#if NK_TARGET_GENOA
+
+#include "numkong/types.h"
+#include "numkong/spatial/genoa.h"  // `nk_substract_bf16x32_genoa_`
+#include "numkong/reduce/skylake.h" // `nk_reduce_add_f32x16_skylake_`
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if defined(__clang__)
+#pragma clang attribute push(                                                                        \
+    __attribute__((target("avx2,avx512f,avx512vl,avx512bw,avx512dq,avx512bf16,f16c,fma,bmi,bmi2"))), \
+    apply_to = function)
+#elif defined(__GNUC__)
+#pragma GCC push_options
+#pragma GCC target("avx2", "avx512f", "avx512vl", "avx512bw", "avx512dq", "avx512bf16", "f16c", "fma", "bmi", "bmi2")
+#endif
+
+NK_PUBLIC void nk_bilinear_bf16_genoa(nk_bf16_t const *a, nk_bf16_t const *b, nk_bf16_t const *c, nk_size_t n,
+                                      nk_f32_t *result) {
+    nk_size_t const tail_length = n % 32;
+    nk_size_t const tail_start = n - tail_length;
+    __mmask32 const tail_mask = (__mmask32)_bzhi_u32(0xFFFFFFFF, tail_length);
+    __m512 sum_f32x16 = _mm512_setzero_ps();
+
+    for (nk_size_t i = 0; i != n; ++i) {
+        nk_f32_t a_f32;
+        nk_bf16_to_f32_serial(a + i, &a_f32);
+        __m512 a_f32x16 = _mm512_set1_ps(a_f32);
+        __m512 cb_j_f32x16 = _mm512_setzero_ps();
+        __m512i b_bf16x32, c_bf16x32;
+        nk_size_t j = 0;
+
+    nk_bilinear_bf16_genoa_cycle:
+        if (j + 32 <= n) {
+            b_bf16x32 = _mm512_loadu_epi16(b + j);
+            c_bf16x32 = _mm512_loadu_epi16(c + i * n + j);
+        }
+        else {
+            b_bf16x32 = _mm512_maskz_loadu_epi16(tail_mask, b + tail_start);
+            c_bf16x32 = _mm512_maskz_loadu_epi16(tail_mask, c + i * n + tail_start);
+        }
+        cb_j_f32x16 = _mm512_dpbf16_ps(cb_j_f32x16, nk_m512bh_from_m512i_(b_bf16x32), nk_m512bh_from_m512i_(c_bf16x32));
+        j += 32;
+        if (j < n) goto nk_bilinear_bf16_genoa_cycle;
+        sum_f32x16 = _mm512_fmadd_ps(a_f32x16, cb_j_f32x16, sum_f32x16);
+    }
+
+    *result = _mm512_reduce_add_ps(sum_f32x16);
+}
+
+NK_PUBLIC void nk_mahalanobis_bf16_genoa(nk_bf16_t const *a, nk_bf16_t const *b, nk_bf16_t const *c, nk_size_t n,
+                                         nk_f32_t *result) {
+    nk_size_t const tail_length = n % 32;
+    nk_size_t const tail_start = n - tail_length;
+    __mmask32 const tail_mask = (__mmask32)_bzhi_u32(0xFFFFFFFF, tail_length);
+    __m512 sum_f32x16 = _mm512_setzero_ps();
+
+    for (nk_size_t i = 0; i != n; ++i) {
+        nk_f32_t a_i, b_i;
+        nk_bf16_to_f32_serial(a + i, &a_i);
+        nk_bf16_to_f32_serial(b + i, &b_i);
+        __m512 diff_i_f32x16 = _mm512_set1_ps(a_i - b_i);
+        __m512 cdiff_j_f32x16 = _mm512_setzero_ps();
+        __m512i a_j_bf16x32, b_j_bf16x32, diff_j_bf16x32, c_bf16x32;
+        nk_size_t j = 0;
+
+        // The nested loop is cleaner to implement with a `goto` in this case:
+    nk_mahalanobis_bf16_genoa_cycle:
+        if (j + 32 <= n) {
+            a_j_bf16x32 = _mm512_loadu_epi16(a + j);
+            b_j_bf16x32 = _mm512_loadu_epi16(b + j);
+            c_bf16x32 = _mm512_loadu_epi16(c + i * n + j);
+        }
+        else {
+            a_j_bf16x32 = _mm512_maskz_loadu_epi16(tail_mask, a + tail_start);
+            b_j_bf16x32 = _mm512_maskz_loadu_epi16(tail_mask, b + tail_start);
+            c_bf16x32 = _mm512_maskz_loadu_epi16(tail_mask, c + i * n + tail_start);
+        }
+        diff_j_bf16x32 = nk_substract_bf16x32_genoa_(a_j_bf16x32, b_j_bf16x32);
+        cdiff_j_f32x16 = _mm512_dpbf16_ps(cdiff_j_f32x16, nk_m512bh_from_m512i_(diff_j_bf16x32),
+                                          nk_m512bh_from_m512i_(c_bf16x32));
+        j += 32;
+        if (j < n) goto nk_mahalanobis_bf16_genoa_cycle;
+        sum_f32x16 = _mm512_fmadd_ps(diff_i_f32x16, cdiff_j_f32x16, sum_f32x16);
+    }
+
+    nk_f32_t quadratic = _mm512_reduce_add_ps(sum_f32x16);
+    *result = nk_f32_sqrt_haswell(quadratic > 0 ? quadratic : 0);
+}
+
+NK_PUBLIC void nk_bilinear_bf16c_genoa(nk_bf16c_t const *a, nk_bf16c_t const *b, nk_bf16c_t const *c, nk_size_t n,
+                                       nk_f32c_t *results) {
+
+    // We take into account, that FMS is the same as FMA with a negative multiplier.
+    // To multiply a floating-point value by -1, we can use the `XOR` instruction to flip the sign bit.
+    // This way we can avoid the shuffling and the need for separate real and imaginary parts.
+    // For the imaginary part of the product, we would need to swap the real and imaginary parts of
+    // one of the vectors.
+    __m512i const sign_flip_i32x16 = _mm512_set1_epi32(0x80000000);
+    __m512i const swap_adjacent_i8x64 = _mm512_set_epi8(                //
+        61, 60, 63, 62, 57, 56, 59, 58, 53, 52, 55, 54, 49, 48, 51, 50, // 4th 128-bit lane
+        45, 44, 47, 46, 41, 40, 43, 42, 37, 36, 39, 38, 33, 32, 35, 34, // 3rd 128-bit lane
+        29, 28, 31, 30, 25, 24, 27, 26, 21, 20, 23, 22, 17, 16, 19, 18, // 2nd 128-bit lane
+        13, 12, 15, 14, 9, 8, 11, 10, 5, 4, 7, 6, 1, 0, 3, 2            // 1st 128-bit lane
+    );
+
+    // Default case for arbitrary size `n`
+    nk_size_t const tail_length = n % 16;
+    nk_size_t const tail_start = n - tail_length;
+    __mmask32 const tail_mask = (__mmask32)_bzhi_u32(0xFFFFFFFF, tail_length * 2);
+    nk_f32_t sum_real = 0;
+    nk_f32_t sum_imag = 0;
+
+    for (nk_size_t i = 0; i != n; ++i) {
+        nk_f32_t a_i_real, a_i_imag;
+        nk_bf16_to_f32_serial(&a[i].real, &a_i_real);
+        nk_bf16_to_f32_serial(&a[i].imag, &a_i_imag);
+        __m512 cb_j_real_f32x16 = _mm512_setzero_ps();
+        __m512 cb_j_imag_f32x16 = _mm512_setzero_ps();
+        __m512i b_bf16x32, c_bf16x32;
+        nk_size_t j = 0;
+
+    nk_bilinear_bf16c_skylake_cycle:
+        if (j + 16 <= n) {
+            b_bf16x32 = _mm512_loadu_epi16((nk_i16_t const *)(b + j));
+            c_bf16x32 = _mm512_loadu_epi16((nk_i16_t const *)(c + i * n + j));
+        }
+        else {
+            b_bf16x32 = _mm512_maskz_loadu_epi16(tail_mask, (nk_i16_t const *)(b + tail_start));
+            c_bf16x32 = _mm512_maskz_loadu_epi16(tail_mask, (nk_i16_t const *)(c + i * n + tail_start));
+        }
+        cb_j_real_f32x16 = _mm512_dpbf16_ps(                                      //
+            cb_j_real_f32x16,                                                     //
+            nk_m512bh_from_m512i_(_mm512_xor_si512(c_bf16x32, sign_flip_i32x16)), //
+            nk_m512bh_from_m512i_(b_bf16x32));
+        cb_j_imag_f32x16 = _mm512_dpbf16_ps(                                            //
+            cb_j_imag_f32x16,                                                           //
+            nk_m512bh_from_m512i_(_mm512_shuffle_epi8(c_bf16x32, swap_adjacent_i8x64)), //
+            nk_m512bh_from_m512i_(b_bf16x32));
+        j += 16;
+        if (j < n) goto nk_bilinear_bf16c_skylake_cycle;
+        // Horizontal sums are the expensive part of the computation:
+        nk_f32_t const cb_j_real = nk_reduce_add_f32x16_skylake_(cb_j_real_f32x16);
+        nk_f32_t const cb_j_imag = nk_reduce_add_f32x16_skylake_(cb_j_imag_f32x16);
+        sum_real += a_i_real * cb_j_real - a_i_imag * cb_j_imag;
+        sum_imag += a_i_real * cb_j_imag + a_i_imag * cb_j_real;
+    }
+
+    // Reduce horizontal sums:
+    results->real = sum_real;
+    results->imag = sum_imag;
+}
+
+#if defined(__clang__)
+#pragma clang attribute pop
+#elif defined(__GNUC__)
+#pragma GCC pop_options
+#endif
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif // NK_TARGET_GENOA
+#endif // NK_TARGET_X86_
+#endif // NK_CURVED_GENOA_H

package/include/numkong/curved/haswell.h

@@ -0,0 +1,276 @@
+/**
+ *  @brief SIMD-accelerated Curved Space Similarity for Haswell.
+ *  @file include/numkong/curved/haswell.h
+ *  @author Ash Vardanian
+ *  @date January 14, 2026
+ *
+ *  @sa include/numkong/curved.h
+ *
+ *  Implements f16 and bf16 bilinear forms using AVX2 with F16C conversion.
+ */
+#ifndef NK_CURVED_HASWELL_H
+#define NK_CURVED_HASWELL_H
+
+#if NK_TARGET_X86_
+#if NK_TARGET_HASWELL
+
+#include "numkong/types.h"
+#include "numkong/reduce/haswell.h"  // `nk_reduce_add_f32x8_haswell_`
+#include "numkong/spatial/haswell.h" // `nk_f32_sqrt_haswell`
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if defined(__clang__)
+#pragma clang attribute push(__attribute__((target("avx2,f16c,fma,bmi,bmi2"))), apply_to = function)
+#elif defined(__GNUC__)
+#pragma GCC push_options
+#pragma GCC target("avx2", "f16c", "fma", "bmi", "bmi2")
+#endif
+
+NK_PUBLIC void nk_bilinear_f32_haswell(nk_f32_t const *a, nk_f32_t const *b, nk_f32_t const *c, nk_size_t n,
+                                       nk_f64_t *result) {
+    nk_size_t const tail_length = n % 4;
+    nk_size_t const tail_start = n - tail_length;
+    __m256d sum_f64x4 = _mm256_setzero_pd();
+
+    for (nk_size_t i = 0; i != n; ++i) {
+        __m256d a_f64x4 = _mm256_set1_pd((nk_f64_t)a[i]);
+        __m256d cb_j_f64x4 = _mm256_setzero_pd();
+        for (nk_size_t j = 0; j + 4 <= n; j += 4) {
+            __m256d b_f64x4 = _mm256_cvtps_pd(_mm_loadu_ps(b + j));
+            __m256d c_f64x4 = _mm256_cvtps_pd(_mm_loadu_ps(c + i * n + j));
+            cb_j_f64x4 = _mm256_fmadd_pd(b_f64x4, c_f64x4, cb_j_f64x4);
+        }
+        sum_f64x4 = _mm256_fmadd_pd(a_f64x4, cb_j_f64x4, sum_f64x4);
+    }
+
+    nk_f64_t sum = nk_reduce_add_f64x4_haswell_(sum_f64x4);
+    if (tail_length) {
+        nk_b128_vec_t b_tail_vec;
+        nk_partial_load_b32x4_haswell_(b + tail_start, &b_tail_vec, tail_length);
+        __m256d b_tail_f64x4 = _mm256_cvtps_pd(b_tail_vec.xmm_ps);
+        for (nk_size_t i = 0; i != n; ++i) {
+            nk_f64_t a_i = (nk_f64_t)a[i];
+            nk_b128_vec_t c_tail_vec;
+            nk_partial_load_b32x4_haswell_(c + i * n + tail_start, &c_tail_vec, tail_length);
+            __m256d c_tail_f64x4 = _mm256_cvtps_pd(c_tail_vec.xmm_ps);
+            sum += a_i * nk_reduce_add_f64x4_haswell_(_mm256_mul_pd(b_tail_f64x4, c_tail_f64x4));
+        }
+    }
+
+    *result = sum;
+}
+
+NK_PUBLIC void nk_mahalanobis_f32_haswell(nk_f32_t const *a, nk_f32_t const *b, nk_f32_t const *c, nk_size_t n,
+                                          nk_f64_t *result) {
+    nk_size_t const tail_length = n % 4;
+    nk_size_t const tail_start = n - tail_length;
+    __m256d sum_f64x4 = _mm256_setzero_pd();
+
+    for (nk_size_t i = 0; i != n; ++i) {
+        __m256d diff_i_f64x4 = _mm256_set1_pd((nk_f64_t)a[i] - (nk_f64_t)b[i]);
+        __m256d cdiff_j_f64x4 = _mm256_setzero_pd();
+        for (nk_size_t j = 0; j + 4 <= n; j += 4) {
+            __m256d diff_j_f64x4 = _mm256_sub_pd( //
+                _mm256_cvtps_pd(_mm_loadu_ps(a + j)), _mm256_cvtps_pd(_mm_loadu_ps(b + j)));
+            __m256d c_f64x4 = _mm256_cvtps_pd(_mm_loadu_ps(c + i * n + j));
+            cdiff_j_f64x4 = _mm256_fmadd_pd(diff_j_f64x4, c_f64x4, cdiff_j_f64x4);
+        }
+        sum_f64x4 = _mm256_fmadd_pd(diff_i_f64x4, cdiff_j_f64x4, sum_f64x4);
+    }
+
+    nk_f64_t sum = nk_reduce_add_f64x4_haswell_(sum_f64x4);
+    if (tail_length) {
+        nk_b128_vec_t a_tail_vec, b_tail_vec;
+        nk_partial_load_b32x4_haswell_(a + tail_start, &a_tail_vec, tail_length);
+        nk_partial_load_b32x4_haswell_(b + tail_start, &b_tail_vec, tail_length);
+        __m256d diff_tail_f64x4 = _mm256_sub_pd(_mm256_cvtps_pd(a_tail_vec.xmm_ps), _mm256_cvtps_pd(b_tail_vec.xmm_ps));
+        for (nk_size_t i = 0; i != n; ++i) {
+            nk_f64_t diff_i = (nk_f64_t)a[i] - (nk_f64_t)b[i];
+            nk_b128_vec_t c_tail_vec;
+            nk_partial_load_b32x4_haswell_(c + i * n + tail_start, &c_tail_vec, tail_length);
+            __m256d c_tail_f64x4 = _mm256_cvtps_pd(c_tail_vec.xmm_ps);
+            sum += diff_i * nk_reduce_add_f64x4_haswell_(_mm256_mul_pd(diff_tail_f64x4, c_tail_f64x4));
+        }
+    }
+
+    *result = nk_f64_sqrt_haswell(sum > 0 ? sum : 0);
+}
+
+NK_PUBLIC void nk_bilinear_f16_haswell(nk_f16_t const *a, nk_f16_t const *b, nk_f16_t const *c, nk_size_t n,
+                                       nk_f32_t *result) {
+    __m256 sum_f32x8 = _mm256_setzero_ps();
+    for (nk_size_t i = 0; i != n; ++i) {
+        __m256 a_f32x8 = _mm256_cvtph_ps(_mm_set1_epi16(*(short const *)(a + i)));
+        __m256 cb_j_f32x8 = _mm256_setzero_ps();
+        for (nk_size_t j = 0; j + 8 <= n; j += 8) {
+            __m256 b_f32x8 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i const *)(b + j)));
+            __m256 c_f32x8 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i const *)(c + i * n + j)));
+            cb_j_f32x8 = _mm256_fmadd_ps(b_f32x8, c_f32x8, cb_j_f32x8);
+        }
+        sum_f32x8 = _mm256_fmadd_ps(a_f32x8, cb_j_f32x8, sum_f32x8);
+    }
+
+    // Handle the tail of every row
+    nk_f32_t sum = nk_reduce_add_f32x8_haswell_(sum_f32x8);
+    nk_size_t const tail_length = n % 8;
+    nk_size_t const tail_start = n - tail_length;
+    if (tail_length) {
+        for (nk_size_t i = 0; i != n; ++i) {
+            nk_f32_t a_i = _mm256_cvtss_f32(_mm256_cvtph_ps(_mm_set1_epi16(*(short const *)(a + i))));
+            nk_b256_vec_t b_vec;
+            nk_partial_load_f16x8_to_f32x8_haswell_(b + tail_start, &b_vec, tail_length);
+            __m256 b_f32x8 = b_vec.ymm_ps;
+            nk_b256_vec_t c_vec;
+            nk_partial_load_f16x8_to_f32x8_haswell_(c + i * n + tail_start, &c_vec, tail_length);
+            __m256 c_f32x8 = c_vec.ymm_ps;
+            nk_f32_t cb_j = nk_reduce_add_f32x8_haswell_(_mm256_mul_ps(b_f32x8, c_f32x8));
+            sum += a_i * cb_j;
+        }
+    }
+
+    *result = sum;
+}
+
+NK_PUBLIC void nk_mahalanobis_f16_haswell(nk_f16_t const *a, nk_f16_t const *b, nk_f16_t const *c, nk_size_t n,
+                                          nk_f32_t *result) {
+    __m256 sum_f32x8 = _mm256_setzero_ps();
+    for (nk_size_t i = 0; i != n; ++i) {
+        __m256 diff_i_f32x8 = _mm256_sub_ps(                          //
+            _mm256_cvtph_ps(_mm_set1_epi16(*(short const *)(a + i))), //
+            _mm256_cvtph_ps(_mm_set1_epi16(*(short const *)(b + i))));
+        __m256 cdiff_j_f32x8 = _mm256_setzero_ps();
+        for (nk_size_t j = 0; j + 8 <= n; j += 8) {
+            __m256 diff_j_f32x8 = _mm256_sub_ps( //
+                _mm256_cvtph_ps(_mm_loadu_si128((__m128i const *)(a + j))),
+                _mm256_cvtph_ps(_mm_loadu_si128((__m128i const *)(b + j))));
+            __m256 c_f32x8 = _mm256_cvtph_ps(_mm_loadu_si128((__m128i const *)(c + i * n + j)));
+            cdiff_j_f32x8 = _mm256_fmadd_ps(diff_j_f32x8, c_f32x8, cdiff_j_f32x8);
+        }
+        sum_f32x8 = _mm256_fmadd_ps(diff_i_f32x8, cdiff_j_f32x8, sum_f32x8);
+    }
+
+    // Handle the tail of every row
+    nk_f32_t sum = nk_reduce_add_f32x8_haswell_(sum_f32x8);
+    nk_size_t const tail_length = n % 8;
+    nk_size_t const tail_start = n - tail_length;
+    if (tail_length) {
+        for (nk_size_t i = 0; i != n; ++i) {
+            nk_f32_t diff_i = _mm256_cvtss_f32(_mm256_sub_ps(             //
+                _mm256_cvtph_ps(_mm_set1_epi16(*(short const *)(a + i))), //
+                _mm256_cvtph_ps(_mm_set1_epi16(*(short const *)(b + i)))));
+            nk_b256_vec_t a_tail_vec, b_tail_vec;
+            nk_partial_load_f16x8_to_f32x8_haswell_(a + tail_start, &a_tail_vec, tail_length);
+            nk_partial_load_f16x8_to_f32x8_haswell_(b + tail_start, &b_tail_vec, tail_length);
+            __m256 diff_j_f32x8 = _mm256_sub_ps(a_tail_vec.ymm_ps, b_tail_vec.ymm_ps);
+            nk_b256_vec_t c_vec;
+            nk_partial_load_f16x8_to_f32x8_haswell_(c + i * n + tail_start, &c_vec, tail_length);
+            __m256 c_f32x8 = c_vec.ymm_ps;
+            nk_f32_t cdiff_j = nk_reduce_add_f32x8_haswell_(_mm256_mul_ps(diff_j_f32x8, c_f32x8));
+            sum += diff_i * cdiff_j;
+        }
+    }
+
+    *result = nk_f32_sqrt_haswell(sum > 0 ? sum : 0);
+}
+
+NK_PUBLIC void nk_bilinear_bf16_haswell(nk_bf16_t const *a, nk_bf16_t const *b, nk_bf16_t const *c, nk_size_t n,
+                                        nk_f32_t *result) {
+    __m256 sum_f32x8 = _mm256_setzero_ps();
+    for (nk_size_t i = 0; i != n; ++i) {
+        // The `nk_bf16_to_f32_serial` is cheaper than `nk_bf16x8_to_f32x8_haswell_`
+        nk_f32_t a_f32;
+        nk_bf16_to_f32_serial(a + i, &a_f32);
+        __m256 a_f32x8 = _mm256_set1_ps(a_f32);
+        __m256 cb_j_f32x8 = _mm256_setzero_ps();
+        for (nk_size_t j = 0; j + 8 <= n; j += 8) {
+            __m256 b_f32x8 = nk_bf16x8_to_f32x8_haswell_(_mm_loadu_si128((__m128i const *)(b + j)));
+            __m256 c_f32x8 = nk_bf16x8_to_f32x8_haswell_(_mm_loadu_si128((__m128i const *)(c + i * n + j)));
+            cb_j_f32x8 = _mm256_fmadd_ps(b_f32x8, c_f32x8, cb_j_f32x8);
+        }
+        sum_f32x8 = _mm256_fmadd_ps(a_f32x8, cb_j_f32x8, sum_f32x8);
+    }
+
+    // Handle the tail of every row
+    nk_f32_t sum = nk_reduce_add_f32x8_haswell_(sum_f32x8);
+    nk_size_t const tail_length = n % 8;
+    nk_size_t const tail_start = n - tail_length;
+    if (tail_length) {
+        for (nk_size_t i = 0; i != n; ++i) {
+            nk_f32_t a_i;
+            nk_bf16_to_f32_serial(a + i, &a_i);
+            nk_b256_vec_t b_vec;
+            nk_partial_load_bf16x8_to_f32x8_haswell_(b + tail_start, &b_vec, tail_length);
+            __m256 b_f32x8 = b_vec.ymm_ps;
+            nk_b256_vec_t c_vec;
+            nk_partial_load_bf16x8_to_f32x8_haswell_(c + i * n + tail_start, &c_vec, tail_length);
+            __m256 c_f32x8 = c_vec.ymm_ps;
+            nk_f32_t cb_j = nk_reduce_add_f32x8_haswell_(_mm256_mul_ps(b_f32x8, c_f32x8));
+            sum += a_i * cb_j;
+        }
+    }
+
+    *result = sum;
+}
+
+NK_PUBLIC void nk_mahalanobis_bf16_haswell(nk_bf16_t const *a, nk_bf16_t const *b, nk_bf16_t const *c, nk_size_t n,
+                                           nk_f32_t *result) {
+    __m256 sum_f32x8 = _mm256_setzero_ps();
+    for (nk_size_t i = 0; i != n; ++i) {
+        nk_f32_t a_i, b_i;
+        nk_bf16_to_f32_serial(a + i, &a_i);
+        nk_bf16_to_f32_serial(b + i, &b_i);
+        __m256 diff_i_f32x8 = _mm256_sub_ps( //
+            _mm256_set1_ps(a_i),             //
+            _mm256_set1_ps(b_i));
+        __m256 cdiff_j_f32x8 = _mm256_setzero_ps();
+        for (nk_size_t j = 0; j + 8 <= n; j += 8) {
+            __m256 diff_j_f32x8 = _mm256_sub_ps(                                        //
+                nk_bf16x8_to_f32x8_haswell_(_mm_loadu_si128((__m128i const *)(a + j))), //
+                nk_bf16x8_to_f32x8_haswell_(_mm_loadu_si128((__m128i const *)(b + j))));
+            __m256 c_f32x8 = nk_bf16x8_to_f32x8_haswell_(_mm_loadu_si128((__m128i const *)(c + i * n + j)));
+            cdiff_j_f32x8 = _mm256_fmadd_ps(diff_j_f32x8, c_f32x8, cdiff_j_f32x8);
+        }
+        sum_f32x8 = _mm256_fmadd_ps(diff_i_f32x8, cdiff_j_f32x8, sum_f32x8);
+    }
+
+    // Handle the tail of every row
+    nk_f32_t sum = nk_reduce_add_f32x8_haswell_(sum_f32x8);
+    nk_size_t const tail_length = n % 8;
+    nk_size_t const tail_start = n - tail_length;
+    if (tail_length) {
+        for (nk_size_t i = 0; i != n; ++i) {
+            nk_f32_t a_i, b_i;
+            nk_bf16_to_f32_serial(a + i, &a_i);
+            nk_bf16_to_f32_serial(b + i, &b_i);
+            nk_f32_t diff_i = a_i - b_i;
+            nk_b256_vec_t a_tail_vec, b_tail_vec;
+            nk_partial_load_bf16x8_to_f32x8_haswell_(a + tail_start, &a_tail_vec, tail_length);
+            nk_partial_load_bf16x8_to_f32x8_haswell_(b + tail_start, &b_tail_vec, tail_length);
+            __m256 diff_j_f32x8 = _mm256_sub_ps(a_tail_vec.ymm_ps, b_tail_vec.ymm_ps);
+            nk_b256_vec_t c_vec;
+            nk_partial_load_bf16x8_to_f32x8_haswell_(c + i * n + tail_start, &c_vec, tail_length);
+            __m256 c_f32x8 = c_vec.ymm_ps;
+            nk_f32_t cdiff_j = nk_reduce_add_f32x8_haswell_(_mm256_mul_ps(diff_j_f32x8, c_f32x8));
+            sum += diff_i * cdiff_j;
+        }
+    }
+
+    *result = nk_f32_sqrt_haswell(sum > 0 ? sum : 0);
+}
+
+#if defined(__clang__)
+#pragma clang attribute pop
+#elif defined(__GNUC__)
+#pragma GCC pop_options
+#endif
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif // NK_TARGET_HASWELL
+#endif // NK_TARGET_X86_
+#endif // NK_CURVED_HASWELL_H

package/include/numkong/curved/neon.h

@@ -0,0 +1,205 @@
+/**
+ *  @brief SIMD-accelerated Curved Space Similarity for NEON.
+ *  @file include/numkong/curved/neon.h
+ *  @author Ash Vardanian
+ *  @date January 14, 2026
+ *
+ *  @sa include/numkong/curved.h
+ *
+ *  Implements f32 bilinear forms and Mahalanobis distance using ARM NEON SIMD.
+ *  Accumulates f32 inputs in f64 precision to avoid catastrophic cancellation.
+ *
+ *  @section neon_curved_instructions Key NEON Instructions
+ *
+ *      Intrinsic         Instruction                   Latency     Throughput
+ *                                                                  A76     M4+/V1+/Oryon
+ *      vfmaq_f64         FMLA (V.2D, V.2D, V.2D)       4cy         2/cy    4/cy
+ *      vcvt_f64_f32      FCVTL (V.2D, V.2S)            3cy         2/cy    2/cy
+ *      vaddvq_f64        FADDP (V.2D to scalar)        3cy         1/cy    1/cy
+ *      vld1_f32          LD1 ({Vt.2S}, [Xn])           4cy         2/cy    2/cy
+ *      vld2_f32          LD2 ({Vt.2S, Vt2.2S}, [Xn])   4cy         1/cy    1/cy
+ *
+ *  For f32 bilinear and Mahalanobis, we upcast to f64 for accumulation to preserve
+ *  precision and avoid catastrophic cancellation in large-magnitude sums.
+ */
+#ifndef NK_CURVED_NEON_H
+#define NK_CURVED_NEON_H
+
+#if NK_TARGET_ARM_
+#if NK_TARGET_NEON
+
+#include "numkong/types.h"
+#include "numkong/spatial/neon.h" // nk_f64_sqrt_neon
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if defined(__clang__)
+#pragma clang attribute push(__attribute__((target("arch=armv8-a+simd"))), apply_to = function)
+#elif defined(__GNUC__)
+#pragma GCC push_options
+#pragma GCC target("arch=armv8-a+simd")
+#endif
+
+NK_PUBLIC void nk_bilinear_f32_neon(nk_f32_t const *a, nk_f32_t const *b, nk_f32_t const *c, nk_size_t n,
+                                    nk_f64_t *result) {
+    nk_f64_t outer_sum_f64 = 0;
+
+    for (nk_size_t i = 0; i != n; ++i) {
+        // Convert a[i] to f64 for precision
+        nk_f64_t a_i_f64 = (nk_f64_t)a[i];
+
+        // Inner loop: accumulate Σⱼ cᵢⱼ × bⱼ in f64
+        float64x2_t inner_sum_f64x2 = vdupq_n_f64(0);
+        nk_size_t j = 0;
+
+        // Vectorized inner loop: process 2 elements at a time
+        for (; j + 2 <= n; j += 2) {
+            // Load b[j:j+2] as f32, upcast to f64
+            float32x2_t b_f32x2 = vld1_f32(b + j);
+            float64x2_t b_f64x2 = vcvt_f64_f32(b_f32x2);
+
+            // Load c[i*n+j : i*n+j+2] as f32, upcast to f64
+            float32x2_t c_f32x2 = vld1_f32(c + i * n + j);
+            float64x2_t c_f64x2 = vcvt_f64_f32(c_f32x2);
+
+            // FMA: inner_sum += c × b
+            inner_sum_f64x2 = vfmaq_f64(inner_sum_f64x2, c_f64x2, b_f64x2);
+        }
+
+        // Reduce the f64x2 accumulator to scalar
+        nk_f64_t inner_sum_f64 = vaddvq_f64(inner_sum_f64x2);
+
+        // Handle tail elements
+        for (; j < n; ++j) { inner_sum_f64 += (nk_f64_t)c[i * n + j] * (nk_f64_t)b[j]; }
+
+        // Outer accumulation: outer_sum += aᵢ × inner_sum
+        outer_sum_f64 += a_i_f64 * inner_sum_f64;
+    }
+
+    *result = outer_sum_f64;
+}
+
+NK_PUBLIC void nk_mahalanobis_f32_neon(nk_f32_t const *a, nk_f32_t const *b, nk_f32_t const *c, nk_size_t n,
+                                       nk_f64_t *result) {
+    nk_f64_t outer_sum_f64 = 0;
+
+    for (nk_size_t i = 0; i != n; ++i) {
+        // Compute difference (aᵢ - bᵢ) in f64 for precision
+        nk_f64_t diff_i_f64 = (nk_f64_t)a[i] - (nk_f64_t)b[i];
+
+        // Inner loop: accumulate Σⱼ cᵢⱼ × (aⱼ - bⱼ) in f64
+        float64x2_t inner_sum_f64x2 = vdupq_n_f64(0);
+        nk_size_t j = 0;
+
+        // Vectorized inner loop: process 2 elements at a time
+        for (; j + 2 <= n; j += 2) {
+            // Load a[j:j+2] and b[j:j+2] as f32
+            float32x2_t a_f32x2 = vld1_f32(a + j);
+            float32x2_t b_f32x2 = vld1_f32(b + j);
+
+            // Compute difference in f32, then upcast to f64
+            float32x2_t diff_f32x2 = vsub_f32(a_f32x2, b_f32x2);
+            float64x2_t diff_f64x2 = vcvt_f64_f32(diff_f32x2);
+
+            // Load c[i*n+j : i*n+j+2] as f32, upcast to f64
+            float32x2_t c_f32x2 = vld1_f32(c + i * n + j);
+            float64x2_t c_f64x2 = vcvt_f64_f32(c_f32x2);
+
+            // FMA: inner_sum += c × diff
+            inner_sum_f64x2 = vfmaq_f64(inner_sum_f64x2, c_f64x2, diff_f64x2);
+        }
+
+        // Reduce the f64x2 accumulator to scalar
+        nk_f64_t inner_sum_f64 = vaddvq_f64(inner_sum_f64x2);
+
+        // Handle tail elements
+        for (; j < n; ++j) {
+            nk_f64_t diff_j_f64 = (nk_f64_t)a[j] - (nk_f64_t)b[j];
+            inner_sum_f64 += (nk_f64_t)c[i * n + j] * diff_j_f64;
+        }
+
+        // Outer accumulation: outer_sum += diff_i × inner_sum
+        outer_sum_f64 += diff_i_f64 * inner_sum_f64;
+    }
+
+    // Take square root of the result (clamp to 0 for numerical stability)
+    *result = nk_f64_sqrt_neon(outer_sum_f64 > 0 ? outer_sum_f64 : 0);
+}
+
+NK_PUBLIC void nk_bilinear_f32c_neon(nk_f32c_t const *a_pairs, nk_f32c_t const *b_pairs, nk_f32c_t const *c_pairs,
+                                     nk_size_t n, nk_f64c_t *results) {
+    // ARMv8.3-A FCMLA (`vcmlaq_f32`) was benchmarked for this complex inner loop.
+    // The deinterleave+4FMA pattern is 2.3x faster on Apple M4 — see `dot/neon.h` comment.
+    nk_f64_t outer_sum_real_f64 = 0;
+    nk_f64_t outer_sum_imag_f64 = 0;
+
+    for (nk_size_t i = 0; i != n; ++i) {
+        // Convert a[i] to f64 for precision
+        nk_f64_t a_real_f64 = (nk_f64_t)a_pairs[i].real;
+        nk_f64_t a_imag_f64 = (nk_f64_t)a_pairs[i].imag;
+
+        // Inner loop: accumulate Σⱼ cᵢⱼ × bⱼ in f64
+        float64x2_t inner_sum_real_f64x2 = vdupq_n_f64(0);
+        float64x2_t inner_sum_imag_f64x2 = vdupq_n_f64(0);
+        nk_size_t j = 0;
+
+        // Vectorized inner loop: process 2 complex elements at a time
+        for (; j + 2 <= n; j += 2) {
+            // Load b[j:j+2] as interleaved complex pairs (real, imag, real, imag)
+            float32x2x2_t b_f32x2x2 = vld2_f32((nk_f32_t const *)(b_pairs + j));
+            float64x2_t b_real_f64x2 = vcvt_f64_f32(b_f32x2x2.val[0]);
+            float64x2_t b_imag_f64x2 = vcvt_f64_f32(b_f32x2x2.val[1]);
+
+            // Load c[i*n+j : i*n+j+2] as interleaved complex pairs
+            float32x2x2_t c_f32x2x2 = vld2_f32((nk_f32_t const *)(c_pairs + i * n + j));
+            float64x2_t c_real_f64x2 = vcvt_f64_f32(c_f32x2x2.val[0]);
+            float64x2_t c_imag_f64x2 = vcvt_f64_f32(c_f32x2x2.val[1]);
+
+            // Complex multiply
+            inner_sum_real_f64x2 = vfmaq_f64(inner_sum_real_f64x2, c_real_f64x2, b_real_f64x2);
+            inner_sum_real_f64x2 = vfmsq_f64(inner_sum_real_f64x2, c_imag_f64x2, b_imag_f64x2);
+
+            // Imaginary part: c_real×b_imag + c_imag×b_real
+            inner_sum_imag_f64x2 = vfmaq_f64(inner_sum_imag_f64x2, c_real_f64x2, b_imag_f64x2);
+            inner_sum_imag_f64x2 = vfmaq_f64(inner_sum_imag_f64x2, c_imag_f64x2, b_real_f64x2);
+        }
+
+        // Reduce the f64x2 accumulators to scalars
+        nk_f64_t inner_sum_real_f64 = vaddvq_f64(inner_sum_real_f64x2);
+        nk_f64_t inner_sum_imag_f64 = vaddvq_f64(inner_sum_imag_f64x2);
+
+        // Handle tail elements
+        for (; j < n; ++j) {
+            nk_f64_t b_real = (nk_f64_t)b_pairs[j].real;
+            nk_f64_t b_imag = (nk_f64_t)b_pairs[j].imag;
+            nk_f64_t c_real = (nk_f64_t)c_pairs[i * n + j].real;
+            nk_f64_t c_imag = (nk_f64_t)c_pairs[i * n + j].imag;
+            // Complex multiply: c × b
+            inner_sum_real_f64 += c_real * b_real - c_imag * b_imag;
+            inner_sum_imag_f64 += c_real * b_imag + c_imag * b_real;
+        }
+
+        // Outer accumulation
+        outer_sum_real_f64 += a_real_f64 * inner_sum_real_f64 - a_imag_f64 * inner_sum_imag_f64;
+        outer_sum_imag_f64 += a_real_f64 * inner_sum_imag_f64 + a_imag_f64 * inner_sum_real_f64;
+    }
+
+    results->real = outer_sum_real_f64;
+    results->imag = outer_sum_imag_f64;
+}
+
+#if defined(__clang__)
+#pragma clang attribute pop
+#elif defined(__GNUC__)
+#pragma GCC pop_options
+#endif
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif // NK_TARGET_NEON
+#endif // NK_TARGET_ARM_
+#endif // NK_CURVED_NEON_H