From b13476d54864209ac75968e513f5fd23d6bbf35a Mon Sep 17 00:00:00 2001
From: Yuxi Hu <darrenyxhu@gmail.com>
Date: Thu, 21 Mar 2019 22:17:58 -0700
Subject: [PATCH] update Alloc in naive storage manager

---
 src/storage/cpu_device_storage.h    | 11 +++++++----
 src/storage/gpu_device_storage.h    | 11 +++++++----
 src/storage/naive_storage_manager.h |  2 +-
 src/storage/pinned_memory_storage.h | 13 ++++++++-----
 4 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/storage/cpu_device_storage.h b/src/storage/cpu_device_storage.h
index ae4a2bcb9331..efd6848b0e4f 100644
--- a/src/storage/cpu_device_storage.h
+++ b/src/storage/cpu_device_storage.h
@@ -42,7 +42,7 @@ class CPUDeviceStorage {
    * \brief Aligned allocation on CPU.
    * \param handle Handle struct.
    */
-  inline static void* Alloc(Storage::Handle* handle);
+  inline static void Alloc(Storage::Handle* handle);
   /*!
    * \brief Deallocation.
    * \param handle Handle struct.
@@ -62,10 +62,13 @@ class CPUDeviceStorage {
 #endif
 };  // class CPUDeviceStorage
 
-inline void* CPUDeviceStorage::Alloc(Storage::Handle* handle) {
+inline void CPUDeviceStorage::Alloc(Storage::Handle* handle) {
   void* ptr = nullptr;
   const size_t size = handle->size;
-  if (size == 0) return ptr;
+  if (size == 0) {
+    handle->dptr = ptr;
+    return;
+  }
 
 #if _MSC_VER
   ptr = _aligned_malloc(size, alignment_);
@@ -74,7 +77,7 @@ inline void* CPUDeviceStorage::Alloc(Storage::Handle* handle) {
   int ret = posix_memalign(&ptr, alignment_, size);
   if (ret != 0) LOG(FATAL) << "Failed to allocate CPU Memory";
 #endif
-  return ptr;
+  handle->dptr = ptr;
 }
 
 inline void CPUDeviceStorage::Free(Storage::Handle handle) {
diff --git a/src/storage/gpu_device_storage.h b/src/storage/gpu_device_storage.h
index ac1160f63431..0c3bed558620 100644
--- a/src/storage/gpu_device_storage.h
+++ b/src/storage/gpu_device_storage.h
@@ -45,7 +45,7 @@ class GPUDeviceStorage {
    * \brief Allocation.
    * \param handle Handle struct.
    */
-  inline static void* Alloc(Storage::Handle* handle);
+  inline static void Alloc(Storage::Handle* handle);
   /*!
    * \brief Deallocation.
    * \param handle Handle struct.
@@ -53,10 +53,13 @@ class GPUDeviceStorage {
   inline static void Free(Storage::Handle handle);
 };  // class GPUDeviceStorage
 
-inline void* GPUDeviceStorage::Alloc(Storage::Handle* handle) {
+inline void GPUDeviceStorage::Alloc(Storage::Handle* handle) {
   void* ret = nullptr;
   const size_t size = handle->size;
-  if (size == 0) return ret;
+  if (size == 0) {
+    handle->dptr = ret;
+    return;
+  }
 
 #if MXNET_USE_CUDA
   mxnet::common::cuda::DeviceStore device_store(handle->ctx.real_dev_id(), true);
@@ -69,7 +72,7 @@ inline void* GPUDeviceStorage::Alloc(Storage::Handle* handle) {
 #else   // MXNET_USE_CUDA
   LOG(FATAL) << "Please compile with CUDA enabled";
 #endif  // MXNET_USE_CUDA
-  return ret;
+  handle->dptr = ret;
 }
 
 inline void GPUDeviceStorage::Free(Storage::Handle handle) {
diff --git a/src/storage/naive_storage_manager.h b/src/storage/naive_storage_manager.h
index 55112b5a82e9..471b015eb32c 100644
--- a/src/storage/naive_storage_manager.h
+++ b/src/storage/naive_storage_manager.h
@@ -58,7 +58,7 @@ class NaiveStorageManager final : public StorageManager {
 
 template <class DeviceStorage>
 void NaiveStorageManager<DeviceStorage>::Alloc(Storage::Handle* handle) {
-  handle->dptr = DeviceStorage::Alloc(handle);
+  DeviceStorage::Alloc(handle);
 }
 
 template <class DeviceStorage>
diff --git a/src/storage/pinned_memory_storage.h b/src/storage/pinned_memory_storage.h
index 61820bfe914a..3fe95f3904d5 100644
--- a/src/storage/pinned_memory_storage.h
+++ b/src/storage/pinned_memory_storage.h
@@ -19,7 +19,7 @@
 
 /*!
  * Copyright (c) 2015 by Contributors
- * \file cpu_device_storage.h
+ * \file pinned_memory_storage.h
  * \brief CPU storage with pinned memory
  */
 #ifndef MXNET_STORAGE_PINNED_MEMORY_STORAGE_H_
@@ -40,7 +40,7 @@ class PinnedMemoryStorage {
    * \brief Allocation.
    * \param handle Handle struct.
    */
-  inline static void* Alloc(Storage::Handle* handle);
+  inline static void Alloc(Storage::Handle* handle);
 
   /*!
    * \brief Deallocation.
@@ -49,10 +49,13 @@ class PinnedMemoryStorage {
   inline static void Free(Storage::Handle handle);
 };
 
-inline void* PinnedMemoryStorage::Alloc(Storage::Handle* handle) {
+inline void PinnedMemoryStorage::Alloc(Storage::Handle* handle) {
   void* ret = nullptr;
   const size_t size = handle->size;
-  if (size == 0) return ret;
+  if (size == 0) {
+    handle->dptr = ret;
+    return;
+  }
 
 #if MXNET_USE_NCCL
   std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
@@ -60,7 +63,7 @@ inline void* PinnedMemoryStorage::Alloc(Storage::Handle* handle) {
   mxnet::common::cuda::DeviceStore device_store(handle->ctx.real_dev_id(), true);
   // make the memory available across all devices
   CUDA_CALL(cudaHostAlloc(&ret, size, cudaHostAllocPortable));
-  return ret;
+  handle->dptr = ret;
 }
 
 inline void PinnedMemoryStorage::Free(Storage::Handle handle) {