free memory for size zero storage handle

apache · Mar 12, 2019 · 79a7f11 · 79a7f11
1 parent ab0ca86
commit 79a7f11
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 11 deletions.
diff --git a/include/mxnet/ndarray.h b/include/mxnet/ndarray.h
@@ -986,9 +986,8 @@ class NDArray {
 #endif
         delay_alloc = false;
       } else if (shandle.size < dbytes) {
-        // free storage if necessary and alloc again
-        if (shandle.size > 0) Storage::Get()->Free(shandle);
-        // init storage
+        // free storage and alloc again
+        Storage::Get()->Free(shandle);
         shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
 #if MXNET_USE_MKLDNN == 1
         mkl_mem_ = nullptr;
@@ -1055,9 +1054,8 @@ class NDArray {
       }
       size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
       if (aux_handles[i].size < aux_bytes) {
-        // free storage if necessary and alloc again
-        if (aux_handles[i].size > 0) Storage::Get()->Free(aux_handles[i]);
-        // init aux storage
+        // free storage and alloc again
+        Storage::Get()->Free(aux_handles[i]);
         aux_handles[i] = Storage::Get()->Alloc(aux_bytes, ctx);
       }
       // init shape

diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
@@ -121,9 +121,9 @@ NDArray::Chunk::~Chunk() {
         CHECK_EQ(mem.mem->GetDataHandle(), mem.h.dptr);
       }
 #endif
-      if (mem.h.size > 0) Storage::Get()->Free(mem.h);
+      Storage::Get()->Free(mem.h);
       for (const auto& aux : mem.aux_h) {
-        if (aux.size > 0) Storage::Get()->Free(aux);
+        Storage::Get()->Free(aux);
       }
     }
   }, shandle.ctx, var);
@@ -134,9 +134,8 @@ void NDArray::Chunk::CheckAndAllocData(const mxnet::TShape &shape, int dtype) {
       << "data is expected to be allocated after aux_data";
   auto dbytes = shape.Size() * mshadow::mshadow_sizeof(dtype);
   if (shandle.size < dbytes) {
-    // free storage if necessary and alloc again
-    if (shandle.size > 0) Storage::Get()->Free(shandle);
-    // init storage
+    // free storage and alloc again
+    Storage::Get()->Free(shandle);
     shandle = Storage::Get()->Alloc(dbytes, ctx);
 #if MXNET_USE_MKLDNN == 1
     mkl_mem_ = nullptr;