upgrade faiss baseline f9f116d30 -> 9d567497e (22 upstream commits) (#1606)

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

Author: alexanderguzhva

thirdparty/faiss/benchs/bench_hnsw.py

@@ -32,7 +32,7 @@
 nq, d = xq.shape
 
 if todo == []:
-    todo = 'hnsw hnsw_sq ivf ivf_hnsw_quantizer kmeans kmeans_hnsw nsg'.split()
+    todo = "hnsw hnsw_sq ivf ivf_hnsw_quantizer kmeans kmeans_hnsw nsg".split()
 
 
 def evaluate(index):
@@ -45,11 +45,13 @@ def evaluate(index):
 
     missing_rate = (I == -1).sum() / float(k * nq)
     recall_at_1 = (I == gt[:, :1]).sum() / float(nq)
-    print("\t %7.3f ms per query, R@1 %.4f, missing rate %.4f" % (
-        (t1 - t0) * 1000.0 / nq, recall_at_1, missing_rate))
+    print(
+        "\t %7.3f ms per query, R@1 %.4f, missing rate %.4f"
+        % ((t1 - t0) * 1000.0 / nq, recall_at_1, missing_rate)
+    )
 
 
-if 'hnsw' in todo:
+if "hnsw" in todo:
 
     print("Testing HNSW Flat")
 
@@ -69,12 +71,12 @@ def evaluate(index):
     print("search")
     for efSearch in 16, 32, 64, 128, 256:
         for bounded_queue in [True, False]:
-            print("efSearch", efSearch, "bounded queue", bounded_queue, end=' ')
+            print("efSearch", efSearch, "bounded queue", bounded_queue, end=" ")
             index.hnsw.search_bounded_queue = bounded_queue
             index.hnsw.efSearch = efSearch
             evaluate(index)
 
-if 'hnsw_sq' in todo:
+if "hnsw_sq" in todo:
 
     print("Testing HNSW with a scalar quantizer")
     # also set M so that the vectors and links both use 128 bytes per
@@ -96,16 +98,16 @@ def evaluate(index):
 
     print("search")
     for efSearch in 16, 32, 64, 128, 256:
-        print("efSearch", efSearch, end=' ')
+        print("efSearch", efSearch, end=" ")
         index.hnsw.efSearch = efSearch
         evaluate(index)
 
-if 'ivf' in todo:
+if "ivf" in todo:
 
     print("Testing IVF Flat (baseline)")
     quantizer = faiss.IndexFlatL2(d)
     index = faiss.IndexIVFFlat(quantizer, d, 16384)
-    index.cp.min_points_per_centroid = 5   # quiet warning
+    index.cp.min_points_per_centroid = 5  # quiet warning
 
     # to see progress
     index.verbose = True
@@ -118,16 +120,16 @@ def evaluate(index):
 
     print("search")
     for nprobe in 1, 4, 16, 64, 256:
-        print("nprobe", nprobe, end=' ')
+        print("nprobe", nprobe, end=" ")
         index.nprobe = nprobe
         evaluate(index)
 
-if 'ivf_hnsw_quantizer' in todo:
+if "ivf_hnsw_quantizer" in todo:
 
     print("Testing IVF Flat with HNSW quantizer")
     quantizer = faiss.IndexHNSWFlat(d, 32)
     index = faiss.IndexIVFFlat(quantizer, d, 16384)
-    index.cp.min_points_per_centroid = 5   # quiet warning
+    index.cp.min_points_per_centroid = 5  # quiet warning
     index.quantizer_trains_alone = 2
 
     # to see progress
@@ -142,13 +144,13 @@ def evaluate(index):
     print("search")
     quantizer.hnsw.efSearch = 64
     for nprobe in 1, 4, 16, 64, 256:
-        print("nprobe", nprobe, end=' ')
+        print("nprobe", nprobe, end=" ")
         index.nprobe = nprobe
         evaluate(index)
 
 # Bonus: 2 kmeans tests
 
-if 'kmeans' in todo:
+if "kmeans" in todo:
     print("Performing kmeans on sift1M database vectors (baseline)")
     clus = faiss.Clustering(d, 16384)
     clus.verbose = True
@@ -157,7 +159,7 @@ def evaluate(index):
     clus.train(xb, index)
 
 
-if 'kmeans_hnsw' in todo:
+if "kmeans_hnsw" in todo:
     print("Performing kmeans on sift1M using HNSW assignment")
     clus = faiss.Clustering(d, 16384)
     clus.verbose = True
@@ -168,7 +170,7 @@ def evaluate(index):
     index.hnsw.efSearch = 128
     clus.train(xb, index)
 
-if 'nsg' in todo:
+if "nsg" in todo:
 
     print("Testing NSG Flat")
 
@@ -186,6 +188,47 @@ def evaluate(index):
 
     print("search")
     for search_L in -1, 16, 32, 64, 128, 256:
-        print("search_L", search_L, end=' ')
+        print("search_L", search_L, end=" ")
         index.nsg.search_L = search_L
         evaluate(index)
+
+
+if "hnsw_locks" in todo:
+
+    ntotal, _ = xb.shape
+    batch_size = ntotal // 100
+    print(
+        f"Testing HNSW Flat: add with {batch_size=}, "
+        "with and without retaining locks"
+    )
+
+    # Unbatched
+    t0 = time.time()
+    index = faiss.IndexHNSWFlat(d, 32)
+    index.add(xb)
+    t1 = time.time()
+    print(
+        f"\t single bulk add(): {index.ntotal} added in {t1 - t0:6.3f}s"
+        f" = {index.ntotal / (t1 - t0):.0f}/s"
+    )
+
+    for retain_locks in [False, True]:
+        index = faiss.IndexHNSWFlat(d, 32)
+        index.retain_locks = retain_locks
+
+        t0 = time.time()
+        t1 = None
+        t2 = None
+        for i in range(0, len(xb), batch_size):
+            t1 = time.time()
+            index.add(xb[i: i + batch_size])
+            t2 = time.time()
+            if i > 2 and t2 - t0 > 2:
+                break
+
+        assert t1 and t2
+        dt = t2 - t0
+        print(
+            f"\t {retain_locks=:1}: {index.ntotal} added in {t2 - t0:6.3f}s"
+            f" = {index.ntotal / (t2 - t0):.0f}/s"
+        )

thirdparty/faiss/conda/faiss-gpu/meta.yaml

@@ -49,7 +49,7 @@ outputs:
         - FAISS_FLATTEN_CONDA_INCLUDES
     requirements:
       build:
-        - {{ compiler('cxx') }} =12.4
+        - {{ compiler('cxx') }} >=13.4,<14
         - sysroot_linux-64 =2.17 # [linux64]
         - llvm-openmp  # [osx]
         - cmake >=3.24.0
@@ -85,7 +85,7 @@ outputs:
       string: "py{{ PY_VER }}_h{{ PKG_HASH }}_{{ number }}_cuda{{ cudatoolkit }}{{ suffix }}"
     requirements:
       build:
-        - {{ compiler('cxx') }} =12.4
+        - {{ compiler('cxx') }} >=13.4,<14
         - sysroot_linux-64 =2.17 # [linux64]
         - swig =4.0
         - cmake >=3.24.0

thirdparty/faiss/faiss/CMakeLists.txt

@@ -117,6 +117,7 @@ set(FAISS_SRC
   impl/IDSelector.cpp
   impl/FaissException.cpp
   impl/HNSW.cpp
+  impl/hnsw/LockVector.cpp
   impl/hnsw/MinimaxHeap.cpp
   impl/NSG.cpp
   impl/PolysemousTraining.cpp
@@ -242,6 +243,7 @@ set(FAISS_HEADERS
   impl/FaissAssert.h
   impl/FaissException.h
   impl/HNSW.h
+  impl/hnsw/LockVector.h
   impl/hnsw/MinimaxHeap.h
   impl/LocalSearchQuantizer.h
   impl/ProductAdditiveQuantizer.h
@@ -280,6 +282,7 @@ set(FAISS_HEADERS
   impl/lattice_Zn.h
   impl/platform_macros.h
   impl/fast_scan/accumulate_loops.h
+  impl/fast_scan/accumulate_loops_512.h
   impl/fast_scan/dispatching.h
   impl/fast_scan/fast_scan.h
   impl/fast_scan/decompose_qbs.h

thirdparty/faiss/faiss/Clustering.cpp

@@ -16,6 +16,7 @@
 #include <cmath>
 #include <cstdio>
 #include <cstring>
+#include <limits>
 
 #include <omp.h>
 
@@ -84,20 +85,27 @@ idx_t subsample_training_set(
 
     const uint64_t actual_seed = get_actual_rng_seed(clus.seed);
 
-    std::vector<int> perm;
+    std::vector<idx_t> perm;
     if (clus.use_faster_subsampling) {
         // use subsampling with splitmix64 rng
         SplitMix64RandomGenerator rng(actual_seed);
 
         const idx_t new_nx = clus.k * clus.max_points_per_centroid;
         perm.resize(new_nx);
         for (idx_t i = 0; i < new_nx; i++) {
-            perm[i] = rng.rand_int(nx);
+            perm[i] = rng.rand_int64() % nx;
         }
     } else {
         // use subsampling with a default std rng
-        perm.resize(nx);
-        rand_perm(perm.data(), nx, actual_seed);
+        FAISS_THROW_IF_NOT_FMT(
+                nx <= static_cast<idx_t>(std::numeric_limits<int>::max()),
+                "Dataset too large (%" PRId64
+                ") for standard subsampling; "
+                "set use_faster_subsampling=true",
+                nx);
+        std::vector<int> int_perm(nx);
+        rand_perm(int_perm.data(), nx, actual_seed);
+        perm.assign(int_perm.begin(), int_perm.end());
     }
 
     nx = clus.k * clus.max_points_per_centroid;
@@ -232,12 +240,27 @@ int split_clusters(
     for (size_t ci = 0; ci < k; ci++) {
         if (hassign[ci] == 0) { /* need to redefine a centroid */
             size_t cj;
-            for (cj = 0; true; cj = (cj + 1) % k) {
-                /* probability to pick this cluster for split */
+            // Try probabilistic selection, with a deterministic fallback
+            // to the largest cluster if too many iterations pass.
+            size_t max_tries = 10 * k;
+            size_t n_tries = 0;
+            bool found = false;
+            for (cj = 0; n_tries < max_tries; cj = (cj + 1) % k) {
                 float p = (hassign[cj] - 1.0) / (float)(n - k);
                 float r = rng.rand_float();
                 if (r < p) {
-                    break; /* found our cluster to be split */
+                    found = true;
+                    break;
+                }
+                n_tries++;
+            }
+            if (!found) {
+                // Deterministic fallback: split the largest cluster.
+                cj = 0;
+                for (size_t j = 1; j < k; j++) {
+                    if (hassign[j] > hassign[cj]) {
+                        cj = j;
+                    }
                 }
             }
             memcpy(centroids + ci * d,
@@ -510,7 +533,7 @@ void Clustering::train_encoded(
 
             // accumulate objective
             obj = 0;
-            for (int j = 0; j < nx; j++) {
+            for (idx_t j = 0; j < nx; j++) {
                 obj += dis[j];
             }