Make: +bf16 flag for sparse SVE2

Author: ashvardanian

include/simsimd/simsimd.h

@@ -1044,8 +1044,10 @@ SIMSIMD_PUBLIC void simsimd_find_metric_punned( //
         if (viable & simsimd_cap_sve2_k)
             switch (kind) {
             case simsimd_metric_intersect_k: *m = (m_t)&simsimd_intersect_u16_sve2, *c = simsimd_cap_sve2_k; return;
-            case simsimd_spdot_counts_u16_k: *m = (m_t)&simsimd_spdot_counts_u16_sve2, *c = simsimd_cap_sve2_k; return;
-            case simsimd_spdot_weights_u16_k:
+            case simsimd_metric_spdot_counts_k:
+                *m = (m_t)&simsimd_spdot_counts_u16_sve2, *c = simsimd_cap_sve2_k;
+                return;
+            case simsimd_metric_spdot_weights_k:
                 *m = (m_t)&simsimd_spdot_weights_u16_sve2, *c = simsimd_cap_sve2_k;
                 return;
             default: break;
@@ -1062,10 +1064,10 @@ SIMSIMD_PUBLIC void simsimd_find_metric_punned( //
         if (viable & simsimd_cap_turin_k)
             switch (kind) {
             case simsimd_metric_intersect_k: *m = (m_t)&simsimd_intersect_u16_turin, *c = simsimd_cap_turin_k; return;
-            case simsimd_spdot_counts_u16_k:
+            case simsimd_metric_spdot_counts_k:
                 *m = (m_t)&simsimd_spdot_counts_u16_turin, *c = simsimd_cap_turin_k;
                 return;
-            case simsimd_spdot_weights_u16_k:
+            case simsimd_metric_spdot_weights_k:
                 *m = (m_t)&simsimd_spdot_weights_u16_turin, *c = simsimd_cap_turin_k;
                 return;
             default: break;

include/simsimd/sparse.h

@@ -1247,17 +1247,17 @@ SIMSIMD_PUBLIC void simsimd_intersect_u32_sve2(simsimd_u32_t const* a, simsimd_u
     *results = c;
 }
 
-SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_sve2(                   //
-    simsimd_u16_t const* a, simsimd_u16_t const* b,                   //
-    simsimd_bf16_t const* a_weights, simsimd_bf16_t const* b_weights, //
-    simsimd_size_t a_length, simsimd_size_t b_length,                 //
+SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_sve2(                  //
+    simsimd_u16_t const* a, simsimd_u16_t const* b,                 //
+    simsimd_i16_t const* a_weights, simsimd_i16_t const* b_weights, //
+    simsimd_size_t a_length, simsimd_size_t b_length,               //
     simsimd_distance_t* results) {
 
     // A single SVE lane is 128 bits wide, so one lane fits 8 values.
     simsimd_size_t const register_size = svcnth();
     simsimd_size_t const lanes_count = register_size / 8;
     simsimd_size_t a_idx = 0, b_idx = 0;
-    svfloat32_t product_vec = svdupq_n_f32(0.f, 0.f, 0.f, 0.f);
+    svint64_t product_vec = svdupq_n_s64(0, 0);
     simsimd_size_t intersection_size = 0;
 
     while (a_idx < a_length && b_idx < b_length) {
@@ -1303,12 +1303,12 @@ SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_sve2(                   //
         simsimd_u64_t b_step = svcntp_b16(b_progress, b_mask);
 
         // Compare `a_vec` with each lane of `b_vec`
-        svbfloat16_t a_weights_vec = svld1_bf16(a_progress, a_weights + a_idx);
-        svbfloat16_t b_weights_vec = svld1_bf16(b_progress, b_weights + b_idx);
+        svint16_t a_weights_vec = svld1_s16(a_progress, a_weights + a_idx);
+        svint16_t b_weights_vec = svld1_s16(b_progress, b_weights + b_idx);
         for (simsimd_size_t i = 0; i < lanes_count; i++) {
             svbool_t equal_mask = svmatch_u16(a_progress, a_vec, b_vec);
-            svbfloat16_t b_equal_weights_vec = svsel_bf16(equal_mask, b_weights_vec, svdup_n_bf16(0.f));
-            product_vec = svbfdot_f32(product_vec, a_weights_vec, b_equal_weights_vec);
+            svint16_t b_equal_weights_vec = svsel_s16(equal_mask, b_weights_vec, svdup_n_s16(0.f));
+            product_vec = svdot_s64(product_vec, a_weights_vec, b_equal_weights_vec);
             b_vec = svext_u16(b_vec, b_vec, 8);
             intersection_size += svcntp_b16(svptrue_b16(), equal_mask);
         }
@@ -1318,20 +1318,29 @@ SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_sve2(                   //
         b_idx += b_step;
     }
     results[0] = (simsimd_distance_t)intersection_size;
-    results[1] = svaddv_f32(svptrue_b32(), product_vec);
+    results[1] = svaddv_s64(svptrue_b64(), product_vec);
 }
 
-SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_sve2(                  //
-    simsimd_u16_t const* a, simsimd_u16_t const* b,                 //
-    simsimd_i16_t const* a_weights, simsimd_i16_t const* b_weights, //
-    simsimd_size_t a_length, simsimd_size_t b_length,               //
+#pragma clang attribute pop
+#pragma GCC pop_options
+#endif // SIMSIMD_TARGET_SVE2
+
+#if SIMSIMD_TARGET_SVE2 && SIMSIMD_TARGET_SVE_BF16
+#pragma GCC push_options
+#pragma GCC target("arch=armv8.6-a+sve+sve2+bf16")
+#pragma clang attribute push(__attribute__((target("arch=armv8.6-a+sve+sve2+bf16"))), apply_to = function)
+
+SIMSIMD_PUBLIC void simsimd_spdot_weights_u16_sve2(                   //
+    simsimd_u16_t const* a, simsimd_u16_t const* b,                   //
+    simsimd_bf16_t const* a_weights, simsimd_bf16_t const* b_weights, //
+    simsimd_size_t a_length, simsimd_size_t b_length,                 //
     simsimd_distance_t* results) {
 
     // A single SVE lane is 128 bits wide, so one lane fits 8 values.
     simsimd_size_t const register_size = svcnth();
     simsimd_size_t const lanes_count = register_size / 8;
     simsimd_size_t a_idx = 0, b_idx = 0;
-    svint64_t product_vec = svdupq_n_s64(0, 0);
+    svfloat32_t product_vec = svdupq_n_f32(0.f, 0.f, 0.f, 0.f);
     simsimd_size_t intersection_size = 0;
 
     while (a_idx < a_length && b_idx < b_length) {
@@ -1377,12 +1386,15 @@ SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_sve2(                  //
         simsimd_u64_t b_step = svcntp_b16(b_progress, b_mask);
 
         // Compare `a_vec` with each lane of `b_vec`
-        svbfloat16_t a_weights_vec = svld1_s16(a_progress, a_weights + a_idx);
-        svbfloat16_t b_weights_vec = svld1_s16(b_progress, b_weights + b_idx);
+        svbfloat16_t a_weights_vec = svld1_bf16(a_progress, a_weights + a_idx);
+        svbfloat16_t b_weights_vec = svld1_bf16(b_progress, b_weights + b_idx);
         for (simsimd_size_t i = 0; i < lanes_count; i++) {
             svbool_t equal_mask = svmatch_u16(a_progress, a_vec, b_vec);
-            svbfloat16_t b_equal_weights_vec = svsel_s16(equal_mask, b_weights_vec, svdup_n_bf16(0.f));
-            product_vec = svdot_s64(product_vec, a_weights_vec, b_equal_weights_vec);
+            //! The `svsel_bf16` intrinsic is broken in many compilers, not returning the correct type.
+            //! So we reinterprete floats as integers and apply `svsel_s16`.
+            svint16_t b_equal_weights_vec =
+                svsel_s16(equal_mask, svreinterpret_s16_bs16(b_weights_vec), svdup_n_s16(0));
+            product_vec = svbfdot_f32(product_vec, a_weights_vec, svreinterpret_bf16_s16(b_equal_weights_vec));
             b_vec = svext_u16(b_vec, b_vec, 8);
             intersection_size += svcntp_b16(svptrue_b16(), equal_mask);
         }
@@ -1392,12 +1404,12 @@ SIMSIMD_PUBLIC void simsimd_spdot_counts_u16_sve2(                  //
         b_idx += b_step;
     }
     results[0] = (simsimd_distance_t)intersection_size;
-    results[1] = svaddv_s64(svptrue_b64(), product_vec);
+    results[1] = svaddv_f32(svptrue_b32(), product_vec);
 }
 
 #pragma clang attribute pop
 #pragma GCC pop_options
-#endif // SIMSIMD_TARGET_SVE2
+#endif // SIMSIMD_TARGET_SVE2 && SIMSIMD_TARGET_SVE_BF16
 #endif // SIMSIMD_TARGET_ARM
 
 #ifdef __cplusplus