add small benchmark for hamming computers

Summary: to measure impact of hamming computer diff

Reviewed By: algoriddle

Differential Revision: D46913890

fbshipit-source-id: 7b9850205885b9b7c5f394f17a79ba222e7b1e2e

Author: mdouze

benchs/bench_hamming_knn.py

@@ -0,0 +1,29 @@
+# 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.
+
+import time
+import numpy as np
+import faiss
+
+if __name__ == "__main__":
+    faiss.omp_set_num_threads(1)
+
+    for d in 4, 8, 16, 13:
+        nq = 10000
+        nb = 30000
+        print('Bits per vector = 8 *', d)
+        xq = faiss.randint((nq, d // 4), seed=1234, vmax=256**4).view('uint8')
+        xb = faiss.randint((nb, d // 4), seed=1234, vmax=256**4).view('uint8')
+        for variant in "hc", "mc":
+            print(f"{variant=:}", end="\t")
+            for k in 1, 4, 16, 64, 256:
+                times = []
+                for _run in range(5):
+                    t0 = time.time()
+                    D, I = faiss.knn_hamming(xq, xb, k, variant=variant)
+                    t1 = time.time()
+                    times.append(t1 - t0)
+                print(f'| {k=:} t={np.mean(times):.3f} s ± {np.std(times):.3f} ', flush=True, end="")
+            print()

faiss/python/__init__.py

@@ -22,7 +22,7 @@
 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
+    merge_knn_results, MapInt64ToInt64, knn_hamming
 
 
 __version__ = "%d.%d.%d" % (FAISS_VERSION_MAJOR,

faiss/python/extra_wrappers.py

@@ -335,10 +335,10 @@ def knn(xq, xb, k, metric=METRIC_L2):
     Parameters
     ----------
     xq : array_like
-        Query vectors, shape (nq, d) where d is appropriate for the index.
+        Query vectors, shape (nq, d) where the dimension d is that same as xb
         `dtype` must be float32.
     xb : array_like
-        Database vectors, shape (nb, d) where d is appropriate for the index.
+        Database vectors, shape (nb, d) where dimension d is the same as xq
         `dtype` must be float32.
     k : int
         Number of nearest neighbors.
@@ -375,6 +375,56 @@ def knn(xq, xb, k, metric=METRIC_L2):
         raise NotImplementedError("only L2 and INNER_PRODUCT are supported")
     return D, I
 
+def knn_hamming(xq, xb, k, variant="hc"):
+    """
+    Compute the k nearest neighbors of a set of vectors without constructing an index.
+
+    Parameters
+    ----------
+    xq : array_like
+        Query vectors, shape (nq, d) where d is the number of bits / 8
+        `dtype` must be uint8.
+    xb : array_like
+        Database vectors, shape (nb, d) where d is the number of bits / 8
+        `dtype` must be uint8.
+    k : int
+        Number of nearest neighbors.
+    variant : string
+        Function variant to use, either "mc" (counter) or "hc" (heap)
+
+    Returns
+    -------
+    D : array_like
+        Distances of the nearest neighbors, shape (nq, k)
+    I : array_like
+        Labels of the nearest neighbors, shape (nq, k)
+    """
+    # other variant is "mc"
+    nq, d = xq.shape
+    nb, d2 = xb.shape
+    assert d == d2
+    D = np.empty((nq, k), dtype='int32')
+    I = np.empty((nq, k), dtype='int64')
+
+    if variant == "hc":
+        heap = faiss.int_maxheap_array_t()
+        heap.k = k
+        heap.nh = nq
+        heap.ids = faiss.swig_ptr(I)
+        heap.val = faiss.swig_ptr(D)
+        faiss.hammings_knn_hc(
+            heap, faiss.swig_ptr(xq), faiss.swig_ptr(xb), nb,
+            d, 1
+        )
+    elif variant == "mc":
+        faiss.hammings_knn_mc(
+            faiss.swig_ptr(xq), faiss.swig_ptr(xb), nq, nb, k, d,
+            faiss.swig_ptr(D), faiss.swig_ptr(I)
+        )
+    else:
+        raise NotImplementedError
+    return D, I
+
 
 ###########################################
 # Kmeans object