faiss 0.1.0 → 0.1.1

data/vendor/faiss/gpu/utils/DeviceMemory.cpp

@@ -0,0 +1,77 @@
+/**
+ * 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/utils/DeviceMemory.h>
+#include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss { namespace gpu {
+
+DeviceMemoryReservation::DeviceMemoryReservation()
+    : state_(NULL),
+      device_(0),
+      data_(NULL),
+      size_(0),
+      stream_(0) {
+}
+
+DeviceMemoryReservation::DeviceMemoryReservation(DeviceMemory* state,
+                                             int device,
+                                             void* p,
+                                             size_t size,
+                                             cudaStream_t stream)
+    : state_(state),
+      device_(device),
+      data_(p),
+      size_(size),
+      stream_(stream) {
+}
+
+DeviceMemoryReservation::DeviceMemoryReservation(
+  DeviceMemoryReservation&& m) noexcept {
+
+  state_ = m.state_;
+  device_ = m.device_;
+  data_ = m.data_;
+  size_ = m.size_;
+  stream_ = m.stream_;
+
+  m.data_ = NULL;
+}
+
+DeviceMemoryReservation::~DeviceMemoryReservation() {
+  if (data_) {
+    FAISS_ASSERT(state_);
+    state_->returnAllocation(*this);
+  }
+
+  data_ = NULL;
+}
+
+DeviceMemoryReservation&
+DeviceMemoryReservation::operator=(DeviceMemoryReservation&& m) {
+  if (data_) {
+    FAISS_ASSERT(state_);
+    state_->returnAllocation(*this);
+  }
+
+  state_ = m.state_;
+  device_ = m.device_;
+  data_ = m.data_;
+  size_ = m.size_;
+  stream_ = m.stream_;
+
+  m.data_ = NULL;
+
+  return *this;
+}
+
+DeviceMemory::~DeviceMemory() {
+}
+
+} } // namespace

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

@@ -0,0 +1,71 @@
+/**
+ * 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 <cuda_runtime.h>
+#include <string>
+
+namespace faiss { namespace gpu {
+
+class DeviceMemory;
+
+class DeviceMemoryReservation {
+ public:
+  DeviceMemoryReservation();
+  DeviceMemoryReservation(DeviceMemory* state,
+                          int device, void* p, size_t size,
+                          cudaStream_t stream);
+  DeviceMemoryReservation(DeviceMemoryReservation&& m) noexcept;
+  ~DeviceMemoryReservation();
+
+  DeviceMemoryReservation& operator=(DeviceMemoryReservation&& m);
+
+  int device() { return device_; }
+  void* get() { return data_; }
+  size_t size() { return size_; }
+  cudaStream_t stream() { return stream_; }
+
+ private:
+  DeviceMemory* state_;
+
+  int device_;
+  void* data_;
+  size_t size_;
+  cudaStream_t stream_;
+};
+
+/// Manages temporary memory allocations on a GPU device
+class DeviceMemory {
+ public:
+  virtual ~DeviceMemory();
+
+  /// Returns the device we are managing memory for
+  virtual int getDevice() const = 0;
+
+  /// Obtains a temporary memory allocation for our device,
+  /// whose usage is ordered with respect to the given stream.
+  virtual DeviceMemoryReservation getMemory(cudaStream_t stream,
+                                            size_t size) = 0;
+
+  /// Returns the current size available without calling cudaMalloc
+  virtual size_t getSizeAvailable() const = 0;
+
+  /// Returns a string containing our current memory manager state
+  virtual std::string toString() const = 0;
+
+  /// Returns the high-water mark of cudaMalloc allocations for our
+  /// device
+  virtual size_t getHighWaterCudaMalloc() const = 0;
+
+ protected:
+  friend class DeviceMemoryReservation;
+  virtual void returnAllocation(DeviceMemoryReservation& m) = 0;
+};
+
+} } // namespace

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

@@ -0,0 +1,185 @@
+/**
+ * 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/impl/FaissAssert.h>
+#include <cuda_runtime.h>
+#include <cublas_v2.h>
+#include <vector>
+
+namespace faiss { namespace gpu {
+
+/// Returns the current thread-local GPU device
+int getCurrentDevice();
+
+/// Sets the current thread-local GPU device
+void setCurrentDevice(int device);
+
+/// Returns the number of available GPU devices
+int getNumDevices();
+
+/// Starts the CUDA profiler (exposed via SWIG)
+void profilerStart();
+
+/// Stops the CUDA profiler (exposed via SWIG)
+void profilerStop();
+
+/// Synchronizes the CPU against all devices (equivalent to
+/// cudaDeviceSynchronize for each device)
+void synchronizeAllDevices();
+
+/// Returns a cached cudaDeviceProp for the given device
+const cudaDeviceProp& getDeviceProperties(int device);
+
+/// Returns the cached cudaDeviceProp for the current device
+const cudaDeviceProp& getCurrentDeviceProperties();
+
+/// Returns the maximum number of threads available for the given GPU
+/// device
+int getMaxThreads(int device);
+
+/// Equivalent to getMaxThreads(getCurrentDevice())
+int getMaxThreadsCurrentDevice();
+
+/// Returns the maximum smem available for the given GPU device
+size_t getMaxSharedMemPerBlock(int device);
+
+/// Equivalent to getMaxSharedMemPerBlock(getCurrentDevice())
+size_t getMaxSharedMemPerBlockCurrentDevice();
+
+/// For a given pointer, returns whether or not it is located on
+/// a device (deviceId >= 0) or the host (-1).
+int getDeviceForAddress(const void* p);
+
+/// Does the given device support full unified memory sharing host
+/// memory?
+bool getFullUnifiedMemSupport(int device);
+
+/// Equivalent to getFullUnifiedMemSupport(getCurrentDevice())
+bool getFullUnifiedMemSupportCurrentDevice();
+
+/// 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
+int getMaxKSelection();
+
+/// RAII object to set the current device, and restore the previous
+/// device upon destruction
+class DeviceScope {
+ public:
+  explicit DeviceScope(int device);
+  ~DeviceScope();
+
+ private:
+  int prevDevice_;
+};
+
+/// RAII object to manage a cublasHandle_t
+class CublasHandleScope {
+ public:
+  CublasHandleScope();
+  ~CublasHandleScope();
+
+  cublasHandle_t get() { return blasHandle_; }
+
+ private:
+  cublasHandle_t blasHandle_;
+};
+
+// RAII object to manage a cudaEvent_t
+class CudaEvent {
+ public:
+  /// Creates an event and records it in this stream
+  explicit CudaEvent(cudaStream_t stream);
+  CudaEvent(const CudaEvent& event) = delete;
+  CudaEvent(CudaEvent&& event) noexcept;
+  ~CudaEvent();
+
+  inline cudaEvent_t get() { return event_; }
+
+  /// Wait on this event in this stream
+  void streamWaitOnEvent(cudaStream_t stream);
+
+  /// Have the CPU wait for the completion of this event
+  void cpuWaitOnEvent();
+
+  CudaEvent& operator=(CudaEvent&& event) noexcept;
+  CudaEvent& operator=(CudaEvent& event) = delete;
+
+ private:
+  cudaEvent_t event_;
+};
+
+/// Wrapper to test return status of CUDA functions
+#define CUDA_VERIFY(X)                                                  \
+  do {                                                                  \
+    auto err__ = (X);                                                   \
+    FAISS_ASSERT_FMT(err__ == cudaSuccess, "CUDA error %d %s",          \
+                     (int) err__, cudaGetErrorString(err__));           \
+  } while (0)
+
+/// Wrapper to synchronously probe for CUDA errors
+// #define FAISS_GPU_SYNC_ERROR 1
+
+#ifdef FAISS_GPU_SYNC_ERROR
+#define CUDA_TEST_ERROR()                       \
+  do {                                          \
+    CUDA_VERIFY(cudaDeviceSynchronize());       \
+  } while (0)
+#else
+#define CUDA_TEST_ERROR()                       \
+  do {                                          \
+    CUDA_VERIFY(cudaGetLastError());            \
+  } while (0)
+#endif
+
+/// Call for a collection of streams to wait on
+template <typename L1, typename L2>
+void streamWaitBase(const L1& listWaiting, const L2& listWaitOn) {
+  // For all the streams we are waiting on, create an event
+  std::vector<cudaEvent_t> events;
+  for (auto& stream : listWaitOn) {
+    cudaEvent_t event;
+    CUDA_VERIFY(cudaEventCreateWithFlags(&event, cudaEventDisableTiming));
+    CUDA_VERIFY(cudaEventRecord(event, stream));
+    events.push_back(event);
+  }
+
+  // For all the streams that are waiting, issue a wait
+  for (auto& stream : listWaiting) {
+    for (auto& event : events) {
+      CUDA_VERIFY(cudaStreamWaitEvent(stream, event, 0));
+    }
+  }
+
+  for (auto& event : events) {
+    CUDA_VERIFY(cudaEventDestroy(event));
+  }
+}
+
+/// These versions allow usage of initializer_list as arguments, since
+/// otherwise {...} doesn't have a type
+template <typename L1>
+void streamWait(const L1& a,
+                const std::initializer_list<cudaStream_t>& b) {
+  streamWaitBase(a, b);
+}
+
+template <typename L2>
+void streamWait(const std::initializer_list<cudaStream_t>& a,
+                const L2& b) {
+  streamWaitBase(a, b);
+}
+
+inline void streamWait(const std::initializer_list<cudaStream_t>& a,
+                       const std::initializer_list<cudaStream_t>& b) {
+  streamWaitBase(a, b);
+}
+
+} } // namespace

data/vendor/faiss/gpu/utils/MemorySpace.cpp

@@ -0,0 +1,89 @@
+/**
+ * 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/utils/MemorySpace.h>
+#include <faiss/impl/FaissAssert.h>
+#include <cuda_runtime.h>
+
+namespace faiss { namespace gpu {
+
+/// Allocates CUDA memory for a given memory space
+void allocMemorySpaceV(MemorySpace space, void** p, size_t size) {
+  switch (space) {
+    case MemorySpace::Device:
+    {
+      auto err = cudaMalloc(p, size);
+
+      // Throw if we fail to allocate
+      FAISS_THROW_IF_NOT_FMT(
+        err == cudaSuccess,
+        "failed to cudaMalloc %zu bytes (error %d %s)",
+        size, (int) err, cudaGetErrorString(err));
+    }
+    break;
+    case MemorySpace::Unified:
+    {
+#ifdef FAISS_UNIFIED_MEM
+      auto err = cudaMallocManaged(p, size);
+
+      // Throw if we fail to allocate
+      FAISS_THROW_IF_NOT_FMT(
+        err == cudaSuccess,
+        "failed to cudaMallocManaged %zu bytes (error %d %s)",
+        size, (int) err, cudaGetErrorString(err));
+#else
+      FAISS_THROW_MSG("Attempting to allocate via cudaMallocManaged "
+                      "without CUDA 8+ support");
+#endif
+    }
+    break;
+    case MemorySpace::HostPinned:
+    {
+      auto err = cudaHostAlloc(p, size, cudaHostAllocDefault);
+
+      // Throw if we fail to allocate
+      FAISS_THROW_IF_NOT_FMT(
+        err == cudaSuccess,
+        "failed to cudaHostAlloc %zu bytes (error %d %s)",
+        size, (int) err, cudaGetErrorString(err));
+    }
+    break;
+    default:
+      FAISS_ASSERT_FMT(false, "unknown MemorySpace %d", (int) space);
+      break;
+  }
+}
+
+// We'll allow allocation to fail, but free should always succeed and be a
+// fatal error if it doesn't free
+void freeMemorySpace(MemorySpace space, void* p) {
+  switch (space) {
+    case MemorySpace::Device:
+    case MemorySpace::Unified:
+    {
+      auto err = cudaFree(p);
+      FAISS_ASSERT_FMT(err == cudaSuccess,
+                       "Failed to cudaFree pointer %p (error %d %s)",
+                       p, (int) err, cudaGetErrorString(err));
+    }
+    break;
+    case MemorySpace::HostPinned:
+    {
+      auto err = cudaFreeHost(p);
+      FAISS_ASSERT_FMT(err == cudaSuccess,
+                       "Failed to cudaFreeHost pointer %p (error %d %s)",
+                       p, (int) err, cudaGetErrorString(err));
+    }
+    break;
+    default:
+      FAISS_ASSERT_FMT(false, "unknown MemorySpace %d", (int) space);
+      break;
+  }
+}
+
+} }

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

@@ -0,0 +1,44 @@
+/**
+ * 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 <cuda.h>
+
+#if CUDA_VERSION >= 8000
+// Whether or not we enable usage of CUDA Unified Memory
+#define FAISS_UNIFIED_MEM 1
+#endif
+
+namespace faiss { namespace gpu {
+
+enum MemorySpace {
+  /// Managed using cudaMalloc/cudaFree
+  Device = 1,
+  /// Managed using cudaMallocManaged/cudaFree
+  Unified = 2,
+  /// Managed using cudaHostAlloc/cudaFreeHost
+  HostPinned = 3,
+};
+
+/// All memory allocations and de-allocations come through these functions
+
+/// Allocates CUDA memory for a given memory space (void pointer)
+/// Throws a FaissException if we are unable to allocate the memory
+void allocMemorySpaceV(MemorySpace space, void** p, size_t size);
+
+template <typename T>
+inline void allocMemorySpace(MemorySpace space, T** p, size_t size) {
+  allocMemorySpaceV(space, (void**)(void*) p, size);
+}
+
+/// Frees CUDA memory for a given memory space
+/// Asserts if we are unable to free the region
+void freeMemorySpace(MemorySpace space, void* p);
+
+} }