update Alloc in naive storage manager

apache · Mar 26, 2019 · b13476d · b13476d
1 parent fbc34b3
commit b13476d
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 14 deletions.
diff --git 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
@@ -45,18 +45,21 @@ 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.
    */
   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
@@ -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
@@ -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,18 +49,21 @@ 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));
 #endif
   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) {