faiss 0.2.0 → 0.2.1

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

@@ -5,210 +5,204 @@
  * 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/StackDeviceMemory.h>
 #include <faiss/gpu/utils/StaticUtils.h>
 #include <faiss/impl/FaissAssert.h>
 #include <algorithm>
 #include <sstream>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 namespace {
 
 size_t adjustStackSize(size_t sz) {
-  if (sz == 0) {
-    return 0;
-  } else {
-    // ensure that we have at least 16 bytes, as all allocations are bumped up
-    // to 16
-    return utils::roundUp(sz, (size_t) 16);
-  }
+    if (sz == 0) {
+        return 0;
+    } else {
+        // ensure that we have at least 16 bytes, as all allocations are bumped
+        // up to 16
+        return utils::roundUp(sz, (size_t)16);
+    }
 }
 
 } // namespace
 
 StackDeviceMemory::Stack::Stack(GpuResources* res, int d, size_t sz)
-    : res_(res),
-      device_(d),
-      alloc_(nullptr),
-      allocSize_(adjustStackSize(sz)),
-      start_(nullptr),
-      end_(nullptr),
-      head_(nullptr),
-      highWaterMemoryUsed_(0) {
-  if (allocSize_ == 0) {
-    return;
-  }
-
-  DeviceScope s(device_);
-  auto req = AllocRequest(AllocType::TemporaryMemoryBuffer,
-                          device_,
-                          MemorySpace::Device,
-                          res_->getDefaultStream(device_),
-                          allocSize_);
-
-  alloc_ = (char*) res_->allocMemory(req);
-  FAISS_ASSERT_FMT(
-    alloc_,
-    "could not reserve temporary memory region of size %zu", allocSize_);
-
-  // In order to disambiguate between our entire region of temporary memory
-  // versus the first allocation in the temporary memory region, ensure that the
-  // first address returned is +16 bytes from the beginning
-  start_ = alloc_ + 16;
-  head_ = start_;
-  end_ = alloc_ + allocSize_;
+        : res_(res),
+          device_(d),
+          alloc_(nullptr),
+          allocSize_(adjustStackSize(sz)),
+          start_(nullptr),
+          end_(nullptr),
+          head_(nullptr),
+          highWaterMemoryUsed_(0) {
+    if (allocSize_ == 0) {
+        return;
+    }
+
+    DeviceScope s(device_);
+    auto req = AllocRequest(
+            AllocType::TemporaryMemoryBuffer,
+            device_,
+            MemorySpace::Device,
+            res_->getDefaultStream(device_),
+            allocSize_);
+
+    alloc_ = (char*)res_->allocMemory(req);
+    FAISS_ASSERT_FMT(
+            alloc_,
+            "could not reserve temporary memory region of size %zu",
+            allocSize_);
+
+    // In order to disambiguate between our entire region of temporary memory
+    // versus the first allocation in the temporary memory region, ensure that
+    // the first address returned is +16 bytes from the beginning
+    start_ = alloc_ + 16;
+    head_ = start_;
+    end_ = alloc_ + allocSize_;
 }
 
 StackDeviceMemory::Stack::~Stack() {
-  DeviceScope s(device_);
+    DeviceScope s(device_);
 
-  // FIXME: make sure there are no outstanding memory allocations?
-  if (alloc_) {
-    res_->deallocMemory(device_, alloc_);
-  }
+    // FIXME: make sure there are no outstanding memory allocations?
+    if (alloc_) {
+        res_->deallocMemory(device_, alloc_);
+    }
 }
 
-size_t
-StackDeviceMemory::Stack::getSizeAvailable() const {
-  return (end_ - head_);
+size_t StackDeviceMemory::Stack::getSizeAvailable() const {
+    return (end_ - head_);
 }
 
-char*
-StackDeviceMemory::Stack::getAlloc(size_t size,
-                                   cudaStream_t stream) {
-  // The user must check to see that the allocation fit within us
-  auto sizeRemaining = getSizeAvailable();
+char* StackDeviceMemory::Stack::getAlloc(size_t size, cudaStream_t stream) {
+    // The user must check to see that the allocation fit within us
+    auto sizeRemaining = getSizeAvailable();
 
-  FAISS_ASSERT(size <= sizeRemaining);
+    FAISS_ASSERT(size <= sizeRemaining);
 
-  // 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;
+    // 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();
+    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);
+        // Because there is a previous user, we must overlap it
+        FAISS_ASSERT(
+                prevUser.start_ <= endAlloc && prevUser.end_ >= startAlloc);
 
-    if (stream != prevUser.stream_) {
-      // Synchronization required
-      streamWait({stream}, {prevUser.stream_});
-    }
+        if (stream != prevUser.stream_) {
+            // Synchronization required
+            streamWait({stream}, {prevUser.stream_});
+        }
 
-    if (endAlloc < prevUser.end_) {
-      // Update the previous user info
-      prevUser.start_ = endAlloc;
+        if (endAlloc < prevUser.end_) {
+            // Update the previous user info
+            prevUser.start_ = endAlloc;
 
-      break;
-    }
+            break;
+        }
 
-    // If we're the exact size of the previous request, then we
-    // don't need to continue
-    bool done = (prevUser.end_ == endAlloc);
+        // 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();
+        lastUsers_.pop_back();
 
-    if (done) {
-      break;
+        if (done) {
+            break;
+        }
     }
-  }
 
-  head_ = endAlloc;
-  FAISS_ASSERT(head_ <= end_);
+    head_ = endAlloc;
+    FAISS_ASSERT(head_ <= end_);
 
-  highWaterMemoryUsed_ = std::max(highWaterMemoryUsed_,
-                                  (size_t) (head_ - start_));
-  FAISS_ASSERT(startAlloc);
-  return startAlloc;
+    highWaterMemoryUsed_ =
+            std::max(highWaterMemoryUsed_, (size_t)(head_ - start_));
+    FAISS_ASSERT(startAlloc);
+    return startAlloc;
 }
 
-void
-StackDeviceMemory::Stack::returnAlloc(char* p,
-                                      size_t size,
-                                      cudaStream_t stream) {
-  // This allocation should be within ourselves
-  FAISS_ASSERT(p >= start_ && p < end_);
+void StackDeviceMemory::Stack::returnAlloc(
+        char* p,
+        size_t size,
+        cudaStream_t stream) {
+    // This allocation should be within ourselves
+    FAISS_ASSERT(p >= start_ && p < end_);
 
-  // All allocations should have been adjusted to a multiple of 16 bytes
-  FAISS_ASSERT(size % 16 == 0);
+    // All allocations should have been adjusted to a multiple of 16 bytes
+    FAISS_ASSERT(size % 16 == 0);
 
-  // This is on our stack
-  // Allocations should be freed in the reverse order they are made
-  if (p + size != head_) {
-    FAISS_ASSERT(p + size == head_);
-  }
+    // This is on our stack
+    // Allocations should be freed in the reverse order they are made
+    if (p + size != head_) {
+        FAISS_ASSERT(p + size == head_);
+    }
 
-  head_ = p;
-  lastUsers_.push_back(Range(p, p + size, stream));
+    head_ = p;
+    lastUsers_.push_back(Range(p, p + size, stream));
 }
 
-std::string
-StackDeviceMemory::Stack::toString() const {
-  std::stringstream s;
+std::string StackDeviceMemory::Stack::toString() const {
+    std::stringstream s;
 
-  s << "SDM device " << device_ << ": Total memory " << allocSize_ << " ["
-    << (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 << "SDM device " << device_ << ": Total memory " << allocSize_ << " ["
+      << (void*)start_ << ", " << (void*)end_ << ")\n";
+    s << "     Available memory " << (size_t)(end_ - head_) << " ["
+      << (void*)head_ << ", " << (void*)end_ << ")\n";
+    s << "     High water temp alloc " << highWaterMemoryUsed_ << "\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";
-  }
+    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();
+    return s.str();
 }
 
-StackDeviceMemory::StackDeviceMemory(GpuResources* res,
-                                     int device,
-                                     size_t allocPerDevice)
-    : device_(device),
-      stack_(res, device, allocPerDevice) {
-}
+StackDeviceMemory::StackDeviceMemory(
+        GpuResources* res,
+        int device,
+        size_t allocPerDevice)
+        : device_(device), stack_(res, device, allocPerDevice) {}
 
-StackDeviceMemory::~StackDeviceMemory() {
-}
+StackDeviceMemory::~StackDeviceMemory() {}
 
-int
-StackDeviceMemory::getDevice() const {
-  return device_;
+int StackDeviceMemory::getDevice() const {
+    return device_;
 }
 
-size_t
-StackDeviceMemory::getSizeAvailable() const {
-  return stack_.getSizeAvailable();
+size_t StackDeviceMemory::getSizeAvailable() const {
+    return stack_.getSizeAvailable();
 }
 
-std::string
-StackDeviceMemory::toString() const {
-  return stack_.toString();
+std::string StackDeviceMemory::toString() const {
+    return stack_.toString();
 }
 
-void*
-StackDeviceMemory::allocMemory(cudaStream_t stream, size_t size) {
-  // All allocations should have been adjusted to a multiple of 16 bytes
-  FAISS_ASSERT(size % 16 == 0);
-  return stack_.getAlloc(size, stream);
+void* StackDeviceMemory::allocMemory(cudaStream_t stream, size_t size) {
+    // All allocations should have been adjusted to a multiple of 16 bytes
+    FAISS_ASSERT(size % 16 == 0);
+    return stack_.getAlloc(size, stream);
 }
 
-void
-StackDeviceMemory::deallocMemory(int device,
-                                 cudaStream_t stream,
-                                 size_t size,
-                                 void* p) {
-  FAISS_ASSERT(p);
-  FAISS_ASSERT(device == device_);
+void StackDeviceMemory::deallocMemory(
+        int device,
+        cudaStream_t stream,
+        size_t size,
+        void* p) {
+    FAISS_ASSERT(p);
+    FAISS_ASSERT(device == device_);
 
-  stack_.returnAlloc((char*) p, size, stream);
+    stack_.returnAlloc((char*)p, size, stream);
 }
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss

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

@@ -5,110 +5,108 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-
 #pragma once
 
-#include <faiss/gpu/GpuResources.h>
 #include <cuda_runtime.h>
+#include <faiss/gpu/GpuResources.h>
 #include <list>
 #include <memory>
-#include <unordered_map>
 #include <tuple>
+#include <unordered_map>
 
-namespace faiss { namespace gpu {
+namespace faiss {
+namespace gpu {
 
 /// Device memory manager that provides temporary memory allocations
 /// out of a region of memory, for a single device
 class StackDeviceMemory {
- public:
-  /// Allocate a new region of memory that we manage
-  StackDeviceMemory(GpuResources* res,
-                    int device,
-                    size_t allocPerDevice);
+   public:
+    /// Allocate a new region of memory that we manage
+    StackDeviceMemory(GpuResources* res, 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);
+    /// Manage a region of memory for a particular device, with or
+    /// without ownership
+    StackDeviceMemory(int device, void* p, size_t size, bool isOwner);
 
-  ~StackDeviceMemory();
+    ~StackDeviceMemory();
 
-  int getDevice() const;
+    int getDevice() const;
 
-  /// All allocations requested should be a multiple of 16 bytes
-  void* allocMemory(cudaStream_t stream, size_t size);
-  void deallocMemory(int device, cudaStream_t, size_t size, void* p);
+    /// All allocations requested should be a multiple of 16 bytes
+    void* allocMemory(cudaStream_t stream, size_t size);
+    void deallocMemory(int device, cudaStream_t, size_t size, void* p);
 
-  size_t getSizeAvailable() const;
-  std::string toString() const;
+    size_t getSizeAvailable() const;
+    std::string toString() const;
 
- 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) {
-    }
+   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_;
-  };
+        // References a memory range [start, end)
+        char* start_;
+        char* end_;
+        cudaStream_t stream_;
+    };
 
-  struct Stack {
-    /// Constructor that allocates memory via cudaMalloc
-    Stack(GpuResources* res, int device, size_t size);
+    struct Stack {
+        /// Constructor that allocates memory via cudaMalloc
+        Stack(GpuResources* res, int device, size_t size);
 
-    ~Stack();
+        ~Stack();
 
-    /// Returns how much size is available for an allocation without
-    /// calling cudaMalloc
-    size_t getSizeAvailable() const;
+        /// 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);
+        /// 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 an allocation
+        void returnAlloc(char* p, size_t size, cudaStream_t stream);
 
-    /// Returns the stack state
-    std::string toString() const;
+        /// Returns the stack state
+        std::string toString() const;
 
-    /// Our GpuResources object
-    GpuResources* res_;
+        /// Our GpuResources object
+        GpuResources* res_;
 
-    /// Device this allocation is on
-    int device_;
+        /// Device this allocation is on
+        int device_;
 
-    /// Where our temporary memory buffer is allocated; we allocate starting 16
-    /// bytes into this
-    char* alloc_;
+        /// Where our temporary memory buffer is allocated; we allocate starting
+        /// 16 bytes into this
+        char* alloc_;
 
-    /// Total size of our allocation
-    size_t allocSize_;
+        /// Total size of our allocation
+        size_t allocSize_;
 
-    /// Our temporary memory region; [start_, end_) is valid
-    char* start_;
-    char* end_;
+        /// Our temporary memory region; [start_, end_) is valid
+        char* start_;
+        char* end_;
 
-    /// Stack head within [start, end)
-    char* head_;
+        /// 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_;
+        /// List of previous last users of allocations on our stack, for
+        /// possible synchronization purposes
+        std::list<Range> lastUsers_;
 
-    /// 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 used from the
+        /// temporary buffer?
+        size_t highWaterMemoryUsed_;
+    };
 
-  /// Our device
-  int device_;
+    /// Our device
+    int device_;
 
-  /// Memory stack
-  Stack stack_;
+    /// Memory stack
+    Stack stack_;
 };
 
-} } // namespace
+} // namespace gpu
+} // namespace faiss