Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

… fix-7346
PaddlePaddle · Jan 9, 2018 · 106d063 · 106d063
2 parents 427e474 + 45e7715
commit 106d063
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 0 deletions.
diff --git a/paddle/memory/memory.cc b/paddle/memory/memory.cc
@@ -114,5 +114,21 @@ void Free<platform::CUDAPlace>(platform::CUDAPlace place, void* p) {
 
 #endif
 
+size_t Usage::operator()(const platform::CPUPlace& cpu) const {
+  return Used(cpu);
+}
+
+size_t Usage::operator()(const platform::CUDAPlace& gpu) const {
+#ifdef PADDLE_WITH_CUDA
+  return Used(gpu);
+#else
+  PADDLE_THROW("'CUDAPlace' is not supported in CPU only device.");
+#endif
+}
+
+size_t memory_usage(const platform::Place& p) {
+  return boost::apply_visitor(Usage(), p);
+}
+
 }  // namespace memory
 }  // namespace paddle
diff --git a/paddle/memory/memory.h b/paddle/memory/memory.h
@@ -54,6 +54,13 @@ void Free(Place place, void* ptr);
 template <typename Place>
 size_t Used(Place place);
 
+struct Usage : public boost::static_visitor<size_t> {
+  size_t operator()(const platform::CPUPlace& cpu) const;
+  size_t operator()(const platform::CUDAPlace& gpu) const;
+};
+
+size_t memory_usage(const platform::Place& p);
+
 /**
  * \brief   Free memory block in one place.
  *

diff --git a/paddle/memory/memory_test.cc b/paddle/memory/memory_test.cc
@@ -44,6 +44,9 @@ TEST(BuddyAllocator, CPUAllocation) {
 
   EXPECT_NE(p, nullptr);
 
+  paddle::platform::Place place = cpu;
+  EXPECT_EQ(paddle::memory::Used(cpu), paddle::memory::memory_usage(place));
+
   paddle::memory::Free(cpu, p);
 }
 
@@ -99,6 +102,9 @@ TEST(BuddyAllocator, GPUAllocation) {
 
   EXPECT_NE(p, nullptr);
 
+  paddle::platform::Place place = gpu;
+  EXPECT_EQ(paddle::memory::Used(gpu), paddle::memory::memory_usage(place));
+
   paddle::memory::Free(gpu, p);
 }