ENH: Add paddle_memory for external usage

paddle/memory/CMakeLists.txt

 add_subdirectory(detail)
+
+cc_library(memory
+    SRCS
+    memory.cc)
+
+cc_library(paddle_memory
+    DEPS
+    memory meta_data
+    meta_cache memory_block
+    buddy_allocator system_allocator)

paddle/memory/detail/CMakeLists.txt

@@ -3,7 +3,13 @@ if(${WITH_GPU})
 else(${WITH_GPU})
   cc_library(system_allocator SRCS system_allocator.cc DEPS gflags)
 endif(${WITH_GPU})
+
 cc_test(system_allocator_test SRCS system_allocator_test.cc DEPS system_allocator)
 
 cc_library(meta_data SRCS meta_data.cc)
+
+cc_library(meta_cache SRCS meta_cache.cc)
+
+cc_library(memory_block SRCS memory_block.cc)
+
 cc_library(buddy_allocator SRCS buddy_allocator.cc)

paddle/memory/detail/buddy_allocator.cc

@@ -20,14 +20,14 @@ namespace memory {
 namespace detail {
 
 BuddyAllocator::BuddyAllocator(SystemAllocator* system_allocator,
-                               size_t min_chunk_size, size_t max_chunk_size) {
+                               size_t min_chunk_size, size_t max_chunk_size)
+    : min_chunk_size_(min_chunk_size),
+      max_chunk_size_(max_chunk_size),
+      cache_(system_allocator->UseGpu()),
+      system_allocator_(std::move(system_allocator)) {
   PADDLE_ASSERT(min_chunk_size > 0);
   PADDLE_ASSERT(max_chunk_size > 0);
   PADDLE_ASSERT(system_allocator != nullptr);
-
-  system_allocator_ = std::move(system_allocator);
-  min_chunk_size_ = min_chunk_size;
-  max_chunk_size_ = max_chunk_size;
 }
 
 inline size_t align(size_t size, size_t alignment) {
@@ -90,7 +90,7 @@ void BuddyAllocator::Free(void* p) {
 
     // Invalidate GPU allocation from cache
     if (system_allocator_->UseGpu()) {
-      cache_.erase(block);
+      cache_.invalidate(block);
     }
     return;
   }

paddle/memory/detail/buddy_allocator.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/memory/detail/meta_cache.h"
 #include "paddle/memory/detail/meta_data.h"
 #include "paddle/memory/detail/system_allocator.h"
 #include "paddle/platform/assert.h"
@@ -80,7 +81,6 @@ class BuddyAllocator {
 
  private:
   // Unify the metadata format between GPU and CPU allocations
-  using MetadataCache = std::unordered_map<const MemoryBlock*, Metadata>;
   MetadataCache cache_;
 
  private:

paddle/memory/detail/memory_block.cc

+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License. */
+
 #include "paddle/memory/detail/memory_block.h"
+#include "paddle/memory/detail/meta_cache.h"
+#include "paddle/memory/detail/meta_data.h"
 #include "paddle/platform/assert.h"
 
 namespace paddle {
@@ -7,8 +23,7 @@ namespace detail {
 
 void MemoryBlock::init(MetadataCache& cache, Type t, size_t index, size_t size,
                        void* left_buddy, void* right_buddy) {
-  cache.store(this,
-              MemoryBlockMetadata(t, index, size - overhead(), size,
+  cache.store(this, Metadata(t, index, size - sizeof(Metadata), size,
                              static_cast<MemoryBlock*>(left_buddy),
                              static_cast<MemoryBlock*>(right_buddy)));
 }
@@ -35,10 +50,10 @@ MemoryBlock* MemoryBlock::right_buddy(MetadataCache& cache) const {
 
 void MemoryBlock::split(MetadataCache& cache, size_t size) {
   // make sure the split fits
-  assert(total_size(cache) >= size);
+  PADDLE_ASSERT(total_size(cache) >= size);
 
   // bail out if there is no room for another partition
-  if (total_size(cache) - size <= overhead()) {
+  if (total_size(cache) - size <= sizeof(Metadata)) {
     return;
   }
 
@@ -53,13 +68,13 @@ void MemoryBlock::split(MetadataCache& cache, size_t size) {
   // Write the metadata for the new block
   auto new_block_right_buddy = metadata.right_buddy;
 
-  cache.store(static_cast<MemoryBlock*>(right_partition),
-              MemoryBlockMetadata(FREE_MEMORY, index(cache),
-                                  remaining_size - overhead(), remaining_size,
-                                  this, new_block_right_buddy));
+  cache.store(
+      static_cast<MemoryBlock*>(right_partition),
+      Metadata(FREE_CHUNK, index(cache), remaining_size - sizeof(Metadata),
+               remaining_size, this, new_block_right_buddy));
 
   metadata.right_buddy = static_cast<MemoryBlock*>(right_partition);
-  metadata.size = size - overhead();
+  metadata.size = size - sizeof(Metadata);
   metadata.total_size = size;
 
   cache.store(this, metadata);
@@ -76,8 +91,8 @@ void MemoryBlock::split(MetadataCache& cache, size_t size) {
 
 void MemoryBlock::merge(MetadataCache& cache, MemoryBlock* right_buddy) {
   // only free blocks can be merged
-  assert(type(cache) == FREE_MEMORY);
-  assert(right_buddy->type(cache) == FREE_MEMORY);
+  PADDLE_ASSERT(type(cache) == FREE_MEMORY);
+  PADDLE_ASSERT(right_buddy->type(cache) == FREE_MEMORY);
 
   auto metadata = cache.load(this);
 
@@ -97,16 +112,15 @@ void MemoryBlock::merge(MetadataCache& cache, MemoryBlock* right_buddy) {
   metadata.total_size += right_buddy->total_size(cache);
 
   cache.store(this, metadata);
-  cache.store(right_buddy,
-              MemoryBlockMetadata(INVALID_MEMORY, 0, 0, 0, nullptr, nullptr));
+  cache.store(right_buddy, Metadata(INVALID_CHUNK, 0, 0, 0, nullptr, nullptr));
 }
 
 void MemoryBlock::mark_as_free(MetadataCache& cache) {
   // check for double free or corruption
-  assert(type(cache) != FREE_MEMORY);
-  assert(type(cache) != INVALID_MEMORY);
+  PADDLE_ASSERT(type(cache) != FREE_CHUNK);
+  PADDLE_ASSERT(type(cache) != INVALID_CHUNK);
 
-  set_type(cache, FREE_MEMORY);
+  set_type(cache, FREE_CHUNK);
 }
 
 void MemoryBlock::set_type(MetadataCache& cache, Type t) {
@@ -130,14 +144,12 @@ size_t MemoryBlock::index(MetadataCache& cache) const {
 }
 
 void* MemoryBlock::data() const {
-  return const_cast<MemoryBlockMetadata*>(
-             reinterpret_cast<const MemoryBlockMetadata*>(this)) +
-         1;
+  return const_cast<Metadata*>(reinterpret_cast<const Metadata*>(this)) + 1;
 }
 
 MemoryBlock* MemoryBlock::metadata() const {
   return const_cast<MemoryBlock*>(reinterpret_cast<const MemoryBlock*>(
-      reinterpret_cast<const MemoryBlockMetadata*>(this) - 1));
+      reinterpret_cast<const Metadata*>(this) - 1));
 }
 
 }  // detail

paddle/memory/detail/memory_block.h

@@ -14,24 +14,18 @@
 
 #pragma once
 
-#include "paddle/memory/detail/meta_data.h"
-
 #include <cstddef>
-#include <unordered_map>
 
 namespace paddle {
 namespace memory {
 namespace detail {
 
-// Forward Declaration
-class Metadata;
+// Forward Declarations
+class MetadataCache;
 
 /*! \brief A class used to interpret the contents of a memory block */
 class MemoryBlock {
  public:
-  // Unify the metadata format between GPU and CPU allocations
-  using MetadataCache = std::unordered_map<const MemoryBlock*, Metadata>;
-
   enum Type {
     FREE_CHUNK,    // memory is free and idle
     ARENA_CHUNK,   // memory is being occupied

paddle/memory/detail/meta_cache.cc

+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License. */
+
+#include "paddle/memory/detail/meta_cache.h"
+#include "paddle/memory/detail/memory_block.h"
+#include "paddle/platform/assert.h"
+
+namespace paddle {
+namespace memory {
+namespace detail {
+
+MetadataCache::MetadataCache(bool uses_gpu) : uses_gpu_(uses_gpu) {}
+
+Metadata MetadataCache::load(const MemoryBlock* block) {
+  if (uses_gpu_) {
+    auto existing_metadata = cache_.find(block);
+    assert(existing_metadata->second.check_guards());
+    return existing_metadata->second;
+  } else {
+    PADDLE_ASSERT(reinterpret_cast<const Metadata*>(block)->check_guards());
+    return *reinterpret_cast<const Metadata*>(block);
+  }
+}
+
+void MetadataCache::store(MemoryBlock* block,
+                          const Metadata& original_metadata) {
+  auto metadata = original_metadata;
+
+  metadata.update_guards();
+
+  if (uses_gpu_) {
+    cache_[block] = metadata;
+  } else {
+    *reinterpret_cast<Metadata*>(block) = metadata;
+  }
+}
+
+void MetadataCache::invalidate(MemoryBlock* block) {
+  if (uses_gpu_) {
+    cache_.erase(block);
+  }
+}
+
+}  // namespace detail
+}  // namespace memory
+}  // namespace paddle

paddle/memory/detail/meta_cache.h

+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License. */
+
+#pragma once
+
+#include "paddle/memory/detail/memory_block.h"
+#include "paddle/memory/detail/meta_data.h"
+
+#include <unordered_map>
+
+namespace paddle {
+namespace memory {
+namespace detail {
+
+/*! A cache for accessing memory block meta-data that may be expensive to access
+   directly.
+
+    Note: this class exists to unify the metadata format between GPU and CPU
+   allocations.
+    It should be removed when the CPU can access all GPU allocations directly
+   via UVM.
+*/
+class MetadataCache {
+ public:
+  MetadataCache(bool uses_gpu);
+
+ public:
+  /*! \brief Load the associated metadata for the specified memory block. */
+  Metadata load(const MemoryBlock*);
+
+  /*! \brief Store the associated metadata for the specified memory block. */
+  void store(MemoryBlock*, const Metadata&);
+
+ public:
+  /*! \brief Acquire any external metadata updates. */
+  void acquire(MemoryBlock*);
+
+  /*! \brief Publish any local updates externally. */
+  void release(MemoryBlock*);
+
+  /*! \brief Indicate that the specified metadata will no longer be used */
+  void invalidate(MemoryBlock*);
+
+ public:
+  MetadataCache(const MetadataCache&) = delete;
+  MetadataCache& operator=(const MetadataCache&) = delete;
+
+ private:
+  bool uses_gpu_;
+
+ private:
+  typedef std::unordered_map<const MemoryBlock*, Metadata> MetadataMap;
+
+ private:
+  MetadataMap cache_;
+};
+
+}  // namespace detail
+}  // namespace memory
+}  // namespace paddle

paddle/memory/detail/meta_data.cc

@@ -29,6 +29,14 @@ Metadata::Metadata(MemoryBlock::Type t, size_t i, size_t s, size_t ts,
       left_buddy(l),
       right_buddy(r) {}
 
+Metadata::Metadata()
+    : type(MemoryBlock::INVALID_CHUNK),
+      index(0),
+      size(0),
+      total_size(0),
+      left_buddy(nullptr),
+      right_buddy(nullptr) {}
+
 template <class T>
 inline void hash_combine(std::size_t& seed, const T& v) {
   std::hash<T> hasher;

paddle/memory/detail/meta_data.h

@@ -26,6 +26,7 @@ class Metadata {
  public:
   Metadata(MemoryBlock::Type t, size_t i, size_t s, size_t ts, MemoryBlock* l,
            MemoryBlock* r);
+  Metadata();
 
  public:
   /*! \brief Update the guards when metadata is changed */