MOD-19169: Use native FP16 arithmetic for SQ8–FP16 inner product

cmake/x86_64InstructionFlags.cmake

@@ -61,6 +61,10 @@ if(CXX_AVX512VL AND CXX_AVX512FP16)
 	add_compile_definitions(OPT_AVX512_FP16_VL)
 endif()
 
+if(CXX_AVX512F AND CXX_AVX512FP16 AND CXX_AVX512BW AND CXX_AVX512VL)
+	add_compile_definitions(OPT_AVX512_FP16_BW_VL)
+endif()
+
 if(CXX_AVX512F)
 	add_compile_definitions(OPT_AVX512F)
 endif()

src/VecSim/spaces/CMakeLists.txt

@@ -32,6 +32,12 @@ if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64)|(AMD64|amd64)|(^i.86$)")
 		list(APPEND OPTIMIZATIONS functions/AVX512FP16_VL.cpp)
 	endif()
 
+	if(CXX_AVX512F AND CXX_AVX512FP16 AND CXX_AVX512BW AND CXX_AVX512VL)
+		message("Building with AVX512F, AVX512FP16, AVX512BW and AVX512VL")
+		set_source_files_properties(functions/AVX512FP16_BW_VL.cpp PROPERTIES COMPILE_FLAGS "-mavx512f -mavx512fp16 -mavx512bw -mavx512vl")
+		list(APPEND OPTIMIZATIONS functions/AVX512FP16_BW_VL.cpp)
+	endif()
+
 	if(CXX_AVX512BW AND CXX_AVX512VBMI2)
 		message("Building with AVX512BW and AVX512VBMI2")
 		set_source_files_properties(functions/AVX512BW_VBMI2.cpp PROPERTIES COMPILE_FLAGS "-mavx512bw  -mavx512vbmi2")

src/VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h

@@ -0,0 +1,127 @@
+/*
+ * Copyright (c) 2006-Present, Redis Ltd.
+ * All rights reserved.
+ *
+ * Licensed under your choice of the Redis Source Available License 2.0
+ * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
+ * GNU Affero General Public License v3 (AGPLv3).
+ */
+#pragma once
+
+#include <cmath>
+#include "VecSim/spaces/space_includes.h"
+#include "VecSim/spaces/functions/AVX512F.h"
+#include "VecSim/types/sq8.h"
+#include "VecSim/types/float16.h"
+#include "VecSim/utils/alignment.h"
+
+static inline __m512h SQ8_FP16_LoadQuantizedValues_AVX512FP16(const uint8_t *quantized_values) {
+    const __m256i bytes = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(quantized_values));
+    return _mm512_cvtepu16_ph(_mm512_cvtepu8_epi16(bytes));
+}
+
+static inline void SQ8_FP16_InnerProductStep_AVX512FP16(const uint8_t *&quantized_values,
+                                                        const vecsim_types::float16 *&query,
+                                                        __m512h &sum) {
+    const __m512h values = SQ8_FP16_LoadQuantizedValues_AVX512FP16(quantized_values);
+    const __m512h query_values = _mm512_loadu_ph(query);
+    sum = _mm512_fmadd_ph(values, query_values, sum);
+    quantized_values += 32;
+    query += 32;
+}
+
+// dim >= 32 keeps the residual's full loads within the vector payloads.
+template <unsigned char residual, bool four_sums_fp32_reduce>
+static inline __m512h SQ8_FP16_AccumulateInnerProduct_AVX512FP16(const void *storage,
+                                                                 const void *query_blob,
+                                                                 size_t dimension) {
+    using float16 = vecsim_types::float16;
+    const auto *quantized_values = static_cast<const uint8_t *>(storage);
+    const auto *query = static_cast<const float16 *>(query_blob);
+    const auto *end = quantized_values + dimension;
+    __m512h sum = _mm512_setzero_ph();
+    if constexpr (residual) {
+        constexpr __mmask32 mask = (1U << residual) - 1;
+        const __m512h values = SQ8_FP16_LoadQuantizedValues_AVX512FP16(quantized_values);
+        const __m512h query_values = _mm512_loadu_ph(query);
+        sum = _mm512_maskz_mul_ph(mask, values, query_values);
+        quantized_values += residual;
+        query += residual;
+    }
+    if constexpr (four_sums_fp32_reduce) {
+        __m512h sum1 = _mm512_setzero_ph();
+        __m512h sum2 = _mm512_setzero_ph();
+        __m512h sum3 = _mm512_setzero_ph();
+        while (end - quantized_values >= 128) {
+            SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum);
+            SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum1);
+            SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum2);
+            SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum3);
+        }
+        const size_t remaining = end - quantized_values;
+        if (remaining >= 32)
+            SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum);
+        if (remaining >= 64)
+            SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum1);
+        if (remaining >= 96)
+            SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum2);
+        sum = _mm512_add_ph(_mm512_add_ph(sum, sum1), _mm512_add_ph(sum2, sum3));
+    } else {
+        do {
+            SQ8_FP16_InnerProductStep_AVX512FP16(quantized_values, query, sum);
+        } while (quantized_values < end);
+    }
+    return sum;
+}
+
+template <bool four_sums_fp32_reduce>
+static inline float SQ8_FP16_ReduceInnerProduct_AVX512FP16(__m512h sum) {
+    if constexpr (four_sums_fp32_reduce) {
+        const __m512i bits = _mm512_castph_si512(sum);
+        const __m512 low = _mm512_cvtph_ps(_mm512_castsi512_si256(bits));
+        const __m512 high = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(bits, 1));
+        return _mm512_reduce_add_ps(_mm512_add_ps(low, high));
+    } else {
+        const _Float16 reduced = _mm512_reduce_add_ph(sum);
+        return static_cast<float>(reduced);
+    }
+}
+
+static inline float SQ8_FP16_ApplyInnerProductCorrection_AVX512FP16(const void *storage,
+                                                                    const void *query_blob,
+                                                                    size_t dimension, float dot) {
+    using sq8 = vecsim_types::sq8;
+    using float16 = vecsim_types::float16;
+    const auto *storage_meta = static_cast<const uint8_t *>(storage) + dimension;
+    const float min_val = load_unaligned<float>(storage_meta + sq8::MIN_VAL * sizeof(float));
+    const float delta = load_unaligned<float>(storage_meta + sq8::DELTA * sizeof(float));
+    const auto *query_meta =
+        reinterpret_cast<const uint8_t *>(static_cast<const float16 *>(query_blob) + dimension);
+    const float query_sum = load_unaligned<float>(query_meta + sq8::SUM_QUERY * sizeof(float));
+    return 1.0f - (min_val * query_sum + delta * dot);
+}
+
+template <unsigned char residual, bool four_sums_fp32_reduce = false>
+float SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL(const void *storage, const void *query_blob,
+                                                   size_t dimension) {
+    // Directed rounding can saturate half overflow, hiding it from the finite check below.
+    if ((_mm_getcsr() & _MM_ROUND_MASK) != _MM_ROUND_NEAREST)
+        return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob,
+                                                                            dimension);
+    const __m512h sum = SQ8_FP16_AccumulateInnerProduct_AVX512FP16<residual, four_sums_fp32_reduce>(
+        storage, query_blob, dimension);
+    const float dot = SQ8_FP16_ReduceInnerProduct_AVX512FP16<four_sums_fp32_reduce>(sum);
+    // Nearest rounding leaves intermediate and reduction overflow nonfinite.
+    if (!std::isfinite(dot))
+        return spaces::Choose_SQ8_FP16_IP_implementation_AVX512F(dimension)(storage, query_blob,
+                                                                            dimension);
+    return SQ8_FP16_ApplyInnerProductCorrection_AVX512FP16(storage, query_blob, dimension, dot);
+}
+
+template <unsigned char residual>
+float SQ8_FP16_InnerProductSIMD32_FourSums_AVX512FP16_BW_VL(const void *storage,
+                                                            const void *query_blob,
+                                                            size_t dimension) {
+    return SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL<residual, true>(storage, query_blob,
+                                                                        dimension);
+}

src/VecSim/spaces/IP_space.cpp

@@ -17,6 +17,7 @@
 #include "VecSim/spaces/functions/SSE.h"
 #include "VecSim/spaces/functions/AVX512BW_VBMI2.h"
 #include "VecSim/spaces/functions/AVX512FP16_VL.h"
+#include "VecSim/spaces/functions/AVX512FP16_BW_VL.h"
 #include "VecSim/spaces/functions/AVX512BF16_VL.h"
 #include "VecSim/spaces/functions/AVX512F_BW_VL_VNNI.h"
 #include "VecSim/spaces/functions/AVX2.h"
@@ -193,7 +194,15 @@ dist_func_t<float> IP_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignment,
         return ret_dist_func;
     }
     // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract.
-    // AVX-512 tier only needs AVX-512F (cvtph_ps is part of AVX-512F, no VNNI/BW/VL required).
+#ifdef OPT_AVX512_FP16_BW_VL
+    if (dim >= spaces::SQ8_FP16_NATIVE_MIN_DIM && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM &&
+        features.avx512f && features.avx512_fp16 && features.avx512bw && features.avx512vl) {
+        if (dim % 32 == 0)
+            *alignment = 32 * sizeof(uint8_t);
+        return Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim);
+    }
+#endif
+    // The mixed-arithmetic tier needs only AVX-512F; cvtph_ps is part of that ISA.
 #ifdef OPT_AVX512F
     if (features.avx512f) {
         if (dim % 16 == 0) // SQ8 chunk = 16 bytes
@@ -271,6 +280,14 @@ dist_func_t<float> Cosine_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignm
     if (dim < 16) {
         return ret_dist_func;
     }
+#ifdef OPT_AVX512_FP16_BW_VL
+    if (dim >= spaces::SQ8_FP16_NATIVE_MIN_DIM && dim <= spaces::FP16_MAX_UNIT_IP_SIMD_DIM &&
+        features.avx512f && features.avx512_fp16 && features.avx512bw && features.avx512vl) {
+        if (dim % 32 == 0)
+            *alignment = 32 * sizeof(uint8_t);
+        return Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(dim);
+    }
+#endif
 #ifdef OPT_AVX512F
     if (features.avx512f) {
         if (dim % 16 == 0)

src/VecSim/spaces/functions/AVX512FP16_BW_VL.cpp

@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2006-Present, Redis Ltd.
+ * All rights reserved.
+ *
+ * Licensed under your choice of the Redis Source Available License 2.0
+ * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
+ * GNU Affero General Public License v3 (AGPLv3).
+ */
+#include "AVX512FP16_BW_VL.h"
+
+#include "VecSim/spaces/IP/IP_AVX512FP16_BW_VL_SQ8_FP16.h"
+
+namespace spaces {
+
+#include "implementation_chooser.h"
+
+dist_func_t<float> Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(size_t dim) {
+    dist_func_t<float> ret_dist_func;
+    if (dim >= spaces::SQ8_FP16_FOUR_SUMS_MIN_DIM) {
+        CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32,
+                              SQ8_FP16_InnerProductSIMD32_FourSums_AVX512FP16_BW_VL);
+    } else {
+        CHOOSE_IMPLEMENTATION(ret_dist_func, dim, 32, SQ8_FP16_InnerProductSIMD32_AVX512FP16_BW_VL);
+    }
+    return ret_dist_func;
+}
+
+dist_func_t<float> Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(size_t dim) {
+    return Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(dim);
+}
+
+#include "implementation_chooser_cleanup.h"
+
+} // namespace spaces

src/VecSim/spaces/functions/AVX512FP16_BW_VL.h

@@ -0,0 +1,18 @@
+/*
+ * Copyright (c) 2006-Present, Redis Ltd.
+ * All rights reserved.
+ *
+ * Licensed under your choice of the Redis Source Available License 2.0
+ * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
+ * GNU Affero General Public License v3 (AGPLv3).
+ */
+#pragma once
+
+#include "VecSim/spaces/spaces.h"
+
+namespace spaces {
+
+dist_func_t<float> Choose_SQ8_FP16_IP_implementation_AVX512FP16_BW_VL(size_t dim);
+dist_func_t<float> Choose_SQ8_FP16_Cosine_implementation_AVX512FP16_BW_VL(size_t dim);
+
+} // namespace spaces

src/VecSim/spaces/spaces.h

@@ -75,6 +75,9 @@ static constexpr size_t UINT8_MAX_EXACT_SIMD_DIM =
 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 constexpr size_t SQ8_FP16_NATIVE_MIN_DIM = 128;
+static constexpr size_t SQ8_FP16_FOUR_SUMS_MIN_DIM = 256;
+
 static inline auto getCpuOptimizationFeatures(const void *arch_opt = nullptr) {
 
 #if defined(CPU_FEATURES_ARCH_AARCH64)