szha commented on a change in pull request #14733: [MXNET-1398] Enable
zero-copy from numpy to MXNet NDArray
URL: https://github.com/apache/incubator-mxnet/pull/14733#discussion_r279419569
##########
File path: python/mxnet/ndarray/ndarray.py
##########
@@ -4115,3 +4115,108 @@ def from_dlpack(dlpack):
# delete the deleter of the old dlpack
ctypes.pythonapi.PyCapsule_SetDestructor(dlpack, None)
return NDArray(handle=handle)
+
+class DLContext(ctypes.Structure):
+ _fields_ = [("device_type", ctypes.c_int),
+ ("device_id", ctypes.c_int)]
+
+
+class DLDataType(ctypes.Structure):
+ _fields_ = [("type_code", ctypes.c_uint8),
+ ("bits", ctypes.c_uint8),
+ ("lanes", ctypes.c_uint16)]
+ TYPE_MAP = {
+ "int32": (0, 32, 1),
+ "int64": (0, 64, 1),
+ "bool": (1, 1, 1),
+ "uint32": (1, 32, 1),
+ "uint64": (1, 64, 1),
+ "float32": (2, 32, 1),
+ "float64": (2, 64, 1),
+ }
+
+
+class DLTensor(ctypes.Structure):
+ _fields_ = [("data", ctypes.c_void_p),
+ ("ctx", DLContext),
+ ("ndim", ctypes.c_int),
+ ("dtype", DLDataType),
+ ("shape", ctypes.POINTER(ctypes.c_int64)),
+ ("strides", ctypes.POINTER(ctypes.c_int64)),
+ ("byte_offset", ctypes.c_uint64)]
+
+class DLManagedTensor(ctypes.Structure):
+ pass
+
+
+DeleterFunc = ctypes.CFUNCTYPE(None, ctypes.POINTER(DLManagedTensor))
+
+
+DLManagedTensor._fields_ = [("dl_tensor", DLTensor), # pylint:
disable=protected-access
+ ("manager_ctx", ctypes.c_void_p),
+ ("deleter", DeleterFunc)]
+
+
+@DeleterFunc
+def dl_managed_tensor_deleter(dl_managed_tensor_handle):
+ void_p = dl_managed_tensor_handle.contents.manager_ctx
+ pyobj = ctypes.cast(void_p, ctypes.py_object)
+ ctypes.pythonapi.Py_DecRef(pyobj)
+
+
+def from_numpy(ndarray, zero_copy=True):
+ """Returns an MXNet's NDArray backed by Numpy's ndarray.
+
+ Parameters
+ ----------
+ ndarray: numpy.ndarray
+ input data
+
+ zero_copy: bool
+ Whether we use DLPack's zero-copy conversion to convert to MXNet's
NDArray.
+ This is only available for c-contiguous arrays, i.e.
array.flags[C_CONTIGUOUS] == True.
+
+ Returns
+ -------
+ NDArray
+ a NDArray backed by a dlpack tensor
+
+ """
+
+ def _make_manager_ctx(obj):
+ pyobj = ctypes.py_object(obj)
+ void_p = ctypes.c_void_p.from_buffer(pyobj)
+ ctypes.pythonapi.Py_IncRef(pyobj)
+ return void_p
+
+ def _make_dl_tensor(array):
+ if str(array.dtype) not in DLDataType.TYPE_MAP:
+ raise ValueError(str(array.dtype) + " is not supported.")
+ dl_tensor = DLTensor()
+ dl_tensor.data = array.ctypes.data_as(ctypes.c_void_p)
+ dl_tensor.ctx = DLContext(1, 0)
+ dl_tensor.ndim = array.ndim
+ dl_tensor.dtype = DLDataType.TYPE_MAP[str(array.dtype)]
+ dl_tensor.shape = array.ctypes.shape_as(ctypes.c_int64)
+ dl_tensor.strides = None
+ dl_tensor.byte_offset = 0
+ return dl_tensor
+
+ def _make_dl_managed_tensor(array):
+ c_obj = DLManagedTensor()
+ c_obj.dl_tensor = _make_dl_tensor(array)
+ c_obj.manager_ctx = _make_manager_ctx(array)
+ c_obj.deleter = dl_managed_tensor_deleter
+ return c_obj
+
+ if not zero_copy:
+ return array(ndarray, dtype=ndarray.dtype)
+
+ if not ndarray.flags['C_CONTIGUOUS']:
+ raise ValueError("Only c-contiguous arrays are supported for
zero-copy")
+ c_obj = _make_dl_managed_tensor(ndarray)
Review comment:
I think completely transferring the ownership from numpy and disabling
access to the original array is preferred. If we are to expose numpy array for
joint use, it should be done in asnumpy as another zero-copy creation feature,
so that the wait_to_read can be done by us
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
