MAINT, BUG: Documentation for DLPack protocol and refcounting bug fixes.

6 files changed, 58 insertions, 25 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index a162a637c..c808f0baf 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -1413,6 +1413,7 @@ _SupportsBuffer = Union[
 
 _T = TypeVar("_T")
 _T_co = TypeVar("_T_co", covariant=True)
+_T_contra = TypeVar("_T_contra", contravariant=True)
 _2Tuple = Tuple[_T, _T]
 _CastingKind = L["no", "equiv", "safe", "same_kind", "unsafe"]
 
@@ -4330,3 +4331,9 @@ class chararray(ndarray[_ShapeType, _CharDType]):
 
 # NOTE: Deprecated
 # class MachAr: ...
+
+class _SupportsDLPack(Protocol[_T_contra]):
+    def __dlpack__(self, *, stream: Optional[int] = ...) -> _PyCapsule: ...
+
+def from_dlpack(__obj: _SupportsDLPack[None]) -> NDArray[Any]: ...
+
diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py
index a4c588a3b..f1a42dffe 100644
--- a/numpy/core/_add_newdocs.py
+++ b/numpy/core/_add_newdocs.py
@@ -1575,7 +1575,15 @@ add_newdoc('numpy.core.multiarray', 'frombuffer',
 
 add_newdoc('numpy.core.multiarray', 'from_dlpack',
     """
-    Create a NumPy array from a DLPack struct.
+    from_dlpack(x, /)
+
+    Create a NumPy array from an object implementing the ``__dlpack__``
+    protocol.
+
+    See Also
+    --------
+    `Array API documentation
+    <https://data-apis.org/array-api/latest/design_topics/data_interchange.html#syntax-for-data-interchange-with-dlpack>`_
     """)
 
 add_newdoc('numpy.core', 'fastCopyAndTranspose',
@@ -2268,6 +2276,15 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('__array_priority__',
 add_newdoc('numpy.core.multiarray', 'ndarray', ('__array_struct__',
     """Array protocol: C-struct side."""))
 
+add_newdoc('numpy.core.multiarray', 'ndarray', ('__dlpack__',
+    """a.__dlpack__(*, stream=None)
+    
+    DLPack Protocol: Part of the Array API."""))
+
+add_newdoc('numpy.core.multiarray', 'ndarray', ('__dlpack_device__',
+    """a.__dlpack_device__()
+    
+    DLPack Protocol: Part of the Array API."""))
 
 add_newdoc('numpy.core.multiarray', 'ndarray', ('base',
     """
diff --git a/numpy/core/src/common/dlpack/dlpack.h b/numpy/core/src/common/dlpack/dlpack.h
index 84afca248..8b19ea2b1 100644
--- a/numpy/core/src/common/dlpack/dlpack.h
+++ b/numpy/core/src/common/dlpack/dlpack.h
@@ -55,11 +55,19 @@ typedef enum {
   /*! \brief ROCm GPUs for AMD GPUs */
   kDLROCM = 10,
   /*!
+   * \brief Pinned ROCm CPU memory allocated by hipMallocHost
+   */
+  kDLROCMHost = 11,
+  /*!
    * \brief Reserved extension device type,
    * used for quickly test extension device
    * The semantics can differ depending on the implementation.
    */
   kDLExtDev = 12,
+  /*!
+   * \brief CUDA managed/unified memory allocated by cudaMallocManaged
+   */
+  kDLCUDAManaged = 13,
 } DLDeviceType;
 
 /*!
@@ -68,7 +76,10 @@ typedef enum {
 typedef struct {
   /*! \brief The device type used in the device. */
   DLDeviceType device_type;
-  /*! \brief The device index */
+  /*!
+   * \brief The device index.
+   * For vanilla CPU memory, pinned memory, or managed memory, this is set to 0.
+   */
   int device_id;
 } DLDevice;
 
diff --git a/numpy/core/src/common/npy_dlpack.h b/numpy/core/src/common/npy_dlpack.h
index 407469058..47559191e 100644
--- a/numpy/core/src/common/npy_dlpack.h
+++ b/numpy/core/src/common/npy_dlpack.h
@@ -4,8 +4,13 @@
 #ifndef NPY_DLPACK_H
 #define NPY_DLPACK_H
 
+// Part of the Array API specification.
 #define NPY_DLPACK_CAPSULE_NAME "dltensor"
 #define NPY_DLPACK_USED_CAPSULE_NAME "used_dltensor"
+
+// Used internally by NumPy to store a base object
+// as it has to release a reference to the original
+// capsule.
 #define NPY_DLPACK_INTERNAL_CAPSULE_NAME "numpy_dltensor"
 
 static void array_dlpack_capsule_deleter(PyObject *self)
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 7414cb9a8..e348a36cb 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -4233,19 +4233,11 @@ _reload_guard(PyObject *NPY_UNUSED(self)) {
     Py_RETURN_NONE;
 }
 
-static void array_dlpack_deleter(DLManagedTensor *self)
-{
-    PyArrayObject *array = (PyArrayObject *)self->manager_ctx;
-    // This will also free the strides as it's one allocation.
-    PyMem_Free(self->dl_tensor.shape);
-    PyMem_Free(self);
-    Py_XDECREF(array);
-}
-
 
 NPY_NO_EXPORT PyObject *
 from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) {
-    PyObject *capsule = PyObject_CallMethod(obj, "__dlpack__", NULL);
+    PyObject *capsule = PyObject_CallMethod((PyObject *)obj->ob_type,
+            "__dlpack__", "O", obj);
     if (capsule == NULL) {
         return NULL;
     }
@@ -4260,7 +4252,7 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) {
     }
 
     const int ndim = managed->dl_tensor.ndim;
-    if (ndim >= NPY_MAXDIMS) {
+    if (ndim > NPY_MAXDIMS) {
         PyErr_SetString(PyExc_RuntimeError,
                 "maxdims of DLPack tensor is higher than the supported "
                 "maxdims.");
@@ -4268,8 +4260,11 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) {
         return NULL;
     }
 
-    if (managed->dl_tensor.device.device_type != kDLCPU &&
-            managed->dl_tensor.device.device_type != kDLCUDAHost) {
+    DLDeviceType device_type = managed->dl_tensor.device.device_type;
+    if (device_type != kDLCPU &&
+            device_type != kDLCUDAHost &&
+            device_type != kDLROCMHost &&
+            device_type != kDLCUDAManaged) {
         PyErr_SetString(PyExc_RuntimeError,
                 "Unsupported device in DLTensor.");
         Py_XDECREF(capsule);
@@ -4340,6 +4335,7 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) {
 
     for (int i = 0; i < ndim; ++i) {
         shape[i] = managed->dl_tensor.shape[i];
+        // DLPack has elements as stride units, NumPy has bytes.
         strides[i] = managed->dl_tensor.strides[i] * itemsize;
     }
 
@@ -4357,25 +4353,20 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) {
     PyObject *new_capsule = PyCapsule_New(managed,
             NPY_DLPACK_INTERNAL_CAPSULE_NAME, array_dlpack_capsule_deleter);
     if (new_capsule == NULL) {
-        Py_XDECREF(descr);
-        Py_XDECREF(ret);
         Py_XDECREF(capsule);
+        Py_XDECREF(ret);
         return NULL;
     }
 
     if (PyArray_SetBaseObject((PyArrayObject *)ret, new_capsule) < 0) {
-        Py_XDECREF(descr);
-        Py_XDECREF(ret);
-        Py_XDECREF(new_capsule);
         Py_XDECREF(capsule);
+        Py_XDECREF(ret);
         return NULL;
     }
 
     if (PyCapsule_SetName(capsule, NPY_DLPACK_USED_CAPSULE_NAME) < 0) {
-        Py_XDECREF(descr);
-        Py_XDECREF(ret);
-        Py_XDECREF(new_capsule);
         Py_XDECREF(capsule);
+        Py_XDECREF(ret);
         return NULL;
     }
 
diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py
index 72fbab0de..9f5cf9192 100644
--- a/numpy/core/tests/test_dlpack.py
+++ b/numpy/core/tests/test_dlpack.py
@@ -2,10 +2,11 @@ import sys
 import pytest
 
 import numpy as np
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_array_equal, IS_PYPY
 
 
 class TestDLPack:
+    @pytest.mark.skipif(IS_PYPY, reason="PyPy can't get refcounts.")
     def test_dunder_dlpack_refcount(self):
         x = np.arange(5)
         y = x.__dlpack__()
@@ -24,6 +25,7 @@ class TestDLPack:
         with pytest.raises(RuntimeError):
             np.from_dlpack(z)
 
+    @pytest.mark.skipif(IS_PYPY, reason="PyPy can't get refcounts.")
     def test_from_dlpack_refcount(self):
         x = np.arange(5)
         y = np.from_dlpack(x)
@@ -35,7 +37,7 @@ class TestDLPack:
         np.int8, np.int16, np.int32, np.int64,
         np.uint8, np.uint16, np.uint32, np.uint64,
         np.float16, np.float32, np.float64,
-        np.complex64, np.complex64
+        np.complex64, np.complex128
     ])
     def test_dtype_passthrough(self, dtype):
         x = np.arange(5, dtype=dtype)