From 02a0fecafa0f82e3e194a3dae37dbe1945a4bc13 Mon Sep 17 00:00:00 2001
From: Yuxi Hu <darrenyxhu@gmail.com>
Date: Wed, 20 Mar 2019 14:57:17 -0700
Subject: [PATCH] skip memory allocation when handle size is 0

---
 src/storage/cpu_device_storage.h         |  8 +++++---
 src/storage/cpu_shared_storage_manager.h |  7 ++++++-
 src/storage/gpu_device_storage.h         |  6 ++++--
 src/storage/pinned_memory_storage.h      |  4 +++-
 src/storage/pooled_storage_manager.h     | 12 ++++++++++++
 5 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/storage/cpu_device_storage.h b/src/storage/cpu_device_storage.h
index 25ad61efb232..0053392c67c9 100644
--- a/src/storage/cpu_device_storage.h
+++ b/src/storage/cpu_device_storage.h
@@ -64,11 +64,13 @@ class CPUDeviceStorage {
 };  // class CPUDeviceStorage
 
 inline void* CPUDeviceStorage::Alloc(Storage::Handle* handle) {
+  void* ptr = nullptr;
   const size_t size = handle->size;
-  void* ptr;
+  if (size == 0) return ptr;
+
 #if _MSC_VER
   ptr = _aligned_malloc(size, alignment_);
-  if (ptr == NULL) LOG(FATAL) << "Failed to allocate CPU Memory";
+  if (ptr == nullptr) LOG(FATAL) << "Failed to allocate CPU Memory";
 #else
   int ret = posix_memalign(&ptr, alignment_, size);
   if (ret != 0) LOG(FATAL) << "Failed to allocate CPU Memory";
@@ -77,7 +79,7 @@ inline void* CPUDeviceStorage::Alloc(Storage::Handle* handle) {
 }
 
 inline void CPUDeviceStorage::Free(Storage::Handle handle) {
-  void * ptr = handle.dptr;
+  void* ptr = handle.dptr;
 #if _MSC_VER
   _aligned_free(ptr);
 #else
diff --git a/src/storage/cpu_shared_storage_manager.h b/src/storage/cpu_shared_storage_manager.h
index a52d779d2318..9c57a4b61eed 100644
--- a/src/storage/cpu_shared_storage_manager.h
+++ b/src/storage/cpu_shared_storage_manager.h
@@ -115,13 +115,18 @@ class CPUSharedStorageManager final : public StorageManager {
 };  // class CPUSharedStorageManager
 
 void CPUSharedStorageManager::Alloc(Storage::Handle* handle) {
+  if (handle->size == 0) {
+    handle->dptr = nullptr;
+    return;
+  }
+
   std::lock_guard<std::recursive_mutex> lock(mutex_);
   std::uniform_int_distribution<> dis(0, std::numeric_limits<int>::max());
   int fid = -1;
   std::string filename;
   bool is_new = false;
   size_t size = handle->size + alignment_;
-  void *ptr = nullptr;
+  void* ptr = nullptr;
 #ifdef _WIN32
   CheckAndRealFree();
   HANDLE map_handle = nullptr;
diff --git a/src/storage/gpu_device_storage.h b/src/storage/gpu_device_storage.h
index 562badb8752e..72b1fb3559fb 100644
--- a/src/storage/gpu_device_storage.h
+++ b/src/storage/gpu_device_storage.h
@@ -55,8 +55,10 @@ class GPUDeviceStorage {
 };  // class GPUDeviceStorage
 
 inline void* GPUDeviceStorage::Alloc(Storage::Handle* handle) {
-  const size_t size = handle->size;
   void* ret = nullptr;
+  const size_t size = handle->size;
+  if (size == 0) return ret;
+
 #if MXNET_USE_CUDA
   mxnet::common::cuda::DeviceStore device_store(handle->ctx.real_dev_id(), true);
 #if MXNET_USE_NCCL
@@ -73,7 +75,7 @@ inline void* GPUDeviceStorage::Alloc(Storage::Handle* handle) {
 
 inline void GPUDeviceStorage::Free(Storage::Handle handle) {
 #if MXNET_USE_CUDA
-  void * ptr = handle.dptr;
+  void* ptr = handle.dptr;
   mxnet::common::cuda::DeviceStore device_store(handle.ctx.real_dev_id(), true);
 #if MXNET_USE_NCCL
   std::lock_guard<std::mutex> l(Storage::Get()->GetMutex(Context::kGPU));
diff --git a/src/storage/pinned_memory_storage.h b/src/storage/pinned_memory_storage.h
index c4ababbdc03a..d84cd0c96554 100644
--- a/src/storage/pinned_memory_storage.h
+++ b/src/storage/pinned_memory_storage.h
@@ -53,6 +53,8 @@ class PinnedMemoryStorage {
 inline void* PinnedMemoryStorage::Alloc(Storage::Handle* handle) {
   void* ret = nullptr;
   const size_t size = handle->size;
+  if (size == 0) return ret;
+
 #if MXNET_USE_NCCL
   std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
 #endif
@@ -63,7 +65,7 @@ inline void* PinnedMemoryStorage::Alloc(Storage::Handle* handle) {
 }
 
 inline void PinnedMemoryStorage::Free(Storage::Handle handle) {
-  void * ptr = handle.dptr;
+  void* ptr = handle.dptr;
 #if MXNET_USE_NCCL
   std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
 #endif
diff --git a/src/storage/pooled_storage_manager.h b/src/storage/pooled_storage_manager.h
index 7c4b070afdd2..4c8ae4eb12dd 100644
--- a/src/storage/pooled_storage_manager.h
+++ b/src/storage/pooled_storage_manager.h
@@ -129,6 +129,12 @@ class GPUPooledStorageManager final : public StorageManager {
 };  // class GPUPooledStorageManager
 
 void GPUPooledStorageManager::Alloc(Storage::Handle* handle) {
+  // Set dptr to nullptr when handle size is 0.
+  if (handle->size == 0) {
+    handle->dptr = nullptr;
+    return;
+  }
+
   std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
   size_t size = RoundAllocSize(handle->size);
   auto&& reuse_it = memory_pool_.find(size);
@@ -290,6 +296,12 @@ class GPUPooledRoundedStorageManager final : public StorageManager {
 };  // class GPUPooledRoundedStorageManager
 
 void GPUPooledRoundedStorageManager::Alloc(Storage::Handle* handle) {
+  // Set dptr to nullptr when handle size is 0.
+  if (handle->size == 0) {
+    handle->dptr = nullptr;
+    return;
+  }
+
   std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
   int bucket = get_bucket(handle->size);
   size_t size = get_size(bucket);