diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h index 0acfde0686d4..772c40068dbe 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 97cfd827c7fe..d912d38930a5 100644 --- a/python/mxnet/ndarray/ndarray.py +++ b/python/mxnet/ndarray/ndarray.py @@ -47,7 +47,7 @@ "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 eddfbcff9ce8..0bfca8c10a1a 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(tensor)); + DLManagedTensor tensor_copy = *tensor; + auto deleter = [tensor_copy](){ + if (tensor_copy.deleter != nullptr) { + tensor_copy.deleter(const_cast(&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 94777677354d..374050668612 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()