From 48937020dc75bd2ede45ff2f1ef9d33386b0934a Mon Sep 17 00:00:00 2001
From: From00 <chenruibiao@baidu.com>
Date: Mon, 20 Dec 2021 19:05:56 +0800
Subject: [PATCH] Skip zero-size Allocation in RecordStream (#38264)

---
 .../memory/allocation/allocator_facade.cc     |  4 ++++
 .../memory/stream_safe_cuda_alloc_test.cu     | 21 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/paddle/fluid/memory/allocation/allocator_facade.cc b/paddle/fluid/memory/allocation/allocator_facade.cc
index 44cd915b168..3e9c2ba269f 100644
--- a/paddle/fluid/memory/allocation/allocator_facade.cc
+++ b/paddle/fluid/memory/allocation/allocator_facade.cc
@@ -311,6 +311,10 @@ class AllocatorFacadePrivate {
 
   void RecordStream(std::shared_ptr<Allocation> allocation,
                     const gpuStream_t& stream) {
+    if (allocation->size() == 0) {
+      return;
+    }
+
     StreamSafeCUDAAllocation* stream_safe_cuda_allocation =
         dynamic_cast<StreamSafeCUDAAllocation*>(allocation.get());
     PADDLE_ENFORCE_NOT_NULL(stream_safe_cuda_allocation,
diff --git a/paddle/fluid/memory/stream_safe_cuda_alloc_test.cu b/paddle/fluid/memory/stream_safe_cuda_alloc_test.cu
index 286dcdba8f2..52c3825053c 100644
--- a/paddle/fluid/memory/stream_safe_cuda_alloc_test.cu
+++ b/paddle/fluid/memory/stream_safe_cuda_alloc_test.cu
@@ -243,6 +243,27 @@ TEST(StreamSafeCUDAAllocInterfaceTest, GetAllocatorInterfaceTest) {
   CheckMemLeak(place);
 }
 
+TEST(StreamSafeCUDAAllocInterfaceTest, ZeroSizeRecordStreamTest) {
+  platform::CUDAPlace place = platform::CUDAPlace();
+  std::shared_ptr<Allocation> zero_size_allocation = AllocShared(place, 0);
+  EXPECT_EQ(zero_size_allocation->ptr(), nullptr);
+
+  gpuStream_t stream;
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamCreate(&stream));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamCreate(&stream));
+#endif
+
+  EXPECT_NO_THROW(RecordStream(zero_size_allocation, stream));
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamDestroy(stream));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipStreamDestroy(stream));
+#endif
+}
+
 TEST(StreamSafeCUDAAllocInterfaceTest, GetStreamInterfaceTest) {
   platform::CUDAPlace place = platform::CUDAPlace();
   size_t alloc_size = 256;
-- 
GitLab