RaBitQ implementation (facebookresearch#4235)

Summary:
This is a reference implementation of the https://arxiv.org/pdf/2405.12497
> Jianyang Gao, Cheng Long, "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search".

The goal is to correctly set up the internals using Faiss.

The following comments for the implementation:
* The code does not include the computations for the symmetric distance, because it is absent in the original article. This can be added later, though.
* The original `RaBitQ` includes random matrix rotation as a part of it, but I've decided to rely on external `faiss::IndexPreTransform` and `faiss::RandomRotationMatrix` facilities.
* Certain features required internal changes in `faiss::IndexIVF`, but I did that as least invasive as possible, without breaking the backward compatibility.
* Not sure about naming convensions, maybe certain classes and structures need to be renamed
* `METRIC_INNER_PRODUCT` is supported as well
* More unit tests are needed?
* I did not bring any hardware-specific optimizations, bcz this is a reference implementation. Certain `simdlib` facilities may be added later, if needed

Here's how to use IndexRaBitQ
```Python
        ds = datasets.SyntheticDataset(...)

        index_rbq = faiss.IndexRaBitQ(ds.d, faiss.METRIC_L2)
        index_rbq.qb = 8

        # wrap with random rotations
        rrot = faiss.RandomRotationMatrix(ds.d, ds.d)
        rrot.init(rrot_seed)

        index_cand = faiss.IndexPreTransform(rrot, index_rbq)
        index_cand.train(ds.get_train())
        index_cand.add(ds.get_database())
```

Here's how to use IndexIVFRaBitQ
```Python
        ds = datasets.SyntheticDataset(...)

        index_flat = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
        index_rbq = faiss.IndexIVFRaBitQ(index_flat, ds.d, nlist, faiss.METRIC_L2)
        index_rbq.qb = 8

        # wrap with random rotations
        rrot = faiss.RandomRotationMatrix(ds.d, ds.d)
        rrot.init(rrot_seed)

        index_cand = faiss.IndexPreTransform(rrot, index_rbq)
        index_cand.train(ds.get_train())
        index_cand.add(ds.get_database())
```

Pull Request resolved: facebookresearch#4235

Test Plan:
Imported from GitHub, without a `Test Plan:` line.

buck run 'fbcode//mode/dev' fbcode//faiss/tests:test_rabitq

Reviewed By: mdouze

Differential Revision: D71638302

Pulled By: junjieqi

fbshipit-source-id: de981a6aed91d296237d8accf337359de04a552e

Author: alexanderguzhva

faiss/CMakeLists.txt

@@ -28,6 +28,7 @@ set(FAISS_SRC
   IndexIVFAdditiveQuantizerFastScan.cpp
   IndexIVFPQFastScan.cpp
   IndexIVFPQR.cpp
+  IndexIVFRaBitQ.cpp
   IndexIVFSpectralHash.cpp
   IndexLSH.cpp
   IndexNNDescent.cpp
@@ -39,6 +40,7 @@ set(FAISS_SRC
   IndexIVFIndependentQuantizer.cpp
   IndexPQFastScan.cpp
   IndexPreTransform.cpp
+  IndexRaBitQ.cpp
   IndexRefine.cpp
   IndexReplicas.cpp
   IndexRowwiseMinMax.cpp
@@ -60,6 +62,7 @@ set(FAISS_SRC
   impl/PolysemousTraining.cpp
   impl/ProductQuantizer.cpp
   impl/AdditiveQuantizer.cpp
+  impl/RaBitQuantizer.cpp
   impl/ResidualQuantizer.cpp
   impl/LocalSearchQuantizer.cpp
   impl/ProductAdditiveQuantizer.cpp
@@ -123,6 +126,7 @@ set(FAISS_HEADERS
   IndexIVFAdditiveQuantizerFastScan.h
   IndexIVFPQFastScan.h
   IndexIVFPQR.h
+  IndexIVFRaBitQ.h
   IndexIVFSpectralHash.h
   IndexLSH.h
   IndexNeuralNetCodec.h
@@ -136,6 +140,7 @@ set(FAISS_HEADERS
   IndexPreTransform.h
   IndexRefine.h
   IndexReplicas.h
+  IndexRaBitQ.h
   IndexRowwiseMinMax.h
   IndexScalarQuantizer.h
   IndexShards.h
@@ -164,6 +169,7 @@ set(FAISS_HEADERS
   impl/ProductQuantizer-inl.h
   impl/ProductQuantizer.h
   impl/Quantizer.h
+  impl/RaBitQuantizer.h
   impl/ResidualQuantizer.h
   impl/ResultHandler.h
   impl/ScalarQuantizer.h

faiss/IndexIVF.cpp

@@ -455,7 +455,7 @@ void IndexIVF::search_preassigned(
 #pragma omp parallel if (do_parallel) reduction(+ : nlistv, ndis, nheap)
     {
         std::unique_ptr<InvertedListScanner> scanner(
-                get_InvertedListScanner(store_pairs, sel));
+                get_InvertedListScanner(store_pairs, sel, params));
 
         /*****************************************************
          * Depending on parallel_mode, there are two possible ways
@@ -796,7 +796,7 @@ void IndexIVF::range_search_preassigned(
     {
         RangeSearchPartialResult pres(result);
         std::unique_ptr<InvertedListScanner> scanner(
-                get_InvertedListScanner(store_pairs, sel));
+                get_InvertedListScanner(store_pairs, sel, params));
         FAISS_THROW_IF_NOT(scanner.get());
         all_pres[omp_get_thread_num()] = &pres;
 
@@ -912,7 +912,8 @@ void IndexIVF::range_search_preassigned(
 
 InvertedListScanner* IndexIVF::get_InvertedListScanner(
         bool /*store_pairs*/,
-        const IDSelector* /* sel */) const {
+        const IDSelector* /* sel */,
+        const IVFSearchParameters* /* params */) const {
     FAISS_THROW_MSG("get_InvertedListScanner not implemented");
 }
 
@@ -1290,6 +1291,14 @@ size_t InvertedListScanner::scan_codes(
 
     if (!keep_max) {
         for (size_t j = 0; j < list_size; j++) {
+            if (sel != nullptr) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                if (!sel->is_member(id)) {
+                    codes += code_size;
+                    continue;
+                }
+            }
+
             float dis = distance_to_code(codes);
             if (dis < simi[0]) {
                 int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
@@ -1300,6 +1309,14 @@ size_t InvertedListScanner::scan_codes(
         }
     } else {
         for (size_t j = 0; j < list_size; j++) {
+            if (sel != nullptr) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                if (!sel->is_member(id)) {
+                    codes += code_size;
+                    continue;
+                }
+            }
+
             float dis = distance_to_code(codes);
             if (dis > simi[0]) {
                 int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];

faiss/IndexIVF.h

@@ -312,11 +312,14 @@ struct IndexIVF : Index, IndexIVFInterface {
 
     /** Get a scanner for this index (store_pairs means ignore labels)
      *
-     * The default search implementation uses this to compute the distances
+     * The default search implementation uses this to compute the distances.
+     * Use sel instead of params->sel, because sel is initialized with
+     * params->sel, but may get overridden by IndexIVF's internal logic.
      */
     virtual InvertedListScanner* get_InvertedListScanner(
             bool store_pairs = false,
-            const IDSelector* sel = nullptr) const;
+            const IDSelector* sel = nullptr,
+            const IVFSearchParameters* params = nullptr) const;
 
     /** reconstruct a vector. Works only if maintain_direct_map is set to 1 or 2
      */

faiss/IndexIVFAdditiveQuantizer.cpp

@@ -253,7 +253,8 @@ struct AQInvertedListScannerLUT : AQInvertedListScanner {
 
 InvertedListScanner* IndexIVFAdditiveQuantizer::get_InvertedListScanner(
         bool store_pairs,
-        const IDSelector* sel) const {
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
     FAISS_THROW_IF_NOT(!sel);
     if (metric_type == METRIC_INNER_PRODUCT) {
         if (aq->search_type == AdditiveQuantizer::ST_decompress) {

faiss/IndexIVFAdditiveQuantizer.h

@@ -52,7 +52,8 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
-            const IDSelector* sel) const override;
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
 
     void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
 

faiss/IndexIVFFlat.cpp

@@ -223,7 +223,8 @@ InvertedListScanner* get_InvertedListScanner1(
 
 InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
         bool store_pairs,
-        const IDSelector* sel) const {
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
     if (sel) {
         return get_InvertedListScanner1<true>(this, store_pairs, sel);
     } else {

faiss/IndexIVFFlat.h

@@ -44,7 +44,8 @@ struct IndexIVFFlat : IndexIVF {
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
-            const IDSelector* sel) const override;
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
 
     void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
             const override;

faiss/IndexIVFPQ.cpp

@@ -1321,7 +1321,8 @@ InvertedListScanner* get_InvertedListScanner2(
 
 InvertedListScanner* IndexIVFPQ::get_InvertedListScanner(
         bool store_pairs,
-        const IDSelector* sel) const {
+        const IDSelector* sel,
+        const IVFSearchParameters*) const {
     if (sel) {
         return get_InvertedListScanner2<true>(*this, store_pairs, sel);
     } else {

faiss/IndexIVFPQ.h

@@ -134,7 +134,8 @@ struct IndexIVFPQ : IndexIVF {
 
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs,
-            const IDSelector* sel) const override;
+            const IDSelector* sel,
+            const IVFSearchParameters* params) const override;
 
     /// build precomputed table
     void precompute_table();