未验证 提交 45e77154 编写于 作者: Q QI JUN 提交者: GitHub

add general memory usage interface for both CPU/CUDA (#7352)

上级 e249ad12
...@@ -114,5 +114,21 @@ void Free<platform::CUDAPlace>(platform::CUDAPlace place, void* p) { ...@@ -114,5 +114,21 @@ void Free<platform::CUDAPlace>(platform::CUDAPlace place, void* p) {
#endif #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 memory
} // namespace paddle } // namespace paddle
...@@ -54,6 +54,13 @@ void Free(Place place, void* ptr); ...@@ -54,6 +54,13 @@ void Free(Place place, void* ptr);
template <typename Place> template <typename Place>
size_t Used(Place 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. * \brief Free memory block in one place.
* *
......
...@@ -44,6 +44,9 @@ TEST(BuddyAllocator, CPUAllocation) { ...@@ -44,6 +44,9 @@ TEST(BuddyAllocator, CPUAllocation) {
EXPECT_NE(p, nullptr); 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); paddle::memory::Free(cpu, p);
} }
...@@ -99,6 +102,9 @@ TEST(BuddyAllocator, GPUAllocation) { ...@@ -99,6 +102,9 @@ TEST(BuddyAllocator, GPUAllocation) {
EXPECT_NE(p, nullptr); 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); paddle::memory::Free(gpu, p);
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册