faiss 0.1.1 → 0.1.2

data/vendor/faiss/gpu/utils/DeviceUtils.h

@@ -64,6 +64,12 @@ bool getFullUnifiedMemSupport(int device);
 /// Equivalent to getFullUnifiedMemSupport(getCurrentDevice())
 bool getFullUnifiedMemSupportCurrentDevice();
 
+/// Does the given device support tensor core operations?
+bool getTensorCoreSupport(int device);
+
+/// Equivalent to getTensorCoreSupport(getCurrentDevice())
+bool getTensorCoreSupportCurrentDevice();
+
 /// Returns the maximum k-selection value supported based on the CUDA SDK that
 /// we were compiled with. .cu files can use DeviceDefs.cuh, but this is for
 /// non-CUDA files

data/vendor/faiss/impl/AuxIndexStructures.cpp

@@ -228,6 +228,23 @@ bool IDSelectorRange::is_member (idx_t id) const
     return id >= imin && id < imax;
 }
 
+/***********************************************************************
+ * IDSelectorArray
+ ***********************************************************************/
+
+IDSelectorArray::IDSelectorArray (size_t n, const idx_t *ids):
+   n (n), ids(ids)
+{
+}
+
+bool IDSelectorArray::is_member (idx_t id) const
+{
+    for (idx_t i = 0; i < n; i++) {
+        if (ids[i] == id) return true;
+    }
+    return false;
+}
+
 
 /***********************************************************************
  * IDSelectorBatch

data/vendor/faiss/impl/AuxIndexStructures.h

@@ -51,9 +51,7 @@ struct RangeSearchResult {
 };
 
 
-/**
-
- Encapsulates a set of ids to remove. */
+/** Encapsulates a set of ids to remove. */
 struct IDSelector {
     typedef Index::idx_t idx_t;
     virtual bool is_member (idx_t id) const = 0;
@@ -71,6 +69,19 @@ struct IDSelectorRange: IDSelector {
     ~IDSelectorRange() override {}
 };
 
+/** simple list of elements to remove
+ *
+ * this is inefficient in most cases, except for IndexIVF with
+ * maintain_direct_map
+ */
+struct IDSelectorArray: IDSelector {
+    size_t n;
+    const idx_t *ids;
+
+    IDSelectorArray (size_t n, const idx_t *ids);
+    bool is_member(idx_t id) const override;
+    ~IDSelectorArray() override {}
+};
 
 /** Remove ids from a set. Repetitions of ids in the indices set
  * passed to the constructor does not hurt performance. The hash

data/vendor/faiss/impl/HNSW.cpp

@@ -15,7 +15,6 @@
 
 namespace faiss {
 
-using idx_t = Index::idx_t;
 
 /**************************************************************
  * HNSW structure implementation

data/vendor/faiss/impl/PolysemousTraining.h

@@ -123,15 +123,15 @@ struct PolysemousTraining: SimulatedAnnealingParameters {
     enum Optimization_type_t {
         OT_None,
         OT_ReproduceDistances_affine,  ///< default
-        OT_Ranking_weighted_diff  /// same as _2, but use rank of y+ - rank of y-
+        OT_Ranking_weighted_diff       ///< same as _2, but use rank of y+ - rank of y-
     };
     Optimization_type_t optimization_type;
 
-    // use 1/4 of the training points for the optimization, with
-    // max. ntrain_permutation. If ntrain_permutation == 0: train on
-    // centroids
+    /** use 1/4 of the training points for the optimization, with
+     * max. ntrain_permutation. If ntrain_permutation == 0: train on
+     * centroids */
     int ntrain_permutation;
-    double dis_weight_factor; // decay of exp that weights distance loss
+    double dis_weight_factor; ///< decay of exp that weights distance loss
 
     // filename pattern for the logging of iterations
     std::string log_pattern;

data/vendor/faiss/impl/ProductQuantizer-inl.h

@@ -0,0 +1,138 @@
+/**
+ * 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.
+ */
+
+namespace faiss {
+
+inline
+PQEncoderGeneric::PQEncoderGeneric(uint8_t *code, int nbits,
+                                                     uint8_t offset)
+    : code(code), offset(offset), nbits(nbits), reg(0)
+{
+    assert(nbits <= 64);
+    if (offset > 0) {
+        reg = (*code & ((1 << offset) - 1));
+    }
+}
+
+inline
+void PQEncoderGeneric::encode(uint64_t x)
+{
+    reg |= (uint8_t)(x << offset);
+    x >>= (8 - offset);
+    if (offset + nbits >= 8) {
+        *code++ = reg;
+
+        for (int i = 0; i < (nbits - (8 - offset)) / 8; ++i) {
+            *code++ = (uint8_t)x;
+            x >>= 8;
+        }
+
+        offset += nbits;
+        offset &= 7;
+        reg = (uint8_t)x;
+    } else {
+        offset += nbits;
+    }
+}
+
+inline
+PQEncoderGeneric::~PQEncoderGeneric()
+{
+    if (offset > 0) {
+        *code = reg;
+    }
+}
+
+
+inline
+PQEncoder8::PQEncoder8(uint8_t *code, int nbits)
+    : code(code) {
+    assert(8 == nbits);
+}
+
+inline
+void PQEncoder8::encode(uint64_t x) {
+    *code++ = (uint8_t)x;
+}
+
+inline
+PQEncoder16::PQEncoder16(uint8_t *code, int nbits)
+    : code((uint16_t *)code) {
+    assert(16 == nbits);
+}
+
+inline
+void PQEncoder16::encode(uint64_t x) {
+    *code++ = (uint16_t)x;
+}
+
+
+inline
+PQDecoderGeneric::PQDecoderGeneric(const uint8_t *code,
+                                                     int nbits)
+    : code(code),
+      offset(0),
+      nbits(nbits),
+      mask((1ull << nbits) - 1),
+      reg(0) {
+    assert(nbits <= 64);
+}
+
+inline
+uint64_t PQDecoderGeneric::decode() {
+    if (offset == 0) {
+        reg = *code;
+    }
+    uint64_t c = (reg >> offset);
+
+    if (offset + nbits >= 8) {
+        uint64_t e = 8 - offset;
+        ++code;
+        for (int i = 0; i < (nbits - (8 - offset)) / 8; ++i) {
+            c |= ((uint64_t)(*code++) << e);
+            e += 8;
+        }
+
+        offset += nbits;
+        offset &= 7;
+        if (offset > 0) {
+            reg = *code;
+            c |= ((uint64_t)reg << e);
+        }
+    } else {
+        offset += nbits;
+    }
+
+    return c & mask;
+}
+
+
+inline
+PQDecoder8::PQDecoder8(const uint8_t *code, int nbits)
+    : code(code) {
+    assert(8 == nbits);
+}
+
+inline
+uint64_t PQDecoder8::decode() {
+    return (uint64_t)(*code++);
+}
+
+
+inline
+PQDecoder16::PQDecoder16(const uint8_t *code, int nbits)
+    : code((uint16_t *)code) {
+     assert(16 == nbits);
+}
+
+inline
+uint64_t PQDecoder16::decode() {
+    return (uint64_t)(*code++);
+}
+
+
+} // namespace faiss

data/vendor/faiss/impl/ProductQuantizer.cpp

@@ -151,7 +151,7 @@ static inline void pq_estimators_from_tables_generic(const ProductQuantizer& pq,
   const size_t M = pq.M;
   const size_t ksub = pq.ksub;
   for (size_t j = 0; j < ncodes; ++j) {
-    faiss::ProductQuantizer::PQDecoderGeneric decoder(
+    PQDecoderGeneric decoder(
       codes + j * pq.code_size, nbits
     );
     float dis = 0;
@@ -760,117 +760,5 @@ void ProductQuantizer::search_sdc (const uint8_t * qcodes,
 }
 
 
-ProductQuantizer::PQEncoderGeneric::PQEncoderGeneric(uint8_t *code, int nbits,
-                                                     uint8_t offset)
-    : code(code), offset(offset), nbits(nbits), reg(0) {
-  assert(nbits <= 64);
-  if (offset > 0) {
-    reg = (*code & ((1 << offset) - 1));
-  }
-}
-
-void ProductQuantizer::PQEncoderGeneric::encode(uint64_t x) {
-  reg |= (uint8_t)(x << offset);
-  x >>= (8 - offset);
-  if (offset + nbits >= 8) {
-    *code++ = reg;
-
-    for (int i = 0; i < (nbits - (8 - offset)) / 8; ++i) {
-      *code++ = (uint8_t)x;
-      x >>= 8;
-    }
-
-    offset += nbits;
-    offset &= 7;
-    reg = (uint8_t)x;
-  } else {
-    offset += nbits;
-  }
-}
-
-ProductQuantizer::PQEncoderGeneric::~PQEncoderGeneric() {
-  if (offset > 0) {
-    *code = reg;
-  }
-}
-
-
-ProductQuantizer::PQEncoder8::PQEncoder8(uint8_t *code, int nbits)
-    : code(code) {
-  assert(8 == nbits);
-}
-
-void ProductQuantizer::PQEncoder8::encode(uint64_t x) {
-  *code++ = (uint8_t)x;
-}
-
-
-ProductQuantizer::PQEncoder16::PQEncoder16(uint8_t *code, int nbits)
-    : code((uint16_t *)code) {
-  assert(16 == nbits);
-}
-
-void ProductQuantizer::PQEncoder16::encode(uint64_t x) {
-  *code++ = (uint16_t)x;
-}
-
-
-ProductQuantizer::PQDecoderGeneric::PQDecoderGeneric(const uint8_t *code,
-                                                     int nbits)
-    : code(code),
-      offset(0),
-      nbits(nbits),
-      mask((1ull << nbits) - 1),
-      reg(0) {
-  assert(nbits <= 64);
-}
-
-uint64_t ProductQuantizer::PQDecoderGeneric::decode() {
-  if (offset == 0) {
-    reg = *code;
-  }
-  uint64_t c = (reg >> offset);
-
-  if (offset + nbits >= 8) {
-    uint64_t e = 8 - offset;
-    ++code;
-    for (int i = 0; i < (nbits - (8 - offset)) / 8; ++i) {
-      c |= ((uint64_t)(*code++) << e);
-      e += 8;
-    }
-
-    offset += nbits;
-    offset &= 7;
-    if (offset > 0) {
-      reg = *code;
-      c |= ((uint64_t)reg << e);
-    }
-  } else {
-    offset += nbits;
-  }
-
-  return c & mask;
-}
-
-
-ProductQuantizer::PQDecoder8::PQDecoder8(const uint8_t *code, int nbits)
-    : code(code) {
-  assert(8 == nbits);
-}
-
-uint64_t ProductQuantizer::PQDecoder8::decode() {
-  return (uint64_t)(*code++);
-}
-
-
-ProductQuantizer::PQDecoder16::PQDecoder16(const uint8_t *code, int nbits)
-    : code((uint16_t *)code) {
-  assert(16 == nbits);
-}
-
-uint64_t ProductQuantizer::PQDecoder16::decode() {
-  return (uint64_t)(*code++);
-}
-
 
 }  // namespace faiss

data/vendor/faiss/impl/ProductQuantizer.h

@@ -173,70 +173,65 @@ struct ProductQuantizer {
                      float_maxheap_array_t * res,
                      bool init_finalize_heap = true) const;
 
-    struct PQEncoderGeneric {
-        uint8_t *code;   ///< code for this vector
-        uint8_t offset;
-        const int nbits; ///< number of bits per subquantizer index
-
-        uint8_t reg;
-
-        PQEncoderGeneric(uint8_t *code, int nbits, uint8_t offset = 0);
-
-        void encode(uint64_t x);
-
-        ~PQEncoderGeneric();
-    };
-
-
-    struct PQEncoder8 {
-        uint8_t *code;
-
-        PQEncoder8(uint8_t *code, int nbits);
-
-        void encode(uint64_t x);
-    };
-
-    struct PQEncoder16 {
-        uint16_t *code;
+};
 
-        PQEncoder16(uint8_t *code, int nbits);
 
-        void encode(uint64_t x);
-    };
+/*************************************************
+ * Objects to encode / decode strings of bits
+ *************************************************/
 
+struct PQEncoderGeneric {
+    uint8_t *code;   ///< code for this vector
+    uint8_t offset;
+    const int nbits; ///< number of bits per subquantizer index
 
-    struct PQDecoderGeneric {
-        const uint8_t *code;
-        uint8_t offset;
-        const int nbits;
-        const uint64_t mask;
-        uint8_t reg;
+    uint8_t reg;
 
-        PQDecoderGeneric(const uint8_t *code, int nbits);
+    PQEncoderGeneric(uint8_t *code, int nbits, uint8_t offset = 0);
 
-        uint64_t decode();
-    };
+    void encode(uint64_t x);
 
-    struct PQDecoder8 {
-        const uint8_t *code;
+    ~PQEncoderGeneric();
+};
 
-        PQDecoder8(const uint8_t *code, int nbits);
 
-        uint64_t decode();
-    };
+struct PQEncoder8 {
+    uint8_t *code;
+    PQEncoder8(uint8_t *code, int nbits);
+    void encode(uint64_t x);
+};
 
-    struct PQDecoder16 {
-        const uint16_t *code;
+struct PQEncoder16 {
+    uint16_t *code;
+    PQEncoder16(uint8_t *code, int nbits);
+    void encode(uint64_t x);
+};
 
-        PQDecoder16(const uint8_t *code, int nbits);
 
-        uint64_t decode();
-    };
+struct PQDecoderGeneric {
+    const uint8_t *code;
+    uint8_t offset;
+    const int nbits;
+    const uint64_t mask;
+    uint8_t reg;
+    PQDecoderGeneric(const uint8_t *code, int nbits);
+    uint64_t decode();
+};
 
+struct PQDecoder8 {
+    const uint8_t *code;
+    PQDecoder8(const uint8_t *code, int nbits);
+    uint64_t decode();
 };
 
+struct PQDecoder16 {
+    const uint16_t *code;
+    PQDecoder16(const uint8_t *code, int nbits);
+    uint64_t decode();
+};
 
 }  // namespace faiss
 
+#include <faiss/impl/ProductQuantizer-inl.h>
 
 #endif