From ed7727e8f04c215f4ff77f486e46347efe0ad3cd Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Tue, 20 Mar 2018 18:17:13 +0800
Subject: [PATCH] Fix bug in system allocator

---
 paddle/fluid/memory/detail/system_allocator.cc | 11 +++++++++++
 paddle/fluid/memory/detail/system_allocator.h  |  3 +++
 paddle/fluid/memory/memory.cc                  |  2 +-
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/memory/detail/system_allocator.cc b/paddle/fluid/memory/detail/system_allocator.cc
index 8ac8978120..9949d80434 100644
--- a/paddle/fluid/memory/detail/system_allocator.cc
+++ b/paddle/fluid/memory/detail/system_allocator.cc
@@ -79,7 +79,18 @@ void* GPUAllocator::Alloc(size_t& index, size_t size) {
   // if size is 0.  We just make sure it does.
   if (size <= 0) return nullptr;
   void* p;
+  int prev_id;
+  cudaGetDevice(&prev_id);
+  if (prev_id != gpu_id_) {
+    cudaSetDevice(gpu_id_);
+  }
+
   cudaError_t result = cudaMalloc(&p, size);
+
+  if (prev_id != gpu_id_) {
+    cudaSetDevice(prev_id);
+  }
+
   if (result == cudaSuccess) {
     index = 0;
     gpu_alloc_size_ += size;
diff --git a/paddle/fluid/memory/detail/system_allocator.h b/paddle/fluid/memory/detail/system_allocator.h
index e93c2c1e32..c103d08640 100644
--- a/paddle/fluid/memory/detail/system_allocator.h
+++ b/paddle/fluid/memory/detail/system_allocator.h
@@ -43,6 +43,8 @@ class CPUAllocator : public SystemAllocator {
 #ifdef PADDLE_WITH_CUDA
 class GPUAllocator : public SystemAllocator {
  public:
+  explicit GPUAllocator(int gpu_id) : gpu_id_(gpu_id) {}
+
   virtual void* Alloc(size_t& index, size_t size);
   virtual void Free(void* p, size_t size, size_t index);
   virtual bool UseGpu() const;
@@ -50,6 +52,7 @@ class GPUAllocator : public SystemAllocator {
  private:
   size_t gpu_alloc_size_ = 0;
   size_t fallback_alloc_size_ = 0;
+  int gpu_id_;
 };
 #endif
 
diff --git a/paddle/fluid/memory/memory.cc b/paddle/fluid/memory/memory.cc
index d07f89439a..1985f1f4e6 100644
--- a/paddle/fluid/memory/memory.cc
+++ b/paddle/fluid/memory/memory.cc
@@ -69,7 +69,7 @@ BuddyAllocator* GetGPUBuddyAllocator(int gpu_id) {
   }
   platform::SetDeviceId(gpu_id);
   if (!as[gpu_id]) {
-    as[gpu_id] = new BuddyAllocator(new detail::GPUAllocator,
+    as[gpu_id] = new BuddyAllocator(new detail::GPUAllocator(gpu_id),
                                     platform::GpuMinChunkSize(),
                                     platform::GpuMaxChunkSize());
     VLOG(10) << "\n\nNOTE: each GPU device use "
-- 
GitLab