Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
Revert "Fix memory leak for size-zero ndarray (#14365)" (#14477)
Browse files Browse the repository at this point in the history
This reverts commit 3ab1dec.
  • Loading branch information
yuxihu authored and TaoLv committed Mar 20, 2019
1 parent 56b7b67 commit f98820c
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 19 deletions.
15 changes: 8 additions & 7 deletions include/mxnet/ndarray.h
Original file line number Diff line number Diff line change
Expand Up @@ -909,6 +909,9 @@ class NDArray {
// aux_handles always reflect the correct number of aux data
for (size_t i = 0; i < aux_shapes.size(); i++) {
CheckAndAllocAuxData(i, aux_shapes[i]);
// this line is needed in case when aux_shapes[i].Size() = 0
// aux_handles[i] will not be updated and take only default value.
aux_handles[i].ctx = ctx;
}
if (!delay_alloc) {
CheckAndAllocData(storage_shape, dtype);
Expand Down Expand Up @@ -983,8 +986,8 @@ class NDArray {
#endif
delay_alloc = false;
} else if (shandle.size < dbytes) {
// free storage
Storage::Get()->Free(shandle);
// free storage if necessary and alloc again
if (shandle.size > 0) Storage::Get()->Free(shandle);
// init storage
shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
#if MXNET_USE_MKLDNN == 1
Expand Down Expand Up @@ -1049,14 +1052,12 @@ class NDArray {
<< "storage type cannot be kDefaultStorage in CheckAndAllocAuxData";
if (aux_handles.size() <= i) {
aux_handles.resize(i + 1);
// set context for the newly created aux handle
aux_handles[i].ctx = ctx;
}
size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
if (aux_handles[i].size < aux_bytes) {
// free storage
Storage::Get()->Free(aux_handles[i]);
// init storage
// free storage if necessary and alloc again
if (aux_handles[i].size > 0) Storage::Get()->Free(aux_handles[i]);
// init aux storage
aux_handles[i] = Storage::Get()->Alloc(aux_bytes, ctx);
}
// init shape
Expand Down
8 changes: 4 additions & 4 deletions src/ndarray/ndarray.cc
Original file line number Diff line number Diff line change
Expand Up @@ -121,9 +121,9 @@ NDArray::Chunk::~Chunk() {
CHECK_EQ(mem.mem->GetDataHandle(), mem.h.dptr);
}
#endif
Storage::Get()->Free(mem.h);
if (mem.h.size > 0) Storage::Get()->Free(mem.h);
for (const auto& aux : mem.aux_h) {
Storage::Get()->Free(aux);
if (aux.size > 0) Storage::Get()->Free(aux);
}
}
}, shandle.ctx, var);
Expand All @@ -134,8 +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
Storage::Get()->Free(shandle);
// free storage if necessary and alloc again
if (shandle.size > 0) Storage::Get()->Free(shandle);
// init storage
shandle = Storage::Get()->Alloc(dbytes, ctx);
#if MXNET_USE_MKLDNN == 1
Expand Down
8 changes: 0 additions & 8 deletions src/storage/pooled_storage_manager.h
Original file line number Diff line number Diff line change
Expand Up @@ -155,10 +155,6 @@ void GPUPooledStorageManager::Alloc(Storage::Handle* handle) {
}

void GPUPooledStorageManager::Free(Storage::Handle handle) {
// Do nothing if dptr is nullptr. Otherwise, nullptr may be reused
// which can cause illegal memory access error.
if (handle.dptr == nullptr) return;

std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
size_t size = RoundAllocSize(handle.size);
auto&& reuse_pool = memory_pool_[size];
Expand Down Expand Up @@ -316,10 +312,6 @@ void GPUPooledRoundedStorageManager::Alloc(Storage::Handle* handle) {
}

void GPUPooledRoundedStorageManager::Free(Storage::Handle handle) {
// Do nothing if dptr is nullptr. Otherwise, nullptr may be reused
// which can cause illegal memory access error.
if (handle.dptr == nullptr) return;

std::lock_guard<std::mutex> lock(Storage::Get()->GetMutex(Context::kGPU));
int bucket = get_bucket(handle.size);
auto&& reuse_pool = memory_pool_[bucket];
Expand Down

0 comments on commit f98820c

Please sign in to comment.