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

[MXNET-1406] [BUG] Fix DLManagedTensor deleter #15016

Merged
merged 3 commits into from
May 22, 2019
Merged

[MXNET-1406] [BUG] Fix DLManagedTensor deleter #15016

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
103e74d
Fix
junrushao May 21, 2019
989bcc6
Fix
junrushao May 21, 2019
414cca8
Retrigger
junrushao May 21, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions include/mxnet/c_api.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -818,10 +818,12 @@ MXNET_DLL int MXNDArrayToDLPack(NDArrayHandle handle,
* The memory is retained until the NDArray went out of scope.
*
* \param dlpack the pointer of the input DLManagedTensor
* \param transient_handle whether the handle will be destructed before calling the deleter
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

before calling?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If it is false, the DLManagedTensor handle won't be freed after the call, and deleter will free it.
If it is true, the DLManagedTensor handle will be freed by the caller.

* \param out_handle pointer holder to get pointer of NDArray
* \return 0 when success, -1 when failure happens
*/
MXNET_DLL int MXNDArrayFromDLPack(DLManagedTensorHandle dlpack,
const bool transient_handle,
NDArrayHandle *out_handle);

/*!
Expand Down
2 changes: 1 addition & 1 deletion include/mxnet/ndarray.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -587,7 +587,7 @@ class NDArray {
*
* \return The created NDArray view.
*/
static NDArray FromDLPack(const DLManagedTensor* tensor);
static NDArray FromDLPack(const DLManagedTensor* tensor, bool transient_handle);

/*!
* \brief Update ndarray chunk storage handles using existing ndarray storage handles
Expand Down
6 changes: 2 additions & 4 deletions python/mxnet/ndarray/ndarray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4156,7 +4156,7 @@ def from_dlpack(dlpack):
assert ctypes.pythonapi.PyCapsule_IsValid(dlpack, _c_str_dltensor), ValueError(
'Invalid DLPack Tensor. DLTensor capsules can be consumed only once.')
dlpack_handle = ctypes.c_void_p(ctypes.pythonapi.PyCapsule_GetPointer(dlpack, _c_str_dltensor))
check_call(_LIB.MXNDArrayFromDLPack(dlpack_handle, ctypes.byref(handle)))
check_call(_LIB.MXNDArrayFromDLPack(dlpack_handle, False, ctypes.byref(handle)))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if transient_handle is false, it means the dlpack_handle won't be free'd any time soon? However, how do we control when the dlpack_handle will be free'd?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It means deleter will free it

# Rename PyCapsule (DLPack)
ctypes.pythonapi.PyCapsule_SetName(dlpack, _c_str_used_dltensor)
# delete the deleter of the old dlpack
Expand Down Expand Up @@ -4262,8 +4262,6 @@ def _make_dl_managed_tensor(array):
if not ndarray.flags['C_CONTIGUOUS']:
raise ValueError("Only c-contiguous arrays are supported for zero-copy")
c_obj = _make_dl_managed_tensor(ndarray)
address = ctypes.addressof(c_obj)
address = ctypes.cast(address, ctypes.c_void_p)
handle = NDArrayHandle()
check_call(_LIB.MXNDArrayFromDLPack(address, ctypes.byref(handle)))
check_call(_LIB.MXNDArrayFromDLPack(ctypes.byref(c_obj), True, ctypes.byref(handle)))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why the handle here is transient handle?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

bc the DLManagedTensor here is freed immediately after this call

return NDArray(handle=handle)
4 changes: 3 additions & 1 deletion src/c_api/c_api.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -562,10 +562,12 @@ int MXNDArrayToDLPack(NDArrayHandle handle,
}

int MXNDArrayFromDLPack(DLManagedTensorHandle dlpack,
const bool transient_handle,
NDArrayHandle *out_handle) {
API_BEGIN();
*out_handle = new NDArray(NDArray::FromDLPack(
static_cast<DLManagedTensor*>(dlpack)));
static_cast<DLManagedTensor*>(dlpack),
transient_handle));
API_END();
}

Expand Down
17 changes: 11 additions & 6 deletions src/ndarray/ndarray.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -355,14 +355,19 @@ DLManagedTensor* NDArray::ToDLPack() const {
return &(dlmanager->tensor);
}

NDArray NDArray::FromDLPack(const DLManagedTensor* tensor) {
DLManagedTensor tensor_copy = *tensor;
auto deleter = [tensor_copy](){
if (tensor_copy.deleter != nullptr) {
tensor_copy.deleter(const_cast<DLManagedTensor*>(&tensor_copy));
NDArray NDArray::FromDLPack(const DLManagedTensor* tensor, bool transient_handle) {
DLManagedTensor *tensor_copy = transient_handle
? new DLManagedTensor(*tensor)
: const_cast<DLManagedTensor*>(tensor);
auto deleter = [tensor_copy, transient_handle](){
if (tensor_copy->deleter != nullptr) {
tensor_copy->deleter(tensor_copy);
}
if (transient_handle) {
delete tensor_copy;
}
};
return NDArray(TBlob(tensor_copy.dl_tensor), tensor_copy.dl_tensor.ctx.device_id, deleter);
return NDArray(TBlob(tensor_copy->dl_tensor), tensor_copy->dl_tensor.ctx.device_id, deleter);
}

bool NDArray::fresh_out_grad() const {
Expand Down