Merge branch 'main' into arm64_failures

Author: mengdilin

.circleci/config.yml

@@ -136,7 +136,6 @@ runs:
               .
         make -k -C build -j$(nproc)
     - name: C++ tests
-      if: inputs.rocm == 'OFF'
       shell: bash
       run: |
         export GTEST_OUTPUT="xml:$(realpath .)/test-results/googletest/"

.github/actions/build_cmake/action.yml

@@ -16,4 +16,10 @@ jobs:
           GH_DEBUG: api
         run: |
           gh run watch ${{ inputs.run_id }} > /dev/null 2>&1
-          gh run rerun ${{ inputs.run_id }} --failed
+
+          # Only retry if there are failed jobs
+          if gh run view ${{ inputs.run_id }} --exit-status; then
+            echo Workflow succeeded - no retry necessary.
+          else
+            gh run rerun ${{ inputs.run_id }} --failed
+          fi

.github/workflows/retry_build.yml

@@ -80,6 +80,11 @@ class DatasetDescriptor:
 
     embedding_column: Optional[str] = None
 
+    sampling_rate: Optional[float] = None
+
+    # sampling column for xdb
+    sampling_column: Optional[str] = None
+
     def __hash__(self):
         return hash(self.get_filename())
 

benchs/bench_fw/descriptors.py

@@ -116,6 +116,21 @@ void IndexIVFAdditiveQuantizer::sa_decode(
     }
 }
 
+void IndexIVFAdditiveQuantizer::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
+    const uint8_t* code = invlists->get_single_code(list_no, offset);
+    aq->decode(code, recons, 1);
+    if (by_residual) {
+        std::vector<float> centroid(d);
+        quantizer->reconstruct(list_no, centroid.data());
+        for (int i = 0; i < d; ++i) {
+            recons[i] += centroid[i];
+        }
+    }
+}
+
 IndexIVFAdditiveQuantizer::~IndexIVFAdditiveQuantizer() = default;
 
 /*********************************************

faiss/IndexIVFAdditiveQuantizer.cpp

@@ -56,6 +56,9 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
 
     void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
 
+    void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
+            const override;
+
     ~IndexIVFAdditiveQuantizer() override;
 };
 

faiss/IndexIVFAdditiveQuantizer.h

@@ -537,6 +537,10 @@ struct InitGpuCompileOptions {
 #ifdef USE_NVIDIA_RAFT
         gpu_compile_options += "NVIDIA_RAFT ";
 #endif
+
+#ifdef USE_AMD_ROCM
+        gpu_compile_options += "AMD_ROCM ";
+#endif
     }
 };
 

faiss/gpu/GpuIndex.cu

@@ -307,13 +307,16 @@ void testMMCodeDistance(faiss::MetricType mt) {
     }
 }
 
+// FIXME: https://github.com/facebookresearch/faiss/issues/3787
+#ifndef USE_AMD_ROCM
 TEST(TestGpuIndexIVFPQ, Query_L2_MMCodeDistance) {
     testMMCodeDistance(faiss::MetricType::METRIC_L2);
 }
 
 TEST(TestGpuIndexIVFPQ, Query_IP_MMCodeDistance) {
     testMMCodeDistance(faiss::MetricType::METRIC_INNER_PRODUCT);
 }
+#endif // USE_AMD_ROCM
 
 TEST(TestGpuIndexIVFPQ, Float16Coarse) {
     Options opt;

faiss/gpu/test/TestGpuIndexIVFPQ.cpp

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <faiss/impl/index_read_utils.h>
 #include <faiss/index_io.h>
 
 #include <faiss/impl/io_macros.h>
@@ -61,7 +62,7 @@ namespace faiss {
  * Read
  **************************************************************/
 
-static void read_index_header(Index* idx, IOReader* f) {
+void read_index_header(Index* idx, IOReader* f) {
     READ1(idx->d);
     READ1(idx->ntotal);
     idx_t dummy;
@@ -230,7 +231,7 @@ InvertedLists* read_InvertedLists(IOReader* f, int io_flags) {
     }
 }
 
-static void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags) {
+void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags) {
     InvertedLists* ils = read_InvertedLists(f, io_flags);
     if (ils) {
         FAISS_THROW_IF_NOT(ils->nlist == ivf->nlist);
@@ -438,7 +439,7 @@ ProductQuantizer* read_ProductQuantizer(IOReader* reader) {
     return pq;
 }
 
-static void read_direct_map(DirectMap* dm, IOReader* f) {
+void read_direct_map(DirectMap* dm, IOReader* f) {
     char maintain_direct_map;
     READ1(maintain_direct_map);
     dm->type = (DirectMap::Type)maintain_direct_map;
@@ -454,10 +455,10 @@ static void read_direct_map(DirectMap* dm, IOReader* f) {
     }
 }
 
-static void read_ivf_header(
+void read_ivf_header(
         IndexIVF* ivf,
         IOReader* f,
-        std::vector<std::vector<idx_t>>* ids = nullptr) {
+        std::vector<std::vector<idx_t>>* ids) {
     read_index_header(ivf, f);
     READ1(ivf->nlist);
     READ1(ivf->nprobe);

faiss/impl/index_read.cpp

@@ -0,0 +1,30 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Utils for index_read
+
+#ifndef FAISS_INDEX_READ_UTILS_H
+#define FAISS_INDEX_READ_UTILS_H
+
+#include <faiss/IndexIVF.h>
+#include <faiss/impl/io.h>
+
+#pragma once
+
+namespace faiss {
+
+void read_index_header(Index* idx, IOReader* f);
+void read_direct_map(DirectMap* dm, IOReader* f);
+void read_ivf_header(
+        IndexIVF* ivf,
+        IOReader* f,
+        std::vector<std::vector<idx_t>>* ids = nullptr);
+void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags);
+
+} // namespace faiss
+
+#endif

faiss/impl/index_read_utils.h

@@ -96,7 +96,7 @@ def replacement_train_encoded(self, x, codec, index, weights=None):
             Index used to decode the vectors. Should have dimension `self.d`.
         index : faiss.Index
             Index used for assignment. The dimension of the index should be `self.d`.
-        weigths : array_like, optional
+        weights : array_like, optional
             Per training sample weight (size n) used when computing the weighted
             average to obtain the centroid (default is 1 for all training vectors).
         """

faiss/python/class_wrappers.py

@@ -345,7 +345,7 @@ def knn(xq, xb, k, metric=METRIC_L2, metric_arg=0.0):
         `dtype` must be float32.
     k : int
         Number of nearest neighbors.
-    distance_type : MetricType, optional
+    metric : MetricType, optional
         distance measure to use (either METRIC_L2 or METRIC_INNER_PRODUCT)
 
     Returns

faiss/python/extra_wrappers.py

@@ -115,10 +115,10 @@ std::string get_compile_options() {
     options += "OPTIMIZE ";
 #endif
 
-#ifdef __AVX2__
-    options += "AVX2 ";
-#elif __AVX512F__
+#ifdef __AVX512F__
     options += "AVX512 ";
+#elif defined(__AVX2__)
+    options += "AVX2 ";
 #elif defined(__ARM_FEATURE_SVE)
     options += "SVE NEON ";
 #elif defined(__aarch64__)

faiss/utils/utils.cpp

@@ -493,6 +493,23 @@ def test_IndexIVFPQ(self):
 
         self.run_search_and_reconstruct(index, xb, xq, eps=1.0)
 
+    def test_IndexIVFRQ(self):
+        d = 32
+        nb = 1000
+        nt = 1500
+        nq = 200
+
+        (xt, xb, xq) = get_dataset(d, nb, nt, nq)
+
+        quantizer = faiss.IndexFlatL2(d)
+        index = faiss.IndexIVFResidualQuantizer(quantizer, d, 32, 8, 8)
+        index.cp.min_points_per_centroid = 5    # quiet warning
+        index.nprobe = 4
+        index.train(xt)
+        index.add(xb)
+
+        self.run_search_and_reconstruct(index, xb, xq, eps=1.0)
+
     def test_MultiIndex(self):
         d = 32
         nb = 1000

tests/test_index.py

@@ -633,14 +633,14 @@ def test_OIVFPQ(self):
         d = ev.d
         quantizer = faiss.IndexFlatL2(d)
         index = faiss.IndexIVFPQ(quantizer, d, ncentroids, M, 8)
-        index.nprobe = 12
+        index.nprobe = 20
 
         res = ev.launch("IVFPQ", index)
         e_ivfpq = ev.evalres(res)
 
         quantizer = faiss.IndexFlatL2(d)
         index_ivfpq = faiss.IndexIVFPQ(quantizer, d, ncentroids, M, 8)
-        index_ivfpq.nprobe = 12
+        index_ivfpq.nprobe = 20
         opq_matrix = faiss.OPQMatrix(d, M)
         opq_matrix.niter = 10
         index = faiss.IndexPreTransform(opq_matrix, index_ivfpq)