Implement optimized FP16 support for ARM architecture - [MOD-9078]

cmake/aarch64InstructionFlags.cmake

@@ -8,6 +8,7 @@ CHECK_CXX_COMPILER_FLAG("-march=armv7-a+neon" CXX_ARMV7_NEON)
 CHECK_CXX_COMPILER_FLAG("-march=armv8-a" CXX_ARMV8A)
 CHECK_CXX_COMPILER_FLAG("-march=armv8-a+sve" CXX_SVE)
 CHECK_CXX_COMPILER_FLAG("-march=armv9-a+sve2" CXX_SVE2)
+CHECK_CXX_COMPILER_FLAG("-march=armv8.2-a+fp16fml" CXX_NEON_HP)
 CHECK_CXX_COMPILER_FLAG("-march=armv8.2-a+bf16" CXX_NEON_BF16)
 CHECK_CXX_COMPILER_FLAG("-march=armv8.2-a+sve+bf16" CXX_SVE_BF16)
 
@@ -17,7 +18,12 @@ if(CXX_SVE2)
   add_compile_definitions(OPT_SVE2)
 endif()
 if (CXX_ARMV8A OR CXX_ARMV7_NEON)
+  message(STATUS "Using ARMv8.0-a with NEON")
   add_compile_definitions(OPT_NEON)
+  endif()
+if (CXX_NEON_HP)
+  message(STATUS "Using ARMv8.2-a with NEON half-percision extension")
+  add_compile_definitions(OPT_NEON_HP)
 endif()
 if (CXX_NEON_BF16)
   add_compile_definitions(OPT_NEON_BF16)

src/VecSim/spaces/CMakeLists.txt

@@ -91,6 +91,13 @@ if (CMAKE_HOST_SYSTEM_PROCESSOR MATCHES "(aarch64)|(arm64)|(ARM64)|(armv.*)")
 		list(APPEND OPTIMIZATIONS functions/NEON.cpp)
 	endif()
 
+	# NEON half-precision support
+	if (CXX_NEON_HP AND CXX_ARMV8A)
+		message("Building with NEON+HP")
+		set_source_files_properties(functions/NEON_HP.cpp PROPERTIES COMPILE_FLAGS "-march=armv8.2-a+fp16fml")
+		list(APPEND OPTIMIZATIONS functions/NEON_HP.cpp)
+  endif()
+
 	# NEON bfloat16 support
 	if (CXX_NEON_BF16)
 		message("Building with NEON + BF16")

src/VecSim/spaces/IP/IP_NEON_FP16.h

@@ -0,0 +1,93 @@
+/*
+ *Copyright Redis Ltd. 2021 - present
+ *Licensed under your choice of the Redis Source Available License 2.0 (RSALv2) or
+ *the Server Side Public License v1 (SSPLv1).
+ */
+
+#include <arm_neon.h>
+
+inline void InnerProduct_Step(const float16_t *&vec1, const float16_t *&vec2, float16x8_t &acc) {
+    // Load half-precision vectors
+    float16x8_t v1 = vld1q_f16(vec1);
+    float16x8_t v2 = vld1q_f16(vec2);
+    vec1 += 8;
+    vec2 += 8;
+
+    // Multiply and accumulate
+    acc = vfmaq_f16(acc, v1, v2);
+}
+
+template <unsigned char residual> // 0..31
+float FP16_InnerProduct_NEON_HP(const void *pVect1v, const void *pVect2v, size_t dimension) {
+    const auto *vec1 = static_cast<const float16_t *>(pVect1v);
+    const auto *vec2 = static_cast<const float16_t *>(pVect2v);
+    const auto *const v1End = vec1 + dimension;
+    float16x8_t acc1 = vdupq_n_f16(0.0f);
+    float16x8_t acc2 = vdupq_n_f16(0.0f);
+    float16x8_t acc3 = vdupq_n_f16(0.0f);
+    float16x8_t acc4 = vdupq_n_f16(0.0f);
+
+    // First, handle the partial chunk residual
+    if constexpr (residual % 8) {
+        auto constexpr chunk_residual = residual % 8;
+        // TODO: spacial cases for some residuals and benchmark if its better
+        constexpr uint16x8_t mask = {
+            0xFFFF,
+            (chunk_residual >= 2) ? 0xFFFF : 0,
+            (chunk_residual >= 3) ? 0xFFFF : 0,
+            (chunk_residual >= 4) ? 0xFFFF : 0,
+            (chunk_residual >= 5) ? 0xFFFF : 0,
+            (chunk_residual >= 6) ? 0xFFFF : 0,
+            (chunk_residual >= 7) ? 0xFFFF : 0,
+            0,
+        };
+
+        // Load partial vectors
+        float16x8_t v1 = vld1q_f16(vec1);
+        float16x8_t v2 = vld1q_f16(vec2);
+
+        // Apply mask to both vectors
+        float16x8_t masked_v1 = vbslq_f16(mask, v1, acc1); // `acc1` should be all zeros here
+        float16x8_t masked_v2 = vbslq_f16(mask, v2, acc2); // `acc2` should be all zeros here
+
+        // Multiply and accumulate
+        acc1 = vfmaq_f16(acc1, masked_v1, masked_v2);
+
+        // Advance pointers
+        vec1 += chunk_residual;
+        vec2 += chunk_residual;
+    }
+
+    // Handle (residual - (residual % 8)) in chunks of 8 float16
+    if constexpr (residual >= 8)
+        InnerProduct_Step(vec1, vec2, acc2);
+    if constexpr (residual >= 16)
+        InnerProduct_Step(vec1, vec2, acc3);
+    if constexpr (residual >= 24)
+        InnerProduct_Step(vec1, vec2, acc4);
+
+    // Process the rest of the vectors (the full chunks part)
+    while (vec1 < v1End) {
+        // TODO: use `vld1q_f16_x4` for quad-loading?
+        InnerProduct_Step(vec1, vec2, acc1);
+        InnerProduct_Step(vec1, vec2, acc2);
+        InnerProduct_Step(vec1, vec2, acc3);
+        InnerProduct_Step(vec1, vec2, acc4);
+    }
+
+    // Accumulate accumulators
+    acc1 = vpaddq_f16(acc1, acc3);
+    acc2 = vpaddq_f16(acc2, acc4);
+    acc1 = vpaddq_f16(acc1, acc2);
+
+    // Horizontal sum of the accumulated values
+    float32x4_t sum_f32 = vcvt_f32_f16(vget_low_f16(acc1));
+    sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc1)));
+
+    // Pairwise add to get horizontal sum
+    float32x2_t sum_2 = vadd_f32(vget_low_f32(sum_f32), vget_high_f32(sum_f32));
+    sum_2 = vpadd_f32(sum_2, sum_2);
+
+    // Extract result
+    return 1.0f - vget_lane_f32(sum_2, 0);
+}

src/VecSim/spaces/IP/IP_SVE_FP16.h

@@ -0,0 +1,72 @@
+/*
+ *Copyright Redis Ltd. 2021 - present
+ *Licensed under your choice of the Redis Source Available License 2.0 (RSALv2) or
+ *the Server Side Public License v1 (SSPLv1).
+ */
+
+#include <arm_sve.h>
+
+inline void InnerProduct_Step(const float16_t *vec1, const float16_t *vec2, svfloat16_t &acc,
+                              size_t &offset, const size_t chunk) {
+    svbool_t all = svptrue_b16();
+
+    // Load half-precision vectors.
+    svfloat16_t v1 = svld1_f16(all, vec1 + offset);
+    svfloat16_t v2 = svld1_f16(all, vec2 + offset);
+    // Compute multiplications and add to the accumulator
+    acc = svmla_f16_x(all, acc, v1, v2);
+
+    // Move to next chunk
+    offset += chunk;
+}
+
+template <bool partial_chunk, unsigned char additional_steps> // [t/f, 0..3]
+float FP16_InnerProduct_SVE(const void *pVect1v, const void *pVect2v, size_t dimension) {
+    const auto *vec1 = static_cast<const float16_t *>(pVect1v);
+    const auto *vec2 = static_cast<const float16_t *>(pVect2v);
+    const size_t chunk = svcnth(); // number of 16-bit elements in a register
+    svbool_t all = svptrue_b16();
+    svfloat16_t acc1 = svdup_f16(0.0f);
+    svfloat16_t acc2 = svdup_f16(0.0f);
+    svfloat16_t acc3 = svdup_f16(0.0f);
+    svfloat16_t acc4 = svdup_f16(0.0f);
+    size_t offset = 0;
+
+    // Process all full vectors
+    const size_t full_iterations = dimension / chunk / 4;
+    for (size_t iter = 0; iter < full_iterations; iter++) {
+        InnerProduct_Step(vec1, vec2, acc1, offset, chunk);
+        InnerProduct_Step(vec1, vec2, acc2, offset, chunk);
+        InnerProduct_Step(vec1, vec2, acc3, offset, chunk);
+        InnerProduct_Step(vec1, vec2, acc4, offset, chunk);
+    }
+
+    // Perform between 0 and 3 additional steps, according to `additional_steps` value
+    if constexpr (additional_steps >= 1)
+        InnerProduct_Step(vec1, vec2, acc1, offset, chunk);
+    if constexpr (additional_steps >= 2)
+        InnerProduct_Step(vec1, vec2, acc2, offset, chunk);
+    if constexpr (additional_steps >= 3)
+        InnerProduct_Step(vec1, vec2, acc3, offset, chunk);
+
+    // Handle the tail with the residual predicate
+    if constexpr (partial_chunk) {
+        svbool_t pg = svwhilelt_b16_u64(offset, dimension);
+
+        // Load half-precision vectors.
+        svfloat16_t v1 = svld1_f16(pg, vec1 + offset);
+        svfloat16_t v2 = svld1_f16(pg, vec2 + offset);
+        // Compute multiplications and add to the accumulator.
+        // use the existing value of `acc` for the inactive elements (by the `m` suffix)
+        acc4 = svmla_f16_m(pg, acc4, v1, v2);
+    }
+
+    // Accumulate accumulators
+    acc1 = svadd_f16_x(all, acc1, acc3);
+    acc2 = svadd_f16_x(all, acc2, acc4);
+    acc1 = svadd_f16_x(all, acc1, acc2);
+
+    // Reduce the accumulated sum.
+    float result = svaddv_f16(all, acc1);
+    return 1.0f - result;
+}

src/VecSim/spaces/IP_space.cpp

@@ -20,6 +20,7 @@
 #include "VecSim/spaces/functions/AVX2.h"
 #include "VecSim/spaces/functions/SSE3.h"
 #include "VecSim/spaces/functions/NEON.h"
+#include "VecSim/spaces/functions/NEON_HP.h"
 #include "VecSim/spaces/functions/NEON_BF16.h"
 #include "VecSim/spaces/functions/SVE.h"
 #include "VecSim/spaces/functions/SVE_BF16.h"
@@ -214,14 +215,33 @@ dist_func_t<float> IP_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con
     if (alignment == nullptr) {
         alignment = &dummy_alignment;
     }
+    auto features = getCpuOptimizationFeatures(arch_opt);
 
     dist_func_t<float> ret_dist_func = FP16_InnerProduct;
+
+#if defined(CPU_FEATURES_ARCH_AARCH64)
+#ifdef OPT_SVE2
+    if (features.sve2) {
+        return Choose_FP16_IP_implementation_SVE2(dim);
+    }
+#endif
+#ifdef OPT_SVE
+    if (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)
+        return Choose_FP16_IP_implementation_NEON_HP(dim);
+    }
+#endif
+#endif
+
+#if defined(CPU_FEATURES_ARCH_X86_64)
     // Optimizations assume at least 32 16FPs. If we have less, we use the naive implementation.
     if (dim < 32) {
         return ret_dist_func;
     }
-#ifdef CPU_FEATURES_ARCH_X86_64
-    auto features = getCpuOptimizationFeatures(arch_opt);
 #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

src/VecSim/spaces/L2/L2_NEON_FP16.h

@@ -0,0 +1,97 @@
+/*
+ *Copyright Redis Ltd. 2021 - present
+ *Licensed under your choice of the Redis Source Available License 2.0 (RSALv2) or
+ *the Server Side Public License v1 (SSPLv1).
+ */
+
+#include <arm_neon.h>
+
+inline void L2Sqr_Step(const float16_t *&vec1, const float16_t *&vec2, float16x8_t &acc) {
+    // Load half-precision vectors
+    float16x8_t v1 = vld1q_f16(vec1);
+    float16x8_t v2 = vld1q_f16(vec2);
+    vec1 += 8;
+    vec2 += 8;
+
+    // Calculate differences
+    float16x8_t diff = vsubq_f16(v1, v2);
+    // Square and accumulate
+    acc = vfmaq_f16(acc, diff, diff);
+}
+
+template <unsigned char residual> // 0..31
+float FP16_L2Sqr_NEON_HP(const void *pVect1v, const void *pVect2v, size_t dimension) {
+    const auto *vec1 = static_cast<const float16_t *>(pVect1v);
+    const auto *vec2 = static_cast<const float16_t *>(pVect2v);
+    const auto *const v1End = vec1 + dimension;
+    float16x8_t acc1 = vdupq_n_f16(0.0f);
+    float16x8_t acc2 = vdupq_n_f16(0.0f);
+    float16x8_t acc3 = vdupq_n_f16(0.0f);
+    float16x8_t acc4 = vdupq_n_f16(0.0f);
+
+    // First, handle the partial chunk residual
+    if constexpr (residual % 8) {
+        auto constexpr chunk_residual = residual % 8;
+        // TODO: spacial cases for some residuals and benchmark if its better
+        constexpr uint16x8_t mask = {
+            0xFFFF,
+            (chunk_residual >= 2) ? 0xFFFF : 0,
+            (chunk_residual >= 3) ? 0xFFFF : 0,
+            (chunk_residual >= 4) ? 0xFFFF : 0,
+            (chunk_residual >= 5) ? 0xFFFF : 0,
+            (chunk_residual >= 6) ? 0xFFFF : 0,
+            (chunk_residual >= 7) ? 0xFFFF : 0,
+            0,
+        };
+
+        // Load partial vectors
+        float16x8_t v1 = vld1q_f16(vec1);
+        float16x8_t v2 = vld1q_f16(vec2);
+
+        // Apply mask to both vectors
+        float16x8_t masked_v1 = vbslq_f16(mask, v1, acc1); // `acc1` should be all zeros here
+        float16x8_t masked_v2 = vbslq_f16(mask, v2, acc2); // `acc2` should be all zeros here
+
+        // Calculate differences
+        float16x8_t diff = vsubq_f16(masked_v1, masked_v2);
+        // Square and accumulate
+        acc1 = vfmaq_f16(acc1, diff, diff);
+
+        // Advance pointers
+        vec1 += chunk_residual;
+        vec2 += chunk_residual;
+    }
+
+    // Handle (residual - (residual % 8)) in chunks of 8 float16
+    if constexpr (residual >= 8)
+        L2Sqr_Step(vec1, vec2, acc2);
+    if constexpr (residual >= 16)
+        L2Sqr_Step(vec1, vec2, acc3);
+    if constexpr (residual >= 24)
+        L2Sqr_Step(vec1, vec2, acc4);
+
+    // Process the rest of the vectors (the full chunks part)
+    while (vec1 < v1End) {
+        // TODO: use `vld1q_f16_x4` for quad-loading?
+        L2Sqr_Step(vec1, vec2, acc1);
+        L2Sqr_Step(vec1, vec2, acc2);
+        L2Sqr_Step(vec1, vec2, acc3);
+        L2Sqr_Step(vec1, vec2, acc4);
+    }
+
+    // Accumulate accumulators
+    acc1 = vpaddq_f16(acc1, acc3);
+    acc2 = vpaddq_f16(acc2, acc4);
+    acc1 = vpaddq_f16(acc1, acc2);
+
+    // Horizontal sum of the accumulated values
+    float32x4_t sum_f32 = vcvt_f32_f16(vget_low_f16(acc1));
+    sum_f32 = vaddq_f32(sum_f32, vcvt_f32_f16(vget_high_f16(acc1)));
+
+    // Pairwise add to get horizontal sum
+    float32x2_t sum_2 = vadd_f32(vget_low_f32(sum_f32), vget_high_f32(sum_f32));
+    sum_2 = vpadd_f32(sum_2, sum_2);
+
+    // Extract result
+    return vget_lane_f32(sum_2, 0);
+}