faiss 0.1.3 → 0.2.0

data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp

@@ -39,8 +39,12 @@ namespace faiss {
  * that hides the template mess.
  ********************************************************************/
 
-#if defined(__F16C__) && defined(__AVX2__)
+#ifdef __AVX2__
+#ifdef __F16C__
 #define USE_F16C
+#else
+#warning "Cannot enable AVX optimizations in scalar quantizer if -mf16c is not set as well"
+#endif
 #endif
 
 
@@ -1220,33 +1224,41 @@ SQDistanceComputer *select_distance_computer (
 
 ScalarQuantizer::ScalarQuantizer
           (size_t d, QuantizerType qtype):
-              qtype (qtype), rangestat(RS_minmax), rangestat_arg(0), d (d)
+              qtype (qtype), rangestat(RS_minmax), rangestat_arg(0), d(d)
+{
+  set_derived_sizes();
+}
+
+ScalarQuantizer::ScalarQuantizer ():
+    qtype(QT_8bit),
+    rangestat(RS_minmax), rangestat_arg(0), d(0), bits(0), code_size(0)
+{}
+
+void ScalarQuantizer::set_derived_sizes ()
 {
     switch (qtype) {
-    case QT_8bit:
-    case QT_8bit_uniform:
-    case QT_8bit_direct:
+      case QT_8bit:
+      case QT_8bit_uniform:
+      case QT_8bit_direct:
         code_size = d;
+        bits = 8;
         break;
-    case QT_4bit:
-    case QT_4bit_uniform:
+      case QT_4bit:
+      case QT_4bit_uniform:
         code_size = (d + 1) / 2;
+        bits = 4;
         break;
-    case QT_6bit:
+      case QT_6bit:
         code_size = (d * 6 + 7) / 8;
+        bits = 6;
         break;
-    case QT_fp16:
+      case QT_fp16:
         code_size = d * 2;
+        bits = 16;
         break;
     }
-
 }
 
-ScalarQuantizer::ScalarQuantizer ():
-    qtype(QT_8bit),
-    rangestat(RS_minmax), rangestat_arg(0), d (0), code_size(0)
-{}
-
 void ScalarQuantizer::train (size_t n, const float *x)
 {
     int bit_per_dim =
@@ -1418,9 +1430,8 @@ struct IVFSQScannerIP: InvertedListScanner {
             float accu = accu0 + dc.query_to_code (codes);
 
             if (accu > simi [0]) {
-                minheap_pop (k, simi, idxi);
                 int64_t id = store_pairs ? (list_no << 32 | j) : ids[j];
-                minheap_push (k, simi, idxi, accu, id);
+                minheap_replace_top (k, simi, idxi, accu, id);
                 nup++;
             }
             codes += code_size;
@@ -1506,9 +1517,8 @@ struct IVFSQScannerL2: InvertedListScanner {
             float dis = dc.query_to_code (codes);
 
             if (dis < simi [0]) {
-                maxheap_pop (k, simi, idxi);
                 int64_t id = store_pairs ? (list_no << 32 | j) : ids[j];
-                maxheap_push (k, simi, idxi, dis, id);
+                maxheap_replace_top (k, simi, idxi, dis, id);
                 nup++;
             }
             codes += code_size;

data/vendor/faiss/faiss/impl/ScalarQuantizer.h

@@ -53,6 +53,9 @@ struct ScalarQuantizer {
     /// dimension of input vectors
     size_t d;
 
+    /// bits per scalar code
+    size_t bits;
+
     /// bytes per vector
     size_t code_size;
 
@@ -62,6 +65,9 @@ struct ScalarQuantizer {
     ScalarQuantizer (size_t d, QuantizerType qtype);
     ScalarQuantizer ();
 
+    /// updates internal values based on qtype and d
+    void set_derived_sizes ();
+
     void train (size_t n, const float *x);
 
     /// Used by an IVF index to train based on the residuals

data/vendor/faiss/faiss/impl/index_read.cpp

@@ -15,15 +15,13 @@
 #include <sys/types.h>
 #include <sys/stat.h>
 
-#ifndef _MSC_VER
-#include <sys/mman.h>
-#endif // !_MSC_VER
-
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/io.h>
 #include <faiss/impl/io_macros.h>
 #include <faiss/utils/hamming.h>
 
+#include <faiss/invlists/InvertedListsIOHook.h>
+
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
 #include <faiss/IndexPreTransform.h>
@@ -39,17 +37,16 @@
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexLattice.h>
+#include <faiss/IndexPQFastScan.h>
+#include <faiss/IndexIVFPQFastScan.h>
+#include <faiss/IndexRefine.h>
+
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexBinaryFromFloat.h>
 #include <faiss/IndexBinaryHNSW.h>
 #include <faiss/IndexBinaryIVF.h>
 #include <faiss/IndexBinaryHash.h>
 
-#ifndef _MSC_VER
-#include <faiss/OnDiskInvertedLists.h>
-#endif // !_MSC_VER
-
-
 namespace faiss {
 
 
@@ -141,7 +138,10 @@ VectorTransform* read_VectorTransform (IOReader *f) {
         }
         vt = itqt;
     } else {
-        FAISS_THROW_MSG("fourcc not recognized");
+        FAISS_THROW_FMT(
+            "fourcc %ud (\"%s\") not recognized",
+            h, fourcc_inv_printable(h).c_str()
+        );
     }
     READ1 (vt->d_in);
     READ1 (vt->d_out);
@@ -167,7 +167,10 @@ static void read_ArrayInvertedLists_sizes (
             sizes[idsizes[j]] = idsizes[j + 1];
         }
     } else {
-        FAISS_THROW_MSG ("invalid list_type");
+        FAISS_THROW_FMT(
+            "list_type %ud (\"%s\") not recognized",
+            list_type, fourcc_inv_printable(list_type).c_str()
+        );
     }
 }
 
@@ -199,11 +202,6 @@ InvertedLists *read_InvertedLists (IOReader *f, int io_flags) {
         }
         return ails;
 
-#ifdef _MSC_VER
-    } else {
-        FAISS_THROW_MSG("Unsupported inverted list format for Windows");
-    }
-#else
     } else if (h == fourcc ("ilar") && (io_flags & IO_FLAG_SKIP_IVF_DATA)) {
         // code is always ilxx where xx is specific to the type of invlists we want
         // so we get the 16 high bits from the io_flag and the 16 low bits as "il"
@@ -218,7 +216,6 @@ InvertedLists *read_InvertedLists (IOReader *f, int io_flags) {
     } else {
         return InvertedListsIOHook::lookup(h)->read(f, io_flags);
     }
-#endif // !_MSC_VER
 
 }
 
@@ -226,8 +223,11 @@ InvertedLists *read_InvertedLists (IOReader *f, int io_flags) {
 static void read_InvertedLists (
         IndexIVF *ivf, IOReader *f, int io_flags) {
     InvertedLists *ils = read_InvertedLists (f, io_flags);
-    FAISS_THROW_IF_NOT (!ils || (ils->nlist == ivf->nlist &&
-                                 ils->code_size == ivf->code_size));
+    if (ils) {
+        FAISS_THROW_IF_NOT (ils->nlist == ivf->nlist);
+        FAISS_THROW_IF_NOT (ils->code_size == InvertedLists::INVALID_CODE_SIZE ||
+                            ils->code_size == ivf->code_size);
+    }
     ivf->invlists = ils;
     ivf->own_invlists = true;
 }
@@ -247,6 +247,7 @@ static void read_ScalarQuantizer (ScalarQuantizer *ivsc, IOReader *f) {
     READ1 (ivsc->d);
     READ1 (ivsc->code_size);
     READVECTOR (ivsc->trained);
+    ivsc->set_derived_sizes ();
 }
 
 
@@ -551,14 +552,20 @@ Index *read_index (IOReader *f, int io_flags) {
         read_ProductQuantizer (&imiq->pq, f);
         idx = imiq;
     } else if(h == fourcc ("IxRF")) {
-        IndexRefineFlat *idxrf = new IndexRefineFlat ();
+        IndexRefine *idxrf = new IndexRefine ();
         read_index_header (idxrf, f);
         idxrf->base_index = read_index(f, io_flags);
-        idxrf->own_fields = true;
-        IndexFlat *rf = dynamic_cast<IndexFlat*> (read_index (f, io_flags));
-        std::swap (*rf, idxrf->refine_index);
-        delete rf;
+        idxrf->refine_index = read_index(f, io_flags);
         READ1 (idxrf->k_factor);
+        if (dynamic_cast<IndexFlat*>(idxrf->refine_index)) {
+            // then make a RefineFlat with it
+            IndexRefine *idxrf_old = idxrf;
+            idxrf = new IndexRefineFlat();
+            *idxrf = *idxrf_old;
+            delete idxrf_old;
+        }
+        idxrf->own_fields = true;
+        idxrf->own_refine_index = true;
         idx = idxrf;
     } else if(h == fourcc ("IxMp") || h == fourcc ("IxM2")) {
         bool is_map2 = h == fourcc ("IxM2");
@@ -598,8 +605,36 @@ Index *read_index (IOReader *f, int io_flags) {
             dynamic_cast<IndexPQ*>(idxhnsw->storage)->pq.compute_sdc_table ();
         }
         idx = idxhnsw;
+    } else if(h == fourcc("IPfs")) {
+        IndexPQFastScan *idxpqfs = new IndexPQFastScan();
+        read_index_header (idxpqfs, f);
+        read_ProductQuantizer (&idxpqfs->pq, f);
+        READ1 (idxpqfs->implem);
+        READ1 (idxpqfs->bbs);
+        READ1 (idxpqfs->qbs);
+        READ1 (idxpqfs->ntotal2);
+        READ1 (idxpqfs->M2);
+        READVECTOR (idxpqfs->codes);
+        idx = idxpqfs;
+
+    } else if (h == fourcc("IwPf")) {
+        IndexIVFPQFastScan *ivpq = new IndexIVFPQFastScan();
+        read_ivf_header (ivpq, f);
+        READ1 (ivpq->by_residual);
+        READ1 (ivpq->code_size);
+        READ1 (ivpq->bbs);
+        READ1 (ivpq->M2);
+        READ1 (ivpq->implem);
+        READ1 (ivpq->qbs2);
+        read_ProductQuantizer (&ivpq->pq, f);
+        read_InvertedLists (ivpq, f, io_flags);
+        ivpq->precompute_table();
+        idx = ivpq;
     } else {
-        FAISS_THROW_FMT("Index type 0x%08x not supported\n", h);
+        FAISS_THROW_FMT(
+            "Index type 0x%08x (\"%s\") not recognized",
+            h, fourcc_inv_printable(h).c_str()
+        );
         idx = nullptr;
     }
     return idx;
@@ -780,7 +815,10 @@ IndexBinary *read_index_binary (IOReader *f, int io_flags) {
         }
         idx = idxmh;
     } else {
-        FAISS_THROW_FMT("Index type 0x%08x not supported\n", h);
+        FAISS_THROW_FMT(
+            "Index type %08x (\"%s\") not recognized",
+            h, fourcc_inv_printable(h).c_str()
+        );
         idx = nullptr;
     }
     return idx;
@@ -797,76 +835,6 @@ IndexBinary *read_index_binary (const char *fname, int io_flags) {
     return idx;
 }
 
-#ifndef _MSC_VER
-
-/**********************************************************
- * InvertedListIOHook's
- **********************************************************/
-
-InvertedListsIOHook::InvertedListsIOHook(
-        const std::string & key, const std::string & classname):
-    key(key), classname(classname)
-{}
-
-namespace {
-
-/// std::vector that deletes its contents
-struct IOHookTable: std::vector<InvertedListsIOHook*> {
-
-    IOHookTable() {
-        push_back(new OnDiskInvertedListsIOHook());
-    }
-
-    ~IOHookTable() {
-        for (auto x: *this) {
-            delete x;
-        }
-    }
-};
-
-static IOHookTable InvertedListsIOHook_table;
-
-} // anonymous namepsace
-
-InvertedListsIOHook* InvertedListsIOHook::lookup(int h)
-{
-    for(const auto & callback: InvertedListsIOHook_table) {
-        if (h == fourcc(callback->key)) {
-            return callback;
-        }
-    }
-    FAISS_THROW_FMT ("read_InvertedLists: could not load ArrayInvertedLists as %04x", h);
-}
-
-InvertedListsIOHook* InvertedListsIOHook::lookup_classname(const std::string & classname)
-{
-    for(const auto & callback: InvertedListsIOHook_table) {
-        if (callback->classname == classname) {
-            return callback;
-        }
-    }
-    FAISS_THROW_FMT ("read_InvertedLists: could not find classname %s", classname.c_str());
-}
-
-void InvertedListsIOHook::add_callback(InvertedListsIOHook *cb)
-{
-    InvertedListsIOHook_table.push_back(cb);
-}
-
-void InvertedListsIOHook::print_callbacks()
-{
-    printf("registered %zd InvertedListsIOHooks:\n",
-          InvertedListsIOHook_table.size());
-    for(const auto & cb: InvertedListsIOHook_table) {
-        printf("%08x %s %s\n",
-               fourcc(cb->key.c_str()),
-               cb->key.c_str(),
-               cb->classname.c_str());
-    }
-}
-
-#endif // !_MSC_VER
-
 
 
 } // namespace faiss

data/vendor/faiss/faiss/impl/index_write.cpp

@@ -15,9 +15,7 @@
 #include <sys/types.h>
 #include <sys/stat.h>
 
-#ifndef _MSC_VER
-#include <sys/mman.h>
-#endif // !_MSC_VER
+#include <faiss/invlists/InvertedListsIOHook.h>
 
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/io.h>
@@ -39,6 +37,9 @@
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexLattice.h>
+#include <faiss/IndexPQFastScan.h>
+#include <faiss/IndexIVFPQFastScan.h>
+#include <faiss/IndexRefine.h>
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexBinaryFromFloat.h>
@@ -46,11 +47,6 @@
 #include <faiss/IndexBinaryIVF.h>
 #include <faiss/IndexBinaryHash.h>
 
-#ifndef _MSC_VER
-#include <faiss/OnDiskInvertedLists.h>
-#endif // !_MSC_VER
-
-
 /*************************************************************
  * The I/O format is the content of the class. For objects that are
  * inherited, like Index, a 4-character-code (fourcc) indicates which
@@ -66,9 +62,6 @@
  * or deprecated fields), the fourcc can be replaced. New code should
  * be able to read the old fourcc and fill in new classes.
  *
- * TODO: serialization to strings for use in Python pickle or Torch
- * serialization.
- *
  * TODO: in this file, the read functions that encouter errors may
  * leak memory.
  **************************************************************/
@@ -216,19 +209,10 @@ void write_InvertedLists (const InvertedLists *ils, IOWriter *f) {
                 WRITEANDCHECK (ails->ids[i].data(), n);
             }
         }
-#ifndef _MSC_VER
-    } else {
 
+    } else {
         InvertedListsIOHook::lookup_classname(
                 typeid(*ils).name())->write(ils, f);
-
-        /*
-        fprintf(stderr, "WARN! write_InvertedLists: unsupported invlist type, "
-                "saving null invlist\n");
-        uint32_t h = fourcc ("il00");
-        WRITE1 (h);
-        */
-#endif // !_MSC_VER
     }
 }
 
@@ -409,13 +393,13 @@ void write_index (const Index *idx, IOWriter *f) {
         WRITE1 (h);
         write_index_header (imiq, f);
         write_ProductQuantizer (&imiq->pq, f);
-    } else if(const IndexRefineFlat * idxrf =
-              dynamic_cast<const IndexRefineFlat *> (idx)) {
+    } else if(const IndexRefine * idxrf =
+              dynamic_cast<const IndexRefine *> (idx)) {
         uint32_t h = fourcc ("IxRF");
         WRITE1 (h);
         write_index_header (idxrf, f);
         write_index (idxrf->base_index, f);
-        write_index (&idxrf->refine_index, f);
+        write_index (idxrf->refine_index, f);
         WRITE1 (idxrf->k_factor);
     } else if(const IndexIDMap * idxmap =
               dynamic_cast<const IndexIDMap *> (idx)) {
@@ -440,8 +424,33 @@ void write_index (const Index *idx, IOWriter *f) {
         write_index_header (idxhnsw, f);
         write_HNSW (&idxhnsw->hnsw, f);
         write_index (idxhnsw->storage, f);
+    } else if (const IndexPQFastScan *idxpqfs =
+               dynamic_cast<const IndexPQFastScan*>(idx)) {
+        uint32_t h = fourcc("IPfs");
+        WRITE1 (h);
+        write_index_header (idxpqfs, f);
+        write_ProductQuantizer (&idxpqfs->pq, f);
+        WRITE1 (idxpqfs->implem);
+        WRITE1 (idxpqfs->bbs);
+        WRITE1 (idxpqfs->qbs);
+        WRITE1 (idxpqfs->ntotal2);
+        WRITE1 (idxpqfs->M2);
+        WRITEVECTOR (idxpqfs->codes);
+    } else if (const IndexIVFPQFastScan * ivpq =
+              dynamic_cast<const IndexIVFPQFastScan *> (idx)) {
+        uint32_t h = fourcc ("IwPf");
+        WRITE1 (h);
+        write_ivf_header (ivpq, f);
+        WRITE1 (ivpq->by_residual);
+        WRITE1 (ivpq->code_size);
+        WRITE1 (ivpq->bbs);
+        WRITE1 (ivpq->M2);
+        WRITE1 (ivpq->implem);
+        WRITE1 (ivpq->qbs2);
+        write_ProductQuantizer (&ivpq->pq, f);
+        write_InvertedLists (ivpq->invlists, f);
     } else {
-      FAISS_THROW_MSG ("don't know how to serialize this type of index");
+        FAISS_THROW_MSG ("don't know how to serialize this type of index");
     }
 }