From 1f4626bb03525f4fd56fb9c8b518b53fe1d99d85 Mon Sep 17 00:00:00 2001
From: Yuxi Hu <darrenyxhu@gmail.com>
Date: Mon, 11 Mar 2019 16:36:10 -0700
Subject: [PATCH] skip adding nullptr into GPU memory pool

---
 src/storage/pooled_storage_manager.h | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/storage/pooled_storage_manager.h b/src/storage/pooled_storage_manager.h
index c407a9f00cb6..7c4b070afdd2 100644
--- a/src/storage/pooled_storage_manager.h
+++ b/src/storage/pooled_storage_manager.h
@@ -155,6 +155,10 @@ void GPUPooledStorageManager::Alloc(Storage::Handle* handle) {
 }
 
 void GPUPooledStorageManager::Free(Storage::Handle handle) {
+  // Do nothing if dptr is nullptr. Otherwise, nullptr may be reused
+  // which can cause illegal memory access error.
+  if (handle.dptr == nullptr) return;
+
   std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
   size_t size = RoundAllocSize(handle.size);
   auto&& reuse_pool = memory_pool_[size];
@@ -312,6 +316,10 @@ void GPUPooledRoundedStorageManager::Alloc(Storage::Handle* handle) {
 }
 
 void GPUPooledRoundedStorageManager::Free(Storage::Handle handle) {
+  // Do nothing if dptr is nullptr. Otherwise, nullptr may be reused
+  // which can cause illegal memory access error.
+  if (handle.dptr == nullptr) return;
+
   std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
   int bucket = get_bucket(handle.size);
   auto&& reuse_pool = memory_pool_[bucket];