From b34933d9ee3b61dbbd642fd02f244c36d0d14550 Mon Sep 17 00:00:00 2001
From: Zeng Jinle <32832641+sneaxiy@users.noreply.github.com>
Date: Mon, 16 Sep 2019 14:26:20 +0800
Subject: [PATCH] fix retry allocator bug, test=develop (#19794)

---
 paddle/fluid/memory/allocation/cuda_allocator.cc    |  1 +
 paddle/fluid/memory/detail/system_allocator.cc      |  1 +
 paddle/fluid/memory/detail/system_allocator_test.cc | 13 +++++++++++++
 3 files changed, 15 insertions(+)

diff --git a/paddle/fluid/memory/allocation/cuda_allocator.cc b/paddle/fluid/memory/allocation/cuda_allocator.cc
index 10e275ab4f6..2ba3b6d0b5b 100644
--- a/paddle/fluid/memory/allocation/cuda_allocator.cc
+++ b/paddle/fluid/memory/allocation/cuda_allocator.cc
@@ -37,6 +37,7 @@ Allocation* CUDAAllocator::AllocateImpl(size_t size) {
   void* ptr;
   auto status = cudaMalloc(&ptr, size);
   if (UNLIKELY(status != cudaSuccess)) {
+    PADDLE_ENFORCE_NE(cudaGetLastError(), cudaSuccess);
     PADDLE_THROW_BAD_ALLOC("Cannot allocate %d on GPU %d, cuda status %d, %s",
                            size, place_.device, status,
                            cudaGetErrorString(status));
diff --git a/paddle/fluid/memory/detail/system_allocator.cc b/paddle/fluid/memory/detail/system_allocator.cc
index c3d40095c74..0842f33a18c 100644
--- a/paddle/fluid/memory/detail/system_allocator.cc
+++ b/paddle/fluid/memory/detail/system_allocator.cc
@@ -118,6 +118,7 @@ void* GPUAllocator::Alloc(size_t* index, size_t size) {
     gpu_alloc_size_ += size;
     return p;
   } else {
+    PADDLE_ENFORCE_NE(cudaGetLastError(), cudaSuccess);
     PADDLE_THROW_BAD_ALLOC(
         "Cannot malloc " + std::to_string(size / 1024.0 / 1024.0) +
         " MB GPU memory. Please shrink "
diff --git a/paddle/fluid/memory/detail/system_allocator_test.cc b/paddle/fluid/memory/detail/system_allocator_test.cc
index 65ac12323b0..ed5022cac69 100644
--- a/paddle/fluid/memory/detail/system_allocator_test.cc
+++ b/paddle/fluid/memory/detail/system_allocator_test.cc
@@ -19,6 +19,7 @@ limitations under the License. */
 
 #include "gflags/gflags.h"
 #include "gtest/gtest.h"
+#include "paddle/fluid/memory/allocation/allocator.h"
 
 DECLARE_bool(use_pinned_memory);
 
@@ -68,4 +69,16 @@ TEST(CUDAPinnedAllocator, Alloc) {
   TestAllocator(&a, 2048);
   TestAllocator(&a, 0);
 }
+
+TEST(GPUAllocator, AllocFailure) {
+  paddle::memory::detail::GPUAllocator allocator(0);
+  size_t index;
+  size_t alloc_size = -1UL;  // very large size
+  try {
+    allocator.Alloc(&index, alloc_size);
+    ASSERT_TRUE(false);
+  } catch (paddle::memory::allocation::BadAlloc&) {
+    PADDLE_ENFORCE_CUDA_SUCCESS(cudaGetLastError());
+  }
+}
 #endif
-- 
GitLab