perf(flat): accelerate row-major search with batched distance kernels

src/ailego/math_batch/inner_product_distance_batch.h

@@ -55,31 +55,67 @@ struct InnerProductDistanceBatchImpl {
   }
 };
 
+template <typename T, size_t BatchSize>
+struct MinusInnerProductDistanceBatchImpl {
+  using ValueType = typename std::remove_cv<T>::type;
+
+  // Keep the sign flip in the baseline-ISA caller. The dispatch translation
+  // unit is compiled for the highest available ISA so it can reference all
+  // optimized kernels; doing this work in an out-of-line specialization there
+  // can make Clang emit AVX-512 instructions before any runtime feature check.
+  static void compute_one_to_many(
+      const ValueType *query, const ValueType **ptrs,
+      std::array<const ValueType *, BatchSize> &prefetch_ptrs, size_t dim,
+      float *sums) {
+    InnerProductDistanceBatchImpl<ValueType, BatchSize>::compute_one_to_many(
+        query, ptrs, prefetch_ptrs, dim, sums);
+    for (size_t j = 0; j < BatchSize; ++j) {
+      sums[j] = -sums[j];
+    }
+  }
+};
+
+template <template <typename, size_t> class Impl, typename ValueType,
+          size_t BatchSize, size_t PrefetchStep>
+static inline void ComputeBatchChunked(const ValueType **vecs,
+                                       const ValueType *query, size_t num_vecs,
+                                       size_t dim, float *results) {
+  size_t i = 0;
+  for (; i + BatchSize <= num_vecs; i += BatchSize) {
+    std::array<const ValueType *, BatchSize> prefetch_ptrs;
+    for (size_t j = 0; j < BatchSize; ++j) {
+      if (i + j + BatchSize * PrefetchStep < num_vecs) {
+        prefetch_ptrs[j] = vecs[i + j + BatchSize * PrefetchStep];
+      } else {
+        prefetch_ptrs[j] = nullptr;
+      }
+    }
+    Impl<ValueType, BatchSize>::compute_one_to_many(
+        query, &vecs[i], prefetch_ptrs, dim, &results[i]);
+  }
+  if constexpr (std::is_same_v<ValueType, float> && BatchSize > 8) {
+    if (i + 8 <= num_vecs) {
+      ComputeBatchChunked<Impl, ValueType, 8, PrefetchStep>(&vecs[i], query, 8,
+                                                            dim, &results[i]);
+      i += 8;
+    }
+  }
+  for (; i < num_vecs; ++i) {  // TODO: unroll by 1, 2, 4, 8, etc.
+    std::array<const ValueType *, 1> prefetch_ptrs{nullptr};
+    Impl<ValueType, 1>::compute_one_to_many(query, &vecs[i], prefetch_ptrs, dim,
+                                            &results[i]);
+  }
+}
+
 template <typename T, size_t BatchSize, size_t PrefetchStep, typename>
 struct InnerProductDistanceBatch {
   using ValueType = typename std::remove_cv<T>::type;
 
   static inline void ComputeBatch(const ValueType **vecs,
                                   const ValueType *query, size_t num_vecs,
                                   size_t dim, float *results) {
-    size_t i = 0;
-    for (; i + BatchSize <= num_vecs; i += BatchSize) {
-      std::array<const ValueType *, BatchSize> prefetch_ptrs;
-      for (size_t j = 0; j < BatchSize; ++j) {
-        if (i + j + BatchSize * PrefetchStep < num_vecs) {
-          prefetch_ptrs[j] = vecs[i + j + BatchSize * PrefetchStep];
-        } else {
-          prefetch_ptrs[j] = nullptr;
-        }
-      }
-      InnerProductDistanceBatchImpl<ValueType, BatchSize>::compute_one_to_many(
-          query, &vecs[i], prefetch_ptrs, dim, &results[i]);
-    }
-    for (; i < num_vecs; ++i) {  // TODO: unroll by 1, 2, 4, 8, etc.
-      std::array<const ValueType *, 1> prefetch_ptrs{nullptr};
-      InnerProductDistanceBatchImpl<ValueType, 1>::compute_one_to_many(
-          query, &vecs[i], prefetch_ptrs, dim, &results[i]);
-    }
+    ComputeBatchChunked<InnerProductDistanceBatchImpl, ValueType, BatchSize,
+                        PrefetchStep>(vecs, query, num_vecs, dim, results);
   }
 
   static DistanceBatchQueryPreprocessFunc GetQueryPreprocessFunc() {
@@ -88,6 +124,19 @@ struct InnerProductDistanceBatch {
   }
 };
 
+template <typename T, size_t BatchSize, size_t PrefetchStep>
+struct MinusInnerProductDistanceBatch {
+  using ValueType = typename std::remove_cv<T>::type;
+
+  static inline void ComputeBatch(const ValueType **vecs,
+                                  const ValueType *query, size_t num_vecs,
+                                  size_t dim, float *results) {
+    ComputeBatchChunked<MinusInnerProductDistanceBatchImpl, ValueType,
+                        BatchSize, PrefetchStep>(vecs, query, num_vecs, dim,
+                                                 results);
+  }
+};
+
 template <>
 struct InnerProductDistanceBatchImpl<ailego::Float16, 1> {
   using ValueType = ailego::Float16;
@@ -132,6 +181,14 @@ struct InnerProductDistanceBatchImpl<float, 12> {
                                   size_t dim, float *sums);
 };
 
+template <>
+struct InnerProductDistanceBatchImpl<float, 8> {
+  using ValueType = float;
+  static void compute_one_to_many(const float *query, const float **ptrs,
+                                  std::array<const float *, 8> &prefetch_ptrs,
+                                  size_t dim, float *sums);
+};
+
 template <>
 struct InnerProductDistanceBatchImpl<int8_t, 12> {
   using ValueType = int8_t;

src/ailego/math_batch/inner_product_distance_batch_dispatch.cc

@@ -50,6 +50,21 @@ void compute_one_to_many_inner_product_avx512fp16_fp16_12(
 #endif  //__AVX512FP16__
 
 #if defined(__AVX512F__)
+void compute_one_to_many_inner_product_avx512f_fp32_1(
+    const float *query, const float **ptrs,
+    std::array<const float *, 1> &prefetch_ptrs, size_t dimensionality,
+    float *results);
+
+void compute_one_to_many_inner_product_avx512f_fp32_12(
+    const float *query, const float **ptrs,
+    std::array<const float *, 12> &prefetch_ptrs, size_t dimensionality,
+    float *results);
+
+void compute_one_to_many_inner_product_avx512f_fp32_8(
+    const float *query, const float **ptrs,
+    std::array<const float *, 8> &prefetch_ptrs, size_t dimensionality,
+    float *results);
+
 void compute_one_to_many_inner_product_avx512f_fp16_1(
     const ailego::Float16 *query, const ailego::Float16 **ptrs,
     std::array<const ailego::Float16 *, 1> &prefetch_ptrs,
@@ -82,6 +97,11 @@ void compute_one_to_many_inner_product_avx2_fp32_12(
     std::array<const float *, 12> &prefetch_ptrs, size_t dimensionality,
     float *results);
 
+void compute_one_to_many_inner_product_avx2_fp32_8(
+    const float *query, const float **ptrs,
+    std::array<const float *, 8> &prefetch_ptrs, size_t dimensionality,
+    float *results);
+
 void compute_one_to_many_inner_product_avx2_fp16_12(
     const ailego::Float16 *query, const ailego::Float16 **ptrs,
     std::array<const ailego::Float16 *, 12> &prefetch_ptrs,
@@ -96,6 +116,12 @@ void compute_one_to_many_inner_product_avx2_int8_12(
 void InnerProductDistanceBatchImpl<float, 1>::compute_one_to_many(
     const ValueType *query, const ValueType **ptrs,
     std::array<const ValueType *, 1> &prefetch_ptrs, size_t dim, float *sums) {
+#if defined(__AVX512F__)
+  if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
+    return compute_one_to_many_inner_product_avx512f_fp32_1(
+        query, ptrs, prefetch_ptrs, dim, sums);
+  }
+#endif
 #if defined(__AVX2__)
   if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
     return compute_one_to_many_inner_product_avx2_fp32_1(
@@ -167,6 +193,12 @@ InnerProductDistanceBatchImpl<int8_t, 1>::GetQueryPreprocessFunc() {
 void InnerProductDistanceBatchImpl<float, 12>::compute_one_to_many(
     const ValueType *query, const ValueType **ptrs,
     std::array<const ValueType *, 12> &prefetch_ptrs, size_t dim, float *sums) {
+#if defined(__AVX512F__)
+  if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
+    return compute_one_to_many_inner_product_avx512f_fp32_12(
+        query, ptrs, prefetch_ptrs, dim, sums);
+  }
+#endif
 #if defined(__AVX2__)
   if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
     return compute_one_to_many_inner_product_avx2_fp32_12(
@@ -177,6 +209,25 @@ void InnerProductDistanceBatchImpl<float, 12>::compute_one_to_many(
                                                     dim, sums);
 }
 
+void InnerProductDistanceBatchImpl<float, 8>::compute_one_to_many(
+    const ValueType *query, const ValueType **ptrs,
+    std::array<const ValueType *, 8> &prefetch_ptrs, size_t dim, float *sums) {
+#if defined(__AVX512F__)
+  if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
+    return compute_one_to_many_inner_product_avx512f_fp32_8(
+        query, ptrs, prefetch_ptrs, dim, sums);
+  }
+#endif
+#if defined(__AVX2__)
+  if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
+    return compute_one_to_many_inner_product_avx2_fp32_8(
+        query, ptrs, prefetch_ptrs, dim, sums);
+  }
+#endif
+  return compute_one_to_many_inner_product_fallback(query, ptrs, prefetch_ptrs,
+                                                    dim, sums);
+}
+
 void InnerProductDistanceBatchImpl<ailego::Float16, 12>::compute_one_to_many(
     const ailego::Float16 *query, const ailego::Float16 **ptrs,
     std::array<const ailego::Float16 *, 12> &prefetch_ptrs, size_t dim,

src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc

@@ -129,6 +129,13 @@ void compute_one_to_many_inner_product_avx2_fp32_12(
       query, ptrs, prefetch_ptrs, dim, sums);
 }
 
+void compute_one_to_many_inner_product_avx2_fp32_8(
+    const float *query, const float **ptrs,
+    std::array<const float *, 8> &prefetch_ptrs, size_t dim, float *sums) {
+  return compute_one_to_many_inner_product_avx2_fp32<float, 8>(
+      query, ptrs, prefetch_ptrs, dim, sums);
+}
+
 #endif
 
-}  // namespace zvec::ailego::DistanceBatch
+}  // namespace zvec::ailego::DistanceBatch

src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx512.cc

@@ -0,0 +1,89 @@
+// Copyright 2025-present the zvec project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <array>
+#include <ailego/math/matrix_utility.i>
+#include <ailego/utility/math_helper.h>
+#include <zvec/ailego/internal/platform.h>
+#include "distance_batch_math.h"
+
+namespace zvec::ailego::DistanceBatch {
+
+#if defined(__AVX512F__)
+
+template <size_t BatchSize>
+static void compute_one_to_many_inner_product_avx512f_fp32(
+    const float *query, const float **ptrs,
+    std::array<const float *, BatchSize> &prefetch_ptrs, size_t dimensionality,
+    float *results) {
+  __m512 accumulators[BatchSize];
+  for (size_t i = 0; i < BatchSize; ++i) {
+    accumulators[i] = _mm512_setzero_ps();
+  }
+
+  size_t dim = 0;
+  for (; dim + 16 <= dimensionality; dim += 16) {
+    const __m512 query_values = _mm512_loadu_ps(query + dim);
+    for (size_t i = 0; i < BatchSize; ++i) {
+      const __m512 vector_values = _mm512_loadu_ps(ptrs[i] + dim);
+      accumulators[i] =
+          _mm512_fmadd_ps(query_values, vector_values, accumulators[i]);
+    }
+    if (prefetch_ptrs[0]) {
+      for (size_t i = 0; i < BatchSize; ++i) {
+        ailego_prefetch(prefetch_ptrs[i] + dim);
+      }
+    }
+  }
+
+  if (dim < dimensionality) {
+    const auto remaining = static_cast<unsigned>(dimensionality - dim);
+    const __mmask16 mask = static_cast<__mmask16>((1u << remaining) - 1u);
+    const __m512 query_values = _mm512_maskz_loadu_ps(mask, query + dim);
+    for (size_t i = 0; i < BatchSize; ++i) {
+      const __m512 vector_values = _mm512_maskz_loadu_ps(mask, ptrs[i] + dim);
+      accumulators[i] =
+          _mm512_fmadd_ps(query_values, vector_values, accumulators[i]);
+    }
+  }
+
+  for (size_t i = 0; i < BatchSize; ++i) {
+    results[i] = HorizontalAdd_FP32_V512(accumulators[i]);
+  }
+}
+
+void compute_one_to_many_inner_product_avx512f_fp32_1(
+    const float *query, const float **ptrs,
+    std::array<const float *, 1> &prefetch_ptrs, size_t dim, float *results) {
+  compute_one_to_many_inner_product_avx512f_fp32<1>(query, ptrs, prefetch_ptrs,
+                                                    dim, results);
+}
+
+void compute_one_to_many_inner_product_avx512f_fp32_12(
+    const float *query, const float **ptrs,
+    std::array<const float *, 12> &prefetch_ptrs, size_t dim, float *results) {
+  compute_one_to_many_inner_product_avx512f_fp32<12>(query, ptrs, prefetch_ptrs,
+                                                     dim, results);
+}
+
+void compute_one_to_many_inner_product_avx512f_fp32_8(
+    const float *query, const float **ptrs,
+    std::array<const float *, 8> &prefetch_ptrs, size_t dim, float *results) {
+  compute_one_to_many_inner_product_avx512f_fp32<8>(query, ptrs, prefetch_ptrs,
+                                                    dim, results);
+}
+
+#endif
+
+}  // namespace zvec::ailego::DistanceBatch

src/core/algorithm/flat/flat_searcher.cc

@@ -102,6 +102,7 @@ int FlatSearcher<BATCH_SIZE>::load(IndexStorage::Pointer cntr,
 
   column_major_order_ = (meta_.major_order() == IndexMeta::MO_COLUMN);
   distance_matrix_.initialize(*measure_);
+  row_batch_distance_ = measure_->batch_distance();
 
   if (column_major_order_) {
     if (!distance_matrix_.is_valid()) {

src/core/algorithm/flat/flat_searcher.h

@@ -49,6 +49,7 @@ class FlatSearcher : public IndexSearcher {
   int unload(void) override {
     container_ = nullptr;
     measure_ = nullptr;
+    row_batch_distance_ = nullptr;
     features_segment_ = nullptr;
     keys_ = nullptr;
     key_id_mapping_.clear();
@@ -150,6 +151,11 @@ class FlatSearcher : public IndexSearcher {
     return distance_matrix_;
   }
 
+  //! Retrieve the row-major batch distance function
+  const IndexMetric::MatrixBatchDistance &row_batch_distance(void) const {
+    return row_batch_distance_;
+  }
+
   //! Clone a features segment
   IndexStorage::Segment::Pointer clone_features_segment(void) const {
     return features_segment_->clone();
@@ -175,6 +181,7 @@ class FlatSearcher : public IndexSearcher {
   mutable std::vector<uint32_t> mapping_{};
   mutable std::mutex mapping_mutex_{};
   FlatDistanceMatrix<BATCH_SIZE> distance_matrix_{};
+  IndexMetric::MatrixBatchDistance row_batch_distance_{};
   IndexSearcher::Stats stats_{};
 };