faiss 0.1.0 → 0.1.1

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

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

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

@@ -0,0 +1,129 @@
+/**
+ * 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 <list>
+#include <memory>
+#include <unordered_map>
+
+namespace faiss { namespace gpu {
+
+/// Device memory manager that provides temporary memory allocations
+/// out of a region of memory
+class StackDeviceMemory : public DeviceMemory {
+ public:
+  /// Allocate a new region of memory that we manage
+  explicit StackDeviceMemory(int device, size_t allocPerDevice);
+
+  /// Manage a region of memory for a particular device, with or
+  /// without ownership
+  StackDeviceMemory(int device, void* p, size_t size, bool isOwner);
+
+  ~StackDeviceMemory() override;
+
+  /// Enable or disable the warning about not having enough temporary memory
+  /// when cudaMalloc gets called
+  void setCudaMallocWarning(bool b);
+
+  int getDevice() const override;
+
+  DeviceMemoryReservation getMemory(cudaStream_t stream,
+                                    size_t size) override;
+
+  size_t getSizeAvailable() const override;
+  std::string toString() const override;
+  size_t getHighWaterCudaMalloc() const override;
+
+ protected:
+  void returnAllocation(DeviceMemoryReservation& m) override;
+
+ protected:
+  /// Previous allocation ranges and the streams for which
+  /// synchronization is required
+  struct Range {
+    inline Range(char* s, char* e, cudaStream_t str) :
+        start_(s), end_(e), stream_(str) {
+    }
+
+    // References a memory range [start, end)
+    char* start_;
+    char* end_;
+    cudaStream_t stream_;
+  };
+
+  struct Stack {
+    /// Constructor that allocates memory via cudaMalloc
+    Stack(int device, size_t size);
+
+    /// Constructor that references a pre-allocated region of memory
+    Stack(int device, void* p, size_t size, bool isOwner);
+    ~Stack();
+
+    /// Returns how much size is available for an allocation without
+    /// calling cudaMalloc
+    size_t getSizeAvailable() const;
+
+    /// Obtains an allocation; all allocations are guaranteed to be 16
+    /// byte aligned
+    char* getAlloc(size_t size, cudaStream_t stream);
+
+    /// Returns an allocation
+    void returnAlloc(char* p, size_t size, cudaStream_t stream);
+
+    /// Returns the stack state
+    std::string toString() const;
+
+    /// Returns the high-water mark of cudaMalloc activity
+    size_t getHighWaterCudaMalloc() const;
+
+    /// Device this allocation is on
+    int device_;
+
+    /// Do we own our region of memory?
+    bool isOwner_;
+
+    /// Where our allocation begins and ends
+    /// [start_, end_) is valid
+    char* start_;
+    char* end_;
+
+    /// Total size end_ - start_
+    size_t size_;
+
+    /// Stack head within [start, end)
+    char* head_;
+
+    /// List of previous last users of allocations on our stack, for
+    /// possible synchronization purposes
+    std::list<Range> lastUsers_;
+
+    /// How much cudaMalloc memory is currently outstanding?
+    size_t mallocCurrent_;
+
+    /// What's the high water mark in terms of memory used from the
+    /// temporary buffer?
+    size_t highWaterMemoryUsed_;
+
+    /// What's the high water mark in terms of memory allocated via
+    /// cudaMalloc?
+    size_t highWaterMalloc_;
+
+    /// Whether or not a warning upon cudaMalloc is generated
+    bool cudaMallocWarning_;
+  };
+
+  /// Our device
+  int device_;
+
+  /// Memory stack
+  Stack stack_;
+};
+
+} } // namespace

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

@@ -0,0 +1,83 @@
+/**
+ * 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>
+
+namespace faiss { namespace gpu { namespace utils {
+
+template <typename U, typename V>
+constexpr __host__ __device__ auto divDown(U a, V b) -> decltype(a + b) {
+  return (a / b);
+}
+
+template <typename U, typename V>
+constexpr __host__ __device__ auto divUp(U a, V b) -> decltype(a + b) {
+  return (a + b - 1) / b;
+}
+
+template <typename U, typename V>
+constexpr __host__ __device__ auto roundDown(U a, V b) -> decltype(a + b) {
+  return divDown(a, b) * b;
+}
+
+template <typename U, typename V>
+constexpr __host__ __device__ auto roundUp(U a, V b) -> decltype(a + b) {
+  return divUp(a, b) * b;
+}
+
+template <class T>
+constexpr __host__ __device__ T pow(T n, T power) {
+  return (power > 0 ? n * pow(n, power - 1) : 1);
+}
+
+template <class T>
+constexpr __host__ __device__ T pow2(T n) {
+  return pow(2, (T) n);
+}
+
+static_assert(pow2(8) == 256, "pow2");
+
+template <typename T>
+constexpr __host__ __device__ int log2(T n, int p = 0) {
+  return (n <= 1) ? p : log2(n / 2, p + 1);
+}
+
+static_assert(log2(2) == 1, "log2");
+static_assert(log2(3) == 1, "log2");
+static_assert(log2(4) == 2, "log2");
+
+template <typename T>
+constexpr __host__ __device__ bool isPowerOf2(T v) {
+  return (v && !(v & (v - 1)));
+}
+
+static_assert(isPowerOf2(2048), "isPowerOf2");
+static_assert(!isPowerOf2(3333), "isPowerOf2");
+
+template <typename T>
+constexpr __host__ __device__ T nextHighestPowerOf2(T v) {
+  return (isPowerOf2(v) ? (T) 2 * v : ((T) 1 << (log2(v) + 1)));
+}
+
+static_assert(nextHighestPowerOf2(1) == 2, "nextHighestPowerOf2");
+static_assert(nextHighestPowerOf2(2) == 4, "nextHighestPowerOf2");
+static_assert(nextHighestPowerOf2(3) == 4, "nextHighestPowerOf2");
+static_assert(nextHighestPowerOf2(4) == 8, "nextHighestPowerOf2");
+
+static_assert(nextHighestPowerOf2(15) == 16, "nextHighestPowerOf2");
+static_assert(nextHighestPowerOf2(16) == 32, "nextHighestPowerOf2");
+static_assert(nextHighestPowerOf2(17) == 32, "nextHighestPowerOf2");
+
+static_assert(nextHighestPowerOf2(1536000000u) == 2147483648u,
+              "nextHighestPowerOf2");
+static_assert(nextHighestPowerOf2((size_t) 2147483648ULL) ==
+              (size_t) 4294967296ULL, "nextHighestPowerOf2");
+
+} } } // namespace