benchmark SSD IndexIVF

Summary: This is some code for benchmakring the SSD reads.

Reviewed By: MDSilber

Differential Revision: D24457715

fbshipit-source-id: 475668e4dc450dc4652ef8828111335c236bfa44

contrib/ivf_tools.py

@@ -0,0 +1,25 @@
+# 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 numpy as np
+import faiss
+
+def get_invlist(invlists, l):
+    """ returns the inverted lists content. """
+    ls = invlists.list_size(l)
+    list_ids = np.zeros(ls, dtype='int64')
+    ids = codes = None
+    try:
+        ids = invlists.get_ids(l)
+        faiss.memcpy(faiss.swig_ptr(list_ids), ids, list_ids.nbytes)
+        codes = invlists.get_codes(l)
+        list_codes = np.zeros((ls, invlists.code_size), dtype='uint8')
+        faiss.memcpy(faiss.swig_ptr(list_codes), codes, list_codes.nbytes)
+    finally:
+        if ids is not None:
+            invlists.release_ids(l, ids)
+        if codes is not None:
+            invlists.release_codes(l, codes)
+    return list_ids, list_codes

faiss/OnDiskInvertedLists.cpp

@@ -678,6 +678,18 @@ void OnDiskInvertedLists::crop_invlists(size_t l0, size_t l1)
     nlist = l1 - l0;
 }
 
+
+void OnDiskInvertedLists::set_all_lists_sizes(const size_t *sizes)
+{
+    size_t ofs = 0;
+    for (size_t i = 0; i < nlist; i++) {
+        lists[i].offset = ofs;
+        lists[i].capacity = lists[i].size = sizes[i];
+        ofs += sizes[i] * (sizeof(idx_t) + code_size);
+    }
+
+}
+
 /*******************************************************
  * I/O support via callbacks
  *******************************************************/
@@ -755,7 +767,9 @@ InvertedLists * OnDiskInvertedListsIOHook::read(IOReader *f, int io_flags) const
 
     }
     READ1(od->totsize);
-    od->do_mmap();
+    if (!(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
+        od->do_mmap();
+    }
     return od;
 }
 

faiss/OnDiskInvertedLists.h

@@ -126,6 +126,9 @@ struct OnDiskInvertedLists: InvertedLists {
     size_t allocate_slot (size_t capacity);
     void free_slot (size_t offset, size_t capacity);
 
+    /// override all list sizes and make a packed storage
+    void set_all_lists_sizes(const size_t *sizes);
+
     // empty constructor for the I/O functions
     OnDiskInvertedLists ();
 };

faiss/index_io.h

@@ -51,7 +51,7 @@ const int IO_FLAG_READ_ONLY = 2;
 const int IO_FLAG_ONDISK_SAME_DIR = 4;
 // don't load IVF data to RAM, only list sizes
 const int IO_FLAG_SKIP_IVF_DATA = 8;
-// try to memmap data (useful for OnDiskInvertedLists)
+// try to memmap data (useful to load an ArrayInvertedLists as an OnDiskInvertedLists)
 const int IO_FLAG_MMAP = IO_FLAG_SKIP_IVF_DATA | 0x646f0000;
 
 

faiss/python/swigfaiss.swig

@@ -174,7 +174,9 @@ namespace std {
         T * data();
         size_t size();
         T at (size_t n) const;
+        T & operator [] (size_t n);
         void resize (size_t n);
+
         void swap (vector<T> & other);
     };
 };