faiss 0.1.0 → 0.1.1

data/vendor/faiss/gpu/GpuIndexIVFScalarQuantizer.h

@@ -0,0 +1,100 @@
+/**
+ * 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.
+ */
+
+
+#pragma once
+
+#include <faiss/gpu/GpuIndexIVF.h>
+#include <faiss/IndexScalarQuantizer.h>
+
+namespace faiss { namespace gpu {
+
+class IVFFlat;
+class GpuIndexFlat;
+
+struct GpuIndexIVFScalarQuantizerConfig : public GpuIndexIVFConfig {
+};
+
+/// 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(
+    GpuResources* resources,
+    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(
+    GpuResources* resources,
+    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();
+
+  void reset() override;
+
+  void train(Index::idx_t n, const float* x) 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;
+
+ private:
+  GpuIndexIVFScalarQuantizerConfig ivfSQConfig_;
+
+  /// Desired inverted list memory reservation
+  size_t reserveMemoryVecs_;
+
+  /// Instance that we own; contains the inverted list
+  IVFFlat* index_;
+};
+
+} } // namespace

data/vendor/faiss/gpu/GpuIndicesOptions.h

@@ -0,0 +1,30 @@
+/**
+ * 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.
+ */
+
+
+#pragma once
+
+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,
+};
+
+} } // namespace

data/vendor/faiss/gpu/GpuResources.cpp

@@ -0,0 +1,52 @@
+/**
+ * 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.
+ */
+
+
+#include <faiss/gpu/GpuResources.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+
+namespace faiss { namespace gpu {
+
+GpuResources::~GpuResources() {
+}
+
+cublasHandle_t
+GpuResources::getBlasHandleCurrentDevice() {
+  return getBlasHandle(getCurrentDevice());
+}
+
+cudaStream_t
+GpuResources::getDefaultStreamCurrentDevice() {
+  return getDefaultStream(getCurrentDevice());
+}
+
+std::vector<cudaStream_t>
+GpuResources::getAlternateStreamsCurrentDevice() {
+  return getAlternateStreams(getCurrentDevice());
+}
+
+DeviceMemory&
+GpuResources::getMemoryManagerCurrentDevice() {
+  return getMemoryManager(getCurrentDevice());
+}
+
+cudaStream_t
+GpuResources::getAsyncCopyStreamCurrentDevice() {
+  return getAsyncCopyStream(getCurrentDevice());
+}
+
+void
+GpuResources::syncDefaultStream(int device) {
+  CUDA_VERIFY(cudaStreamSynchronize(getDefaultStream(device)));
+}
+
+void
+GpuResources::syncDefaultStreamCurrentDevice() {
+  syncDefaultStream(getCurrentDevice());
+}
+
+} } // namespace

data/vendor/faiss/gpu/GpuResources.h

@@ -0,0 +1,73 @@
+/**
+ * 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.
+ */
+
+
+#pragma once
+
+#include <faiss/gpu/utils/DeviceMemory.h>
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <utility>
+#include <vector>
+
+namespace faiss { namespace gpu {
+
+/// Base class of GPU-side resource provider; hides provision of
+/// cuBLAS handles, CUDA streams and a temporary memory manager
+class GpuResources {
+ public:
+  virtual ~GpuResources();
+
+  /// Call to pre-allocate resources for a particular device. If this is
+  /// not called, then resources will be allocated at the first time
+  /// of demand
+  virtual void initializeForDevice(int device) = 0;
+
+  /// Returns the cuBLAS handle that we use for the given device
+  virtual cublasHandle_t getBlasHandle(int device) = 0;
+
+  /// Returns the stream that we order all computation on for the
+  /// given device
+  virtual cudaStream_t getDefaultStream(int device) = 0;
+
+  /// Returns the set of alternative streams that we use for the given device
+  virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0;
+
+  /// Returns the temporary memory manager for the given device
+  virtual DeviceMemory& getMemoryManager(int device) = 0;
+
+  /// Returns the available CPU pinned memory buffer
+  virtual std::pair<void*, size_t> getPinnedMemory() = 0;
+
+  /// Returns the stream on which we perform async CPU <-> GPU copies
+  virtual cudaStream_t getAsyncCopyStream(int device) = 0;
+
+  /// Calls getBlasHandle with the current device
+  cublasHandle_t getBlasHandleCurrentDevice();
+
+  /// Calls getDefaultStream with the current device
+  cudaStream_t getDefaultStreamCurrentDevice();
+
+  /// Synchronizes the CPU with respect to the default stream for the
+  /// given device
+  // equivalent to cudaDeviceSynchronize(getDefaultStream(device))
+  void syncDefaultStream(int device);
+
+  /// Calls syncDefaultStream for the current device
+  void syncDefaultStreamCurrentDevice();
+
+  /// Calls getAlternateStreams for the current device
+  std::vector<cudaStream_t> getAlternateStreamsCurrentDevice();
+
+  /// Calls getMemoryManager for the current device
+  DeviceMemory& getMemoryManagerCurrentDevice();
+
+  /// Calls getAsyncCopyStream for the current device
+  cudaStream_t getAsyncCopyStreamCurrentDevice();
+};
+
+} } // namespace

data/vendor/faiss/gpu/StandardGpuResources.cpp

@@ -0,0 +1,295 @@
+/**
+ * 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.
+ */
+
+
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/utils/MemorySpace.h>
+#include <faiss/impl/FaissAssert.h>
+#include <limits>
+
+namespace faiss { namespace gpu {
+
+namespace {
+
+// How many streams per device we allocate by default (for multi-streaming)
+constexpr int kNumStreams = 2;
+
+// Use 256 MiB of pinned memory for async CPU <-> GPU copies by default
+constexpr size_t kDefaultPinnedMemoryAllocation = (size_t) 256 * 1024 * 1024;
+
+// Default temporary memory allocation for <= 4 GiB memory GPUs
+constexpr size_t k4GiBTempMem = (size_t) 512 * 1024 * 1024;
+
+// Default temporary memory allocation for <= 8 GiB memory GPUs
+constexpr size_t k8GiBTempMem = (size_t) 1024 * 1024 * 1024;
+
+// Maximum temporary memory allocation for all GPUs
+constexpr size_t kMaxTempMem = (size_t) 1536 * 1024 * 1024;
+
+}
+
+StandardGpuResources::StandardGpuResources() :
+    pinnedMemAlloc_(nullptr),
+    pinnedMemAllocSize_(0),
+    // let the adjustment function determine the memory size for us by passing
+    // in a huge value that will then be adjusted
+    tempMemSize_(getDefaultTempMemForGPU(-1,
+                                         std::numeric_limits<size_t>::max())),
+    pinnedMemSize_(kDefaultPinnedMemoryAllocation),
+    cudaMallocWarning_(true) {
+}
+
+StandardGpuResources::~StandardGpuResources() {
+  for (auto& entry : defaultStreams_) {
+    DeviceScope scope(entry.first);
+
+    auto it = userDefaultStreams_.find(entry.first);
+    if (it == userDefaultStreams_.end()) {
+      // The user did not specify this stream, thus we are the ones
+      // who have created it
+      CUDA_VERIFY(cudaStreamDestroy(entry.second));
+    }
+  }
+
+  for (auto& entry : alternateStreams_) {
+    DeviceScope scope(entry.first);
+
+    for (auto stream : entry.second) {
+      CUDA_VERIFY(cudaStreamDestroy(stream));
+    }
+  }
+
+  for (auto& entry : asyncCopyStreams_) {
+    DeviceScope scope(entry.first);
+
+    CUDA_VERIFY(cudaStreamDestroy(entry.second));
+  }
+
+  for (auto& entry : blasHandles_) {
+    DeviceScope scope(entry.first);
+
+    auto blasStatus = cublasDestroy(entry.second);
+    FAISS_ASSERT(blasStatus == CUBLAS_STATUS_SUCCESS);
+  }
+
+  if (pinnedMemAlloc_) {
+    freeMemorySpace(MemorySpace::HostPinned, pinnedMemAlloc_);
+  }
+}
+
+size_t
+StandardGpuResources::getDefaultTempMemForGPU(int device,
+                                              size_t requested) {
+  auto totalMem = device != -1 ?
+    getDeviceProperties(device).totalGlobalMem :
+    std::numeric_limits<size_t>::max();
+
+  if (totalMem <= (size_t) 4 * 1024 * 1024 * 1024) {
+    // If the GPU has <= 4 GiB of memory, reserve 512 MiB
+
+    if (requested > k4GiBTempMem) {
+      return k4GiBTempMem;
+    }
+  } else if (totalMem <= (size_t) 8 * 1024 * 1024 * 1024) {
+    // If the GPU has <= 8 GiB of memory, reserve 1 GiB
+
+    if (requested > k8GiBTempMem) {
+      return k8GiBTempMem;
+    }
+  } else {
+    // Never use more than 1.5 GiB
+    if (requested > kMaxTempMem) {
+      return kMaxTempMem;
+    }
+  }
+
+  // use whatever lower limit the user requested
+  return requested;
+}
+
+void
+StandardGpuResources::noTempMemory() {
+  setTempMemory(0);
+  setCudaMallocWarning(false);
+}
+
+void
+StandardGpuResources::setTempMemory(size_t size) {
+  if (tempMemSize_ != size) {
+    // adjust based on general limits
+    tempMemSize_ = getDefaultTempMemForGPU(-1, size);
+
+    // We need to re-initialize memory resources for all current devices that
+    // have been initialized.
+    // This should be safe to do, even if we are currently running work, because
+    // the cudaFree call that this implies will force-synchronize all GPUs with
+    // the CPU
+    for (auto& p : memory_) {
+      int device = p.first;
+      // Free the existing memory first
+      p.second.reset();
+
+      // Allocate new
+      p.second = std::unique_ptr<StackDeviceMemory>(
+        new StackDeviceMemory(p.first,
+                              // adjust for this specific device
+                              getDefaultTempMemForGPU(device, tempMemSize_)));
+    }
+  }
+}
+
+void
+StandardGpuResources::setPinnedMemory(size_t size) {
+  // Should not call this after devices have been initialized
+  FAISS_ASSERT(defaultStreams_.size() == 0);
+  FAISS_ASSERT(!pinnedMemAlloc_);
+
+  pinnedMemSize_ = size;
+}
+
+void
+StandardGpuResources::setDefaultStream(int device, cudaStream_t stream) {
+  auto it = defaultStreams_.find(device);
+  if (it != defaultStreams_.end()) {
+    // Replace this stream with the user stream
+    CUDA_VERIFY(cudaStreamDestroy(it->second));
+    it->second = stream;
+  }
+
+  userDefaultStreams_[device] = stream;
+}
+
+void
+StandardGpuResources::setDefaultNullStreamAllDevices() {
+  for (int dev = 0; dev < getNumDevices(); ++dev) {
+    setDefaultStream(dev, nullptr);
+  }
+}
+
+void
+StandardGpuResources::setCudaMallocWarning(bool b) {
+  cudaMallocWarning_ = b;
+
+  for (auto& v : memory_) {
+    v.second->setCudaMallocWarning(b);
+  }
+}
+
+bool
+StandardGpuResources::isInitialized(int device) const {
+  // Use default streams as a marker for whether or not a certain
+  // device has been initialized
+  return defaultStreams_.count(device) != 0;
+}
+
+void
+StandardGpuResources::initializeForDevice(int device) {
+  if (isInitialized(device)) {
+    return;
+  }
+
+  // If this is the first device that we're initializing, create our
+  // pinned memory allocation
+  if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
+    allocMemorySpace(MemorySpace::HostPinned, &pinnedMemAlloc_, pinnedMemSize_);
+    pinnedMemAllocSize_ = pinnedMemSize_;
+  }
+
+  FAISS_ASSERT(device < getNumDevices());
+  DeviceScope scope(device);
+
+  // Make sure that device properties for all devices are cached
+  auto& prop = getDeviceProperties(device);
+
+  // Also check to make sure we meet our minimum compute capability (3.0)
+  FAISS_ASSERT_FMT(prop.major >= 3,
+                   "Device id %d with CC %d.%d not supported, "
+                   "need 3.0+ compute capability",
+                   device, prop.major, prop.minor);
+
+  // Create streams
+  cudaStream_t defaultStream = 0;
+  auto it = userDefaultStreams_.find(device);
+  if (it != userDefaultStreams_.end()) {
+    // We already have a stream provided by the user
+    defaultStream = it->second;
+  } else {
+    CUDA_VERIFY(cudaStreamCreateWithFlags(&defaultStream,
+                                          cudaStreamNonBlocking));
+  }
+
+  defaultStreams_[device] = defaultStream;
+
+  cudaStream_t asyncCopyStream = 0;
+  CUDA_VERIFY(cudaStreamCreateWithFlags(&asyncCopyStream,
+                                        cudaStreamNonBlocking));
+
+  asyncCopyStreams_[device] = asyncCopyStream;
+
+  std::vector<cudaStream_t> deviceStreams;
+  for (int j = 0; j < kNumStreams; ++j) {
+    cudaStream_t stream = 0;
+    CUDA_VERIFY(cudaStreamCreateWithFlags(&stream,
+                                          cudaStreamNonBlocking));
+
+    deviceStreams.push_back(stream);
+  }
+
+  alternateStreams_[device] = std::move(deviceStreams);
+
+  // Create cuBLAS handle
+  cublasHandle_t blasHandle = 0;
+  auto blasStatus = cublasCreate(&blasHandle);
+  FAISS_ASSERT(blasStatus == CUBLAS_STATUS_SUCCESS);
+  blasHandles_[device] = blasHandle;
+
+  FAISS_ASSERT(memory_.count(device) == 0);
+
+  auto mem = std::unique_ptr<StackDeviceMemory>(
+    new StackDeviceMemory(device,
+                          // adjust for this specific device
+                          getDefaultTempMemForGPU(device, tempMemSize_)));
+  mem->setCudaMallocWarning(cudaMallocWarning_);
+
+  memory_.emplace(device, std::move(mem));
+}
+
+cublasHandle_t
+StandardGpuResources::getBlasHandle(int device) {
+  initializeForDevice(device);
+  return blasHandles_[device];
+}
+
+cudaStream_t
+StandardGpuResources::getDefaultStream(int device) {
+  initializeForDevice(device);
+  return defaultStreams_[device];
+}
+
+std::vector<cudaStream_t>
+StandardGpuResources::getAlternateStreams(int device) {
+  initializeForDevice(device);
+  return alternateStreams_[device];
+}
+
+DeviceMemory& StandardGpuResources::getMemoryManager(int device) {
+  initializeForDevice(device);
+  return *memory_[device];
+}
+
+std::pair<void*, size_t>
+StandardGpuResources::getPinnedMemory() {
+  return std::make_pair(pinnedMemAlloc_, pinnedMemAllocSize_);
+}
+
+cudaStream_t
+StandardGpuResources::getAsyncCopyStream(int device) {
+  initializeForDevice(device);
+  return asyncCopyStreams_[device];
+}
+
+} } // namespace