[MXNET-779]Add DLPack Transformation API

include/mxnet/c_api.h

@@ -93,6 +93,10 @@ typedef void *CudaModuleHandle;
 typedef void *CudaKernelHandle;
 /*! \brief handle to a Profile object (domain, duration, counter, etc.) */
 typedef void *ProfileHandle;
+/*! \brief handle to DLManagedTensor*/
+typedef void *DLManagedTensorHandle;
+/*! \brief handle to PyObject*/
+typedef void *PyObjectHandle;
 
 typedef void (*ExecutorMonitorCallback)(const char*,
                                         NDArrayHandle,
@@ -737,6 +741,47 @@ MXNET_DLL int MXNDArrayGetShape(NDArrayHandle handle,
  */
 MXNET_DLL int MXNDArrayGetData(NDArrayHandle handle,
                                void **out_pdata);
+/*!
+* \brief Create a reference view of NDArray that
+*  represents as DLManagedTensor
+*  Notice: MXNet uses asynchronous execution. Please call MXNDArrayWaitToRead or
+*          MXNDArrayWaitToWrite before calling MXNDArrayToDLPack.
+* \param handle the handle to the ndarray
+* \param out_dlpack pointer holder to get pointer of DLManagedTensor
+* \return 0 when success, -1 when failure happens
+*/
+MXNET_DLL int MXNDArrayToDLPack(NDArrayHandle handle,
+                                       DLManagedTensorHandle *out_dlpack);
+
+/*!
+* \brief Create a NDArray backed by a dlpack tensor.
+*
+* This allows us to create a NDArray using the memory
+* allocated by an external deep learning framework
+* that is DLPack compatible.
+*
+* The memory is retained until the NDArray went out of scope.
+*
+* \param dlpack the pointer of the input DLManagedTensor
+* \param out_handle pointer holder to get pointer of NDArray
+* \return 0 when success, -1 when failure happens
+*/
+MXNET_DLL int MXNDArrayFromDLPack(DLManagedTensorHandle dlpack,
+                                  NDArrayHandle *out_handle);
+/*!
+ * \brief Delete a dlpack tensor
+ * \param dlpack the pointer of the input DLManagedTensor
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXNDArrayCallDLPackDeleter(DLManagedTensorHandle dlpack);
+
+/*!
+ * \brief Delete a dlpack tensor
+ * \param dlpack_capsule the pointer of a PyCapsule storing DLManagedTensor
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL void MXNDArrayCallDLPackCapsuleDeleter(PyObjectHandle dlpack_capsule);
+
 /*!
  * \brief get the type of the data in NDArray
  * \param handle the handle to the narray

include/mxnet/ndarray.h

@@ -519,6 +519,26 @@ class NDArray {
     return ret;
   }
 
+  /*!
+   * \brief Create a reference view of NDArray that
+   *  represents as DLManagedTensor.
+   * \return A DLManagedTensor
+   */
+  DLManagedTensor* ToDLPack() const;
+
+  /*!
+   * \brief Create a NDArray backed by a dlpack tensor.
+   *
+   * This allows us to create a NDArray using the memory
+   * allocated by an external deep learning framework
+   * that is DLPack compatible.
+   *
+   * The memory is retained until the NDArray went out of scope.
+   *
+   * \return The created NDArray view.
+   */
+  static NDArray FromDLPack(DLManagedTensor* tensor);
+
   /*!
    * \brief Update ndarray chunk storage handles using existing ndarray storage handles
    * Also update the aux_handle, aux_shapes and aux_types.

include/mxnet/tensor_blob.h

@@ -104,6 +104,37 @@ class TBlob {
       : dptr_(dptr), shape_(shape), type_flag_(type_flag) {
     SetDLTensor(dev_mask, dev_id);
   }
+  /*!
+   * \brief constructor that construct TBlob from DLTensor
+   * \param DLTensor Object
+   */
+  explicit TBlob(const DLTensor &dltensor) : dptr_(dltensor.data),
+    shape_(TShape(dltensor.shape, dltensor.shape + dltensor.ndim)),
+    type_flag_(DLDataTypeTransform(dltensor.dtype)), dltensor_(dltensor) {
+    // compactness check for DLTensor
+    if (dltensor.strides != nullptr) {
+      // check strides
+      const int &ndim = dltensor.ndim;
+      const int64_t *shape = dltensor.shape;
+      const int64_t *strides = dltensor.strides;
+      if (ndim >= 1) {
+        bool err = false;
+        if (strides[ndim - 1] != 1) {
+          err = true;
+        } else {
+          for (int i = ndim - 2; i >= 0; --i) {
+            if (strides[i] != shape[i + 1] * strides[i + 1]) {
+              err = true;
+              break;
+            }
+          }
+        }
+        if (err) {
+          LOG(FATAL) << "Unsupported DLPack because MXNet only support compact tensor now";
+        }
+      }
+    }
+  }
   /*!
    * \brief constructor from tensor
    * \param src source tensor
@@ -336,14 +367,51 @@ class TBlob {
       }
     }
   }
+  static int DLDataTypeTransform(DLDataType dldata_type) {
+    if (dldata_type.lanes != 1) {
+      LOG(FATAL) << "Unsupported DLDataType whose lanes != 1";
+    }
+    switch (dldata_type.code) {
+      case kDLFloat:
+        switch (dldata_type.bits) {
+          case 16:
+            return mshadow::kFloat16;
+          case 32:
+            return mshadow::kFloat32;
+          case 64:
+            return mshadow::kFloat64;
+        }
+        break;
+      case kDLUInt:
+        switch (dldata_type.bits) {
+          case 8:
+            return mshadow::kUint8;
+        }
+        break;
+      case kDLInt:
+        switch (dldata_type.bits) {
+          case 8:
+            return mshadow::kInt8;
+          case 32:
+            return mshadow::kInt32;
+          case 64:
+            return mshadow::kInt64;
+        }
+        break;
+    }
+    LOG(FATAL) << "Unknown DLDataType{" << dldata_type.code
+               << ", " << dldata_type.bits
+               << ", " << dldata_type.lanes << "}";
+    return mshadow::kFloat32;
+  }
 
   inline void SetDLTensor(int dev_mask, int dev_id) {
     dltensor_.data = dptr_;
     dltensor_.ctx = DLContext{static_cast<DLDeviceType>(dev_mask), dev_id};
     dltensor_.ndim = shape_.ndim();
     dltensor_.dtype = DTypeTransform(type_flag_);
     dltensor_.shape = shape_.data();
-    dltensor_.strides = NULL;
+    dltensor_.strides = nullptr;
     dltensor_.byte_offset = 0;
   }
 

python/mxnet/_ctypes/ndarray.py

@@ -31,21 +31,24 @@
 
 class NDArrayBase(object):
     """Base data structure for ndarray"""
-    __slots__ = ["handle", "writable"]
+    __slots__ = ["handle", "writable", "dlpack"]
     # pylint: disable= no-member
 
-    def __init__(self, handle, writable=True):
+    def __init__(self, handle, writable=True, dlpack=None):
         """initialize a new NDArray
 
         Parameters
         ----------
         handle : NDArrayHandle
             NDArray handle of C API
+        dlpack : PyCapsule (DLPack)
+            DLPack Object
         """
         if handle is not None:
             assert isinstance(handle, NDArrayHandle)
         self.handle = handle
         self.writable = writable
+        self.dlpack = dlpack
 
     def __del__(self):
         check_call(_LIB.MXNDArrayFree(self.handle))

python/mxnet/base.py

@@ -235,6 +235,7 @@ def _load_lib():
 CudaModuleHandle = ctypes.c_void_p
 CudaKernelHandle = ctypes.c_void_p
 ProfileHandle = ctypes.c_void_p
+DLPackHandle = ctypes.c_void_p
 
 
 #----------------------------
@@ -729,3 +730,16 @@ def write_all_str(module_file, module_all_list):
     module_op_file.close()
     write_all_str(module_internal_file, module_internal_all)
     module_internal_file.close()
+
+ctypes.pythonapi.PyCapsule_New.restype = ctypes.py_object
+ctypes.pythonapi.PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p,
+                                           ctypes.c_void_p]
+
+ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p
+ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
+
+ctypes.pythonapi.PyCapsule_SetName.restype = ctypes.c_int
+ctypes.pythonapi.PyCapsule_SetName.argtypes = [ctypes.py_object, ctypes.c_char_p]
+
+_LIB.MXNDArrayCallDLPackCapsuleDeleter.restype = None
+_LIB.MXNDArrayCallDLPackCapsuleDeleter.argtypes = [ctypes.c_void_p]

python/mxnet/cython/ndarray.pyx

@@ -30,6 +30,7 @@ cdef class NDArrayBase:
     # handle for symbolic operator.
     cdef NDArrayHandle chandle
     cdef int cwritable
+    cdef object dlpack
 
     cdef _set_handle(self, handle):
         cdef unsigned long long ptr
@@ -52,9 +53,10 @@ cdef class NDArrayBase:
         def __get__(self):
             return bool(self.cwritable)
 
-    def __init__(self, handle, writable=True):
+    def __init__(self, handle, writable=True, dlpack=None):
         self._set_handle(handle)
         self.cwritable = writable
+        self.dlpack = dlpack
 
     def __dealloc__(self):
         CALL(MXNDArrayFree(self.chandle))