This is an automated email from the ASF dual-hosted git repository.
wkcn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new c18381d [MXNET-1398] Enable zero-copy from numpy to MXNet NDArray
(#14733)
c18381d is described below
commit c18381d336e1d2f950c6d11841db15c94062cce4
Author: Junru Shao <[email protected]>
AuthorDate: Sun Apr 28 16:12:01 2019 -0700
[MXNET-1398] Enable zero-copy from numpy to MXNet NDArray (#14733)
* Enable zero-copy from numpy to MXNet NDArray
* should work
* make lint happy
* fix stupid typos
* wip to address comments
* Address comments
* Address comments
* Remove redundant code
---
include/mxnet/c_api.h | 1 +
python/mxnet/ndarray/ndarray.py | 107 +++++++++++++++++++++++++++++++++-
src/ndarray/ndarray.cc | 14 ++---
tests/python/unittest/test_ndarray.py | 31 ++++++++++
4 files changed, 145 insertions(+), 8 deletions(-)
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
index 21018e3..d3c6794 100644
--- a/include/mxnet/c_api.h
+++ b/include/mxnet/c_api.h
@@ -823,6 +823,7 @@ MXNET_DLL int MXNDArrayToDLPack(NDArrayHandle handle,
*/
MXNET_DLL int MXNDArrayFromDLPack(DLManagedTensorHandle dlpack,
NDArrayHandle *out_handle);
+
/*!
* \brief Delete a dlpack tensor
* \param dlpack the pointer of the input DLManagedTensor
diff --git a/python/mxnet/ndarray/ndarray.py b/python/mxnet/ndarray/ndarray.py
index 97cfd82..d912d38 100644
--- a/python/mxnet/ndarray/ndarray.py
+++ b/python/mxnet/ndarray/ndarray.py
@@ -47,7 +47,7 @@ __all__ = ["NDArray", "concatenate", "_DTYPE_NP_TO_MX",
"_DTYPE_MX_TO_NP", "_GRA
"imdecode", "lesser", "lesser_equal", "logical_and", "logical_or",
"logical_xor",
"maximum", "minimum", "moveaxis", "modulo", "multiply",
"not_equal", "onehot_encode",
"power", "subtract", "true_divide", "waitall", "_new_empty_handle",
"histogram",
- "split_v2", "to_dlpack_for_read", "to_dlpack_for_write",
"from_dlpack"]
+ "split_v2", "to_dlpack_for_read", "to_dlpack_for_write",
"from_dlpack", "from_numpy"]
_STORAGE_TYPE_UNDEFINED = -1
_STORAGE_TYPE_DEFAULT = 0
@@ -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)
+ address = ctypes.addressof(c_obj)
+ address = ctypes.cast(address, ctypes.c_void_p)
+ handle = NDArrayHandle()
+ check_call(_LIB.MXNDArrayFromDLPack(address, ctypes.byref(handle)))
+ return NDArray(handle=handle)
diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
index eddfbcf..0bfca8c 100644
--- a/src/ndarray/ndarray.cc
+++ b/src/ndarray/ndarray.cc
@@ -339,8 +339,8 @@ NDArray NDArray::data_ndarray() const {
}
struct NDArrayDLManager {
- NDArray handle; // ref NDArray
- DLManagedTensor tensor;
+ NDArray handle; // ref NDArray
+ DLManagedTensor tensor;
};
DLManagedTensor* NDArray::ToDLPack() const {
@@ -356,13 +356,13 @@ DLManagedTensor* NDArray::ToDLPack() const {
}
NDArray NDArray::FromDLPack(const DLManagedTensor* tensor) {
- const DLTensor &dl_tensor = tensor->dl_tensor;
- auto deleter = [tensor](){
- if (tensor->deleter != nullptr) {
- tensor->deleter(const_cast<DLManagedTensor*>(tensor));
+ DLManagedTensor tensor_copy = *tensor;
+ auto deleter = [tensor_copy](){
+ if (tensor_copy.deleter != nullptr) {
+ tensor_copy.deleter(const_cast<DLManagedTensor*>(&tensor_copy));
}
};
- return NDArray(TBlob(dl_tensor), 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 {
diff --git a/tests/python/unittest/test_ndarray.py
b/tests/python/unittest/test_ndarray.py
index 9477767..3740506 100644
--- a/tests/python/unittest/test_ndarray.py
+++ b/tests/python/unittest/test_ndarray.py
@@ -1653,6 +1653,37 @@ def test_ndarray_nan_comparison():
for i in (np.isnan(data1_grad))[1][0].flatten():
assert i == True
+
+def test_zero_from_numpy():
+ # Test zero_copy
+ arrays = [
+ # ordinary numpy array
+ np.array([[1, 2], [3, 4], [5, 6]], dtype="float32"),
+ # 0-dim
+ np.array((1, )).reshape(()),
+ # 0-size
+ np.array(()).reshape((1, 0, 2)),
+ ]
+ for zero_copy in [False, True]:
+ for np_array in arrays:
+ mx_array = mx.nd.from_numpy(np_array, zero_copy=zero_copy)
+ mx.test_utils.assert_almost_equal(np_array, mx_array.asnumpy())
+ np_array = arrays[0]
+ mx_array = mx.nd.from_numpy(np_array)
+ np_array[2, 1] = 0
+ mx.test_utils.assert_almost_equal(np_array, mx_array.asnumpy())
+ mx_array[2, 1] = 100
+ mx.test_utils.assert_almost_equal(np_array, mx_array.asnumpy())
+ np_array = np.array([[1, 2], [3, 4], [5, 6]]).transpose()
+ assert not np_array.flags["C_CONTIGUOUS"]
+ try:
+ mx_array = mx.nd.from_numpy(np_array)
+ except ValueError:
+ pass
+ else:
+ assert False
+
+
if __name__ == '__main__':
import nose
nose.runmodule()
