faiss 0.1.5 → 0.2.2

data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h

@@ -5,124 +5,122 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/gpu/GpuIndexIVF.h>
 #include <faiss/IndexScalarQuantizer.h>
+#include <faiss/gpu/GpuIndexIVF.h>
 #include <memory>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 class IVFFlat;
 class GpuIndexFlat;
 
 struct GpuIndexIVFScalarQuantizerConfig : public GpuIndexIVFConfig {
-  inline GpuIndexIVFScalarQuantizerConfig()
-      : interleavedLayout(true) {
-  }
+    inline GpuIndexIVFScalarQuantizerConfig() : interleavedLayout(true) {}
 
-  /// Use the alternative memory layout for the IVF lists
-  /// (currently the default)
-  bool interleavedLayout;
+    /// Use the alternative memory layout for the IVF lists
+    /// (currently the default)
+    bool interleavedLayout;
 };
 
 /// Wrapper around the GPU implementation that looks like
 /// faiss::IndexIVFScalarQuantizer
 class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
- public:
-  /// Construct from a pre-existing faiss::IndexIVFScalarQuantizer instance,
-  /// copying data over to the given GPU, if the input index is trained.
-  GpuIndexIVFScalarQuantizer(
-    GpuResourcesProvider* provider,
-    const faiss::IndexIVFScalarQuantizer* index,
-    GpuIndexIVFScalarQuantizerConfig config =
-    GpuIndexIVFScalarQuantizerConfig());
-
-  /// Constructs a new instance with an empty flat quantizer; the user
-  /// provides the number of lists desired.
-  GpuIndexIVFScalarQuantizer(
-    GpuResourcesProvider* provider,
-    int dims,
-    int nlist,
-    faiss::ScalarQuantizer::QuantizerType qtype,
-    faiss::MetricType metric = MetricType::METRIC_L2,
-    bool encodeResidual = true,
-    GpuIndexIVFScalarQuantizerConfig config =
-    GpuIndexIVFScalarQuantizerConfig());
-
-  ~GpuIndexIVFScalarQuantizer() override;
-
-  /// Reserve GPU memory in our inverted lists for this number of vectors
-  void reserveMemory(size_t numVecs);
-
-  /// Initialize ourselves from the given CPU index; will overwrite
-  /// all data in ourselves
-  void copyFrom(const faiss::IndexIVFScalarQuantizer* index);
-
-  /// Copy ourselves to the given CPU index; will overwrite all data
-  /// in the index instance
-  void copyTo(faiss::IndexIVFScalarQuantizer* index) const;
-
-  /// After adding vectors, one can call this to reclaim device memory
-  /// to exactly the amount needed. Returns space reclaimed in bytes
-  size_t reclaimMemory();
-
-  /// Clears out all inverted lists, but retains the coarse and scalar quantizer
-  /// information
-  void reset() override;
-
-  /// Trains the coarse and scalar quantizer based on the given vector data
-  void train(Index::idx_t n, const float* x) override;
-
-  /// Returns the number of vectors present in a particular inverted list
-  int getListLength(int listId) const override;
-
-  /// Return the encoded vector data contained in a particular inverted list,
-  /// for debugging purposes.
-  /// If gpuFormat is true, the data is returned as it is encoded in the
-  /// GPU-side representation.
-  /// Otherwise, it is converted to the CPU format.
-  /// compliant format, while the native GPU format may differ.
-  std::vector<uint8_t>
-  getListVectorData(int listId, bool gpuFormat = false) const override;
-
-  /// Return the vector indices contained in a particular inverted list, for
-  /// debugging purposes.
-  std::vector<Index::idx_t> getListIndices(int listId) const override;
-
- protected:
-  /// Called from GpuIndex for add/add_with_ids
-  void addImpl_(int n,
-                const float* x,
-                const Index::idx_t* ids) override;
-
-  /// Called from GpuIndex for search
-  void searchImpl_(int n,
-                   const float* x,
-                   int k,
-                   float* distances,
-                   Index::idx_t* labels) const override;
-
-  /// Called from train to handle SQ residual training
-  void trainResiduals_(Index::idx_t n, const float* x);
-
- public:
-  /// Exposed like the CPU version
-  faiss::ScalarQuantizer sq;
-
-  /// Exposed like the CPU version
-  bool by_residual;
-
-protected:
-  /// Our configuration options
-  const GpuIndexIVFScalarQuantizerConfig ivfSQConfig_;
-
-  /// Desired inverted list memory reservation
-  size_t reserveMemoryVecs_;
-
-  /// Instance that we own; contains the inverted list
-  std::unique_ptr<IVFFlat> index_;
+   public:
+    /// Construct from a pre-existing faiss::IndexIVFScalarQuantizer instance,
+    /// copying data over to the given GPU, if the input index is trained.
+    GpuIndexIVFScalarQuantizer(
+            GpuResourcesProvider* provider,
+            const faiss::IndexIVFScalarQuantizer* index,
+            GpuIndexIVFScalarQuantizerConfig config =
+                    GpuIndexIVFScalarQuantizerConfig());
+
+    /// Constructs a new instance with an empty flat quantizer; the user
+    /// provides the number of lists desired.
+    GpuIndexIVFScalarQuantizer(
+            GpuResourcesProvider* provider,
+            int dims,
+            int nlist,
+            faiss::ScalarQuantizer::QuantizerType qtype,
+            faiss::MetricType metric = MetricType::METRIC_L2,
+            bool encodeResidual = true,
+            GpuIndexIVFScalarQuantizerConfig config =
+                    GpuIndexIVFScalarQuantizerConfig());
+
+    ~GpuIndexIVFScalarQuantizer() override;
+
+    /// Reserve GPU memory in our inverted lists for this number of vectors
+    void reserveMemory(size_t numVecs);
+
+    /// Initialize ourselves from the given CPU index; will overwrite
+    /// all data in ourselves
+    void copyFrom(const faiss::IndexIVFScalarQuantizer* index);
+
+    /// Copy ourselves to the given CPU index; will overwrite all data
+    /// in the index instance
+    void copyTo(faiss::IndexIVFScalarQuantizer* index) const;
+
+    /// After adding vectors, one can call this to reclaim device memory
+    /// to exactly the amount needed. Returns space reclaimed in bytes
+    size_t reclaimMemory();
+
+    /// Clears out all inverted lists, but retains the coarse and scalar
+    /// quantizer information
+    void reset() override;
+
+    /// Trains the coarse and scalar quantizer based on the given vector data
+    void train(Index::idx_t n, const float* x) override;
+
+    /// Returns the number of vectors present in a particular inverted list
+    int getListLength(int listId) const override;
+
+    /// Return the encoded vector data contained in a particular inverted list,
+    /// for debugging purposes.
+    /// If gpuFormat is true, the data is returned as it is encoded in the
+    /// GPU-side representation.
+    /// Otherwise, it is converted to the CPU format.
+    /// compliant format, while the native GPU format may differ.
+    std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false)
+            const override;
+
+    /// Return the vector indices contained in a particular inverted list, for
+    /// debugging purposes.
+    std::vector<Index::idx_t> getListIndices(int listId) const override;
+
+   protected:
+    /// Called from GpuIndex for add/add_with_ids
+    void addImpl_(int n, const float* x, const Index::idx_t* ids) override;
+
+    /// Called from GpuIndex for search
+    void searchImpl_(
+            int n,
+            const float* x,
+            int k,
+            float* distances,
+            Index::idx_t* labels) const override;
+
+    /// Called from train to handle SQ residual training
+    void trainResiduals_(Index::idx_t n, const float* x);
+
+   public:
+    /// Exposed like the CPU version
+    faiss::ScalarQuantizer sq;
+
+    /// Exposed like the CPU version
+    bool by_residual;
+
+   protected:
+    /// Our configuration options
+    const GpuIndexIVFScalarQuantizerConfig ivfSQConfig_;
+
+    /// Desired inverted list memory reservation
+    size_t reserveMemoryVecs_;
+
+    /// Instance that we own; contains the inverted list
+    std::unique_ptr<IVFFlat> index_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h

@@ -5,26 +5,27 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// How user vector index data is stored on the GPU
 enum IndicesOptions {
-  /// The user indices are only stored on the CPU; the GPU returns
-  /// (inverted list, offset) to the CPU which is then translated to
-  /// the real user index.
-  INDICES_CPU = 0,
-  /// The indices are not stored at all, on either the CPU or
-  /// GPU. Only (inverted list, offset) is returned to the user as the
-  /// index.
-  INDICES_IVF = 1,
-  /// Indices are stored as 32 bit integers on the GPU, but returned
-  /// as 64 bit integers
-  INDICES_32_BIT = 2,
-  /// Indices are stored as 64 bit integers on the GPU
-  INDICES_64_BIT = 3,
+    /// The user indices are only stored on the CPU; the GPU returns
+    /// (inverted list, offset) to the CPU which is then translated to
+    /// the real user index.
+    INDICES_CPU = 0,
+    /// The indices are not stored at all, on either the CPU or
+    /// GPU. Only (inverted list, offset) is returned to the user as the
+    /// index.
+    INDICES_IVF = 1,
+    /// Indices are stored as 32 bit integers on the GPU, but returned
+    /// as 64 bit integers
+    INDICES_32_BIT = 2,
+    /// Indices are stored as 64 bit integers on the GPU
+    INDICES_64_BIT = 3,
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

data/vendor/faiss/faiss/gpu/GpuResources.cpp

@@ -5,76 +5,72 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #include <faiss/gpu/GpuResources.h>
 #include <faiss/gpu/utils/DeviceUtils.h>
 #include <sstream>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 std::string allocTypeToString(AllocType t) {
-  switch (t) {
-    case AllocType::Other:
-      return "Other";
-    case AllocType::FlatData:
-      return "FlatData";
-    case AllocType::IVFLists:
-      return "IVFLists";
-    case AllocType::Quantizer:
-      return "Quantizer";
-    case AllocType::QuantizerPrecomputedCodes:
-      return "QuantizerPrecomputedCodes";
-    case AllocType::TemporaryMemoryBuffer:
-      return "TemporaryMemoryBuffer";
-    case AllocType::TemporaryMemoryOverflow:
-      return "TemporaryMemoryOverflow";
-    default:
-      return "Unknown";
-  }
+    switch (t) {
+        case AllocType::Other:
+            return "Other";
+        case AllocType::FlatData:
+            return "FlatData";
+        case AllocType::IVFLists:
+            return "IVFLists";
+        case AllocType::Quantizer:
+            return "Quantizer";
+        case AllocType::QuantizerPrecomputedCodes:
+            return "QuantizerPrecomputedCodes";
+        case AllocType::TemporaryMemoryBuffer:
+            return "TemporaryMemoryBuffer";
+        case AllocType::TemporaryMemoryOverflow:
+            return "TemporaryMemoryOverflow";
+        default:
+            return "Unknown";
+    }
 }
 
 std::string memorySpaceToString(MemorySpace s) {
-  switch (s) {
-    case MemorySpace::Temporary:
-      return "Temporary";
-    case MemorySpace::Device:
-      return "Device";
-    case MemorySpace::Unified:
-      return "Unified";
-    default:
-      return "Unknown";
-  }
+    switch (s) {
+        case MemorySpace::Temporary:
+            return "Temporary";
+        case MemorySpace::Device:
+            return "Device";
+        case MemorySpace::Unified:
+            return "Unified";
+        default:
+            return "Unknown";
+    }
 }
 
-std::string
-AllocInfo::toString() const {
-  std::stringstream ss;
-  ss << "type " << allocTypeToString(type)
-     << " dev " << device
-     << " space " << memorySpaceToString(space)
-     << " stream " << (void*) stream;
+std::string AllocInfo::toString() const {
+    std::stringstream ss;
+    ss << "type " << allocTypeToString(type) << " dev " << device << " space "
+       << memorySpaceToString(space) << " stream " << (void*)stream;
 
-  return ss.str();
+    return ss.str();
 }
 
-std::string
-AllocRequest::toString() const {
-  std::stringstream ss;
-  ss << AllocInfo::toString() << " size " << size << " bytes";
+std::string AllocRequest::toString() const {
+    std::stringstream ss;
+    ss << AllocInfo::toString() << " size " << size << " bytes";
 
-  return ss.str();
+    return ss.str();
 }
 
 AllocInfo makeDevAlloc(AllocType at, cudaStream_t st) {
-  return AllocInfo(at, getCurrentDevice(), MemorySpace::Device, st);
+    return AllocInfo(at, getCurrentDevice(), MemorySpace::Device, st);
 }
 
 AllocInfo makeTempAlloc(AllocType at, cudaStream_t st) {
-  return AllocInfo(at, getCurrentDevice(), MemorySpace::Temporary, st);
+    return AllocInfo(at, getCurrentDevice(), MemorySpace::Temporary, st);
 }
 
 AllocInfo makeSpaceAlloc(AllocType at, MemorySpace sp, cudaStream_t st) {
-  return AllocInfo(at, getCurrentDevice(), sp, st);
+    return AllocInfo(at, getCurrentDevice(), sp, st);
 }
 
 //
@@ -82,119 +78,111 @@ AllocInfo makeSpaceAlloc(AllocType at, MemorySpace sp, cudaStream_t st) {
 //
 
 GpuMemoryReservation::GpuMemoryReservation()
-    : res(nullptr),
-      device(0),
-      stream(nullptr),
-      data(nullptr),
-      size(0) {
-}
+        : res(nullptr), device(0), stream(nullptr), data(nullptr), size(0) {}
 
-GpuMemoryReservation::GpuMemoryReservation(GpuResources* r,
-                                           int dev,
-                                           cudaStream_t str,
-                                           void* p,
-                                           size_t sz)
-    : res(r),
-      device(dev),
-      stream(str),
-      data(p),
-      size(sz) {
-}
+GpuMemoryReservation::GpuMemoryReservation(
+        GpuResources* r,
+        int dev,
+        cudaStream_t str,
+        void* p,
+        size_t sz)
+        : res(r), device(dev), stream(str), data(p), size(sz) {}
 
 GpuMemoryReservation::GpuMemoryReservation(GpuMemoryReservation&& m) noexcept {
-  res = m.res; m.res = nullptr;
-  device = m.device; m.device = 0;
-  stream = m.stream; m.stream = nullptr;
-  data = m.data; m.data = nullptr;
-  size = m.size; m.size = 0;
-}
-
-GpuMemoryReservation&
-GpuMemoryReservation::operator=(GpuMemoryReservation&& m) {
-  // Can't be both a valid allocation and the same allocation
-  FAISS_ASSERT(!(res && res == m.res && device == m.device && data == m.data));
-
-  release();
-  res = m.res; m.res = nullptr;
-  device = m.device; m.device = 0;
-  stream = m.stream; m.stream = nullptr;
-  data = m.data; m.data = nullptr;
-  size = m.size; m.size = 0;
-
-  return *this;
-}
-
-void
-GpuMemoryReservation::release() {
-  if (res) {
-    res->deallocMemory(device, data);
-    res = nullptr;
-    device = 0;
-    stream = nullptr;
-    data = nullptr;
-    size = 0;
-  }
+    res = m.res;
+    m.res = nullptr;
+    device = m.device;
+    m.device = 0;
+    stream = m.stream;
+    m.stream = nullptr;
+    data = m.data;
+    m.data = nullptr;
+    size = m.size;
+    m.size = 0;
+}
+
+GpuMemoryReservation& GpuMemoryReservation::operator=(
+        GpuMemoryReservation&& m) {
+    // Can't be both a valid allocation and the same allocation
+    FAISS_ASSERT(
+            !(res && res == m.res && device == m.device && data == m.data));
+
+    release();
+    res = m.res;
+    m.res = nullptr;
+    device = m.device;
+    m.device = 0;
+    stream = m.stream;
+    m.stream = nullptr;
+    data = m.data;
+    m.data = nullptr;
+    size = m.size;
+    m.size = 0;
+
+    return *this;
+}
+
+void GpuMemoryReservation::release() {
+    if (res) {
+        res->deallocMemory(device, data);
+        res = nullptr;
+        device = 0;
+        stream = nullptr;
+        data = nullptr;
+        size = 0;
+    }
 }
 
 GpuMemoryReservation::~GpuMemoryReservation() {
-  if (res) {
-    res->deallocMemory(device, data);
-  }
+    if (res) {
+        res->deallocMemory(device, data);
+    }
 }
 
 //
 // GpuResources
 //
 
-GpuResources::~GpuResources() {
-}
+GpuResources::~GpuResources() {}
 
-cublasHandle_t
-GpuResources::getBlasHandleCurrentDevice() {
-  return getBlasHandle(getCurrentDevice());
+cublasHandle_t GpuResources::getBlasHandleCurrentDevice() {
+    return getBlasHandle(getCurrentDevice());
 }
 
-cudaStream_t
-GpuResources::getDefaultStreamCurrentDevice() {
-  return getDefaultStream(getCurrentDevice());
+cudaStream_t GpuResources::getDefaultStreamCurrentDevice() {
+    return getDefaultStream(getCurrentDevice());
 }
 
-std::vector<cudaStream_t>
-GpuResources::getAlternateStreamsCurrentDevice() {
-  return getAlternateStreams(getCurrentDevice());
+std::vector<cudaStream_t> GpuResources::getAlternateStreamsCurrentDevice() {
+    return getAlternateStreams(getCurrentDevice());
 }
 
-cudaStream_t
-GpuResources::getAsyncCopyStreamCurrentDevice() {
-  return getAsyncCopyStream(getCurrentDevice());
+cudaStream_t GpuResources::getAsyncCopyStreamCurrentDevice() {
+    return getAsyncCopyStream(getCurrentDevice());
 }
 
-void
-GpuResources::syncDefaultStream(int device) {
-  CUDA_VERIFY(cudaStreamSynchronize(getDefaultStream(device)));
+void GpuResources::syncDefaultStream(int device) {
+    CUDA_VERIFY(cudaStreamSynchronize(getDefaultStream(device)));
 }
 
-void
-GpuResources::syncDefaultStreamCurrentDevice() {
-  syncDefaultStream(getCurrentDevice());
+void GpuResources::syncDefaultStreamCurrentDevice() {
+    syncDefaultStream(getCurrentDevice());
 }
 
-GpuMemoryReservation
-GpuResources::allocMemoryHandle(const AllocRequest& req) {
-  return GpuMemoryReservation(
-    this, req.device, req.stream, allocMemory(req), req.size);
+GpuMemoryReservation GpuResources::allocMemoryHandle(const AllocRequest& req) {
+    return GpuMemoryReservation(
+            this, req.device, req.stream, allocMemory(req), req.size);
 }
 
-size_t
-GpuResources::getTempMemoryAvailableCurrentDevice() const {
-  return getTempMemoryAvailable(getCurrentDevice());
+size_t GpuResources::getTempMemoryAvailableCurrentDevice() const {
+    return getTempMemoryAvailable(getCurrentDevice());
 }
 
 //
 // GpuResourcesProvider
 //
 
-GpuResourcesProvider::~GpuResourcesProvider() {
-}
+GpuResourcesProvider::~GpuResourcesProvider() {}
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss