faiss 0.1.2 → 0.1.3

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

@@ -1,71 +0,0 @@
-/**
- * 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/MemorySpace.cpp

@@ -1,89 +0,0 @@
-/**
- * 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

@@ -1,44 +0,0 @@
-/**
- * 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);
-
-} }

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

@@ -1,239 +0,0 @@
-/**
- * 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/StackDeviceMemory.h>
-#include <faiss/gpu/utils/DeviceUtils.h>
-#include <faiss/gpu/utils/MemorySpace.h>
-#include <faiss/gpu/utils/StaticUtils.h>
-#include <faiss/impl/FaissAssert.h>
-#include <stdio.h>
-#include <sstream>
-
-namespace faiss { namespace gpu {
-
-StackDeviceMemory::Stack::Stack(int d, size_t sz)
-    : device_(d),
-      isOwner_(true),
-      start_(nullptr),
-      end_(nullptr),
-      size_(sz),
-      head_(nullptr),
-      mallocCurrent_(0),
-      highWaterMemoryUsed_(0),
-      highWaterMalloc_(0),
-      cudaMallocWarning_(true) {
-  DeviceScope s(device_);
-
-  allocMemorySpace(MemorySpace::Device, &start_, size_);
-
-  head_ = start_;
-  end_ = start_ + size_;
-}
-
-StackDeviceMemory::Stack::Stack(int d, void* p, size_t sz, bool isOwner)
-    : device_(d),
-      isOwner_(isOwner),
-      start_((char*) p),
-      end_(((char*) p) + sz),
-      size_(sz),
-      head_((char*) p),
-      mallocCurrent_(0),
-      highWaterMemoryUsed_(0),
-      highWaterMalloc_(0),
-      cudaMallocWarning_(true) {
-}
-
-StackDeviceMemory::Stack::~Stack() {
-  if (isOwner_) {
-    DeviceScope s(device_);
-
-    freeMemorySpace(MemorySpace::Device, start_);
-  }
-}
-
-size_t
-StackDeviceMemory::Stack::getSizeAvailable() const {
-  return (end_ - head_);
-}
-
-char*
-StackDeviceMemory::Stack::getAlloc(size_t size,
-                                   cudaStream_t stream) {
-  if (size > (end_ - head_)) {
-    // Too large for our stack
-    DeviceScope s(device_);
-
-    if (cudaMallocWarning_) {
-      // Print our requested size before we attempt the allocation
-      fprintf(stderr, "WARN: increase temp memory to avoid cudaMalloc, "
-              "or decrease query/add size (alloc %zu B, highwater %zu B)\n",
-              size, highWaterMalloc_);
-    }
-
-    char* p = nullptr;
-    allocMemorySpace(MemorySpace::Device, &p, size);
-
-    mallocCurrent_ += size;
-    highWaterMalloc_ = std::max(highWaterMalloc_, mallocCurrent_);
-
-    return p;
-  } else {
-    // We can make the allocation out of our stack
-    // Find all the ranges that we overlap that may have been
-    // previously allocated; our allocation will be [head, endAlloc)
-    char* startAlloc = head_;
-    char* endAlloc = head_ + size;
-
-    while (lastUsers_.size() > 0) {
-      auto& prevUser = lastUsers_.back();
-
-      // Because there is a previous user, we must overlap it
-      FAISS_ASSERT(prevUser.start_ <= endAlloc && prevUser.end_ >= startAlloc);
-
-      if (stream != prevUser.stream_) {
-        // Synchronization required
-        // FIXME
-        FAISS_ASSERT(false);
-      }
-
-      if (endAlloc < prevUser.end_) {
-        // Update the previous user info
-        prevUser.start_ = endAlloc;
-
-        break;
-      }
-
-      // If we're the exact size of the previous request, then we
-      // don't need to continue
-      bool done = (prevUser.end_ == endAlloc);
-
-      lastUsers_.pop_back();
-
-      if (done) {
-        break;
-      }
-    }
-
-    head_ = endAlloc;
-    FAISS_ASSERT(head_ <= end_);
-
-    highWaterMemoryUsed_ = std::max(highWaterMemoryUsed_,
-                                    (size_t) (head_ - start_));
-    return startAlloc;
-  }
-}
-
-void
-StackDeviceMemory::Stack::returnAlloc(char* p,
-                                      size_t size,
-                                      cudaStream_t stream) {
-  if (p < start_ || p >= end_) {
-    // This is not on our stack; it was a one-off allocation
-    DeviceScope s(device_);
-
-    freeMemorySpace(MemorySpace::Device, p);
-
-    FAISS_ASSERT(mallocCurrent_ >= size);
-    mallocCurrent_ -= size;
-  } else {
-    // This is on our stack
-    // Allocations should be freed in the reverse order they are made
-    FAISS_ASSERT(p + size == head_);
-
-    head_ = p;
-    lastUsers_.push_back(Range(p, p + size, stream));
-  }
-}
-
-std::string
-StackDeviceMemory::Stack::toString() const {
-  std::stringstream s;
-
-  s << "SDM device " << device_ << ": Total memory " << size_ << " ["
-    << (void*) start_ << ", " << (void*) end_ << ")\n";
-  s << "     Available memory " << (size_t) (end_ - head_)
-    << " [" << (void*) head_ << ", " << (void*) end_ << ")\n";
-  s << "     High water temp alloc " << highWaterMemoryUsed_ << "\n";
-  s << "     High water cudaMalloc " << highWaterMalloc_ << "\n";
-
-  int i = lastUsers_.size();
-  for (auto it = lastUsers_.rbegin(); it != lastUsers_.rend(); ++it) {
-    s << i-- << ": size " << (size_t) (it->end_ - it->start_)
-      << " stream " << it->stream_
-      << " [" << (void*) it->start_ << ", " << (void*) it->end_ << ")\n";
-  }
-
-  return s.str();
-}
-
-size_t
-StackDeviceMemory::Stack::getHighWaterCudaMalloc() const {
-  return highWaterMalloc_;
-}
-
-StackDeviceMemory::StackDeviceMemory(int device, size_t allocPerDevice)
-    : device_(device),
-      stack_(device, allocPerDevice) {
-}
-
-StackDeviceMemory::StackDeviceMemory(int device,
-                                     void* p, size_t size, bool isOwner)
-    : device_(device),
-      stack_(device, p, size, isOwner) {
-}
-
-StackDeviceMemory::~StackDeviceMemory() {
-}
-
-void
-StackDeviceMemory::setCudaMallocWarning(bool b) {
-  stack_.cudaMallocWarning_ = b;
-}
-
-int
-StackDeviceMemory::getDevice() const {
-  return device_;
-}
-
-DeviceMemoryReservation
-StackDeviceMemory::getMemory(cudaStream_t stream, size_t size) {
-  // We guarantee 16 byte alignment for allocations, so bump up `size`
-  // to the next highest multiple of 16
-  size = utils::roundUp(size, (size_t) 16);
-
-  return DeviceMemoryReservation(this,
-                                 device_,
-                                 stack_.getAlloc(size, stream),
-                                 size,
-                                 stream);
-}
-
-size_t
-StackDeviceMemory::getSizeAvailable() const {
-  return stack_.getSizeAvailable();
-}
-
-std::string
-StackDeviceMemory::toString() const {
-  return stack_.toString();
-}
-
-size_t
-StackDeviceMemory::getHighWaterCudaMalloc() const {
-  return stack_.getHighWaterCudaMalloc();
-}
-
-void
-StackDeviceMemory::returnAllocation(DeviceMemoryReservation& m) {
-  FAISS_ASSERT(m.get());
-  FAISS_ASSERT(device_ == m.device());
-
-  stack_.returnAlloc((char*) m.get(), m.size(), m.stream());
-}
-
-} } // namespace