[MOD-17916] Bound native FP16 dispatch without hot-path overhead

src/VecSim/spaces/IP/IP_AVX512FP16_VL_FP16.h

@@ -22,30 +22,39 @@ static void InnerProductStep(float16 *&pVect1, float16 *&pVect2, __m512h &sum) {
     pVect2 += 32;
 }
 
-template <unsigned short residual> // 0..31
+template <unsigned short residual> // 0..63
 float FP16_InnerProductSIMD32_AVX512FP16_VL(const void *pVect1v, const void *pVect2v,
                                             size_t dimension) {
     auto *pVect1 = (float16 *)pVect1v;
     auto *pVect2 = (float16 *)pVect2v;
 
     const float16 *pEnd1 = pVect1 + dimension;
 
-    __m512h sum = _mm512_setzero_ph();
+    // Two accumulators break the FMA dependency chain, letting more FMAs be in flight at once.
+    __m512h sum0 = _mm512_setzero_ph();
+    __m512h sum1 = _mm512_setzero_ph();
 
-    if constexpr (residual) {
-        constexpr __mmask32 mask = (1LU << residual) - 1;
+    if constexpr (residual % 32) {
+        constexpr __mmask32 mask = (1LU << (residual % 32)) - 1;
         __m512h v1 = _mm512_loadu_ph(pVect1);
-        pVect1 += residual;
+        pVect1 += residual % 32;
         __m512h v2 = _mm512_loadu_ph(pVect2);
-        pVect2 += residual;
-        sum = _mm512_maskz_mul_ph(mask, v1, v2);
+        pVect2 += residual % 32;
+        sum0 = _mm512_maskz_mul_ph(mask, v1, v2);
     }
 
-    // We dealt with the residual part. We are left with some multiple of 32 16-bit floats.
-    do {
-        InnerProductStep(pVect1, pVect2, sum);
-    } while (pVect1 < pEnd1);
+    if constexpr (residual >= 32) {
+        InnerProductStep(pVect1, pVect2, sum1);
+    }
+
+    // We dealt with the residual part. We are left with some multiple of 64 16-bit floats.
+    while (pVect1 < pEnd1) {
+        InnerProductStep(pVect1, pVect2, sum0);
+        InnerProductStep(pVect1, pVect2, sum1);
+    }
 
-    _Float16 res = _mm512_reduce_add_ph(sum);
-    return _Float16(1) - res;
+    const __m512h sum = _mm512_add_ph(sum0, sum1);
+    const _Float16 reduced = _mm512_reduce_add_ph(sum);
+    // Subtract in fp32 so distances close to 1.0 are not rounded back to fp16.
+    return 1.0f - static_cast<float>(reduced);
 }

src/VecSim/spaces/IP_space.cpp

@@ -632,17 +632,18 @@ dist_func_t<float> IP_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con
 
 #if defined(CPU_FEATURES_ARCH_AARCH64)
 #ifdef OPT_SVE2
-    if (features.sve2) {
+    if (dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.sve2) {
         return Choose_FP16_IP_implementation_SVE2(dim);
     }
 #endif
 #ifdef OPT_SVE
-    if (features.sve) {
+    if (dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.sve) {
         return Choose_FP16_IP_implementation_SVE(dim);
     }
 #endif
 #ifdef OPT_NEON_HP
-    if (features.asimdhp && dim >= 8) { // Optimization assumes at least 8 16FPs (full chunk)
+    if (dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.asimdhp &&
+        dim >= 8) { // Optimization assumes at least 8 16FPs (full chunk)
         return Choose_FP16_IP_implementation_NEON_HP(dim);
     }
 #endif
@@ -655,7 +656,8 @@ dist_func_t<float> IP_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con
 #ifdef OPT_AVX512_FP16_VL
     // More details about the dimension limitation can be found in this PR's description:
     // https://github.com/RedisAI/VectorSimilarity/pull/477
-    if (dim >= 32 && features.avx512_fp16 && features.avx512vl) {
+    if (dim >= 32 && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM && features.avx512_fp16 &&
+        features.avx512vl) {
         if (dim % 32 == 0) // no point in aligning if we have an offsetting residual
             *alignment = 32 * sizeof(float16); // handles 32 floats
         return Choose_FP16_IP_implementation_AVX512FP16_VL(dim);

src/VecSim/spaces/L2_space.cpp

@@ -368,17 +368,18 @@ dist_func_t<float> L2_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con
 
 #if defined(CPU_FEATURES_ARCH_AARCH64)
 #ifdef OPT_SVE2
-    if (features.sve2) {
+    if (dim <= spaces::FP16_MAX_UNIT_L2_SIMD_DIM && features.sve2) {
         return Choose_FP16_L2_implementation_SVE2(dim);
     }
 #endif
 #ifdef OPT_SVE
-    if (features.sve) {
+    if (dim <= spaces::FP16_MAX_UNIT_L2_SIMD_DIM && features.sve) {
         return Choose_FP16_L2_implementation_SVE(dim);
     }
 #endif
 #ifdef OPT_NEON_HP
-    if (features.asimdhp && dim >= 8) { // Optimization assumes at least 8 16FPs (full chunk)
+    if (dim <= spaces::FP16_MAX_UNIT_L2_SIMD_DIM && features.asimdhp &&
+        dim >= 8) { // Optimization assumes at least 8 16FPs (full chunk)
         return Choose_FP16_L2_implementation_NEON_HP(dim);
     }
 #endif
@@ -391,7 +392,8 @@ dist_func_t<float> L2_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con
 #ifdef OPT_AVX512_FP16_VL
     // More details about the dimension limitation can be found in this PR's description:
     // https://github.com/RedisAI/VectorSimilarity/pull/477
-    if (dim >= 32 && features.avx512_fp16 && features.avx512vl) {
+    if (dim >= 32 && dim <= spaces::FP16_MAX_UNIT_L2_SIMD_DIM && features.avx512_fp16 &&
+        features.avx512vl) {
         if (dim % 32 == 0) // no point in aligning if we have an offsetting residual
             *alignment = 32 * sizeof(float16); // handles 32 floats
         return Choose_FP16_L2_implementation_AVX512FP16_VL(dim);

src/VecSim/spaces/functions/AVX512FP16_VL.cpp

@@ -17,7 +17,7 @@ namespace spaces {
 
 dist_func_t<float> Choose_FP16_IP_implementation_AVX512FP16_VL(size_t dim) {
     dist_func_t<float> ret_dist_func;
-    CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, FP16_InnerProductSIMD32_AVX512FP16_VL);
+    CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 64, FP16_InnerProductSIMD32_AVX512FP16_VL);
     return ret_dist_func;
 }
 

src/VecSim/spaces/spaces.h

@@ -65,6 +65,16 @@ static constexpr size_t UINT8_MAX_EXACT_SIMD_DIM =
     std::numeric_limits<int32_t>::max() /
     (std::numeric_limits<uint8_t>::max() * std::numeric_limits<uint8_t>::max());
 
+// Native-fp16 accumulation is a throughput optimization and may overflow for sufficiently large
+// values at any dimension. These limits are conservative dispatch guardrails, not an input-range
+// contract or a general overflow guarantee. They retain the native path at ordinary embedding
+// dimensions while avoiding dimensions where unit-scale components alone can exceed fp16's largest
+// finite value (65,504): IP contributes at most 1 per component, and an L2 difference of at most 2
+// contributes at most 4. The chooser pays this check once when an index is created; the native
+// accumulation/reduction loops pay no additional instructions.
+static constexpr size_t FP16_MAX_UNIT_IP_SIMD_DIM = 65504;
+static constexpr size_t FP16_MAX_UNIT_L2_SIMD_DIM = FP16_MAX_UNIT_IP_SIMD_DIM / 4;
+
 static inline auto getCpuOptimizationFeatures(const void *arch_opt = nullptr) {
 
 #if defined(CPU_FEATURES_ARCH_AARCH64)

tests/flow/common.py

@@ -119,6 +119,55 @@ def get_ground_truth_results(dist_func, query, vectors, k):
 
     return results, keys
 
+# Native-fp16 SIMD kernels intentionally perform arithmetic in half precision. Their unit tests
+# bound the resulting error at 1%, so flow tests must validate score/order stability within the
+# same contract instead of assuming scalar-fp32-identical distances on every runner CPU.
+FLOAT16_NATIVE_RTOL = 1e-2
+
+
+def get_distances_by_label(dist_func, query, vectors):
+    distances = {}
+    for label, vector in vectors:
+        distance = dist_func(query, vector)
+        distances[label] = min(distance, distances.get(label, distance))
+    return distances
+
+
+def assert_float16_l2_scores(labels, distances, exact_distances):
+    for label, distance in zip(labels, distances):
+        assert math.isclose(float(distance), exact_distances[int(label)],
+                            rel_tol=FLOAT16_NATIVE_RTOL, abs_tol=0)
+
+
+def assert_float16_l2_knn(labels, distances, exact_distances, k):
+    assert len(labels) == k
+    returned = set(map(int, labels))
+    assert len(returned) == k
+
+    cutoff = sorted(exact_distances.values())[k - 1]
+    mandatory = {label for label, distance in exact_distances.items()
+                 if distance < cutoff * (1 - FLOAT16_NATIVE_RTOL)}
+    allowed = {label for label, distance in exact_distances.items()
+               if distance <= cutoff * (1 + FLOAT16_NATIVE_RTOL)}
+    assert mandatory.issubset(returned)
+    assert returned.issubset(allowed)
+    assert_float16_l2_scores(labels, distances, exact_distances)
+
+
+def assert_float16_l2_range(labels, distances, exact_distances, radius, require_inner=True):
+    returned = set(map(int, labels))
+    assert len(returned) == len(labels)
+
+    inner = {label for label, distance in exact_distances.items()
+             if distance <= radius * (1 - FLOAT16_NATIVE_RTOL)}
+    allowed = {label for label, distance in exact_distances.items()
+               if distance <= radius * (1 + FLOAT16_NATIVE_RTOL)}
+    if require_inner:
+        assert inner.issubset(returned)
+    assert returned.issubset(allowed)
+    assert_float16_l2_scores(labels, distances, exact_distances)
+
+
 def fp32_expand_and_calc_cosine_dist(a, b):
     # stupid numpy doesn't make any intermediate conversions when handling small types
     # so we might get overflow. We need to convert to float32 ourselves.

tests/flow/test_bruteforce.py

@@ -443,10 +443,10 @@ class TestFloat16():
     def test_bf_float16_L2(self, test_logger):
         k = 10
 
-        keys, dists = self.data.measure_dists(k)
         bf_labels, bf_distances = self.data.index.knn_query(self.data.query, k=k)
-        assert_allclose(bf_labels, [keys],  rtol=1e-5, atol=0)
-        assert_allclose(bf_distances, [dists],  rtol=1e-5, atol=0)
+        exact_distances = get_distances_by_label(spatial.distance.sqeuclidean,
+                                                 self.data.query.flat, self.data.vectors)
+        assert_float16_l2_knn(bf_labels[0], bf_distances[0], exact_distances, k)
         test_logger.info(f"sanity test for {self.data.metric} and {self.data.type} pass")
 
     def test_bf_float16_batch_iterator(self, test_logger):
@@ -461,8 +461,8 @@ def test_bf_float16_batch_iterator(self, test_logger):
 
         _, distances_second_batch = batch_iterator.get_next_results(10, BY_SCORE)
         for i, dist in enumerate(distances_second_batch[0][:-1]):
-            # assert sorting by score
-            assert(distances_second_batch[0][i] < distances_second_batch[0][i+1])
+            # Native fp16 scores can tie after rounding; they must remain nondecreasing.
+            assert(distances_second_batch[0][i] <= distances_second_batch[0][i+1])
             # assert that every distance in the second batch is higher than any distance of the first batch
             assert(len(distances_first_batch[0][np.where(distances_first_batch[0] > dist)]) == 0)
 
@@ -494,14 +494,10 @@ def test_bf_float16_range_query(self, test_logger):
         res_num = len(bf_labels[0])
         test_logger.info(f'lookup time for {self.num_labels} vectors with dim={self.dim} took {end - start} seconds, got {res_num} results')
 
-        # Verify that we got exactly all vectors within the range
-        results, keys = get_ground_truth_results(spatial.distance.sqeuclidean, query_data.flat, self.data.vectors, res_num)
-
-        assert_allclose(max(bf_distances[0]), results[res_num-1]["dist"], rtol=1e-05)
-        assert np.array_equal(np.array(bf_labels[0]), np.array(keys))
+        exact_distances = get_distances_by_label(spatial.distance.sqeuclidean,
+                                                 query_data.flat, self.data.vectors)
+        assert_float16_l2_range(bf_labels[0], bf_distances[0], exact_distances, radius)
         assert max(bf_distances[0]) <= radius
-        # Verify that the next closest vector that hasn't returned is not within the range
-        assert results[res_num]["dist"] > radius
 
         # Expect zero results for radius==0
         bf_labels, bf_distances = bfindex.range_query(query_data, radius=0)
@@ -519,27 +515,16 @@ def test_bf_float16_multivalue(test_logger):
     k=10
 
     query_data = data.query
-    dists = {}
-    for key, vec in data.vectors:
-        # Setting or updating the score for each label.
-        # If it's the first time we calculate a score for a label dists.get(key, dist)
-        # will return dist so we will choose the actual score the first time.
-        dist = spatial.distance.sqeuclidean(query_data.flat, vec)
-        dists[key] = min(dist, dists.get(key, dist))
-
-    dists = list(dists.items())
-    dists = sorted(dists, key=lambda pair: pair[1])[:k]
-    keys = [key for key, _ in dists[:k]]
-    dists = [dist for _, dist in dists[:k]]
 
     start = time.time()
     bf_labels, bf_distances = data.index.knn_query(query_data, k=10)
     end = time.time()
 
     test_logger.info(f'lookup time for {num_elements} vectors ({num_labels} labels and {num_per_label} vectors per label) with dim={dim} took {end - start} seconds')
 
-    assert_allclose(bf_labels, [keys],  rtol=1e-5, atol=0)
-    assert_allclose(bf_distances, [dists],  rtol=1e-5, atol=0)
+    exact_distances = get_distances_by_label(spatial.distance.sqeuclidean,
+                                             query_data.flat, data.vectors)
+    assert_float16_l2_knn(bf_labels[0], bf_distances[0], exact_distances, k)
 
 '''
 A Class to run common tests for BF index

tests/flow/test_hnsw.py

@@ -761,8 +761,8 @@ def test_batch_iterator(self):
         labels_second_batch, distances_second_batch = batch_iterator.get_next_results(10, BY_SCORE)
         should_have_return_in_first_batch = []
         for i, dist in enumerate(distances_second_batch[0][:-1]):
-            # Assert sorting by score
-            assert (distances_second_batch[0][i] < distances_second_batch[0][i + 1])
+            # Native fp16 scores can tie after rounding; they must remain nondecreasing.
+            assert (distances_second_batch[0][i] <= distances_second_batch[0][i + 1])
             # Assert that every distance in the second batch is higher than any distance of the first batch
             if len(distances_first_batch[0][np.where(distances_first_batch[0] > dist)]) != 0:
                 should_have_return_in_first_batch.append(dist)
@@ -790,18 +790,23 @@ def test_range_query(self, test_logger):
             end = time.time()
             res_num = len(hnsw_labels[0])
 
-            dists = sorted([(key, spatial.distance.sqeuclidean(self.query_data[0], vec)) for key, vec in self.vectors])
-            actual_results = [(key, dist) for key, dist in dists if dist <= radius]
+            exact_distances = get_distances_by_label(spatial.distance.sqeuclidean,
+                                                     self.query_data[0], self.vectors)
+            actual_labels = {label for label, distance in exact_distances.items()
+                             if distance <= radius}
 
             test_logger.info(
                 f'lookup time for {self.num_elements} vectors with dim={self.dim} took {end - start} seconds with epsilon={epsilon_rt},'
-                f' got {res_num} results, which are {res_num / len(actual_results)} of the entire results in the range.')
+                f' got {res_num} results, which are {res_num / len(actual_labels)} of the entire results in the range.')
 
-            # Compare the number of vectors that are actually within the range to the returned results.
-            assert np.all(np.isin(hnsw_labels, np.array([label for label, _ in actual_results])))
+            # HNSW is approximate, so validate returned labels and scores without requiring every
+            # vector in the exact inner range to be present.
+            assert_float16_l2_range(hnsw_labels[0], hnsw_distances[0], exact_distances,
+                                    radius, require_inner=False)
 
             assert max(hnsw_distances[0]) <= radius
-            recalls[epsilon_rt] = res_num / len(actual_results)
+            returned_labels = set(map(int, hnsw_labels[0]))
+            recalls[epsilon_rt] = len(returned_labels.intersection(actual_labels)) / len(actual_labels)
 
         # Expect higher recalls for higher epsilon values.
         assert recalls[0.001] <= recalls[0.01] <= recalls[0.1]