From 89110fd2660098bc949a1f13f7b53515e0c931a3 Mon Sep 17 00:00:00 2001
From: liaogang <liaogang@baidu.com>
Date: Mon, 3 Jul 2017 19:51:32 +0800
Subject: [PATCH] ENH: Add useGpu in system allocator

---
 paddle/memory/detail/system_allocator.cc | 4 ++++
 paddle/memory/detail/system_allocator.h  | 4 +++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/paddle/memory/detail/system_allocator.cc b/paddle/memory/detail/system_allocator.cc
index 2b0fbfa87e2..75a2c91ef93 100644
--- a/paddle/memory/detail/system_allocator.cc
+++ b/paddle/memory/detail/system_allocator.cc
@@ -60,6 +60,8 @@ void CPUAllocator::Free(void* p, size_t size, size_t index) {
   free(p);
 }
 
+bool CPUAllocator::UseGpu() { return false; }
+
 #ifndef PADDLE_ONLY_CPU
 
 void* GPUAllocator::Alloc(size_t& index, size_t size) {
@@ -131,6 +133,8 @@ void GPUAllocator::Free(void* p, size_t size, size_t index) {
   }
 }
 
+bool GPUAllocator::UseGpu() { return true; }
+
 #endif  // PADDLE_ONLY_CPU
 
 }  // namespace detail
diff --git a/paddle/memory/detail/system_allocator.h b/paddle/memory/detail/system_allocator.h
index 7093c429671..f3bbfef8435 100644
--- a/paddle/memory/detail/system_allocator.h
+++ b/paddle/memory/detail/system_allocator.h
@@ -32,12 +32,14 @@ class SystemAllocator {
   virtual ~SystemAllocator() {}
   virtual void* Alloc(size_t& index, size_t size) = 0;
   virtual void Free(void* p, size_t size, size_t index) = 0;
+  virtual bool UseGpu() = 0;
 };
 
 class CPUAllocator : public SystemAllocator {
  public:
   virtual void* Alloc(size_t& index, size_t size);
   virtual void Free(void* p, size_t size, size_t index);
+  virtual bool UseGpu();
 };
 
 #ifndef PADDLE_ONLY_CPU
@@ -45,7 +47,7 @@ class GPUAllocator : public SystemAllocator {
  public:
   virtual void* Alloc(size_t& index, size_t size);
   virtual void Free(void* p, size_t size, size_t index);
-
+  virtual bool UseGpu();
  private:
   size_t gpu_alloc_size_ = 0;
   size_t fallback_alloc_size_ = 0;
-- 
GitLab