Search and return codes (#3143)

Summary:
This PR adds a functionality where an IVF index can be searched and the corresponding codes be returned. It also adds a few functions to compress int arrays into a bit-compact representation.

Pull Request resolved: #3143

Test Plan:
```
buck test //faiss/tests/:test_index_composite -- TestSearchAndReconstruct

buck test //faiss/tests/:test_standalone_codec -- test_arrays
```

Reviewed By: algoriddle

Differential Revision: D51544613

Pulled By: mdouze

fbshipit-source-id: 875f72d0f9140096851592422570efa0f65431fc

benchs/bench_hybrid_cpu_gpu.py

@@ -530,14 +530,7 @@ def aa(*args, **kwargs):
         raise RuntimeError()
 
     totex = op.num_experiments()
-    rs = np.random.RandomState(123)
-    if totex < args.n_autotune:
-        experiments = rs.permutation(totex - 2) + 1
-    else:
-        experiments = rs.randint(
-            totex - 2, size=args.n_autotune - 2, replace=False)
-
-    experiments = [0, totex - 1] + list(experiments)
+    experiments = op.sample_experiments()
     print(f"total nb experiments {totex}, running {len(experiments)}")
 
     print("perform search")

contrib/evaluation.py

@@ -380,7 +380,23 @@ def do_nothing_key(self):
         return np.zeros(len(self.ranges), dtype=int)
 
     def num_experiments(self):
-        return np.prod([len(values) for name, values in self.ranges])
+        return int(np.prod([len(values) for name, values in self.ranges]))
+
+    def sample_experiments(self, n_autotune, rs=np.random):
+        """ sample a set of experiments of max size n_autotune
+        (run all experiments in random order if n_autotune is 0)
+        """
+        assert n_autotune == 0 or n_autotune >= 2
+        totex = self.num_experiments()
+        rs = np.random.RandomState(123)
+        if n_autotune == 0 or totex < n_autotune:
+            experiments = rs.permutation(totex - 2)
+        else:
+            experiments = rs.choice(
+                totex - 2, size=n_autotune - 2, replace=False)
+
+        experiments = [0, totex - 1] + [int(cno) + 1 for cno in experiments]
+        return experiments
 
     def cno_to_key(self, cno):
         """Convert a sequential experiment number to a key"""

faiss/IndexIVF.cpp

@@ -977,44 +977,107 @@ void IndexIVF::search_and_reconstruct(
             std::min(nlist, params ? params->nprobe : this->nprobe);
     FAISS_THROW_IF_NOT(nprobe > 0);
 
-    idx_t* idx = new idx_t[n * nprobe];
-    ScopeDeleter<idx_t> del(idx);
-    float* coarse_dis = new float[n * nprobe];
-    ScopeDeleter<float> del2(coarse_dis);
+    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+    std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
 
-    quantizer->search(n, x, nprobe, coarse_dis, idx);
+    quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
 
-    invlists->prefetch_lists(idx, n * nprobe);
+    invlists->prefetch_lists(idx.get(), n * nprobe);
 
     // search_preassigned() with `store_pairs` enabled to obtain the list_no
     // and offset into `codes` for reconstruction
     search_preassigned(
             n,
             x,
             k,
-            idx,
-            coarse_dis,
+            idx.get(),
+            coarse_dis.get(),
             distances,
             labels,
             true /* store_pairs */,
             params);
-    for (idx_t i = 0; i < n; ++i) {
-        for (idx_t j = 0; j < k; ++j) {
-            idx_t ij = i * k + j;
-            idx_t key = labels[ij];
-            float* reconstructed = recons + ij * d;
-            if (key < 0) {
-                // Fill with NaNs
-                memset(reconstructed, -1, sizeof(*reconstructed) * d);
-            } else {
-                int list_no = lo_listno(key);
-                int offset = lo_offset(key);
+#pragma omp parallel for if (n * k > 1000)
+    for (idx_t ij = 0; ij < n * k; ij++) {
+        idx_t key = labels[ij];
+        float* reconstructed = recons + ij * d;
+        if (key < 0) {
+            // Fill with NaNs
+            memset(reconstructed, -1, sizeof(*reconstructed) * d);
+        } else {
+            int list_no = lo_listno(key);
+            int offset = lo_offset(key);
+
+            // Update label to the actual id
+            labels[ij] = invlists->get_single_id(list_no, offset);
+
+            reconstruct_from_offset(list_no, offset, reconstructed);
+        }
+    }
+}
+
+void IndexIVF::search_and_return_codes(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels,
+        uint8_t* codes,
+        bool include_listno,
+        const SearchParameters* params_in) const {
+    const IVFSearchParameters* params = nullptr;
+    if (params_in) {
+        params = dynamic_cast<const IVFSearchParameters*>(params_in);
+        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
+    }
+    const size_t nprobe =
+            std::min(nlist, params ? params->nprobe : this->nprobe);
+    FAISS_THROW_IF_NOT(nprobe > 0);
+
+    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
+    std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
+
+    quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get());
+
+    invlists->prefetch_lists(idx.get(), n * nprobe);
+
+    // search_preassigned() with `store_pairs` enabled to obtain the list_no
+    // and offset into `codes` for reconstruction
+    search_preassigned(
+            n,
+            x,
+            k,
+            idx.get(),
+            coarse_dis.get(),
+            distances,
+            labels,
+            true /* store_pairs */,
+            params);
+
+    size_t code_size_1 = code_size;
+    if (include_listno) {
+        code_size_1 += coarse_code_size();
+    }
+
+#pragma omp parallel for if (n * k > 1000)
+    for (idx_t ij = 0; ij < n * k; ij++) {
+        idx_t key = labels[ij];
+        uint8_t* code1 = codes + ij * code_size_1;
+
+        if (key < 0) {
+            // Fill with 0xff
+            memset(code1, -1, code_size_1);
+        } else {
+            int list_no = lo_listno(key);
+            int offset = lo_offset(key);
+            const uint8_t* cc = invlists->get_single_code(list_no, offset);
 
-                // Update label to the actual id
-                labels[ij] = invlists->get_single_id(list_no, offset);
+            labels[ij] = invlists->get_single_id(list_no, offset);
 
-                reconstruct_from_offset(list_no, offset, reconstructed);
+            if (include_listno) {
+                encode_listno(list_no, code1);
+                code1 += code_size_1 - code_size;
             }
+            memcpy(code1, cc, code_size);
         }
     }
 }

faiss/IndexIVF.h

@@ -357,6 +357,24 @@ struct IndexIVF : Index, IndexIVFInterface {
             float* recons,
             const SearchParameters* params = nullptr) const override;
 
+    /** Similar to search, but also returns the codes corresponding to the
+     * stored vectors for the search results.
+     *
+     * @param codes      codes (n, k, code_size)
+     * @param include_listno
+     *                   include the list ids in the code (in this case add
+     *                   ceil(log8(nlist)) to the code size)
+     */
+    void search_and_return_codes(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            uint8_t* recons,
+            bool include_listno = false,
+            const SearchParameters* params = nullptr) const;
+
     /** Reconstruct a vector given the location in terms of (inv list index +
      * inv list offset) instead of the id.
      *

faiss/IndexIVFAdditiveQuantizer.cpp

@@ -149,6 +149,7 @@ struct AQInvertedListScanner : InvertedListScanner {
     const float* q;
     /// following codes come from this inverted list
     void set_list(idx_t list_no, float coarse_dis) override {
+        this->list_no = list_no;
         if (ia.metric_type == METRIC_L2 && ia.by_residual) {
             ia.quantizer->compute_residual(q0, tmp.data(), list_no);
             q = tmp.data();

faiss/impl/AdditiveQuantizer.cpp

@@ -261,7 +261,7 @@ void AdditiveQuantizer::decode(const uint8_t* code, float* x, size_t n) const {
             is_trained, "The additive quantizer is not trained yet.");
 
     // standard additive quantizer decoding
-#pragma omp parallel for if (n > 1000)
+#pragma omp parallel for if (n > 100)
     for (int64_t i = 0; i < n; i++) {
         BitstringReader bsr(code + i * code_size, code_size);
         float* xi = x + i * d;

faiss/impl/ProductQuantizer.cpp

@@ -306,7 +306,8 @@ void ProductQuantizer::decode(const uint8_t* code, float* x) const {
 }
 
 void ProductQuantizer::decode(const uint8_t* code, float* x, size_t n) const {
-    for (size_t i = 0; i < n; i++) {
+#pragma omp parallel for if (n > 100)
+    for (int64_t i = 0; i < n; i++) {
         this->decode(code + code_size * i, x + d * i);
     }
 }

faiss/python/__init__.py

@@ -22,7 +22,8 @@
 from faiss.extra_wrappers import kmin, kmax, pairwise_distances, rand, randint, \
     lrand, randn, rand_smooth_vectors, eval_intersection, normalize_L2, \
     ResultHeap, knn, Kmeans, checksum, matrix_bucket_sort_inplace, bucket_sort, \
-    merge_knn_results, MapInt64ToInt64, knn_hamming
+    merge_knn_results, MapInt64ToInt64, knn_hamming, \
+    pack_bitstrings, unpack_bitstrings
 
 
 __version__ = "%d.%d.%d" % (FAISS_VERSION_MAJOR,

faiss/python/class_wrappers.py

@@ -402,6 +402,74 @@ def replacement_search_and_reconstruct(self, x, k, *, params=None, D=None, I=Non
         )
         return D, I, R
 
+    def replacement_search_and_return_codes(
+            self, x, k, *,
+            include_listnos=False, params=None, D=None, I=None, codes=None):
+        """Find the k nearest neighbors of the set of vectors x in the index,
+        and return the codes stored for these vectors
+
+        Parameters
+        ----------
+        x : array_like
+            Query vectors, shape (n, d) where d is appropriate for the index.
+            `dtype` must be float32.
+        k : int
+            Number of nearest neighbors.
+        params : SearchParameters
+            Search parameters of the current search (overrides the class-level params)
+        include_listnos : bool, optional
+            whether to include the list ids in the first bytes of each code
+        D : array_like, optional
+            Distance array to store the result.
+        I : array_like, optional
+            Labels array to store the result.
+        codes : array_like, optional
+            codes array to store
+
+        Returns
+        -------
+        D : array_like
+            Distances of the nearest neighbors, shape (n, k). When not enough results are found
+            the label is set to +Inf or -Inf.
+        I : array_like
+            Labels of the nearest neighbors, shape (n, k). When not enough results are found,
+            the label is set to -1
+        R : array_like
+            Approximate (reconstructed) nearest neighbor vectors, shape (n, k, d).
+        """
+        n, d = x.shape
+        assert d == self.d
+        x = np.ascontiguousarray(x, dtype='float32')
+
+        assert k > 0
+
+        if D is None:
+            D = np.empty((n, k), dtype=np.float32)
+        else:
+            assert D.shape == (n, k)
+
+        if I is None:
+            I = np.empty((n, k), dtype=np.int64)
+        else:
+            assert I.shape == (n, k)
+
+        code_size_1 = self.code_size
+        if include_listnos:
+            code_size_1 += self.coarse_code_size()
+
+        if codes is None:
+            codes = np.empty((n, k, code_size_1), dtype=np.uint8)
+        else:
+            assert codes.shape == (n, k, code_size_1)
+
+        self.search_and_return_codes_c(
+            n, swig_ptr(x),
+            k, swig_ptr(D),
+            swig_ptr(I), swig_ptr(codes), include_listnos,
+            params
+        )
+        return D, I, codes
+
     def replacement_remove_ids(self, x):
         """Remove some ids from the index.
         This is a O(ntotal) operation by default, so could be expensive.
@@ -734,6 +802,8 @@ def replacement_permute_entries(self, perm):
                    ignore_missing=True)
     replace_method(the_class, 'search_and_reconstruct',
                    replacement_search_and_reconstruct, ignore_missing=True)
+    replace_method(the_class, 'search_and_return_codes',
+                   replacement_search_and_return_codes, ignore_missing=True)
 
     # these ones are IVF-specific
     replace_method(the_class, 'search_preassigned',

faiss/python/extra_wrappers.py

@@ -14,6 +14,9 @@
 
 import faiss
 
+import collections.abc
+
+
 ###########################################
 # Wrapper for a few functions
 ###########################################
@@ -579,3 +582,72 @@ def assign(self, x):
         self.index.add(self.centroids)
         D, I = self.index.search(x, 1)
         return D.ravel(), I.ravel()
+
+
+###########################################
+# Packing and unpacking bistrings
+###########################################
+
+def is_sequence(x):
+    return isinstance(x, collections.abc.Sequence)
+
+pack_bitstrings_c = pack_bitstrings
+
+def pack_bitstrings(a, nbit):
+    """
+    Pack a set integers (i, j) where i=0:n and j=0:M into
+    n bitstrings.
+    Output is an uint8 array of size (n, code_size), where code_size is
+    such that at most 7 bits per code are wasted.
+
+    If nbit is an integer: all entries takes nbit bits.
+    If nbit is an array: entry (i, j) takes nbit[j] bits.
+    """
+    n, M = a.shape
+    a = np.ascontiguousarray(a, dtype='int32')
+    if is_sequence(nbit):
+        nbit = np.ascontiguousarray(nbit, dtype='int32')
+        assert nbit.shape == (M,)
+        code_size = int((nbit.sum() + 7) // 8)
+        b = np.empty((n, code_size), dtype='uint8')
+        pack_bitstrings_c(
+            n, M, swig_ptr(nbit), swig_ptr(a), swig_ptr(b), code_size)
+    else:
+        code_size = (M * nbit + 7) // 8
+        b = np.empty((n, code_size), dtype='uint8')
+        pack_bitstrings_c(n, M, nbit, swig_ptr(a), swig_ptr(b), code_size)
+    return b
+
+unpack_bitstrings_c = unpack_bitstrings
+
+def unpack_bitstrings(b, M_or_nbits, nbit=None):
+    """
+    Unpack a set integers (i, j) where i=0:n and j=0:M from
+    n bitstrings (encoded as uint8s).
+    Input is an uint8 array of size (n, code_size), where code_size is
+    such that at most 7 bits per code are wasted.
+
+    Two forms:
+    - when called with (array, M, nbit): there are M entries of size
+      nbit per row
+    - when called with (array, nbits): element (i, j) is encoded in
+      nbits[j] bits
+    """
+    n, code_size = b.shape
+    if nbit is None:
+        nbit = np.ascontiguousarray(M_or_nbits, dtype='int32')
+        M = len(nbit)
+        min_code_size = int((nbit.sum() + 7) // 8)
+        assert code_size >= min_code_size
+        a = np.empty((n, M), dtype='int32')
+        unpack_bitstrings_c(
+            n, M, swig_ptr(nbit),
+            swig_ptr(b), code_size, swig_ptr(a))
+    else:
+        M = M_or_nbits
+        min_code_size = (M * nbit + 7) // 8
+        assert code_size >= min_code_size
+        a = np.empty((n, M), dtype='int32')
+        unpack_bitstrings_c(
+            n, M, nbit, swig_ptr(b), code_size, swig_ptr(a))
+    return a