numkong 7.0.0 → 7.4.2

package/include/numkong/spatial/loongsonasx.h

@@ -0,0 +1,483 @@
+/**
+ *  @brief SIMD-accelerated Spatial Similarity Measures for LoongArch LASX (256-bit).
+ *  @file include/numkong/spatial/loongsonasx.h
+ *  @author Ash Vardanian
+ *  @date March 23, 2026
+ *
+ *  @sa include/numkong/spatial.h
+ *
+ *  @section spatial_loongsonasx_instructions Key LASX Spatial Instructions
+ *
+ *  LASX provides 256-bit SIMD operations using __m256i as the universal vector type.
+ *  All intrinsics are prefixed with __lasx_. Float operations reinterpret __m256i as
+ *  f32x8 or f64x4. Integer widening multiply-accumulate chains handle i8/u8 distances.
+ *
+ *  For F32 spatial distances, upcasting to F64 and downcasting back is faster than stable
+ *  summation algorithms. For F64 angular we use the Dot2 algorithm (Ogita-Rump-Oishi, 2005)
+ *  for the cross-product accumulation, while self-products use simple FMA since all terms
+ *  are non-negative and don't suffer from cancellation.
+ */
+#ifndef NK_SPATIAL_LOONGSONASX_H
+#define NK_SPATIAL_LOONGSONASX_H
+
+#if NK_TARGET_LOONGARCH_
+#if NK_TARGET_LOONGSONASX
+
+#include "numkong/types.h"
+#include "numkong/spatial/serial.h"
+#include "numkong/dot/loongsonasx.h"    //
+#include "numkong/cast/loongsonasx.h"   // `nk_bf16x8_to_f32x8_loongsonasx_`
+#include "numkong/scalar/loongsonasx.h" // `nk_f32_sqrt_loongsonasx`, `nk_f64_sqrt_loongsonasx`
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#pragma region Angular Normalize Helpers
+
+NK_INTERNAL nk_f64_t nk_angular_normalize_f64_loongsonasx_(nk_f64_t ab, nk_f64_t a2, nk_f64_t b2) {
+    if (a2 == 0 && b2 == 0) return 0;
+    else if (ab == 0) return 1;
+    nk_f64_t result = 1 - ab / (nk_f64_sqrt_loongsonasx(a2) * nk_f64_sqrt_loongsonasx(b2));
+    return result > 0 ? result : 0;
+}
+
+NK_INTERNAL nk_f32_t nk_angular_normalize_i32_loongsonasx_(nk_i32_t ab, nk_i32_t a2, nk_i32_t b2) {
+    if (a2 == 0 && b2 == 0) return 0;
+    else if (ab == 0) return 1;
+    nk_f32_t result = 1.0f -
+                      (nk_f32_t)ab * nk_f32_rsqrt_loongsonasx((nk_f32_t)a2) * nk_f32_rsqrt_loongsonasx((nk_f32_t)b2);
+    return result > 0 ? result : 0;
+}
+
+#pragma endregion Angular Normalize Helpers
+
+#pragma region I8 and U8 Integers
+
+NK_PUBLIC void nk_sqeuclidean_i8_loongsonasx(nk_i8_t const *a, nk_i8_t const *b, nk_size_t n, nk_u32_t *result) {
+    __m256i sum_i32x8 = __lasx_xvreplgr2vr_w(0);
+    nk_size_t i = 0;
+    for (; i + 32 <= n; i += 32) {
+        __m256i a_i8x32 = __lasx_xvld(a + i, 0);
+        __m256i b_i8x32 = __lasx_xvld(b + i, 0);
+        __m256i diff_i8x32 = __lasx_xvsub_b(a_i8x32, b_i8x32);
+        __m256i sq_i16x16 = __lasx_xvreplgr2vr_h(0);
+        sq_i16x16 = __lasx_xvmaddwev_h_b(sq_i16x16, diff_i8x32, diff_i8x32);
+        sq_i16x16 = __lasx_xvmaddwod_h_b(sq_i16x16, diff_i8x32, diff_i8x32);
+        sum_i32x8 = __lasx_xvadd_w(sum_i32x8, __lasx_xvhaddw_w_h(sq_i16x16, sq_i16x16));
+    }
+    nk_i32_t sum = nk_reduce_add_i32x8_loongsonasx_(sum_i32x8);
+    for (; i < n; ++i) {
+        nk_i32_t diff = (nk_i32_t)a[i] - b[i];
+        sum += diff * diff;
+    }
+    *result = (nk_u32_t)sum;
+}
+
+NK_PUBLIC void nk_euclidean_i8_loongsonasx(nk_i8_t const *a, nk_i8_t const *b, nk_size_t n, nk_f32_t *result) {
+    nk_u32_t distance_sq_u32;
+    nk_sqeuclidean_i8_loongsonasx(a, b, n, &distance_sq_u32);
+    *result = nk_f32_sqrt_loongsonasx((nk_f32_t)distance_sq_u32);
+}
+
+NK_PUBLIC void nk_angular_i8_loongsonasx(nk_i8_t const *a, nk_i8_t const *b, nk_size_t n, nk_f32_t *result) {
+    __m256i dot_i32x8 = __lasx_xvreplgr2vr_w(0);
+    __m256i a_sq_i32x8 = __lasx_xvreplgr2vr_w(0);
+    __m256i b_sq_i32x8 = __lasx_xvreplgr2vr_w(0);
+    nk_size_t i = 0;
+    for (; i + 32 <= n; i += 32) {
+        __m256i a_i8x32 = __lasx_xvld(a + i, 0);
+        __m256i b_i8x32 = __lasx_xvld(b + i, 0);
+        // dot(a, b)
+        __m256i ab_i16x16 = __lasx_xvreplgr2vr_h(0);
+        ab_i16x16 = __lasx_xvmaddwev_h_b(ab_i16x16, a_i8x32, b_i8x32);
+        ab_i16x16 = __lasx_xvmaddwod_h_b(ab_i16x16, a_i8x32, b_i8x32);
+        dot_i32x8 = __lasx_xvadd_w(dot_i32x8, __lasx_xvhaddw_w_h(ab_i16x16, ab_i16x16));
+        // norm_sq(a)
+        __m256i aa_i16x16 = __lasx_xvreplgr2vr_h(0);
+        aa_i16x16 = __lasx_xvmaddwev_h_b(aa_i16x16, a_i8x32, a_i8x32);
+        aa_i16x16 = __lasx_xvmaddwod_h_b(aa_i16x16, a_i8x32, a_i8x32);
+        a_sq_i32x8 = __lasx_xvadd_w(a_sq_i32x8, __lasx_xvhaddw_w_h(aa_i16x16, aa_i16x16));
+        // norm_sq(b)
+        __m256i bb_i16x16 = __lasx_xvreplgr2vr_h(0);
+        bb_i16x16 = __lasx_xvmaddwev_h_b(bb_i16x16, b_i8x32, b_i8x32);
+        bb_i16x16 = __lasx_xvmaddwod_h_b(bb_i16x16, b_i8x32, b_i8x32);
+        b_sq_i32x8 = __lasx_xvadd_w(b_sq_i32x8, __lasx_xvhaddw_w_h(bb_i16x16, bb_i16x16));
+    }
+    nk_i32_t dot = nk_reduce_add_i32x8_loongsonasx_(dot_i32x8);
+    nk_i32_t a_sq = nk_reduce_add_i32x8_loongsonasx_(a_sq_i32x8);
+    nk_i32_t b_sq = nk_reduce_add_i32x8_loongsonasx_(b_sq_i32x8);
+    for (; i < n; ++i) {
+        nk_i32_t a_val = a[i], b_val = b[i];
+        dot += a_val * b_val;
+        a_sq += a_val * a_val;
+        b_sq += b_val * b_val;
+    }
+    *result = nk_angular_normalize_i32_loongsonasx_(dot, a_sq, b_sq);
+}
+
+NK_PUBLIC void nk_sqeuclidean_u8_loongsonasx(nk_u8_t const *a, nk_u8_t const *b, nk_size_t n, nk_u32_t *result) {
+    __m256i sum_i32x8 = __lasx_xvreplgr2vr_w(0);
+    __m256i zeros_i8x32 = __lasx_xvreplgr2vr_b(0);
+    nk_size_t i = 0;
+    for (; i + 32 <= n; i += 32) {
+        __m256i a_u8x32 = __lasx_xvld(a + i, 0);
+        __m256i b_u8x32 = __lasx_xvld(b + i, 0);
+        __m256i a_low_u16x16 = __lasx_xvilvl_b(zeros_i8x32, a_u8x32);
+        __m256i a_high_u16x16 = __lasx_xvilvh_b(zeros_i8x32, a_u8x32);
+        __m256i b_low_u16x16 = __lasx_xvilvl_b(zeros_i8x32, b_u8x32);
+        __m256i b_high_u16x16 = __lasx_xvilvh_b(zeros_i8x32, b_u8x32);
+        __m256i diff_low_i16x16 = __lasx_xvsub_h(a_low_u16x16, b_low_u16x16);
+        __m256i diff_high_i16x16 = __lasx_xvsub_h(a_high_u16x16, b_high_u16x16);
+        __m256i sq_ev_low_i32x8 = __lasx_xvmulwev_w_h(diff_low_i16x16, diff_low_i16x16);
+        __m256i sq_od_low_i32x8 = __lasx_xvmulwod_w_h(diff_low_i16x16, diff_low_i16x16);
+        __m256i sq_ev_high_i32x8 = __lasx_xvmulwev_w_h(diff_high_i16x16, diff_high_i16x16);
+        __m256i sq_od_high_i32x8 = __lasx_xvmulwod_w_h(diff_high_i16x16, diff_high_i16x16);
+        sum_i32x8 = __lasx_xvadd_w(sum_i32x8, __lasx_xvadd_w(sq_ev_low_i32x8, sq_od_low_i32x8));
+        sum_i32x8 = __lasx_xvadd_w(sum_i32x8, __lasx_xvadd_w(sq_ev_high_i32x8, sq_od_high_i32x8));
+    }
+    nk_i32_t sum = nk_reduce_add_i32x8_loongsonasx_(sum_i32x8);
+    for (; i < n; ++i) {
+        nk_i32_t diff = (nk_i32_t)a[i] - b[i];
+        sum += diff * diff;
+    }
+    *result = (nk_u32_t)sum;
+}
+
+NK_PUBLIC void nk_euclidean_u8_loongsonasx(nk_u8_t const *a, nk_u8_t const *b, nk_size_t n, nk_f32_t *result) {
+    nk_u32_t distance_sq_u32;
+    nk_sqeuclidean_u8_loongsonasx(a, b, n, &distance_sq_u32);
+    *result = nk_f32_sqrt_loongsonasx((nk_f32_t)distance_sq_u32);
+}
+
+NK_PUBLIC void nk_angular_u8_loongsonasx(nk_u8_t const *a, nk_u8_t const *b, nk_size_t n, nk_f32_t *result) {
+    __m256i dot_i32x8 = __lasx_xvreplgr2vr_w(0);
+    __m256i a_sq_i32x8 = __lasx_xvreplgr2vr_w(0);
+    __m256i b_sq_i32x8 = __lasx_xvreplgr2vr_w(0);
+    __m256i zeros_i8x32 = __lasx_xvreplgr2vr_b(0);
+    nk_size_t i = 0;
+    for (; i + 32 <= n; i += 32) {
+        __m256i a_u8x32 = __lasx_xvld(a + i, 0);
+        __m256i b_u8x32 = __lasx_xvld(b + i, 0);
+        __m256i a_low_u16x16 = __lasx_xvilvl_b(zeros_i8x32, a_u8x32);
+        __m256i a_high_u16x16 = __lasx_xvilvh_b(zeros_i8x32, a_u8x32);
+        __m256i b_low_u16x16 = __lasx_xvilvl_b(zeros_i8x32, b_u8x32);
+        __m256i b_high_u16x16 = __lasx_xvilvh_b(zeros_i8x32, b_u8x32);
+        // dot(a, b)
+        __m256i ab_ev_low_i32x8 = __lasx_xvmulwev_w_h(a_low_u16x16, b_low_u16x16);
+        __m256i ab_od_low_i32x8 = __lasx_xvmulwod_w_h(a_low_u16x16, b_low_u16x16);
+        __m256i ab_ev_high_i32x8 = __lasx_xvmulwev_w_h(a_high_u16x16, b_high_u16x16);
+        __m256i ab_od_high_i32x8 = __lasx_xvmulwod_w_h(a_high_u16x16, b_high_u16x16);
+        dot_i32x8 = __lasx_xvadd_w(dot_i32x8, __lasx_xvadd_w(ab_ev_low_i32x8, ab_od_low_i32x8));
+        dot_i32x8 = __lasx_xvadd_w(dot_i32x8, __lasx_xvadd_w(ab_ev_high_i32x8, ab_od_high_i32x8));
+        // norm_sq(a)
+        __m256i aa_ev_low_i32x8 = __lasx_xvmulwev_w_h(a_low_u16x16, a_low_u16x16);
+        __m256i aa_od_low_i32x8 = __lasx_xvmulwod_w_h(a_low_u16x16, a_low_u16x16);
+        __m256i aa_ev_high_i32x8 = __lasx_xvmulwev_w_h(a_high_u16x16, a_high_u16x16);
+        __m256i aa_od_high_i32x8 = __lasx_xvmulwod_w_h(a_high_u16x16, a_high_u16x16);
+        a_sq_i32x8 = __lasx_xvadd_w(a_sq_i32x8, __lasx_xvadd_w(aa_ev_low_i32x8, aa_od_low_i32x8));
+        a_sq_i32x8 = __lasx_xvadd_w(a_sq_i32x8, __lasx_xvadd_w(aa_ev_high_i32x8, aa_od_high_i32x8));
+        // norm_sq(b)
+        __m256i bb_ev_low_i32x8 = __lasx_xvmulwev_w_h(b_low_u16x16, b_low_u16x16);
+        __m256i bb_od_low_i32x8 = __lasx_xvmulwod_w_h(b_low_u16x16, b_low_u16x16);
+        __m256i bb_ev_high_i32x8 = __lasx_xvmulwev_w_h(b_high_u16x16, b_high_u16x16);
+        __m256i bb_od_high_i32x8 = __lasx_xvmulwod_w_h(b_high_u16x16, b_high_u16x16);
+        b_sq_i32x8 = __lasx_xvadd_w(b_sq_i32x8, __lasx_xvadd_w(bb_ev_low_i32x8, bb_od_low_i32x8));
+        b_sq_i32x8 = __lasx_xvadd_w(b_sq_i32x8, __lasx_xvadd_w(bb_ev_high_i32x8, bb_od_high_i32x8));
+    }
+    nk_i32_t dot = nk_reduce_add_i32x8_loongsonasx_(dot_i32x8);
+    nk_i32_t a_sq = nk_reduce_add_i32x8_loongsonasx_(a_sq_i32x8);
+    nk_i32_t b_sq = nk_reduce_add_i32x8_loongsonasx_(b_sq_i32x8);
+    for (; i < n; ++i) {
+        nk_i32_t a_val = a[i], b_val = b[i];
+        dot += a_val * b_val;
+        a_sq += a_val * a_val;
+        b_sq += b_val * b_val;
+    }
+    *result = nk_angular_normalize_i32_loongsonasx_(dot, a_sq, b_sq);
+}
+
+#pragma endregion I8 and U8 Integers
+
+#pragma region F32 and F64 Floats
+
+NK_PUBLIC void nk_sqeuclidean_f32_loongsonasx(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result) {
+    __m256d sum_f64x4_low = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d sum_f64x4_high = (__m256d)__lasx_xvreplgr2vr_d(0);
+    nk_size_t i = 0;
+    for (; i + 8 <= n; i += 8) {
+        __m256i a_f32x8 = __lasx_xvld(a + i, 0);
+        __m256i b_f32x8 = __lasx_xvld(b + i, 0);
+        __m256d a_low_f64x4 = __lasx_xvfcvtl_d_s((__m256)a_f32x8);
+        __m256d b_low_f64x4 = __lasx_xvfcvtl_d_s((__m256)b_f32x8);
+        __m256d a_high_f64x4 = __lasx_xvfcvth_d_s((__m256)a_f32x8);
+        __m256d b_high_f64x4 = __lasx_xvfcvth_d_s((__m256)b_f32x8);
+        __m256d diff_low_f64x4 = __lasx_xvfsub_d(a_low_f64x4, b_low_f64x4);
+        __m256d diff_high_f64x4 = __lasx_xvfsub_d(a_high_f64x4, b_high_f64x4);
+        sum_f64x4_low = __lasx_xvfmadd_d(diff_low_f64x4, diff_low_f64x4, sum_f64x4_low);
+        sum_f64x4_high = __lasx_xvfmadd_d(diff_high_f64x4, diff_high_f64x4, sum_f64x4_high);
+    }
+    __m256d combined_f64x4 = __lasx_xvfadd_d(sum_f64x4_low, sum_f64x4_high);
+    nk_f64_t sum = nk_reduce_add_f64x4_loongsonasx_(combined_f64x4);
+    for (; i < n; ++i) {
+        nk_f64_t diff = (nk_f64_t)a[i] - b[i];
+        sum += diff * diff;
+    }
+    *result = sum;
+}
+
+NK_PUBLIC void nk_euclidean_f32_loongsonasx(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result) {
+    nk_sqeuclidean_f32_loongsonasx(a, b, n, result);
+    *result = nk_f64_sqrt_loongsonasx(*result);
+}
+
+NK_PUBLIC void nk_angular_f32_loongsonasx(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result) {
+    __m256d dot_f64x4_low = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d dot_f64x4_high = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d a_sq_f64x4_low = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d a_sq_f64x4_high = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d b_sq_f64x4_low = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d b_sq_f64x4_high = (__m256d)__lasx_xvreplgr2vr_d(0);
+    nk_size_t i = 0;
+    for (; i + 8 <= n; i += 8) {
+        __m256i a_f32x8 = __lasx_xvld(a + i, 0);
+        __m256i b_f32x8 = __lasx_xvld(b + i, 0);
+        __m256d a_low_f64x4 = __lasx_xvfcvtl_d_s((__m256)a_f32x8);
+        __m256d b_low_f64x4 = __lasx_xvfcvtl_d_s((__m256)b_f32x8);
+        __m256d a_high_f64x4 = __lasx_xvfcvth_d_s((__m256)a_f32x8);
+        __m256d b_high_f64x4 = __lasx_xvfcvth_d_s((__m256)b_f32x8);
+        dot_f64x4_low = __lasx_xvfmadd_d(a_low_f64x4, b_low_f64x4, dot_f64x4_low);
+        dot_f64x4_high = __lasx_xvfmadd_d(a_high_f64x4, b_high_f64x4, dot_f64x4_high);
+        a_sq_f64x4_low = __lasx_xvfmadd_d(a_low_f64x4, a_low_f64x4, a_sq_f64x4_low);
+        a_sq_f64x4_high = __lasx_xvfmadd_d(a_high_f64x4, a_high_f64x4, a_sq_f64x4_high);
+        b_sq_f64x4_low = __lasx_xvfmadd_d(b_low_f64x4, b_low_f64x4, b_sq_f64x4_low);
+        b_sq_f64x4_high = __lasx_xvfmadd_d(b_high_f64x4, b_high_f64x4, b_sq_f64x4_high);
+    }
+    nk_f64_t dot = nk_reduce_add_f64x4_loongsonasx_(__lasx_xvfadd_d(dot_f64x4_low, dot_f64x4_high));
+    nk_f64_t a_sq = nk_reduce_add_f64x4_loongsonasx_(__lasx_xvfadd_d(a_sq_f64x4_low, a_sq_f64x4_high));
+    nk_f64_t b_sq = nk_reduce_add_f64x4_loongsonasx_(__lasx_xvfadd_d(b_sq_f64x4_low, b_sq_f64x4_high));
+    for (; i < n; ++i) {
+        nk_f64_t a_val = a[i], b_val = b[i];
+        dot += a_val * b_val;
+        a_sq += a_val * a_val;
+        b_sq += b_val * b_val;
+    }
+    *result = nk_angular_normalize_f64_loongsonasx_(dot, a_sq, b_sq);
+}
+
+NK_PUBLIC void nk_sqeuclidean_f64_loongsonasx(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result) {
+    __m256d sum_f64x4 = (__m256d)__lasx_xvreplgr2vr_d(0);
+    nk_size_t i = 0;
+    for (; i + 4 <= n; i += 4) {
+        __m256d a_f64x4 = (__m256d)__lasx_xvld(a + i, 0);
+        __m256d b_f64x4 = (__m256d)__lasx_xvld(b + i, 0);
+        __m256d diff_f64x4 = __lasx_xvfsub_d(a_f64x4, b_f64x4);
+        sum_f64x4 = __lasx_xvfmadd_d(diff_f64x4, diff_f64x4, sum_f64x4);
+    }
+    nk_f64_t sum = nk_reduce_add_f64x4_loongsonasx_(sum_f64x4);
+    for (; i < n; ++i) {
+        nk_f64_t diff = a[i] - b[i];
+        sum += diff * diff;
+    }
+    *result = sum;
+}
+
+NK_PUBLIC void nk_euclidean_f64_loongsonasx(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result) {
+    nk_sqeuclidean_f64_loongsonasx(a, b, n, result);
+    *result = nk_f64_sqrt_loongsonasx(*result);
+}
+
+NK_PUBLIC void nk_angular_f64_loongsonasx(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result) {
+    __m256d dot_sum_f64x4 = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d dot_compensation_f64x4 = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d a_norm_sq_f64x4 = (__m256d)__lasx_xvreplgr2vr_d(0);
+    __m256d b_norm_sq_f64x4 = (__m256d)__lasx_xvreplgr2vr_d(0);
+    nk_size_t i = 0;
+    for (; i + 4 <= n; i += 4) {
+        __m256d a_f64x4 = (__m256d)__lasx_xvld(a + i, 0);
+        __m256d b_f64x4 = (__m256d)__lasx_xvld(b + i, 0);
+
+        __m256d product_f64x4 = __lasx_xvfmul_d(a_f64x4, b_f64x4);
+        __m256d product_error_f64x4 = __lasx_xvfmsub_d(a_f64x4, b_f64x4, product_f64x4);
+
+        __m256d tentative_sum_f64x4 = __lasx_xvfadd_d(dot_sum_f64x4, product_f64x4);
+        __m256d virtual_addend_f64x4 = __lasx_xvfsub_d(tentative_sum_f64x4, dot_sum_f64x4);
+        __m256d sum_error_f64x4 = __lasx_xvfadd_d(
+            __lasx_xvfsub_d(dot_sum_f64x4, __lasx_xvfsub_d(tentative_sum_f64x4, virtual_addend_f64x4)),
+            __lasx_xvfsub_d(product_f64x4, virtual_addend_f64x4));
+
+        dot_sum_f64x4 = tentative_sum_f64x4;
+        dot_compensation_f64x4 = __lasx_xvfadd_d(dot_compensation_f64x4,
+                                                 __lasx_xvfadd_d(sum_error_f64x4, product_error_f64x4));
+
+        a_norm_sq_f64x4 = __lasx_xvfmadd_d(a_f64x4, a_f64x4, a_norm_sq_f64x4);
+        b_norm_sq_f64x4 = __lasx_xvfmadd_d(b_f64x4, b_f64x4, b_norm_sq_f64x4);
+    }
+
+    nk_f64_t dot = nk_dot_stable_sum_f64x4_loongsonasx_(dot_sum_f64x4, dot_compensation_f64x4);
+    nk_f64_t a_sq = nk_reduce_add_f64x4_loongsonasx_(a_norm_sq_f64x4);
+    nk_f64_t b_sq = nk_reduce_add_f64x4_loongsonasx_(b_norm_sq_f64x4);
+    for (; i < n; ++i) {
+        nk_f64_t a_val = a[i], b_val = b[i];
+        dot += a_val * b_val;
+        a_sq += a_val * a_val;
+        b_sq += b_val * b_val;
+    }
+    *result = nk_angular_normalize_f64_loongsonasx_(dot, a_sq, b_sq);
+}
+
+#pragma endregion F32 and F64 Floats
+
+#pragma region F16 and BF16 Floats
+
+NK_INTERNAL nk_f32_t nk_angular_normalize_f32_loongsonasx_(nk_f32_t ab, nk_f32_t a2, nk_f32_t b2) {
+    if (a2 == 0.0f && b2 == 0.0f) return 0.0f;
+    else if (ab == 0.0f) return 1.0f;
+    nk_f32_t result = 1.0f - ab * nk_f32_rsqrt_loongsonasx(a2) * nk_f32_rsqrt_loongsonasx(b2);
+    return result > 0.0f ? result : 0.0f;
+}
+
+/** @brief Horizontal sum of 8 × f32 lanes in a 256-bit LASX register. */
+NK_INTERNAL nk_f32_t nk_reduce_add_f32x8_loongsonasx_(__m256 sum_f32x8) {
+    // Add high 128-bit lane to low 128-bit lane
+    __m256 high_f32x4 = (__m256)__lasx_xvpermi_q((__m256i)sum_f32x8, (__m256i)sum_f32x8, 0x11);
+    __m256 sum_f32x4 = __lasx_xvfadd_s(sum_f32x8, high_f32x4);
+    __m256 swapped_f32x4 = (__m256)__lasx_xvshuf4i_w((__m256i)sum_f32x4, 0b01001110);
+    __m256 reduced_f32x4 = __lasx_xvfadd_s(sum_f32x4, swapped_f32x4);
+    __m256 swapped_f32x2 = (__m256)__lasx_xvshuf4i_w((__m256i)reduced_f32x4, 0b10110001);
+    __m256 reduced_f32x2 = __lasx_xvfadd_s(reduced_f32x4, swapped_f32x2);
+    nk_fui32_t c;
+    c.u = (nk_u32_t)__lasx_xvpickve2gr_w((__m256i)reduced_f32x2, 0);
+    return c.f;
+}
+
+NK_PUBLIC void nk_sqeuclidean_bf16_loongsonasx(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result) {
+    __m256 sum_f32x8 = (__m256)__lasx_xvreplgr2vr_w(0);
+    __m256i mask_high_u32x8 = __lasx_xvreplgr2vr_w((int)0xFFFF0000);
+    nk_size_t i = 0;
+    for (; i + 16 <= n; i += 16) {
+        __m256i a_bf16x16 = __lasx_xvld(a + i, 0);
+        __m256i b_bf16x16 = __lasx_xvld(b + i, 0);
+        __m256 a_even_f32x8 = (__m256)__lasx_xvslli_w(a_bf16x16, 16);
+        __m256 b_even_f32x8 = (__m256)__lasx_xvslli_w(b_bf16x16, 16);
+        __m256 diff_even_f32x8 = __lasx_xvfsub_s(a_even_f32x8, b_even_f32x8);
+        sum_f32x8 = __lasx_xvfmadd_s(diff_even_f32x8, diff_even_f32x8, sum_f32x8);
+        __m256 a_odd_f32x8 = (__m256)__lasx_xvand_v(a_bf16x16, mask_high_u32x8);
+        __m256 b_odd_f32x8 = (__m256)__lasx_xvand_v(b_bf16x16, mask_high_u32x8);
+        __m256 diff_odd_f32x8 = __lasx_xvfsub_s(a_odd_f32x8, b_odd_f32x8);
+        sum_f32x8 = __lasx_xvfmadd_s(diff_odd_f32x8, diff_odd_f32x8, sum_f32x8);
+    }
+    nk_f32_t sum = nk_reduce_add_f32x8_loongsonasx_(sum_f32x8);
+    for (; i < n; ++i) {
+        nk_f32_t a_val, b_val;
+        nk_bf16_to_f32_serial(&a[i], &a_val);
+        nk_bf16_to_f32_serial(&b[i], &b_val);
+        nk_f32_t diff = a_val - b_val;
+        sum += diff * diff;
+    }
+    *result = sum;
+}
+
+NK_PUBLIC void nk_euclidean_bf16_loongsonasx(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result) {
+    nk_sqeuclidean_bf16_loongsonasx(a, b, n, result);
+    *result = nk_f32_sqrt_loongsonasx(*result);
+}
+
+NK_PUBLIC void nk_angular_bf16_loongsonasx(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result) {
+    __m256 dot_f32x8 = (__m256)__lasx_xvreplgr2vr_w(0);
+    __m256 a_sq_f32x8 = (__m256)__lasx_xvreplgr2vr_w(0);
+    __m256 b_sq_f32x8 = (__m256)__lasx_xvreplgr2vr_w(0);
+    __m256i mask_high_u32x8 = __lasx_xvreplgr2vr_w((int)0xFFFF0000);
+    nk_size_t i = 0;
+    for (; i + 16 <= n; i += 16) {
+        __m256i a_bf16x16 = __lasx_xvld(a + i, 0);
+        __m256i b_bf16x16 = __lasx_xvld(b + i, 0);
+        __m256 a_even_f32x8 = (__m256)__lasx_xvslli_w(a_bf16x16, 16);
+        __m256 b_even_f32x8 = (__m256)__lasx_xvslli_w(b_bf16x16, 16);
+        dot_f32x8 = __lasx_xvfmadd_s(a_even_f32x8, b_even_f32x8, dot_f32x8);
+        a_sq_f32x8 = __lasx_xvfmadd_s(a_even_f32x8, a_even_f32x8, a_sq_f32x8);
+        b_sq_f32x8 = __lasx_xvfmadd_s(b_even_f32x8, b_even_f32x8, b_sq_f32x8);
+        __m256 a_odd_f32x8 = (__m256)__lasx_xvand_v(a_bf16x16, mask_high_u32x8);
+        __m256 b_odd_f32x8 = (__m256)__lasx_xvand_v(b_bf16x16, mask_high_u32x8);
+        dot_f32x8 = __lasx_xvfmadd_s(a_odd_f32x8, b_odd_f32x8, dot_f32x8);
+        a_sq_f32x8 = __lasx_xvfmadd_s(a_odd_f32x8, a_odd_f32x8, a_sq_f32x8);
+        b_sq_f32x8 = __lasx_xvfmadd_s(b_odd_f32x8, b_odd_f32x8, b_sq_f32x8);
+    }
+    nk_f32_t dot = nk_reduce_add_f32x8_loongsonasx_(dot_f32x8);
+    nk_f32_t a_sq = nk_reduce_add_f32x8_loongsonasx_(a_sq_f32x8);
+    nk_f32_t b_sq = nk_reduce_add_f32x8_loongsonasx_(b_sq_f32x8);
+    for (; i < n; ++i) {
+        nk_f32_t a_val, b_val;
+        nk_bf16_to_f32_serial(&a[i], &a_val);
+        nk_bf16_to_f32_serial(&b[i], &b_val);
+        dot += a_val * b_val;
+        a_sq += a_val * a_val;
+        b_sq += b_val * b_val;
+    }
+    *result = nk_angular_normalize_f32_loongsonasx_(dot, a_sq, b_sq);
+}
+
+NK_PUBLIC void nk_sqeuclidean_f16_loongsonasx(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
+    __m256 sum_f32x8 = (__m256)__lasx_xvreplgr2vr_w(0);
+    nk_size_t i = 0;
+    for (; i + 8 <= n; i += 8) {
+        __m128i a_f16x8 = __lsx_vld(a + i, 0);
+        __m128i b_f16x8 = __lsx_vld(b + i, 0);
+        __m256 a_f32x8 = (__m256)nk_f16x8_to_f32x8_loongsonasx_(a_f16x8);
+        __m256 b_f32x8 = (__m256)nk_f16x8_to_f32x8_loongsonasx_(b_f16x8);
+        __m256 diff_f32x8 = __lasx_xvfsub_s(a_f32x8, b_f32x8);
+        sum_f32x8 = __lasx_xvfmadd_s(diff_f32x8, diff_f32x8, sum_f32x8);
+    }
+    nk_f32_t sum = nk_reduce_add_f32x8_loongsonasx_(sum_f32x8);
+    for (; i < n; ++i) {
+        nk_f32_t a_val, b_val;
+        nk_f16_to_f32_serial(&a[i], &a_val);
+        nk_f16_to_f32_serial(&b[i], &b_val);
+        nk_f32_t diff = a_val - b_val;
+        sum += diff * diff;
+    }
+    *result = sum;
+}
+
+NK_PUBLIC void nk_euclidean_f16_loongsonasx(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
+    nk_sqeuclidean_f16_loongsonasx(a, b, n, result);
+    *result = nk_f32_sqrt_loongsonasx(*result);
+}
+
+NK_PUBLIC void nk_angular_f16_loongsonasx(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
+    __m256 dot_f32x8 = (__m256)__lasx_xvreplgr2vr_w(0);
+    __m256 a_sq_f32x8 = (__m256)__lasx_xvreplgr2vr_w(0);
+    __m256 b_sq_f32x8 = (__m256)__lasx_xvreplgr2vr_w(0);
+    nk_size_t i = 0;
+    for (; i + 8 <= n; i += 8) {
+        __m128i a_f16x8 = __lsx_vld(a + i, 0);
+        __m128i b_f16x8 = __lsx_vld(b + i, 0);
+        __m256 a_f32x8 = (__m256)nk_f16x8_to_f32x8_loongsonasx_(a_f16x8);
+        __m256 b_f32x8 = (__m256)nk_f16x8_to_f32x8_loongsonasx_(b_f16x8);
+        dot_f32x8 = __lasx_xvfmadd_s(a_f32x8, b_f32x8, dot_f32x8);
+        a_sq_f32x8 = __lasx_xvfmadd_s(a_f32x8, a_f32x8, a_sq_f32x8);
+        b_sq_f32x8 = __lasx_xvfmadd_s(b_f32x8, b_f32x8, b_sq_f32x8);
+    }
+    nk_f32_t dot = nk_reduce_add_f32x8_loongsonasx_(dot_f32x8);
+    nk_f32_t a_sq = nk_reduce_add_f32x8_loongsonasx_(a_sq_f32x8);
+    nk_f32_t b_sq = nk_reduce_add_f32x8_loongsonasx_(b_sq_f32x8);
+    for (; i < n; ++i) {
+        nk_f32_t a_val, b_val;
+        nk_f16_to_f32_serial(&a[i], &a_val);
+        nk_f16_to_f32_serial(&b[i], &b_val);
+        dot += a_val * b_val;
+        a_sq += a_val * a_val;
+        b_sq += b_val * b_val;
+    }
+    *result = nk_angular_normalize_f32_loongsonasx_(dot, a_sq, b_sq);
+}
+
+#pragma endregion F16 and BF16 Floats
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif // NK_TARGET_LOONGSONASX
+#endif // NK_TARGET_LOONGARCH_
+#endif // NK_SPATIAL_LOONGSONASX_H