diff --git a/paddle/memory/detail/system_allocator.cc b/paddle/memory/detail/system_allocator.cc index 2b0fbfa87e22ad5697b1062d1d7b80f5eb251399..75a2c91ef938e417a34de27d10b1a82cd4d85011 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 7093c42967162c6737e1fe68176cf3993a18ba09..f3bbfef843542fed6f304b1649f9092f8cdce66e 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;