Speed up polysemous training with AVX-512

Add AVX-512 implementations of the compute_cost and cost_update hot
loops for both ReproduceWithHammingObjective and
ReproduceDistancesObjective. The vectorized paths use 512-bit packed
double FMA, masked blends for branchless swap handling, and a portable
popcnt_512 helper that uses _mm512_popcnt_epi64 when AVX512VPOPCNTDQ
is available or falls back to a nibble-lookup approach.

Dispatch is guarded by COMPILE_SIMD_AVX512 and the SIMD dynamic
dispatch level, falling back to the existing scalar code with zero
overhead on non-AVX-512 systems.

Benchmarks of the training phase on SIFT1M (bench_polysemous_sift1m.py)
show ~1.09x speedup over the scalar path on Sapphire Rapids.

Signed-off-by: Mulugeta Mammo <[email protected]>

Author: mulugetam

benchs/bench_polysemous_sift1m.py

@@ -6,24 +6,39 @@
 
 from __future__ import print_function
 
+import time
+import numpy as np
 import faiss
 from datasets import load_sift1M, evaluate
 
+NUM_TRAIN_RUNS = 5
 
 print("load data")
 xb, xq, xt, gt = load_sift1M()
 nq, d = xq.shape
 
-# index with 16 subquantizers, 8 bit each
-index = faiss.IndexPQ(d, 16, 8)
-index.do_polysemous_training = True
-index.verbose = True
+train_times = []
+for run in range(NUM_TRAIN_RUNS):
+    index = faiss.IndexPQ(d, 16, 8)
+    index.do_polysemous_training = True
+    index.verbose = (run == 0)
 
-print("train")
+    print("train run %d/%d" % (run + 1, NUM_TRAIN_RUNS))
 
-index.train(xt)
+    t0 = time.time()
+    index.train(xt)
+    t1 = time.time()
+    elapsed = t1 - t0
+    train_times.append(elapsed)
+    print("  Training time: %.2f s" % elapsed)
 
-print("add vectors to index")
+times = np.array(train_times)
+print("\nTraining time over %d runs: "
+      "median %.2f s, mean %.2f s, std %.2f s, min %.2f s, max %.2f s"
+      % (NUM_TRAIN_RUNS, np.median(times), np.mean(times),
+         np.std(times), np.min(times), np.max(times)))
+
+print("\nadd vectors to index")
 
 index.add(xb)
 
@@ -42,3 +57,4 @@
     index.polysemous_ht = ht
     t, r = evaluate(index, xq, gt, 1)
     print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1]))
+

faiss/CMakeLists.txt

@@ -26,6 +26,7 @@ set(FAISS_SIMD_AVX512_SRC
   impl/fast_scan/impl-avx512.cpp
   impl/hnsw/avx512.cpp
   impl/pq_code_distance/avx512.cpp
+  impl/polysemous_training/avx512.cpp
   impl/scalar_quantizer/sq-avx512.cpp
   impl/binary_hamming/avx512.cpp
   utils/simd_impl/distances_avx512.cpp
@@ -282,6 +283,7 @@ set(FAISS_HEADERS
   impl/PanoramaStats.h
   impl/PdxLayout.h
   impl/PolysemousTraining.h
+  impl/polysemous_training/avx512.h
   impl/ProductQuantizer-inl.h
   impl/ProductQuantizer.h
   impl/Quantizer.h

faiss/impl/PolysemousTraining.cpp

@@ -18,6 +18,7 @@
 #include <cstring>
 #include <memory>
 
+#include <faiss/impl/polysemous_training/avx512.h>
 #include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
@@ -168,17 +169,17 @@ static inline int hamming_dis(uint64_t a, uint64_t b) {
     return popcount64(a ^ b);
 }
 
+static inline double sqr(double x) {
+    return x * x;
+}
+
 namespace {
 
 /// optimize permutation to reproduce a distance table with Hamming distances
 struct ReproduceWithHammingObjective : PermutationObjective {
     int nbits;
     double dis_weight_factor;
 
-    static double sqr(double x) {
-        return x * x;
-    }
-
     // weighting of distances: it is more important to reproduce small
     // distances well
     double dis_weight(double x) const {
@@ -190,6 +191,13 @@ struct ReproduceWithHammingObjective : PermutationObjective {
 
     // cost = quadratic difference between actual distance and Hamming distance
     double compute_cost(const int* perm) const override {
+#ifdef COMPILE_SIMD_AVX512
+        if (SIMDConfig::level == SIMDLevel::AVX512 ||
+            SIMDConfig::level == SIMDLevel::AVX512_SPR) {
+            return polysemous_avx512::hamming_compute_cost_avx512(
+                    n, perm, target_dis.data(), weights.data());
+        }
+#endif
         double cost = 0;
         for (int i = 0; i < n; i++) {
             for (int j = 0; j < n; j++) {
@@ -205,6 +213,13 @@ struct ReproduceWithHammingObjective : PermutationObjective {
     // what would the cost update be if iw and jw were swapped?
     // computed in O(n) instead of O(n^2) for the full re-computation
     double cost_update(const int* perm, int iw, int jw) const override {
+#ifdef COMPILE_SIMD_AVX512
+        if (SIMDConfig::level == SIMDLevel::AVX512 ||
+            SIMDConfig::level == SIMDLevel::AVX512_SPR) {
+            return polysemous_avx512::hamming_cost_update_avx512(
+                    n, perm, iw, jw, target_dis.data(), weights.data());
+        }
+#endif
         double delta_cost = 0;
 
         for (int i = 0; i < n; i++) {
@@ -308,6 +323,12 @@ double ReproduceDistancesObjective::get_source_dis(int i, int j) const {
 
 // cost = quadratic difference between actual distance and Hamming distance
 double ReproduceDistancesObjective::compute_cost(const int* perm) const {
+#ifdef COMPILE_SIMD_AVX512
+    if (SIMDConfig::level == SIMDLevel::AVX512 ||
+        SIMDConfig::level == SIMDLevel::AVX512_SPR) {
+        return polysemous_avx512::distances_compute_cost_avx512(*this, perm);
+    }
+#endif
     double cost = 0;
     for (int i = 0; i < n; i++) {
         for (int j = 0; j < n; j++) {
@@ -324,6 +345,13 @@ double ReproduceDistancesObjective::compute_cost(const int* perm) const {
 // computed in O(n) instead of O(n^2) for the full re-computation
 double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
         const {
+#ifdef COMPILE_SIMD_AVX512
+    if (SIMDConfig::level == SIMDLevel::AVX512 ||
+        SIMDConfig::level == SIMDLevel::AVX512_SPR) {
+        return polysemous_avx512::distances_cost_update_avx512(
+                *this, perm, iw, jw);
+    }
+#endif
     double delta_cost = 0;
     for (int i = 0; i < n; i++) {
         if (i == iw) {