提交 89110fd2 编写于 作者: L liaogang

ENH: Add useGpu in system allocator

上级 275e5b7d
......@@ -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
......
......@@ -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;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册