skip memory allocation when handle size is 0

apache · Mar 20, 2019 · 02a0fec · 02a0fec
1 parent 00d59af
commit 02a0fec
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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);