distance-rvv: Add support for f16 distance functions

Author: afonso360

src/distance-rvv.c

@@ -20,45 +20,79 @@ extern const char *distance_backend_name;
 // MARK: - UTILS -
 
 // Reduces a vector by summing all of it's elements into a single scalar float
-float float32_sum_vector_f32m8(vfloat32m8_t vec, size_t vl) {
+static inline float float32_sum_vector_f32m8(vfloat32m8_t vec, size_t vl) {
     vfloat32m1_t acc = __riscv_vfmv_v_f_f32m1(0.0f, 1);
     vl = __riscv_vsetvl_e32m8(vl);
     acc = __riscv_vfredusum_vs_f32m8_f32m1(vec, acc, vl);
     return __riscv_vfmv_f_s_f32m1_f32(acc);
 }
 
 // Reduces a vector by summing all of it's elements into a single scalar float
-float float32_sum_vector_f32m4(vfloat32m4_t vec, size_t vl) {
+static inline float float32_sum_vector_f32m4(vfloat32m4_t vec, size_t vl) {
     vfloat32m1_t acc = __riscv_vfmv_v_f_f32m1(0.0f, 1);
     vl = __riscv_vsetvl_e32m4(vl);
     acc = __riscv_vfredusum_vs_f32m4_f32m1(vec, acc, vl);
     return __riscv_vfmv_f_s_f32m1_f32(acc);
 }
 
+// Reduces a vector by summing all of it's elements into a single scalar double
+static inline double float64_sum_vector_f64m4(vfloat64m4_t vec, size_t vl) {
+    vfloat64m1_t acc = __riscv_vfmv_v_f_f64m1(0.0, 1);
+    vl = __riscv_vsetvl_e64m4(vl);
+    acc = __riscv_vfredusum_vs_f64m4_f64m1(vec, acc, vl);
+    return __riscv_vfmv_f_s_f64m1_f64(acc);
+}
+
 // Reduces a vector by summing all of it's elements into a single scalar integer
-uint64_t uint64_sum_vector_u64m8(vuint64m8_t vec, size_t vl) {
+static inline uint64_t uint64_sum_vector_u64m8(vuint64m8_t vec, size_t vl) {
     vuint64m1_t acc = __riscv_vmv_s_x_u64m1(0, 1);
     vl = __riscv_vsetvl_e64m8(vl);
     acc = __riscv_vredsum_vs_u64m8_u64m1(vec, acc, vl);
     return __riscv_vmv_x_s_u64m1_u64(acc);
 }
 
 // Reduces a vector by summing all of it's elements into a single scalar integer
-uint32_t uint32_sum_vector_u32m8(vuint32m8_t vec, size_t vl) {
+static inline uint32_t uint32_sum_vector_u32m8(vuint32m8_t vec, size_t vl) {
     vuint32m1_t acc = __riscv_vmv_s_x_u32m1(0, 1);
     vl = __riscv_vsetvl_e32m8(vl);
     acc = __riscv_vredsum_vs_u32m8_u32m1(vec, acc, vl);
     return __riscv_vmv_x_s_u32m1_u32(acc);
 }
 
 // Reduces a vector by summing all of it's elements into a single scalar integer
-int32_t int32_sum_vector_i32m8(vint32m8_t vec, size_t vl) {
+static inline int32_t int32_sum_vector_i32m8(vint32m8_t vec, size_t vl) {
     vint32m1_t acc = __riscv_vmv_s_x_i32m1(0, 1);
     vl = __riscv_vsetvl_e32m8(vl);
     acc = __riscv_vredsum_vs_i32m8_i32m1(vec, acc, vl);
     return __riscv_vmv_x_s_i32m1_i32(acc);
 }
 
+// Scalar-load fp16 payloads, convert to fp32, and pack as an f32m2 vector.
+static inline vfloat32m2_t rvv_load_f16_as_f32m2(const uint16_t *src, size_t n) {
+    size_t vl = __riscv_vsetvl_e32m2(n);
+    float lanes[vl];
+    for (size_t i = 0; i < vl; ++i) lanes[i] = float16_to_float32(src[i]);
+    return __riscv_vle32_v_f32m2(lanes, vl);
+}
+
+// Returns true if any lane has an fp16-style infinity mismatch:
+// one side is Inf and the other is not, or both are Inf with different signs.
+static inline bool rvv_has_f16_inf_mismatch_f64m4(vfloat64m4_t va, vfloat64m4_t vb, size_t vl) {
+    vuint64m4_t a_class = __riscv_vfclass_v_u64m4(va, vl);
+    vuint64m4_t b_class = __riscv_vfclass_v_u64m4(vb, vl);
+    vuint64m4_t a_inf_bits = __riscv_vand_vx_u64m4(a_class, 0x81u, vl);
+    vuint64m4_t b_inf_bits = __riscv_vand_vx_u64m4(b_class, 0x81u, vl);
+    vbool16_t inf_mismatch = __riscv_vmsne_vv_u64m4_b16(a_inf_bits, b_inf_bits, vl);
+    return __riscv_vfirst_m_b16(inf_mismatch, vl) >= 0;
+}
+
+// Returns mask of lanes where both vectors are not NaN.
+static inline vbool16_t rvv_both_not_nan_f64m4(vfloat64m4_t va, vfloat64m4_t vb, size_t vl) {
+    vbool16_t a_not_nan = __riscv_vmfeq_vv_f64m4_b16(va, va, vl);
+    vbool16_t b_not_nan = __riscv_vmfeq_vv_f64m4_b16(vb, vb, vl);
+    return __riscv_vmand_mm_b16(a_not_nan, b_not_nan, vl);
+}
+
 
 // MARK: - FLOAT32 -
 
@@ -213,34 +247,182 @@ float float32_distance_cosine_rvv (const void *v1, const void *v2, int n) {
 
 // MARK: - FLOAT16 -
 
+static inline float float16_distance_l2_impl_rvv(const void *v1, const void *v2, int n, bool use_sqrt) {
+    const uint16_t *a = (const uint16_t *)v1;
+    const uint16_t *b = (const uint16_t *)v2;
+
+    size_t vl = __riscv_vsetvlmax_e64m4();
+    vfloat64m4_t vsum = __riscv_vfmv_v_f_f64m4(0.0, vl);
+
+    for (size_t i = n; i > 0;) {
+        // Scalar-load fp16, convert to f32, then widen to f64.
+        vl = __riscv_vsetvl_e32m2(i);
+        vfloat32m2_t va32 = rvv_load_f16_as_f32m2(a, vl);
+        vfloat32m2_t vb32 = rvv_load_f16_as_f32m2(b, vl);
+        vfloat64m4_t va = __riscv_vfwcvt_f_f_v_f64m4(va32, vl);
+        vfloat64m4_t vb = __riscv_vfwcvt_f_f_v_f64m4(vb32, vl);
+
+        vl = __riscv_vsetvl_e64m4(vl);
+
+        // Return +Inf if there is an infinity mismatch.
+        if (rvv_has_f16_inf_mismatch_f64m4(va, vb, vl)) return INFINITY;
+
+        // Skip NaN lanes in accumulation path.
+        vbool16_t not_nan = rvv_both_not_nan_f64m4(va, vb, vl);
+
+        vfloat64m4_t vdiff = __riscv_vfsub_vv_f64m4(va, vb, vl);
+        vsum = __riscv_vfmacc_vv_f64m4_m(not_nan, vsum, vdiff, vdiff, vl);
+
+        a += vl;
+        b += vl;
+        i -= vl;
+    }
+
+    double l2sq = float64_sum_vector_f64m4(vsum, n);
+    return use_sqrt ? sqrtf((float)l2sq) : (float)l2sq;
+}
+
 float float16_distance_l2_rvv (const void *v1, const void *v2, int n) {
-    printf("float16_distance_l2_rvv: unimplemented\n");
-    abort();
-    return 0.0f;
+    return float16_distance_l2_impl_rvv(v1, v2, n, true);
 }
 
 float float16_distance_l2_squared_rvv (const void *v1, const void *v2, int n) {
-    printf("float16_distance_l2_squared_rvv: unimplemented\n");
-    abort();
-    return 0.0f;
+    return float16_distance_l2_impl_rvv(v1, v2, n, false);
 }
 
 float float16_distance_l1_rvv (const void *v1, const void *v2, int n) {
-    printf("float16_distance_l1_rvv: unimplemented\n");
-    abort();
-    return 0.0f;
+    const uint16_t *a = (const uint16_t *)v1;
+    const uint16_t *b = (const uint16_t *)v2;
+
+    size_t vl = __riscv_vsetvlmax_e64m4();
+    vfloat64m4_t vsum = __riscv_vfmv_v_f_f64m4(0.0, vl);
+
+    for (size_t i = n; i > 0;) {
+        // Scalar-load fp16, convert to f32, then widen to f64.
+        vl = __riscv_vsetvl_e32m2(i);
+        vfloat32m2_t va32 = rvv_load_f16_as_f32m2(a, vl);
+        vfloat32m2_t vb32 = rvv_load_f16_as_f32m2(b, vl);
+        vfloat64m4_t va = __riscv_vfwcvt_f_f_v_f64m4(va32, vl);
+        vfloat64m4_t vb = __riscv_vfwcvt_f_f_v_f64m4(vb32, vl);
+
+        vl = __riscv_vsetvl_e64m4(vl);
+
+        // Return +Inf if there is an infinity mismatch.
+        if (rvv_has_f16_inf_mismatch_f64m4(va, vb, vl)) return INFINITY;
+
+        // Skip NaN lanes in accumulation path.
+        vbool16_t not_nan = rvv_both_not_nan_f64m4(va, vb, vl);
+
+        vfloat64m4_t vdiff = __riscv_vfsub_vv_f64m4(va, vb, vl);
+        vfloat64m4_t vabs = __riscv_vfabs_v_f64m4(vdiff, vl);
+        vsum = __riscv_vfadd_vv_f64m4_m(not_nan, vsum, vabs, vl);
+
+        a += vl;
+        b += vl;
+        i -= vl;
+    }
+
+    return (float)float64_sum_vector_f64m4(vsum, n);
 }
 
 float float16_distance_dot_rvv (const void *v1, const void *v2, int n) {
-    printf("float16_distance_dot_rvv: unimplemented\n");
-    abort();
-    return 0.0f;
+    const uint16_t *a = (const uint16_t *)v1;
+    const uint16_t *b = (const uint16_t *)v2;
+
+    // Keep accumulation vectorized while preserving CPU NaN/Inf semantics.
+    size_t vl = __riscv_vsetvlmax_e64m4();
+    vfloat64m4_t vdot = __riscv_vfmv_v_f_f64m4(0.0, vl);
+
+    for (size_t i = n; i > 0;) {
+        // Scalar-load fp16, convert to f32, then widen to f64.
+        vl = __riscv_vsetvl_e32m2(i);
+        vfloat32m2_t va32 = rvv_load_f16_as_f32m2(a, vl);
+        vfloat32m2_t vb32 = rvv_load_f16_as_f32m2(b, vl);
+        vfloat64m4_t va = __riscv_vfwcvt_f_f_v_f64m4(va32, vl);
+        vfloat64m4_t vb = __riscv_vfwcvt_f_f_v_f64m4(vb32, vl);
+        
+        vl = __riscv_vsetvl_e64m4(vl);
+
+        // not_nan = lanes where both sides are not NaN.
+        vbool16_t not_nan = rvv_both_not_nan_f64m4(va, vb, vl);
+
+        // Multiply once, then classify the product only.
+        vfloat64m4_t vprod = __riscv_vfmul_vv_f64m4(va, vb, vl);
+
+        // Try to find infinite values, if there are any, exit early
+        vuint64m4_t p_class = __riscv_vfclass_v_u64m4(vprod, vl);
+        vbool16_t inf_pos = __riscv_vmsne_vx_u64m4_b16_m(not_nan, __riscv_vand_vx_u64m4_m(not_nan, p_class, 0x80u, vl), 0u, vl);
+        vbool16_t inf_neg = __riscv_vmsne_vx_u64m4_b16_m(not_nan, __riscv_vand_vx_u64m4_m(not_nan, p_class, 0x01u, vl), 0u, vl);
+        long first_pos = __riscv_vfirst_m_b16(inf_pos, vl);
+        long first_neg = __riscv_vfirst_m_b16(inf_neg, vl);
+        if (first_pos >= 0 || first_neg >= 0) {
+            if (first_pos >= 0 && (first_neg < 0 || first_pos < first_neg)) return -INFINITY;
+            return INFINITY;
+        }
+
+        // Accumulate only valid lanes; NaN lanes are skipped.
+        vdot = __riscv_vfadd_vv_f64m4_m(not_nan, vdot, vprod, vl);
+
+        a += vl;
+        b += vl;
+        i -= vl;
+    }
+
+    double dot = float64_sum_vector_f64m4(vdot, n);
+    return (float)(-dot);
 }
 
 float float16_distance_cosine_rvv (const void *v1, const void *v2, int n) {
-    printf("float16_distance_cosine_rvv: unimplemented\n");
-    abort();
-    return 0.0f;
+    const uint16_t *a = (const uint16_t *)v1;
+    const uint16_t *b = (const uint16_t *)v2;
+
+    size_t vl = __riscv_vsetvlmax_e64m4();
+    vfloat64m4_t vdot = __riscv_vfmv_v_f_f64m4(0.0, vl);
+    vfloat64m4_t vnx = __riscv_vfmv_v_f_f64m4(0.0, vl);
+    vfloat64m4_t vny = __riscv_vfmv_v_f_f64m4(0.0, vl);
+
+    for (size_t i = n; i > 0;) {
+        // Scalar-load fp16, convert to f32, then widen to f64.
+        vl = __riscv_vsetvl_e32m2(i);
+        vfloat32m2_t va32 = rvv_load_f16_as_f32m2(a, vl);
+        vfloat32m2_t vb32 = rvv_load_f16_as_f32m2(b, vl);
+        vfloat64m4_t va = __riscv_vfwcvt_f_f_v_f64m4(va32, vl);
+        vfloat64m4_t vb = __riscv_vfwcvt_f_f_v_f64m4(vb32, vl);
+
+        vl = __riscv_vsetvl_e64m4(vl);
+
+        // Keep only lanes where both values are not NaN.
+        vbool16_t not_nan = rvv_both_not_nan_f64m4(va, vb, vl);
+
+        // Any infinity on a valid lane returns 1.0f.
+        vuint64m4_t a_class = __riscv_vfclass_v_u64m4(va, vl);
+        vuint64m4_t b_class = __riscv_vfclass_v_u64m4(vb, vl);
+        vuint64m4_t ab_class = __riscv_vor_vv_u64m4(a_class, b_class, vl);
+        vbool16_t ab_inf = __riscv_vmsne_vx_u64m4_b16(__riscv_vand_vx_u64m4(ab_class, 0x81u, vl), 0u, vl);
+        vbool16_t any_inf = __riscv_vmand_mm_b16(not_nan, ab_inf, vl);
+        if (__riscv_vfirst_m_b16(any_inf, vl) >= 0) return 1.0f;
+
+        // Accumulate dot and squared norms on valid lanes.
+        vfloat64m4_t vprod = __riscv_vfmul_vv_f64m4(va, vb, vl);
+        vdot = __riscv_vfadd_vv_f64m4_m(not_nan, vdot, vprod, vl);
+        vnx = __riscv_vfmacc_vv_f64m4_m(not_nan, vnx, va, va, vl);
+        vny = __riscv_vfmacc_vv_f64m4_m(not_nan, vny, vb, vb, vl);
+
+        a += vl;
+        b += vl;
+        i -= vl;
+    }
+
+    double dot = float64_sum_vector_f64m4(vdot, n);
+    double nx = float64_sum_vector_f64m4(vnx, n);
+    double ny = float64_sum_vector_f64m4(vny, n);
+    double denom = sqrt(nx) * sqrt(ny);
+    if (!(denom > 0.0) || !isfinite(denom) || !isfinite(dot)) return 1.0f;
+
+    double cosv = dot / denom;
+    if (cosv > 1.0) cosv = 1.0;
+    if (cosv < -1.0) cosv = -1.0;
+    return (float)(1.0 - cosv);
 }
 
 // MARK: - BFLOAT16 -
@@ -691,31 +873,31 @@ float bit1_distance_hamming_rvv (const void *v1, const void *v2, int n) {
 void init_distance_functions_rvv (void) {
 #if defined(__riscv_v_intrinsic)
     dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F32] = float32_distance_l2_rvv;
-    // dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F16] = float16_distance_l2_rvv;
+    dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F16] = float16_distance_l2_rvv;
     // dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_BF16] = bfloat16_distance_l2_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_U8] = uint8_distance_l2_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_I8] = int8_distance_l2_rvv;
 
     dispatch_distance_table[VECTOR_DISTANCE_SQUARED_L2][VECTOR_TYPE_F32] = float32_distance_l2_squared_rvv;
-    // dispatch_distance_table[VECTOR_DISTANCE_SQUARED_L2][VECTOR_TYPE_F16] = float16_distance_l2_squared_rvv;
+    dispatch_distance_table[VECTOR_DISTANCE_SQUARED_L2][VECTOR_TYPE_F16] = float16_distance_l2_squared_rvv;
     // dispatch_distance_table[VECTOR_DISTANCE_SQUARED_L2][VECTOR_TYPE_BF16] = bfloat16_distance_l2_squared_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_SQUARED_L2][VECTOR_TYPE_U8] = uint8_distance_l2_squared_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_SQUARED_L2][VECTOR_TYPE_I8] = int8_distance_l2_squared_rvv;
 
     dispatch_distance_table[VECTOR_DISTANCE_COSINE][VECTOR_TYPE_F32] = float32_distance_cosine_rvv;
-    // dispatch_distance_table[VECTOR_DISTANCE_COSINE][VECTOR_TYPE_F16] = float16_distance_cosine_rvv;
+    dispatch_distance_table[VECTOR_DISTANCE_COSINE][VECTOR_TYPE_F16] = float16_distance_cosine_rvv;
     // dispatch_distance_table[VECTOR_DISTANCE_COSINE][VECTOR_TYPE_BF16] = bfloat16_distance_cosine_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_COSINE][VECTOR_TYPE_U8] = uint8_distance_cosine_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_COSINE][VECTOR_TYPE_I8] = int8_distance_cosine_rvv;
 
     dispatch_distance_table[VECTOR_DISTANCE_DOT][VECTOR_TYPE_F32] = float32_distance_dot_rvv;
-    // dispatch_distance_table[VECTOR_DISTANCE_DOT][VECTOR_TYPE_F16] = float16_distance_dot_rvv;
+    dispatch_distance_table[VECTOR_DISTANCE_DOT][VECTOR_TYPE_F16] = float16_distance_dot_rvv;
     // dispatch_distance_table[VECTOR_DISTANCE_DOT][VECTOR_TYPE_BF16] = bfloat16_distance_dot_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_DOT][VECTOR_TYPE_U8] = uint8_distance_dot_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_DOT][VECTOR_TYPE_I8] = int8_distance_dot_rvv;
 
     dispatch_distance_table[VECTOR_DISTANCE_L1][VECTOR_TYPE_F32] = float32_distance_l1_rvv;
-    // dispatch_distance_table[VECTOR_DISTANCE_L1][VECTOR_TYPE_F16] = float16_distance_l1_rvv;
+    dispatch_distance_table[VECTOR_DISTANCE_L1][VECTOR_TYPE_F16] = float16_distance_l1_rvv;
     // dispatch_distance_table[VECTOR_DISTANCE_L1][VECTOR_TYPE_BF16] = bfloat16_distance_l1_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_L1][VECTOR_TYPE_U8] = uint8_distance_l1_rvv;
     dispatch_distance_table[VECTOR_DISTANCE_L1][VECTOR_TYPE_I8] = int8_distance_l1_rvv;