From a888ae827d34a47eb8fab7a01353175ab39bba76 Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Mon, 24 May 2021 15:22:01 +0200 Subject: [PATCH 01/18] ENH: Add the __dlpack__ and __dlpack_device__ methods to ndarray. --- numpy/__init__.pyi | 10 ++ numpy/core/src/common/dlpack/dlpack.h | 188 ++++++++++++++++++++++++++ numpy/core/src/multiarray/methods.c | 163 ++++++++++++++++++++++ 3 files changed, 361 insertions(+) create mode 100644 numpy/core/src/common/dlpack/dlpack.h diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 637f8578a3a5..a162a637c528 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -1432,6 +1432,10 @@ _ArrayTD64_co = NDArray[Union[bool_, integer[Any], timedelta64]] # Introduce an alias for `dtype` to avoid naming conflicts. _dtype = dtype +# `builtins.PyCapsule` unfortunately lacks annotations as of the moment; +# use `Any` as a stopgap measure +_PyCapsule = Any + class _SupportsItem(Protocol[_T_co]): def item(self, args: Any, /) -> _T_co: ... @@ -2439,6 +2443,12 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]): def __ior__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... @overload def __ior__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + @overload + def __ior__(self: NDArray[_ScalarType], other: _RecursiveSequence) -> NDArray[_ScalarType]: ... + @overload + def __dlpack__(self: NDArray[number[Any]], *, stream: None = ...) -> _PyCapsule: ... + @overload + def __dlpack_device__(self) -> Tuple[L[1], L[0]]: ... # Keep `dtype` at the bottom to avoid name conflicts with `np.dtype` @property diff --git a/numpy/core/src/common/dlpack/dlpack.h b/numpy/core/src/common/dlpack/dlpack.h new file mode 100644 index 000000000000..84afca248292 --- /dev/null +++ b/numpy/core/src/common/dlpack/dlpack.h @@ -0,0 +1,188 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file dlpack.h + * \brief The common header of DLPack. + */ +#ifndef DLPACK_DLPACK_H_ +#define DLPACK_DLPACK_H_ + +#ifdef __cplusplus +#define DLPACK_EXTERN_C extern "C" +#else +#define DLPACK_EXTERN_C +#endif + +/*! \brief The current version of dlpack */ +#define DLPACK_VERSION 050 + +/*! \brief DLPACK_DLL prefix for windows */ +#ifdef _WIN32 +#ifdef DLPACK_EXPORTS +#define DLPACK_DLL __declspec(dllexport) +#else +#define DLPACK_DLL __declspec(dllimport) +#endif +#else +#define DLPACK_DLL +#endif + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif +/*! + * \brief The device type in DLDevice. + */ +typedef enum { + /*! \brief CPU device */ + kDLCPU = 1, + /*! \brief CUDA GPU device */ + kDLCUDA = 2, + /*! + * \brief Pinned CUDA CPU memory by cudaMallocHost + */ + kDLCUDAHost = 3, + /*! \brief OpenCL devices. */ + kDLOpenCL = 4, + /*! \brief Vulkan buffer for next generation graphics. */ + kDLVulkan = 7, + /*! \brief Metal for Apple GPU. */ + kDLMetal = 8, + /*! \brief Verilog simulator buffer */ + kDLVPI = 9, + /*! \brief ROCm GPUs for AMD GPUs */ + kDLROCM = 10, + /*! + * \brief Reserved extension device type, + * used for quickly test extension device + * The semantics can differ depending on the implementation. + */ + kDLExtDev = 12, +} DLDeviceType; + +/*! + * \brief A Device for Tensor and operator. + */ +typedef struct { + /*! \brief The device type used in the device. */ + DLDeviceType device_type; + /*! \brief The device index */ + int device_id; +} DLDevice; + +/*! + * \brief The type code options DLDataType. + */ +typedef enum { + /*! \brief signed integer */ + kDLInt = 0U, + /*! \brief unsigned integer */ + kDLUInt = 1U, + /*! \brief IEEE floating point */ + kDLFloat = 2U, + /*! + * \brief Opaque handle type, reserved for testing purposes. + * Frameworks need to agree on the handle data type for the exchange to be well-defined. + */ + kDLOpaqueHandle = 3U, + /*! \brief bfloat16 */ + kDLBfloat = 4U, + /*! + * \brief complex number + * (C/C++/Python layout: compact struct per complex number) + */ + kDLComplex = 5U, +} DLDataTypeCode; + +/*! + * \brief The data type the tensor can hold. + * + * Examples + * - float: type_code = 2, bits = 32, lanes=1 + * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes=4 + * - int8: type_code = 0, bits = 8, lanes=1 + * - std::complex: type_code = 5, bits = 64, lanes = 1 + */ +typedef struct { + /*! + * \brief Type code of base types. + * We keep it uint8_t instead of DLDataTypeCode for minimal memory + * footprint, but the value should be one of DLDataTypeCode enum values. + * */ + uint8_t code; + /*! + * \brief Number of bits, common choices are 8, 16, 32. + */ + uint8_t bits; + /*! \brief Number of lanes in the type, used for vector types. */ + uint16_t lanes; +} DLDataType; + +/*! + * \brief Plain C Tensor object, does not manage memory. + */ +typedef struct { + /*! + * \brief The opaque data pointer points to the allocated data. This will be + * CUDA device pointer or cl_mem handle in OpenCL. This pointer is always + * aligned to 256 bytes as in CUDA. + * + * For given DLTensor, the size of memory required to store the contents of + * data is calculated as follows: + * + * \code{.c} + * static inline size_t GetDataSize(const DLTensor* t) { + * size_t size = 1; + * for (tvm_index_t i = 0; i < t->ndim; ++i) { + * size *= t->shape[i]; + * } + * size *= (t->dtype.bits * t->dtype.lanes + 7) / 8; + * return size; + * } + * \endcode + */ + void* data; + /*! \brief The device of the tensor */ + DLDevice device; + /*! \brief Number of dimensions */ + int ndim; + /*! \brief The data type of the pointer*/ + DLDataType dtype; + /*! \brief The shape of the tensor */ + int64_t* shape; + /*! + * \brief strides of the tensor (in number of elements, not bytes) + * can be NULL, indicating tensor is compact and row-majored. + */ + int64_t* strides; + /*! \brief The offset in bytes to the beginning pointer to data */ + uint64_t byte_offset; +} DLTensor; + +/*! + * \brief C Tensor object, manage memory of DLTensor. This data structure is + * intended to facilitate the borrowing of DLTensor by another framework. It is + * not meant to transfer the tensor. When the borrowing framework doesn't need + * the tensor, it should call the deleter to notify the host that the resource + * is no longer needed. + */ +typedef struct DLManagedTensor { + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; + /*! \brief the context of the original host framework of DLManagedTensor in + * which DLManagedTensor is used in the framework. It can also be NULL. + */ + void * manager_ctx; + /*! \brief Destructor signature void (*)(void*) - this should be called + * to destruct manager_ctx which holds the DLManagedTensor. It can be NULL + * if there is no way for the caller to provide a reasonable destructor. + * The destructors deletes the argument self as well. + */ + void (*deleter)(struct DLManagedTensor * self); +} DLManagedTensor; +#ifdef __cplusplus +} // DLPACK_EXTERN_C +#endif +#endif // DLPACK_DLPACK_H_ diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 2d66c77dc4c1..2251d4b69d89 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -31,6 +31,7 @@ #include "alloc.h" #include +#include "common/dlpack/dlpack.h" /* NpyArg_ParseKeywords @@ -2762,6 +2763,158 @@ array_class_getitem(PyObject *cls, PyObject *args) generic_alias = NULL; #endif return generic_alias; + +#define NPY_DLPACK_CAPSULE_NAME "dltensor" +#define NPY_DLPACK_USED_CAPSULE_NAME "used_dltensor" + +static void array_dlpack_capsule_deleter(PyObject *self) +{ + if (!PyCapsule_IsValid(self, NPY_DLPACK_CAPSULE_NAME)) { + if (!PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { + PyErr_SetString(PyExc_RuntimeError, "Invalid capsule name."); + } + return; + } + + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_CAPSULE_NAME); + managed->deleter(managed); +} + +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); + + PyArray_XDECREF(array); +} + +static PyObject * +array_dlpack(PyArrayObject *self, + PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) +{ + PyObject *stream = Py_None; + NPY_PREPARE_ARGPARSER; + if (npy_parse_arguments("__dlpack__", args, len_args, kwnames, + "$stream", NULL, &stream, NULL, NULL, NULL)) { + return NULL; + } + + if (stream != Py_None) { + PyErr_SetString(PyExc_RuntimeError, "NumPy only supports " + "stream=None."); + return NULL; + } + + npy_intp itemsize = PyArray_ITEMSIZE(self); + int ndim = PyArray_NDIM(self); + npy_intp *strides = PyArray_STRIDES(self); + npy_intp *shape = PyArray_SHAPE(self); + + for (int i = 0; i < ndim; ++i) { + if (strides[i] % itemsize != 0) { + PyErr_SetString(PyExc_RuntimeError, + "DLPack only supports strides which are a multiple of " + "itemsize."); + return NULL; + } + } + + DLDataType managed_dtype; + PyArray_Descr *dtype = PyArray_DESCR(self); + + if (PyDataType_ISBYTESWAPPED(dtype)) { + PyErr_SetString(PyExc_TypeError, "DLPack only supports native " + "byte swapping."); + return NULL; + } + + managed_dtype.bits = 8 * itemsize; + managed_dtype.lanes = 1; + + if (PyDataType_ISSIGNED(dtype)) { + managed_dtype.code = kDLInt; + } else if (PyDataType_ISUNSIGNED(dtype)) { + managed_dtype.code = kDLUInt; + } else if (PyDataType_ISFLOAT(dtype)) { + // We can't be sure that the dtype is + // IEEE or padded. + if (itemsize > 8) { + PyErr_SetString(PyExc_TypeError, "DLPack only supports IEEE " + "floating point types without padding."); + return NULL; + } + managed_dtype.code = kDLFloat; + } else if (PyDataType_ISCOMPLEX(dtype)) { + // We can't be sure that the dtype is + // IEEE or padded. + if (itemsize > 16) { + PyErr_SetString(PyExc_TypeError, "DLPack only supports IEEE " + "complex point types without padding."); + return NULL; + } + managed_dtype.code = kDLComplex; + } else { + PyErr_SetString(PyExc_TypeError, + "DLPack only supports signed/unsigned integers, float " + "and complex dtypes."); + return NULL; + } + + DLManagedTensor *managed = PyMem_Malloc(sizeof(DLManagedTensor)); + if (managed == NULL) { + PyErr_NoMemory(); + return NULL; + } + + managed->dl_tensor.data = PyArray_DATA(self); + managed->dl_tensor.device.device_type = kDLCPU; + managed->dl_tensor.device.device_id = 0; + managed->dl_tensor.dtype = managed_dtype; + + + int64_t *managed_shape_strides = PyMem_Malloc(sizeof(int64_t) * ndim * 2); + if (managed_shape_strides == NULL) { + PyErr_NoMemory(); + PyMem_Free(managed); + return NULL; + } + + int64_t *managed_shape = managed_shape_strides; + int64_t *managed_strides = managed_shape_strides + ndim; + for (int i = 0; i < ndim; ++i) { + managed_shape[i] = shape[i]; + // Strides in DLPack are items; in NumPy are bytes. + managed_strides[i] = strides[i] / itemsize; + } + + managed->dl_tensor.ndim = ndim; + managed->dl_tensor.shape = managed_shape; + managed->dl_tensor.strides = managed_strides; + managed->dl_tensor.byte_offset = 0; + managed->manager_ctx = self; + managed->deleter = array_dlpack_deleter; + + PyObject *capsule = PyCapsule_New(managed, NPY_DLPACK_CAPSULE_NAME, + array_dlpack_capsule_deleter); + if (capsule == NULL) { + PyMem_Free(managed); + PyMem_Free(managed_shape_strides); + return NULL; + } + + // the capsule holds a reference + PyArray_INCREF(self); + return capsule; +} + +static PyObject * +array_dlpack_device(PyArrayObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) +{ + return Py_BuildValue("ii", kDLCPU, 0); +>>>>>>> ENH: Add the __dlpack__ and __dlpack_device__ methods to ndarray. } NPY_NO_EXPORT PyMethodDef array_methods[] = { @@ -2989,5 +3142,15 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = { {"view", (PyCFunction)array_view, METH_FASTCALL | METH_KEYWORDS, NULL}, + + // For data interchange between libraries + {"__dlpack__", + (PyCFunction)array_dlpack, + METH_FASTCALL | METH_KEYWORDS, NULL}, + + {"__dlpack_device__", + (PyCFunction)array_dlpack_device, + METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL} /* sentinel */ }; From 3163d57bd1550591e5f6ff1b063d4423d8d2123e Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Thu, 27 May 2021 12:03:55 +0200 Subject: [PATCH 02/18] ENH, TST: Add the from_dlpack method and test DLPack. --- numpy/core/_add_newdocs.py | 5 + numpy/core/multiarray.py | 9 +- numpy/core/numeric.py | 6 +- numpy/core/setup.py | 2 + numpy/core/src/common/npy_dlpack.h | 25 +++ numpy/core/src/multiarray/methods.c | 25 ++- numpy/core/src/multiarray/multiarraymodule.c | 153 +++++++++++++++++++ numpy/core/tests/test_dlpack.py | 76 +++++++++ 8 files changed, 287 insertions(+), 14 deletions(-) create mode 100644 numpy/core/src/common/npy_dlpack.h create mode 100644 numpy/core/tests/test_dlpack.py diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py index c8a24db0cf1a..a4c588a3b068 100644 --- a/numpy/core/_add_newdocs.py +++ b/numpy/core/_add_newdocs.py @@ -1573,6 +1573,11 @@ array_function_like_doc, )) +add_newdoc('numpy.core.multiarray', 'from_dlpack', + """ + Create a NumPy array from a DLPack struct. + """) + add_newdoc('numpy.core', 'fastCopyAndTranspose', """_fastCopyAndTranspose(a)""") diff --git a/numpy/core/multiarray.py b/numpy/core/multiarray.py index 351cd3a1bbb5..9f431f01b7c6 100644 --- a/numpy/core/multiarray.py +++ b/numpy/core/multiarray.py @@ -31,10 +31,10 @@ 'count_nonzero', 'c_einsum', 'datetime_as_string', 'datetime_data', 'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype', 'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat', - 'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'get_handler_name', - 'inner', 'interp', 'interp_complex', 'is_busday', 'lexsort', - 'matmul', 'may_share_memory', 'min_scalar_type', 'ndarray', 'nditer', - 'nested_iters', 'normalize_axis_index', 'packbits', + 'frombuffer', 'from_dlpack', 'fromfile', 'fromiter', 'fromstring', + 'get_handler_name', 'inner', 'interp', 'interp_complex', 'is_busday', + 'lexsort', 'matmul', 'may_share_memory', 'min_scalar_type', 'ndarray', + 'nditer', 'nested_iters', 'normalize_axis_index', 'packbits', 'promote_types', 'putmask', 'ravel_multi_index', 'result_type', 'scalar', 'set_datetimeparse_function', 'set_legacy_print_mode', 'set_numeric_ops', 'set_string_function', 'set_typeDict', 'shares_memory', @@ -55,6 +55,7 @@ datetime_data.__module__ = 'numpy' empty.__module__ = 'numpy' frombuffer.__module__ = 'numpy' +from_dlpack.__module__ = 'numpy' fromfile.__module__ = 'numpy' fromiter.__module__ = 'numpy' frompyfunc.__module__ = 'numpy' diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index 1654e83646d9..987470f9239e 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -13,8 +13,8 @@ WRAP, arange, array, asarray, asanyarray, ascontiguousarray, asfortranarray, broadcast, can_cast, compare_chararrays, concatenate, copyto, dot, dtype, empty, - empty_like, flatiter, frombuffer, fromfile, fromiter, fromstring, - inner, lexsort, matmul, may_share_memory, + empty_like, flatiter, frombuffer, from_dlpack, fromfile, fromiter, + fromstring, inner, lexsort, matmul, may_share_memory, min_scalar_type, ndarray, nditer, nested_iters, promote_types, putmask, result_type, set_numeric_ops, shares_memory, vdot, where, zeros, normalize_axis_index) @@ -41,7 +41,7 @@ 'newaxis', 'ndarray', 'flatiter', 'nditer', 'nested_iters', 'ufunc', 'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'zeros', 'count_nonzero', 'empty', 'broadcast', 'dtype', - 'fromstring', 'fromfile', 'frombuffer', 'where', + 'fromstring', 'fromfile', 'frombuffer', 'from_dlpack', 'where', 'argwhere', 'copyto', 'concatenate', 'fastCopyAndTranspose', 'lexsort', 'set_numeric_ops', 'can_cast', 'promote_types', 'min_scalar_type', 'result_type', 'isfortran', 'empty_like', 'zeros_like', 'ones_like', diff --git a/numpy/core/setup.py b/numpy/core/setup.py index 3e1ed4c9bc0c..5c6fd4dcf8d0 100644 --- a/numpy/core/setup.py +++ b/numpy/core/setup.py @@ -740,6 +740,7 @@ def gl_if_msvc(build_cmd): ####################################################################### common_deps = [ + join('src', 'common', 'dlpack', 'dlpack.h'), join('src', 'common', 'array_assign.h'), join('src', 'common', 'binop_override.h'), join('src', 'common', 'cblasfuncs.h'), @@ -749,6 +750,7 @@ def gl_if_msvc(build_cmd): join('src', 'common', 'npy_cblas.h'), join('src', 'common', 'npy_config.h'), join('src', 'common', 'npy_ctypes.h'), + join('src', 'common', 'npy_dlpack.h'), join('src', 'common', 'npy_extint128.h'), join('src', 'common', 'npy_import.h'), join('src', 'common', 'npy_hashtable.h'), diff --git a/numpy/core/src/common/npy_dlpack.h b/numpy/core/src/common/npy_dlpack.h new file mode 100644 index 000000000000..407469058cf1 --- /dev/null +++ b/numpy/core/src/common/npy_dlpack.h @@ -0,0 +1,25 @@ +#include "Python.h" +#include "dlpack/dlpack.h" + +#ifndef NPY_DLPACK_H +#define NPY_DLPACK_H + +#define NPY_DLPACK_CAPSULE_NAME "dltensor" +#define NPY_DLPACK_USED_CAPSULE_NAME "used_dltensor" +#define NPY_DLPACK_INTERNAL_CAPSULE_NAME "numpy_dltensor" + +static void array_dlpack_capsule_deleter(PyObject *self) +{ + if (!PyCapsule_IsValid(self, NPY_DLPACK_CAPSULE_NAME) && + !PyCapsule_IsValid(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { + if (!PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { + PyErr_SetString(PyExc_RuntimeError, "Invalid capsule name."); + } + return; + } + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(self, PyCapsule_GetName(self)); + managed->deleter(managed); +} + +#endif \ No newline at end of file diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 2251d4b69d89..1f1a5d0190bc 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -31,7 +31,9 @@ #include "alloc.h" #include + #include "common/dlpack/dlpack.h" +#include "common/npy_dlpack.h" /* NpyArg_ParseKeywords @@ -2582,8 +2584,6 @@ array_round(PyArrayObject *self, PyObject *args, PyObject *kwds) } } - - static PyObject * array_setflags(PyArrayObject *self, PyObject *args, PyObject *kwds) { @@ -2781,16 +2781,17 @@ static void array_dlpack_capsule_deleter(PyObject *self) managed->deleter(managed); } -static void array_dlpack_deleter(DLManagedTensor *self) +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); - - PyArray_XDECREF(array); + Py_XDECREF(array); } + static PyObject * array_dlpack(PyArrayObject *self, PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) @@ -2906,13 +2907,23 @@ array_dlpack(PyArrayObject *self, } // the capsule holds a reference - PyArray_INCREF(self); + Py_INCREF(self); return capsule; } static PyObject * -array_dlpack_device(PyArrayObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) +array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)) { + PyObject *base = PyArray_BASE(self); + if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { + DLManagedTensor *managed = PyCapsule_GetPointer(base, + NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + return NULL; + } + return Py_BuildValue("ii", managed->dl_tensor.device.device_type, + managed->dl_tensor.device.device_id); + } return Py_BuildValue("ii", kDLCPU, 0); >>>>>>> ENH: Add the __dlpack__ and __dlpack_device__ methods to ndarray. } diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index dea828ed95e0..7414cb9a8d8e 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -70,6 +70,8 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0; #include "get_attr_string.h" #include "experimental_public_dtype_api.h" /* _get_experimental_dtype_api */ +#include "common/npy_dlpack.h" + /* ***************************************************************************** ** INCLUDE GENERATED CODE ** @@ -4231,6 +4233,155 @@ _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); + if (capsule == NULL) { + return NULL; + } + + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(capsule, + NPY_DLPACK_CAPSULE_NAME); + + if (managed == NULL) { + Py_XDECREF(capsule); + return NULL; + } + + const int ndim = managed->dl_tensor.ndim; + if (ndim >= NPY_MAXDIMS) { + PyErr_SetString(PyExc_RuntimeError, + "maxdims of DLPack tensor is higher than the supported " + "maxdims."); + Py_XDECREF(capsule); + return NULL; + } + + if (managed->dl_tensor.device.device_type != kDLCPU && + managed->dl_tensor.device.device_type != kDLCUDAHost) { + PyErr_SetString(PyExc_RuntimeError, + "Unsupported device in DLTensor."); + Py_XDECREF(capsule); + return NULL; + } + + if (managed->dl_tensor.dtype.lanes != 1) { + PyErr_SetString(PyExc_RuntimeError, + "Unsupported lanes in DLTensor dtype."); + Py_XDECREF(capsule); + return NULL; + } + + int typenum = -1; + const uint8_t bits = managed->dl_tensor.dtype.bits; + const npy_intp itemsize = bits / 8; + switch(managed->dl_tensor.dtype.code) { + case kDLInt: + switch (bits) + { + case 8: typenum = NPY_INT8; break; + case 16: typenum = NPY_INT16; break; + case 32: typenum = NPY_INT32; break; + case 64: typenum = NPY_INT64; break; + } + break; + case kDLUInt: + switch (bits) + { + case 8: typenum = NPY_UINT8; break; + case 16: typenum = NPY_UINT16; break; + case 32: typenum = NPY_UINT32; break; + case 64: typenum = NPY_UINT64; break; + } + break; + case kDLFloat: + switch (bits) + { + case 16: typenum = NPY_FLOAT16; break; + case 32: typenum = NPY_FLOAT32; break; + case 64: typenum = NPY_FLOAT64; break; + } + break; + case kDLComplex: + switch (bits) + { + case 64: typenum = NPY_COMPLEX64; break; + case 128: typenum = NPY_COMPLEX128; break; + } + break; + } + + if (typenum == -1) { + PyErr_SetString(PyExc_RuntimeError, + "Unsupported dtype in DLTensor."); + Py_XDECREF(capsule); + return NULL; + } + + PyArray_Descr *descr = PyArray_DescrFromType(typenum); + if (descr == NULL) { + Py_XDECREF(capsule); + return NULL; + } + + npy_intp shape[NPY_MAXDIMS]; + npy_intp strides[NPY_MAXDIMS]; + + for (int i = 0; i < ndim; ++i) { + shape[i] = managed->dl_tensor.shape[i]; + strides[i] = managed->dl_tensor.strides[i] * itemsize; + } + + char *data = (char *)managed->dl_tensor.data + + managed->dl_tensor.byte_offset; + + PyObject *ret = PyArray_NewFromDescr(&PyArray_Type, descr, ndim, shape, + strides, data, 0, NULL); + if (ret == NULL) { + Py_XDECREF(capsule); + Py_XDECREF(descr); + return NULL; + } + + 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); + return NULL; + } + + if (PyArray_SetBaseObject((PyArrayObject *)ret, new_capsule) < 0) { + Py_XDECREF(descr); + Py_XDECREF(ret); + Py_XDECREF(new_capsule); + Py_XDECREF(capsule); + 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); + return NULL; + } + + Py_XDECREF(capsule); + return ret; +} static struct PyMethodDef array_module_methods[] = { {"_get_implementing_args", @@ -4445,6 +4596,8 @@ static struct PyMethodDef array_module_methods[] = { {"_reload_guard", (PyCFunction)_reload_guard, METH_NOARGS, "Give a warning on reload and big warning in sub-interpreters."}, + {"from_dlpack", (PyCFunction)from_dlpack, + METH_O, NULL}, {NULL, NULL, 0, NULL} /* sentinel */ }; diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py new file mode 100644 index 000000000000..72fbab0de250 --- /dev/null +++ b/numpy/core/tests/test_dlpack.py @@ -0,0 +1,76 @@ +import sys +import pytest + +import numpy as np +from numpy.testing import assert_array_equal + + +class TestDLPack: + def test_dunder_dlpack_refcount(self): + x = np.arange(5) + y = x.__dlpack__() + assert sys.getrefcount(x) == 3 + del y + assert sys.getrefcount(x) == 2 + + def test_dunder_dlpack_stream(self): + x = np.arange(5) + x.__dlpack__(stream=None) + + dt = np.dtype([('int', np.int32), ('char', np.int8)]) + y = np.zeros((5,), dtype=dt) + z = y['int'] + + with pytest.raises(RuntimeError): + np.from_dlpack(z) + + def test_from_dlpack_refcount(self): + x = np.arange(5) + y = np.from_dlpack(x) + assert sys.getrefcount(x) == 3 + del y + assert sys.getrefcount(x) == 2 + + @pytest.mark.parametrize("dtype", [ + 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 + ]) + def test_dtype_passthrough(self, dtype): + x = np.arange(5, dtype=dtype) + y = np.from_dlpack(x) + + assert y.dtype == x.dtype + assert_array_equal(x, y) + + def test_invalid_dtype(self): + x = np.asarray(np.datetime64('2021-05-27')) + + with pytest.raises(TypeError): + np.from_dlpack(x) + + def test_invalid_byte_swapping(self): + dt = np.dtype('=i8').newbyteorder() + x = np.arange(5, dtype=dt) + + with pytest.raises(TypeError): + np.from_dlpack(x) + + def test_non_contiguous(self): + x = np.arange(25).reshape((5, 5)) + + y1 = x[0] + assert_array_equal(y1, np.from_dlpack(y1)) + + y2 = x[:, 0] + assert_array_equal(y2, np.from_dlpack(y2)) + + y3 = x[1, :] + assert_array_equal(y3, np.from_dlpack(y3)) + + y4 = x[1] + assert_array_equal(y4, np.from_dlpack(y4)) + + y5 = np.diagonal(x) + assert_array_equal(y5, np.from_dlpack(y5)) From eb6ee260c94e14481325a0d5ef5db89bbd64349f Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Thu, 27 May 2021 17:26:40 +0200 Subject: [PATCH 03/18] MAINT, BUG: Documentation for DLPack protocol and refcounting bug fixes. --- numpy/__init__.pyi | 7 +++++ numpy/core/_add_newdocs.py | 19 ++++++++++- numpy/core/src/common/dlpack/dlpack.h | 13 +++++++- numpy/core/src/common/npy_dlpack.h | 5 +++ numpy/core/src/multiarray/multiarraymodule.c | 33 +++++++------------- numpy/core/tests/test_dlpack.py | 6 ++-- 6 files changed, 58 insertions(+), 25 deletions(-) diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index a162a637c528..c808f0bafadb 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 a4c588a3b068..f1a42dffeb54 100644 --- a/numpy/core/_add_newdocs.py +++ b/numpy/core/_add_newdocs.py @@ -1575,7 +1575,15 @@ 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 + `_ """) add_newdoc('numpy.core', 'fastCopyAndTranspose', @@ -2268,6 +2276,15 @@ 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 84afca248292..8b19ea2b1143 100644 --- a/numpy/core/src/common/dlpack/dlpack.h +++ b/numpy/core/src/common/dlpack/dlpack.h @@ -54,12 +54,20 @@ typedef enum { kDLVPI = 9, /*! \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 407469058cf1..47559191e355 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 7414cb9a8d8e..e348a36cb7ef 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 72fbab0de250..9f5cf91926c8 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -2,10 +2,11 @@ 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 @@ def test_dunder_dlpack_stream(self): 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 @@ def test_from_dlpack_refcount(self): 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) From 3e5b274164ee8b44ea881922019d61d4905b790f Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Thu, 27 May 2021 18:52:43 +0200 Subject: [PATCH 04/18] MAINT: Add URL to DLPack GitHub. --- numpy/core/src/common/dlpack/dlpack.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/numpy/core/src/common/dlpack/dlpack.h b/numpy/core/src/common/dlpack/dlpack.h index 8b19ea2b1143..29209aee12ab 100644 --- a/numpy/core/src/common/dlpack/dlpack.h +++ b/numpy/core/src/common/dlpack/dlpack.h @@ -1,3 +1,5 @@ +// Taken from: +// https://github.com/dmlc/dlpack/blob/9b6176fdecb55e9bf39b16f08b96913ed3f275b4/include/dlpack/dlpack.h /*! * Copyright (c) 2017 by Contributors * \file dlpack.h From e12695d85bc98193014259ef4cf3ce46a750820f Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Thu, 27 May 2021 19:00:52 +0200 Subject: [PATCH 05/18] MAINT: Split up capsule deleter. --- numpy/core/src/common/npy_dlpack.h | 23 +++++++++++++++----- numpy/core/src/multiarray/multiarraymodule.c | 3 ++- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/numpy/core/src/common/npy_dlpack.h b/numpy/core/src/common/npy_dlpack.h index 47559191e355..ef2c7c6fd2d1 100644 --- a/numpy/core/src/common/npy_dlpack.h +++ b/numpy/core/src/common/npy_dlpack.h @@ -13,17 +13,28 @@ // capsule. #define NPY_DLPACK_INTERNAL_CAPSULE_NAME "numpy_dltensor" +/* This is exactly as mandated by dlpack */ static void array_dlpack_capsule_deleter(PyObject *self) { - if (!PyCapsule_IsValid(self, NPY_DLPACK_CAPSULE_NAME) && - !PyCapsule_IsValid(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { - if (!PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { - PyErr_SetString(PyExc_RuntimeError, "Invalid capsule name."); - } + if (PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { return; } DLManagedTensor *managed = - (DLManagedTensor *)PyCapsule_GetPointer(self, PyCapsule_GetName(self)); + (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_CAPSULE_NAME); + if (managed == NULL) { + return; + } + managed->deleter(managed); +} + +/* used internally */ +static void array_dlpack_internal_capsule_deleter(PyObject *self) +{ + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + return; + } managed->deleter(managed); } diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index e348a36cb7ef..ad77177cdcc6 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -4351,7 +4351,8 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { } PyObject *new_capsule = PyCapsule_New(managed, - NPY_DLPACK_INTERNAL_CAPSULE_NAME, array_dlpack_capsule_deleter); + NPY_DLPACK_INTERNAL_CAPSULE_NAME, + array_dlpack_internal_capsule_deleter); if (new_capsule == NULL) { Py_XDECREF(capsule); Py_XDECREF(ret); From 094e0aac120519216a68244d241a5d3bc9497d52 Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Thu, 27 May 2021 19:20:37 +0200 Subject: [PATCH 06/18] MAINT: Move around code so that there are no more unused warnings. --- numpy/core/src/common/npy_dlpack.h | 25 -------------------- numpy/core/src/multiarray/methods.c | 13 ++++++++++ numpy/core/src/multiarray/multiarraymodule.c | 11 +++++++++ 3 files changed, 24 insertions(+), 25 deletions(-) diff --git a/numpy/core/src/common/npy_dlpack.h b/numpy/core/src/common/npy_dlpack.h index ef2c7c6fd2d1..446d44afaf9c 100644 --- a/numpy/core/src/common/npy_dlpack.h +++ b/numpy/core/src/common/npy_dlpack.h @@ -13,29 +13,4 @@ // capsule. #define NPY_DLPACK_INTERNAL_CAPSULE_NAME "numpy_dltensor" -/* This is exactly as mandated by dlpack */ -static void array_dlpack_capsule_deleter(PyObject *self) -{ - if (PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { - return; - } - DLManagedTensor *managed = - (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_CAPSULE_NAME); - if (managed == NULL) { - return; - } - managed->deleter(managed); -} - -/* used internally */ -static void array_dlpack_internal_capsule_deleter(PyObject *self) -{ - DLManagedTensor *managed = - (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME); - if (managed == NULL) { - return; - } - managed->deleter(managed); -} - #endif \ No newline at end of file diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 1f1a5d0190bc..0cc8605141d6 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -2791,6 +2791,19 @@ array_dlpack_deleter(DLManagedTensor *self) Py_XDECREF(array); } +/* This is exactly as mandated by dlpack */ +static void array_dlpack_capsule_deleter(PyObject *self) +{ + if (PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { + return; + } + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_CAPSULE_NAME); + if (managed == NULL) { + return; + } + managed->deleter(managed); +} static PyObject * array_dlpack(PyArrayObject *self, diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index ad77177cdcc6..517234c888f3 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -4233,6 +4233,17 @@ _reload_guard(PyObject *NPY_UNUSED(self)) { Py_RETURN_NONE; } +/* used internally */ +static void array_dlpack_internal_capsule_deleter(PyObject *self) +{ + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + return; + } + managed->deleter(managed); +} + NPY_NO_EXPORT PyObject * from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { From 158c7283b5de5ec7b832407c15fafacc95cb8ee6 Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Fri, 28 May 2021 09:17:08 +0200 Subject: [PATCH 07/18] TST: Improve testing coverage for DLPack. --- numpy/core/tests/test_dlpack.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py index 9f5cf91926c8..4553c943db7f 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -18,6 +18,10 @@ def test_dunder_dlpack_stream(self): x = np.arange(5) x.__dlpack__(stream=None) + with pytest.raises(RuntimeError): + x.__dlpack__(stream=1) + + def test_strides_not_multiple_of_itemsize(self): dt = np.dtype([('int', np.int32), ('char', np.int8)]) y = np.zeros((5,), dtype=dt) z = y['int'] @@ -76,3 +80,10 @@ def test_non_contiguous(self): y5 = np.diagonal(x) assert_array_equal(y5, np.from_dlpack(y5)) + + @pytest.mark.parametrize("ndim", range(33)) + def test_higher_dims(self, ndim): + shape = (1,) * ndim + x = np.zeros(shape, dtype=np.float64, order='C') + + assert shape == np.from_dlpack(x).shape From 9ebee26c6c5cd89623d531608eed25a770d01fff Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Fri, 28 May 2021 18:21:55 +0200 Subject: [PATCH 08/18] BUG: Fix handling of C-contiguous and 1-element arrays. --- numpy/core/src/multiarray/methods.c | 19 ++++++++++++------- numpy/core/src/multiarray/multiarraymodule.c | 7 +++++-- numpy/core/tests/test_dlpack.py | 2 +- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 0cc8605141d6..42464014c10c 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -2827,12 +2827,14 @@ array_dlpack(PyArrayObject *self, npy_intp *strides = PyArray_STRIDES(self); npy_intp *shape = PyArray_SHAPE(self); - for (int i = 0; i < ndim; ++i) { - if (strides[i] % itemsize != 0) { - PyErr_SetString(PyExc_RuntimeError, - "DLPack only supports strides which are a multiple of " - "itemsize."); - return NULL; + if (!PyArray_IS_C_CONTIGUOUS(self) && PyArray_SIZE(self) != 1) { + for (int i = 0; i < ndim; ++i) { + if (strides[i] % itemsize != 0) { + PyErr_SetString(PyExc_RuntimeError, + "DLPack only supports strides which are a multiple of " + "itemsize."); + return NULL; + } } } @@ -2906,7 +2908,10 @@ array_dlpack(PyArrayObject *self, managed->dl_tensor.ndim = ndim; managed->dl_tensor.shape = managed_shape; - managed->dl_tensor.strides = managed_strides; + managed->dl_tensor.strides = NULL; + if (PyArray_SIZE(self) != 1 && !PyArray_IS_C_CONTIGUOUS(self)) { + managed->dl_tensor.strides = managed_strides; + } managed->dl_tensor.byte_offset = 0; managed->manager_ctx = self; managed->deleter = array_dlpack_deleter; diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 517234c888f3..b04debe101ff 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -4347,14 +4347,17 @@ 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; + if (managed->dl_tensor.strides != NULL) + { + strides[i] = managed->dl_tensor.strides[i] * itemsize; + } } char *data = (char *)managed->dl_tensor.data + managed->dl_tensor.byte_offset; PyObject *ret = PyArray_NewFromDescr(&PyArray_Type, descr, ndim, shape, - strides, data, 0, NULL); + managed->dl_tensor.strides != NULL ? strides : NULL, data, 0, NULL); if (ret == NULL) { Py_XDECREF(capsule); Py_XDECREF(descr); diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py index 4553c943db7f..19f4e92815c4 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -84,6 +84,6 @@ def test_non_contiguous(self): @pytest.mark.parametrize("ndim", range(33)) def test_higher_dims(self, ndim): shape = (1,) * ndim - x = np.zeros(shape, dtype=np.float64, order='C') + x = np.zeros(shape, dtype=np.float64) assert shape == np.from_dlpack(x).shape From e83b8d8044763d36c42d3ab103a5437893fb09d8 Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Sun, 30 May 2021 11:57:46 +0200 Subject: [PATCH 09/18] BUG, TST: Device bugfix and test __dl_device__. --- numpy/__init__.pyi | 2 +- numpy/core/src/multiarray/methods.c | 39 +++++++++++++++++++---------- numpy/core/tests/test_dlpack.py | 5 ++++ 3 files changed, 32 insertions(+), 14 deletions(-) diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index c808f0bafadb..63e723a35471 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -2449,7 +2449,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]): @overload def __dlpack__(self: NDArray[number[Any]], *, stream: None = ...) -> _PyCapsule: ... @overload - def __dlpack_device__(self) -> Tuple[L[1], L[0]]: ... + def __dlpack_device__(self) -> Tuple[int, L[0]]: ... # Keep `dtype` at the bottom to avoid name conflicts with `np.dtype` @property diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 42464014c10c..29372fe2f7e2 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -2805,6 +2805,23 @@ static void array_dlpack_capsule_deleter(PyObject *self) managed->deleter(managed); } +static DLDevice +array_get_dl_device(PyArrayObject *self) { + DLDevice ret; + ret.device_type = kDLCPU; + ret.device_id = 0; + PyObject *base = PyArray_BASE(self); + if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { + DLManagedTensor *managed = PyCapsule_GetPointer( + base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + return ret; + } + return managed->dl_tensor.device; + } + return ret; +} + static PyObject * array_dlpack(PyArrayObject *self, PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) @@ -2886,8 +2903,11 @@ array_dlpack(PyArrayObject *self, } managed->dl_tensor.data = PyArray_DATA(self); - managed->dl_tensor.device.device_type = kDLCPU; - managed->dl_tensor.device.device_id = 0; + managed->dl_tensor.device = array_get_dl_device(self); + if (PyErr_Occurred()) { + PyMem_Free(managed); + return NULL; + } managed->dl_tensor.dtype = managed_dtype; @@ -2932,18 +2952,11 @@ array_dlpack(PyArrayObject *self, static PyObject * array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)) { - PyObject *base = PyArray_BASE(self); - if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { - DLManagedTensor *managed = PyCapsule_GetPointer(base, - NPY_DLPACK_INTERNAL_CAPSULE_NAME); - if (managed == NULL) { - return NULL; - } - return Py_BuildValue("ii", managed->dl_tensor.device.device_type, - managed->dl_tensor.device.device_id); + DLDevice device = array_get_dl_device(self); + if (PyErr_Occurred()) { + return NULL; } - return Py_BuildValue("ii", kDLCPU, 0); ->>>>>>> ENH: Add the __dlpack__ and __dlpack_device__ methods to ndarray. + return Py_BuildValue("ii", device.device_type, device.device_id); } NPY_NO_EXPORT PyMethodDef array_methods[] = { diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py index 19f4e92815c4..926668c594e7 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -87,3 +87,8 @@ def test_higher_dims(self, ndim): x = np.zeros(shape, dtype=np.float64) assert shape == np.from_dlpack(x).shape + + def test_dlpack_device(self): + x = np.arange(5) + assert x.__dlpack_device__() == (1, 0) + assert np.from_dlpack(x).__dlpack_device__() == (1, 0) From e167da747b53baed13a5148750c9c82746ffcb30 Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Mon, 31 May 2021 09:46:20 +0200 Subject: [PATCH 10/18] MAINT: Robustify dlpack_capsule_deleter and add comments. --- numpy/core/src/multiarray/methods.c | 33 ++++++++++++++++++++++------- numpy/core/tests/test_dlpack.py | 9 ++++++++ 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 29372fe2f7e2..9a16adecd252 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -2792,25 +2792,42 @@ array_dlpack_deleter(DLManagedTensor *self) } /* This is exactly as mandated by dlpack */ -static void array_dlpack_capsule_deleter(PyObject *self) -{ - if (PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { +static void dlpack_capsule_deleter(PyObject *self){ + if (PyCapsule_IsValid(self, "used_dltensor")) { return; } - DLManagedTensor *managed = - (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_CAPSULE_NAME); + + /* an exception may be in-flight, we must save it in case we create another one */ + PyObject *type, *value, *traceback; + PyErr_Fetch(&type, &value, &traceback); + + DLManagedTensor *managed = (DLManagedTensor *)PyCapsule_GetPointer(self, "dltensor"); if (managed == NULL) { - return; + PyErr_WriteUnraisable(self); + goto done; } - managed->deleter(managed); + /* the spec says the deleter can be NULL if there is no way for the caller to provide a reasonable destructor. */ + if (managed->deleter) { + managed->deleter(managed); + /* TODO: is the deleter allowed to set a python exception? */ + assert(!PyErr_Occurred()); + } + +done: + PyErr_Restore(type, value, traceback); } +// This function cannot return NULL, but it can fail, +// So call PyErr_Occurred to check if it failed after +// calling it. static DLDevice array_get_dl_device(PyArrayObject *self) { DLDevice ret; ret.device_type = kDLCPU; ret.device_id = 0; PyObject *base = PyArray_BASE(self); + // The outer if is due to the fact that NumPy arrays are on the CPU + // by default. if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { DLManagedTensor *managed = PyCapsule_GetPointer( base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); @@ -2937,7 +2954,7 @@ array_dlpack(PyArrayObject *self, managed->deleter = array_dlpack_deleter; PyObject *capsule = PyCapsule_New(managed, NPY_DLPACK_CAPSULE_NAME, - array_dlpack_capsule_deleter); + dlpack_capsule_deleter); if (capsule == NULL) { PyMem_Free(managed); PyMem_Free(managed_shape_strides); diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py index 926668c594e7..2561991e85ec 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -92,3 +92,12 @@ def test_dlpack_device(self): x = np.arange(5) assert x.__dlpack_device__() == (1, 0) assert np.from_dlpack(x).__dlpack_device__() == (1, 0) + + def dlpack_deleter_exception(self): + x = np.arange(5) + _ = x.__dlpack__() + raise RuntimeError + + def test_dlpack_destructor_exception(self): + with pytest.raises(RuntimeError): + self.dlpack_deleter_exception() From a4cc97abdc4e347b293886f4847d1ddd5a9bfaa7 Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Mon, 31 May 2021 14:32:48 +0200 Subject: [PATCH 11/18] BUG: Offset not properly stored/computed. --- numpy/core/src/multiarray/methods.c | 37 +++++++++++++++----- numpy/core/src/multiarray/multiarraymodule.c | 2 +- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 9a16adecd252..d3162e54948e 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -2792,7 +2792,7 @@ array_dlpack_deleter(DLManagedTensor *self) } /* This is exactly as mandated by dlpack */ -static void dlpack_capsule_deleter(PyObject *self){ +static void dlpack_capsule_deleter(PyObject *self) { if (PyCapsule_IsValid(self, "used_dltensor")) { return; } @@ -2827,7 +2827,7 @@ array_get_dl_device(PyArrayObject *self) { ret.device_id = 0; PyObject *base = PyArray_BASE(self); // The outer if is due to the fact that NumPy arrays are on the CPU - // by default. + // by default (if not created from DLPack). if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { DLManagedTensor *managed = PyCapsule_GetPointer( base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); @@ -2839,6 +2839,20 @@ array_get_dl_device(PyArrayObject *self) { return ret; } +static char * +array_get_dl_data(PyArrayObject *self) { + PyObject *base = PyArray_BASE(self); + if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { + DLManagedTensor *managed = PyCapsule_GetPointer( + base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + return NULL; + } + return managed->dl_tensor.data; + } + return PyArray_DATA(self); +} + static PyObject * array_dlpack(PyArrayObject *self, PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) @@ -2913,18 +2927,23 @@ array_dlpack(PyArrayObject *self, return NULL; } + DLDevice device = array_get_dl_device(self); + if (PyErr_Occurred()) { + return NULL; + } + char *data = array_get_dl_data(self); + if (data == NULL) { + return NULL; + } + DLManagedTensor *managed = PyMem_Malloc(sizeof(DLManagedTensor)); if (managed == NULL) { PyErr_NoMemory(); return NULL; } - managed->dl_tensor.data = PyArray_DATA(self); - managed->dl_tensor.device = array_get_dl_device(self); - if (PyErr_Occurred()) { - PyMem_Free(managed); - return NULL; - } + managed->dl_tensor.data = data; + managed->dl_tensor.device = device; managed->dl_tensor.dtype = managed_dtype; @@ -2949,7 +2968,7 @@ array_dlpack(PyArrayObject *self, if (PyArray_SIZE(self) != 1 && !PyArray_IS_C_CONTIGUOUS(self)) { managed->dl_tensor.strides = managed_strides; } - managed->dl_tensor.byte_offset = 0; + managed->dl_tensor.byte_offset = (char *)PyArray_DATA(self) - data; managed->manager_ctx = self; managed->deleter = array_dlpack_deleter; diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index b04debe101ff..60e3875fc615 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -4384,7 +4384,7 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { Py_XDECREF(ret); return NULL; } - + Py_XDECREF(capsule); return ret; } From f96aaa2ac484cab4e9f40b36902544e1174e0513 Mon Sep 17 00:00:00 2001 From: mattip Date: Mon, 1 Nov 2021 09:09:33 +0200 Subject: [PATCH 12/18] move dlpack functions to a new file --- numpy/core/code_generators/genapi.py | 1 + numpy/core/setup.py | 1 + numpy/core/src/common/npy_dlpack.h | 14 +- numpy/core/src/multiarray/dlpack.c | 382 +++++++++++++++++++ numpy/core/src/multiarray/dlpack.h | 18 + numpy/core/src/multiarray/methods.c | 238 +----------- numpy/core/src/multiarray/multiarraymodule.c | 158 +------- 7 files changed, 419 insertions(+), 393 deletions(-) create mode 100644 numpy/core/src/multiarray/dlpack.c create mode 100644 numpy/core/src/multiarray/dlpack.h diff --git a/numpy/core/code_generators/genapi.py b/numpy/core/code_generators/genapi.py index c2458c2b5d80..b401ee6a581e 100644 --- a/numpy/core/code_generators/genapi.py +++ b/numpy/core/code_generators/genapi.py @@ -41,6 +41,7 @@ join('multiarray', 'datetime_busdaycal.c'), join('multiarray', 'datetime_strings.c'), join('multiarray', 'descriptor.c'), + join('multiarray', 'dlpack.c'), join('multiarray', 'dtypemeta.c'), join('multiarray', 'einsum.c.src'), join('multiarray', 'flagsobject.c'), diff --git a/numpy/core/setup.py b/numpy/core/setup.py index 5c6fd4dcf8d0..2c99060ec59d 100644 --- a/numpy/core/setup.py +++ b/numpy/core/setup.py @@ -883,6 +883,7 @@ def gl_if_msvc(build_cmd): join('src', 'multiarray', 'datetime_busday.c'), join('src', 'multiarray', 'datetime_busdaycal.c'), join('src', 'multiarray', 'descriptor.c'), + join('src', 'multiarray', 'dlpack.c'), join('src', 'multiarray', 'dtypemeta.c'), join('src', 'multiarray', 'dragon4.c'), join('src', 'multiarray', 'dtype_transfer.c'), diff --git a/numpy/core/src/common/npy_dlpack.h b/numpy/core/src/common/npy_dlpack.h index 446d44afaf9c..cb926a26271d 100644 --- a/numpy/core/src/common/npy_dlpack.h +++ b/numpy/core/src/common/npy_dlpack.h @@ -13,4 +13,16 @@ // capsule. #define NPY_DLPACK_INTERNAL_CAPSULE_NAME "numpy_dltensor" -#endif \ No newline at end of file +PyObject * +array_dlpack(PyArrayObject *self, PyObject *const *args, Py_ssize_t len_args, + PyObject *kwnames); + + +PyObject * +array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)); + + +NPY_NO_EXPORT PyObject * +from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj); + +#endif diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c new file mode 100644 index 000000000000..591eddfafdf1 --- /dev/null +++ b/numpy/core/src/multiarray/dlpack.c @@ -0,0 +1,382 @@ +#define NPY_NO_DEPRECATED_API NPY_API_VERSION +#define _MULTIARRAYMODULE + +#define PY_SSIZE_T_CLEAN +#include + +#include "numpy/arrayobject.h" +#include "common/npy_argparse.h" + +#include "common/dlpack/dlpack.h" +#include "common/npy_dlpack.h" + +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); +} + +/* This is exactly as mandated by dlpack */ +static void dlpack_capsule_deleter(PyObject *self) { + if (PyCapsule_IsValid(self, "used_dltensor")) { + return; + } + + /* an exception may be in-flight, we must save it in case we create another one */ + PyObject *type, *value, *traceback; + PyErr_Fetch(&type, &value, &traceback); + + DLManagedTensor *managed = (DLManagedTensor *)PyCapsule_GetPointer(self, "dltensor"); + if (managed == NULL) { + PyErr_WriteUnraisable(self); + goto done; + } + /* the spec says the deleter can be NULL if there is no way for the caller to provide a reasonable destructor. */ + if (managed->deleter) { + managed->deleter(managed); + /* TODO: is the deleter allowed to set a python exception? */ + assert(!PyErr_Occurred()); + } + +done: + PyErr_Restore(type, value, traceback); +} + +// This function cannot return NULL, but it can fail, +// So call PyErr_Occurred to check if it failed after +// calling it. +static DLDevice +array_get_dl_device(PyArrayObject *self) { + DLDevice ret; + ret.device_type = kDLCPU; + ret.device_id = 0; + PyObject *base = PyArray_BASE(self); + // The outer if is due to the fact that NumPy arrays are on the CPU + // by default (if not created from DLPack). + if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { + DLManagedTensor *managed = PyCapsule_GetPointer( + base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + return ret; + } + return managed->dl_tensor.device; + } + return ret; +} + +static char * +array_get_dl_data(PyArrayObject *self) { + PyObject *base = PyArray_BASE(self); + if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { + DLManagedTensor *managed = PyCapsule_GetPointer( + base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + return NULL; + } + return managed->dl_tensor.data; + } + return PyArray_DATA(self); +} + +/* used internally */ +static void array_dlpack_internal_capsule_deleter(PyObject *self) +{ + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + return; + } + managed->deleter(managed); +} + +PyObject * +array_dlpack(PyArrayObject *self, + PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) +{ + PyObject *stream = Py_None; + NPY_PREPARE_ARGPARSER; + if (npy_parse_arguments("__dlpack__", args, len_args, kwnames, + "$stream", NULL, &stream, NULL, NULL, NULL)) { + return NULL; + } + + if (stream != Py_None) { + PyErr_SetString(PyExc_RuntimeError, "NumPy only supports " + "stream=None."); + return NULL; + } + + npy_intp itemsize = PyArray_ITEMSIZE(self); + int ndim = PyArray_NDIM(self); + npy_intp *strides = PyArray_STRIDES(self); + npy_intp *shape = PyArray_SHAPE(self); + + if (!PyArray_IS_C_CONTIGUOUS(self) && PyArray_SIZE(self) != 1) { + for (int i = 0; i < ndim; ++i) { + if (strides[i] % itemsize != 0) { + PyErr_SetString(PyExc_RuntimeError, + "DLPack only supports strides which are a multiple of " + "itemsize."); + return NULL; + } + } + } + + DLDataType managed_dtype; + PyArray_Descr *dtype = PyArray_DESCR(self); + + if (PyDataType_ISBYTESWAPPED(dtype)) { + PyErr_SetString(PyExc_TypeError, "DLPack only supports native " + "byte swapping."); + return NULL; + } + + managed_dtype.bits = 8 * itemsize; + managed_dtype.lanes = 1; + + if (PyDataType_ISSIGNED(dtype)) { + managed_dtype.code = kDLInt; + } else if (PyDataType_ISUNSIGNED(dtype)) { + managed_dtype.code = kDLUInt; + } else if (PyDataType_ISFLOAT(dtype)) { + // We can't be sure that the dtype is + // IEEE or padded. + if (itemsize > 8) { + PyErr_SetString(PyExc_TypeError, "DLPack only supports IEEE " + "floating point types without padding."); + return NULL; + } + managed_dtype.code = kDLFloat; + } else if (PyDataType_ISCOMPLEX(dtype)) { + // We can't be sure that the dtype is + // IEEE or padded. + if (itemsize > 16) { + PyErr_SetString(PyExc_TypeError, "DLPack only supports IEEE " + "complex point types without padding."); + return NULL; + } + managed_dtype.code = kDLComplex; + } else { + PyErr_SetString(PyExc_TypeError, + "DLPack only supports signed/unsigned integers, float " + "and complex dtypes."); + return NULL; + } + + DLDevice device = array_get_dl_device(self); + if (PyErr_Occurred()) { + return NULL; + } + char *data = array_get_dl_data(self); + if (data == NULL) { + return NULL; + } + + DLManagedTensor *managed = PyMem_Malloc(sizeof(DLManagedTensor)); + if (managed == NULL) { + PyErr_NoMemory(); + return NULL; + } + + managed->dl_tensor.data = data; + managed->dl_tensor.device = device; + managed->dl_tensor.dtype = managed_dtype; + + + int64_t *managed_shape_strides = PyMem_Malloc(sizeof(int64_t) * ndim * 2); + if (managed_shape_strides == NULL) { + PyErr_NoMemory(); + PyMem_Free(managed); + return NULL; + } + + int64_t *managed_shape = managed_shape_strides; + int64_t *managed_strides = managed_shape_strides + ndim; + for (int i = 0; i < ndim; ++i) { + managed_shape[i] = shape[i]; + // Strides in DLPack are items; in NumPy are bytes. + managed_strides[i] = strides[i] / itemsize; + } + + managed->dl_tensor.ndim = ndim; + managed->dl_tensor.shape = managed_shape; + managed->dl_tensor.strides = NULL; + if (PyArray_SIZE(self) != 1 && !PyArray_IS_C_CONTIGUOUS(self)) { + managed->dl_tensor.strides = managed_strides; + } + managed->dl_tensor.byte_offset = (char *)PyArray_DATA(self) - data; + managed->manager_ctx = self; + managed->deleter = array_dlpack_deleter; + + PyObject *capsule = PyCapsule_New(managed, NPY_DLPACK_CAPSULE_NAME, + dlpack_capsule_deleter); + if (capsule == NULL) { + PyMem_Free(managed); + PyMem_Free(managed_shape_strides); + return NULL; + } + + // the capsule holds a reference + Py_INCREF(self); + return capsule; +} + +PyObject * +array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)) +{ + DLDevice device = array_get_dl_device(self); + if (PyErr_Occurred()) { + return NULL; + } + return Py_BuildValue("ii", device.device_type, device.device_id); +} + +NPY_NO_EXPORT PyObject * +from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { + PyObject *capsule = PyObject_CallMethod((PyObject *)obj->ob_type, + "__dlpack__", "O", obj); + if (capsule == NULL) { + return NULL; + } + + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(capsule, + NPY_DLPACK_CAPSULE_NAME); + + if (managed == NULL) { + Py_XDECREF(capsule); + return NULL; + } + + const int ndim = managed->dl_tensor.ndim; + if (ndim > NPY_MAXDIMS) { + PyErr_SetString(PyExc_RuntimeError, + "maxdims of DLPack tensor is higher than the supported " + "maxdims."); + Py_XDECREF(capsule); + return NULL; + } + + 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); + return NULL; + } + + if (managed->dl_tensor.dtype.lanes != 1) { + PyErr_SetString(PyExc_RuntimeError, + "Unsupported lanes in DLTensor dtype."); + Py_XDECREF(capsule); + return NULL; + } + + int typenum = -1; + const uint8_t bits = managed->dl_tensor.dtype.bits; + const npy_intp itemsize = bits / 8; + switch(managed->dl_tensor.dtype.code) { + case kDLInt: + switch (bits) + { + case 8: typenum = NPY_INT8; break; + case 16: typenum = NPY_INT16; break; + case 32: typenum = NPY_INT32; break; + case 64: typenum = NPY_INT64; break; + } + break; + case kDLUInt: + switch (bits) + { + case 8: typenum = NPY_UINT8; break; + case 16: typenum = NPY_UINT16; break; + case 32: typenum = NPY_UINT32; break; + case 64: typenum = NPY_UINT64; break; + } + break; + case kDLFloat: + switch (bits) + { + case 16: typenum = NPY_FLOAT16; break; + case 32: typenum = NPY_FLOAT32; break; + case 64: typenum = NPY_FLOAT64; break; + } + break; + case kDLComplex: + switch (bits) + { + case 64: typenum = NPY_COMPLEX64; break; + case 128: typenum = NPY_COMPLEX128; break; + } + break; + } + + if (typenum == -1) { + PyErr_SetString(PyExc_RuntimeError, + "Unsupported dtype in DLTensor."); + Py_XDECREF(capsule); + return NULL; + } + + PyArray_Descr *descr = PyArray_DescrFromType(typenum); + if (descr == NULL) { + Py_XDECREF(capsule); + return NULL; + } + + npy_intp shape[NPY_MAXDIMS]; + npy_intp strides[NPY_MAXDIMS]; + + for (int i = 0; i < ndim; ++i) { + shape[i] = managed->dl_tensor.shape[i]; + // DLPack has elements as stride units, NumPy has bytes. + if (managed->dl_tensor.strides != NULL) + { + strides[i] = managed->dl_tensor.strides[i] * itemsize; + } + } + + char *data = (char *)managed->dl_tensor.data + + managed->dl_tensor.byte_offset; + + PyObject *ret = PyArray_NewFromDescr(&PyArray_Type, descr, ndim, shape, + managed->dl_tensor.strides != NULL ? strides : NULL, data, 0, NULL); + if (ret == NULL) { + Py_XDECREF(capsule); + Py_XDECREF(descr); + return NULL; + } + + PyObject *new_capsule = PyCapsule_New(managed, + NPY_DLPACK_INTERNAL_CAPSULE_NAME, + array_dlpack_internal_capsule_deleter); + if (new_capsule == NULL) { + Py_XDECREF(capsule); + Py_XDECREF(ret); + return NULL; + } + + if (PyArray_SetBaseObject((PyArrayObject *)ret, new_capsule) < 0) { + Py_XDECREF(capsule); + Py_XDECREF(ret); + return NULL; + } + + if (PyCapsule_SetName(capsule, NPY_DLPACK_USED_CAPSULE_NAME) < 0) { + Py_XDECREF(capsule); + Py_XDECREF(ret); + return NULL; + } + + Py_XDECREF(capsule); + return ret; +} + + diff --git a/numpy/core/src/multiarray/dlpack.h b/numpy/core/src/multiarray/dlpack.h new file mode 100644 index 000000000000..7d4a289d9747 --- /dev/null +++ b/numpy/core/src/multiarray/dlpack.h @@ -0,0 +1,18 @@ +#define NPY_NO_DEPRECATED_API NPY_API_VERSION +#define _MULTIARRAYMODULE + +#define PY_SSIZE_T_CLEAN +#include +#include "numpy/arrayobject.h" + +static PyObject * +array_dlpack(PyArrayObject *self, PyObject *const *args, Py_ssize_t len_args, + PyObject *kwnames); + + +static PyObject * +array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)); + + +NPY_NO_EXPORT PyObject * +from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj); diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index d3162e54948e..f40db492ca67 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -26,15 +26,13 @@ #include "shape.h" #include "strfuncs.h" #include "array_assign.h" +#include "npy_dlpack.h" #include "methods.h" #include "alloc.h" #include -#include "common/dlpack/dlpack.h" -#include "common/npy_dlpack.h" - /* NpyArg_ParseKeywords * @@ -2584,6 +2582,8 @@ array_round(PyArrayObject *self, PyObject *args, PyObject *kwds) } } + + static PyObject * array_setflags(PyArrayObject *self, PyObject *args, PyObject *kwds) { @@ -2763,236 +2763,6 @@ array_class_getitem(PyObject *cls, PyObject *args) generic_alias = NULL; #endif return generic_alias; - -#define NPY_DLPACK_CAPSULE_NAME "dltensor" -#define NPY_DLPACK_USED_CAPSULE_NAME "used_dltensor" - -static void array_dlpack_capsule_deleter(PyObject *self) -{ - if (!PyCapsule_IsValid(self, NPY_DLPACK_CAPSULE_NAME)) { - if (!PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { - PyErr_SetString(PyExc_RuntimeError, "Invalid capsule name."); - } - return; - } - - DLManagedTensor *managed = - (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_CAPSULE_NAME); - managed->deleter(managed); -} - -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); -} - -/* This is exactly as mandated by dlpack */ -static void dlpack_capsule_deleter(PyObject *self) { - if (PyCapsule_IsValid(self, "used_dltensor")) { - return; - } - - /* an exception may be in-flight, we must save it in case we create another one */ - PyObject *type, *value, *traceback; - PyErr_Fetch(&type, &value, &traceback); - - DLManagedTensor *managed = (DLManagedTensor *)PyCapsule_GetPointer(self, "dltensor"); - if (managed == NULL) { - PyErr_WriteUnraisable(self); - goto done; - } - /* the spec says the deleter can be NULL if there is no way for the caller to provide a reasonable destructor. */ - if (managed->deleter) { - managed->deleter(managed); - /* TODO: is the deleter allowed to set a python exception? */ - assert(!PyErr_Occurred()); - } - -done: - PyErr_Restore(type, value, traceback); -} - -// This function cannot return NULL, but it can fail, -// So call PyErr_Occurred to check if it failed after -// calling it. -static DLDevice -array_get_dl_device(PyArrayObject *self) { - DLDevice ret; - ret.device_type = kDLCPU; - ret.device_id = 0; - PyObject *base = PyArray_BASE(self); - // The outer if is due to the fact that NumPy arrays are on the CPU - // by default (if not created from DLPack). - if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { - DLManagedTensor *managed = PyCapsule_GetPointer( - base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); - if (managed == NULL) { - return ret; - } - return managed->dl_tensor.device; - } - return ret; -} - -static char * -array_get_dl_data(PyArrayObject *self) { - PyObject *base = PyArray_BASE(self); - if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { - DLManagedTensor *managed = PyCapsule_GetPointer( - base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); - if (managed == NULL) { - return NULL; - } - return managed->dl_tensor.data; - } - return PyArray_DATA(self); -} - -static PyObject * -array_dlpack(PyArrayObject *self, - PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) -{ - PyObject *stream = Py_None; - NPY_PREPARE_ARGPARSER; - if (npy_parse_arguments("__dlpack__", args, len_args, kwnames, - "$stream", NULL, &stream, NULL, NULL, NULL)) { - return NULL; - } - - if (stream != Py_None) { - PyErr_SetString(PyExc_RuntimeError, "NumPy only supports " - "stream=None."); - return NULL; - } - - npy_intp itemsize = PyArray_ITEMSIZE(self); - int ndim = PyArray_NDIM(self); - npy_intp *strides = PyArray_STRIDES(self); - npy_intp *shape = PyArray_SHAPE(self); - - if (!PyArray_IS_C_CONTIGUOUS(self) && PyArray_SIZE(self) != 1) { - for (int i = 0; i < ndim; ++i) { - if (strides[i] % itemsize != 0) { - PyErr_SetString(PyExc_RuntimeError, - "DLPack only supports strides which are a multiple of " - "itemsize."); - return NULL; - } - } - } - - DLDataType managed_dtype; - PyArray_Descr *dtype = PyArray_DESCR(self); - - if (PyDataType_ISBYTESWAPPED(dtype)) { - PyErr_SetString(PyExc_TypeError, "DLPack only supports native " - "byte swapping."); - return NULL; - } - - managed_dtype.bits = 8 * itemsize; - managed_dtype.lanes = 1; - - if (PyDataType_ISSIGNED(dtype)) { - managed_dtype.code = kDLInt; - } else if (PyDataType_ISUNSIGNED(dtype)) { - managed_dtype.code = kDLUInt; - } else if (PyDataType_ISFLOAT(dtype)) { - // We can't be sure that the dtype is - // IEEE or padded. - if (itemsize > 8) { - PyErr_SetString(PyExc_TypeError, "DLPack only supports IEEE " - "floating point types without padding."); - return NULL; - } - managed_dtype.code = kDLFloat; - } else if (PyDataType_ISCOMPLEX(dtype)) { - // We can't be sure that the dtype is - // IEEE or padded. - if (itemsize > 16) { - PyErr_SetString(PyExc_TypeError, "DLPack only supports IEEE " - "complex point types without padding."); - return NULL; - } - managed_dtype.code = kDLComplex; - } else { - PyErr_SetString(PyExc_TypeError, - "DLPack only supports signed/unsigned integers, float " - "and complex dtypes."); - return NULL; - } - - DLDevice device = array_get_dl_device(self); - if (PyErr_Occurred()) { - return NULL; - } - char *data = array_get_dl_data(self); - if (data == NULL) { - return NULL; - } - - DLManagedTensor *managed = PyMem_Malloc(sizeof(DLManagedTensor)); - if (managed == NULL) { - PyErr_NoMemory(); - return NULL; - } - - managed->dl_tensor.data = data; - managed->dl_tensor.device = device; - managed->dl_tensor.dtype = managed_dtype; - - - int64_t *managed_shape_strides = PyMem_Malloc(sizeof(int64_t) * ndim * 2); - if (managed_shape_strides == NULL) { - PyErr_NoMemory(); - PyMem_Free(managed); - return NULL; - } - - int64_t *managed_shape = managed_shape_strides; - int64_t *managed_strides = managed_shape_strides + ndim; - for (int i = 0; i < ndim; ++i) { - managed_shape[i] = shape[i]; - // Strides in DLPack are items; in NumPy are bytes. - managed_strides[i] = strides[i] / itemsize; - } - - managed->dl_tensor.ndim = ndim; - managed->dl_tensor.shape = managed_shape; - managed->dl_tensor.strides = NULL; - if (PyArray_SIZE(self) != 1 && !PyArray_IS_C_CONTIGUOUS(self)) { - managed->dl_tensor.strides = managed_strides; - } - managed->dl_tensor.byte_offset = (char *)PyArray_DATA(self) - data; - managed->manager_ctx = self; - managed->deleter = array_dlpack_deleter; - - PyObject *capsule = PyCapsule_New(managed, NPY_DLPACK_CAPSULE_NAME, - dlpack_capsule_deleter); - if (capsule == NULL) { - PyMem_Free(managed); - PyMem_Free(managed_shape_strides); - return NULL; - } - - // the capsule holds a reference - Py_INCREF(self); - return capsule; -} - -static PyObject * -array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)) -{ - DLDevice device = array_get_dl_device(self); - if (PyErr_Occurred()) { - return NULL; - } - return Py_BuildValue("ii", device.device_type, device.device_id); } NPY_NO_EXPORT PyMethodDef array_methods[] = { @@ -3220,7 +2990,6 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = { {"view", (PyCFunction)array_view, METH_FASTCALL | METH_KEYWORDS, NULL}, - // For data interchange between libraries {"__dlpack__", (PyCFunction)array_dlpack, @@ -3229,6 +2998,5 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = { {"__dlpack_device__", (PyCFunction)array_dlpack_device, METH_NOARGS, NULL}, - {NULL, NULL, 0, NULL} /* sentinel */ }; diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 60e3875fc615..ac4c1ae0b99d 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -70,7 +70,7 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0; #include "get_attr_string.h" #include "experimental_public_dtype_api.h" /* _get_experimental_dtype_api */ -#include "common/npy_dlpack.h" +#include "npy_dlpack.h" /* ***************************************************************************** @@ -4233,162 +4233,6 @@ _reload_guard(PyObject *NPY_UNUSED(self)) { Py_RETURN_NONE; } -/* used internally */ -static void array_dlpack_internal_capsule_deleter(PyObject *self) -{ - DLManagedTensor *managed = - (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME); - if (managed == NULL) { - return; - } - managed->deleter(managed); -} - - -NPY_NO_EXPORT PyObject * -from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { - PyObject *capsule = PyObject_CallMethod((PyObject *)obj->ob_type, - "__dlpack__", "O", obj); - if (capsule == NULL) { - return NULL; - } - - DLManagedTensor *managed = - (DLManagedTensor *)PyCapsule_GetPointer(capsule, - NPY_DLPACK_CAPSULE_NAME); - - if (managed == NULL) { - Py_XDECREF(capsule); - return NULL; - } - - const int ndim = managed->dl_tensor.ndim; - if (ndim > NPY_MAXDIMS) { - PyErr_SetString(PyExc_RuntimeError, - "maxdims of DLPack tensor is higher than the supported " - "maxdims."); - Py_XDECREF(capsule); - return NULL; - } - - 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); - return NULL; - } - - if (managed->dl_tensor.dtype.lanes != 1) { - PyErr_SetString(PyExc_RuntimeError, - "Unsupported lanes in DLTensor dtype."); - Py_XDECREF(capsule); - return NULL; - } - - int typenum = -1; - const uint8_t bits = managed->dl_tensor.dtype.bits; - const npy_intp itemsize = bits / 8; - switch(managed->dl_tensor.dtype.code) { - case kDLInt: - switch (bits) - { - case 8: typenum = NPY_INT8; break; - case 16: typenum = NPY_INT16; break; - case 32: typenum = NPY_INT32; break; - case 64: typenum = NPY_INT64; break; - } - break; - case kDLUInt: - switch (bits) - { - case 8: typenum = NPY_UINT8; break; - case 16: typenum = NPY_UINT16; break; - case 32: typenum = NPY_UINT32; break; - case 64: typenum = NPY_UINT64; break; - } - break; - case kDLFloat: - switch (bits) - { - case 16: typenum = NPY_FLOAT16; break; - case 32: typenum = NPY_FLOAT32; break; - case 64: typenum = NPY_FLOAT64; break; - } - break; - case kDLComplex: - switch (bits) - { - case 64: typenum = NPY_COMPLEX64; break; - case 128: typenum = NPY_COMPLEX128; break; - } - break; - } - - if (typenum == -1) { - PyErr_SetString(PyExc_RuntimeError, - "Unsupported dtype in DLTensor."); - Py_XDECREF(capsule); - return NULL; - } - - PyArray_Descr *descr = PyArray_DescrFromType(typenum); - if (descr == NULL) { - Py_XDECREF(capsule); - return NULL; - } - - npy_intp shape[NPY_MAXDIMS]; - npy_intp strides[NPY_MAXDIMS]; - - for (int i = 0; i < ndim; ++i) { - shape[i] = managed->dl_tensor.shape[i]; - // DLPack has elements as stride units, NumPy has bytes. - if (managed->dl_tensor.strides != NULL) - { - strides[i] = managed->dl_tensor.strides[i] * itemsize; - } - } - - char *data = (char *)managed->dl_tensor.data + - managed->dl_tensor.byte_offset; - - PyObject *ret = PyArray_NewFromDescr(&PyArray_Type, descr, ndim, shape, - managed->dl_tensor.strides != NULL ? strides : NULL, data, 0, NULL); - if (ret == NULL) { - Py_XDECREF(capsule); - Py_XDECREF(descr); - return NULL; - } - - PyObject *new_capsule = PyCapsule_New(managed, - NPY_DLPACK_INTERNAL_CAPSULE_NAME, - array_dlpack_internal_capsule_deleter); - if (new_capsule == NULL) { - Py_XDECREF(capsule); - Py_XDECREF(ret); - return NULL; - } - - if (PyArray_SetBaseObject((PyArrayObject *)ret, new_capsule) < 0) { - Py_XDECREF(capsule); - Py_XDECREF(ret); - return NULL; - } - - if (PyCapsule_SetName(capsule, NPY_DLPACK_USED_CAPSULE_NAME) < 0) { - Py_XDECREF(capsule); - Py_XDECREF(ret); - return NULL; - } - - Py_XDECREF(capsule); - return ret; -} - static struct PyMethodDef array_module_methods[] = { {"_get_implementing_args", (PyCFunction)array__get_implementing_args, From 0c992dca1cc23f12af14a8d66101166ef6c92355 Mon Sep 17 00:00:00 2001 From: mattip Date: Mon, 1 Nov 2021 09:15:18 +0200 Subject: [PATCH 13/18] BUG: fixes from review --- numpy/core/src/multiarray/dlpack.c | 104 ++++++++++++++++++----------- 1 file changed, 64 insertions(+), 40 deletions(-) diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c index 591eddfafdf1..f1591bb1fb4e 100644 --- a/numpy/core/src/multiarray/dlpack.c +++ b/numpy/core/src/multiarray/dlpack.c @@ -22,7 +22,7 @@ array_dlpack_deleter(DLManagedTensor *self) /* This is exactly as mandated by dlpack */ static void dlpack_capsule_deleter(PyObject *self) { - if (PyCapsule_IsValid(self, "used_dltensor")) { + if (PyCapsule_IsValid(self, NPY_DLPACK_USED_CAPSULE_NAME)) { return; } @@ -30,12 +30,16 @@ static void dlpack_capsule_deleter(PyObject *self) { PyObject *type, *value, *traceback; PyErr_Fetch(&type, &value, &traceback); - DLManagedTensor *managed = (DLManagedTensor *)PyCapsule_GetPointer(self, "dltensor"); + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_CAPSULE_NAME); if (managed == NULL) { PyErr_WriteUnraisable(self); goto done; } - /* the spec says the deleter can be NULL if there is no way for the caller to provide a reasonable destructor. */ + /* + * the spec says the deleter can be NULL if there is no way for the caller + * to provide a reasonable destructor. + */ if (managed->deleter) { managed->deleter(managed); /* TODO: is the deleter allowed to set a python exception? */ @@ -46,6 +50,34 @@ static void dlpack_capsule_deleter(PyObject *self) { PyErr_Restore(type, value, traceback); } +/* used internally, almost identical to dlpack_capsule_deleter() */ +static void array_dlpack_internal_capsule_deleter(PyObject *self) +{ + /* an exception may be in-flight, we must save it in case we create another one */ + PyObject *type, *value, *traceback; + PyErr_Fetch(&type, &value, &traceback); + + DLManagedTensor *managed = + (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME); + if (managed == NULL) { + PyErr_WriteUnraisable(self); + goto done; + } + /* + * the spec says the deleter can be NULL if there is no way for the caller + * to provide a reasonable destructor. + */ + if (managed->deleter) { + managed->deleter(managed); + /* TODO: is the deleter allowed to set a python exception? */ + assert(!PyErr_Occurred()); + } + +done: + PyErr_Restore(type, value, traceback); +} + + // This function cannot return NULL, but it can fail, // So call PyErr_Occurred to check if it failed after // calling it. @@ -82,17 +114,6 @@ array_get_dl_data(PyArrayObject *self) { return PyArray_DATA(self); } -/* used internally */ -static void array_dlpack_internal_capsule_deleter(PyObject *self) -{ - DLManagedTensor *managed = - (DLManagedTensor *)PyCapsule_GetPointer(self, NPY_DLPACK_INTERNAL_CAPSULE_NAME); - if (managed == NULL) { - return; - } - managed->deleter(managed); -} - PyObject * array_dlpack(PyArrayObject *self, PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) @@ -140,9 +161,11 @@ array_dlpack(PyArrayObject *self, if (PyDataType_ISSIGNED(dtype)) { managed_dtype.code = kDLInt; - } else if (PyDataType_ISUNSIGNED(dtype)) { + } + else if (PyDataType_ISUNSIGNED(dtype)) { managed_dtype.code = kDLUInt; - } else if (PyDataType_ISFLOAT(dtype)) { + } + else if (PyDataType_ISFLOAT(dtype)) { // We can't be sure that the dtype is // IEEE or padded. if (itemsize > 8) { @@ -151,7 +174,8 @@ array_dlpack(PyArrayObject *self, return NULL; } managed_dtype.code = kDLFloat; - } else if (PyDataType_ISCOMPLEX(dtype)) { + } + else if (PyDataType_ISCOMPLEX(dtype)) { // We can't be sure that the dtype is // IEEE or padded. if (itemsize > 16) { @@ -160,7 +184,8 @@ array_dlpack(PyArrayObject *self, return NULL; } managed_dtype.code = kDLComplex; - } else { + } + else { PyErr_SetString(PyExc_TypeError, "DLPack only supports signed/unsigned integers, float " "and complex dtypes."); @@ -243,12 +268,12 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { return NULL; } - DLManagedTensor *managed = + DLManagedTensor *managed = (DLManagedTensor *)PyCapsule_GetPointer(capsule, NPY_DLPACK_CAPSULE_NAME); if (managed == NULL) { - Py_XDECREF(capsule); + Py_DECREF(capsule); return NULL; } @@ -257,7 +282,7 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { PyErr_SetString(PyExc_RuntimeError, "maxdims of DLPack tensor is higher than the supported " "maxdims."); - Py_XDECREF(capsule); + Py_DECREF(capsule); return NULL; } @@ -268,14 +293,14 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { device_type != kDLCUDAManaged) { PyErr_SetString(PyExc_RuntimeError, "Unsupported device in DLTensor."); - Py_XDECREF(capsule); + Py_DECREF(capsule); return NULL; } if (managed->dl_tensor.dtype.lanes != 1) { PyErr_SetString(PyExc_RuntimeError, "Unsupported lanes in DLTensor dtype."); - Py_XDECREF(capsule); + Py_DECREF(capsule); return NULL; } @@ -321,13 +346,7 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { if (typenum == -1) { PyErr_SetString(PyExc_RuntimeError, "Unsupported dtype in DLTensor."); - Py_XDECREF(capsule); - return NULL; - } - - PyArray_Descr *descr = PyArray_DescrFromType(typenum); - if (descr == NULL) { - Py_XDECREF(capsule); + Py_DECREF(capsule); return NULL; } @@ -346,11 +365,16 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { char *data = (char *)managed->dl_tensor.data + managed->dl_tensor.byte_offset; + PyArray_Descr *descr = PyArray_DescrFromType(typenum); + if (descr == NULL) { + Py_DECREF(capsule); + return NULL; + } + PyObject *ret = PyArray_NewFromDescr(&PyArray_Type, descr, ndim, shape, managed->dl_tensor.strides != NULL ? strides : NULL, data, 0, NULL); if (ret == NULL) { - Py_XDECREF(capsule); - Py_XDECREF(descr); + Py_DECREF(capsule); return NULL; } @@ -358,24 +382,24 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { NPY_DLPACK_INTERNAL_CAPSULE_NAME, array_dlpack_internal_capsule_deleter); if (new_capsule == NULL) { - Py_XDECREF(capsule); - Py_XDECREF(ret); + Py_DECREF(capsule); + Py_DECREF(ret); return NULL; } if (PyArray_SetBaseObject((PyArrayObject *)ret, new_capsule) < 0) { - Py_XDECREF(capsule); - Py_XDECREF(ret); + Py_DECREF(capsule); + Py_DECREF(ret); return NULL; } if (PyCapsule_SetName(capsule, NPY_DLPACK_USED_CAPSULE_NAME) < 0) { - Py_XDECREF(capsule); - Py_XDECREF(ret); + Py_DECREF(capsule); + Py_DECREF(ret); return NULL; } - - Py_XDECREF(capsule); + + Py_DECREF(capsule); return ret; } From ef04f59d1a5cc7320866c48d75059aa52d390003 Mon Sep 17 00:00:00 2001 From: Matti Picus Date: Tue, 2 Nov 2021 09:50:33 +0200 Subject: [PATCH 14/18] Updates Co-authored-by: Sebastian Berg Co-authored-by: Bas van Beek <43369155+BvB93@users.noreply.github.com> --- numpy/__init__.pyi | 2 +- numpy/core/src/multiarray/dlpack.c | 6 ++---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 63e723a35471..1562ce89ed98 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -4333,7 +4333,7 @@ class chararray(ndarray[_ShapeType, _CharDType]): # class MachAr: ... class _SupportsDLPack(Protocol[_T_contra]): - def __dlpack__(self, *, stream: Optional[int] = ...) -> _PyCapsule: ... + def __dlpack__(self, *, stream: None | _T_contra = ...) -> _PyCapsule: ... def from_dlpack(__obj: _SupportsDLPack[None]) -> NDArray[Any]: ... diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c index f1591bb1fb4e..9de3043795f7 100644 --- a/numpy/core/src/multiarray/dlpack.c +++ b/numpy/core/src/multiarray/dlpack.c @@ -211,7 +211,6 @@ array_dlpack(PyArrayObject *self, managed->dl_tensor.device = device; managed->dl_tensor.dtype = managed_dtype; - int64_t *managed_shape_strides = PyMem_Malloc(sizeof(int64_t) * ndim * 2); if (managed_shape_strides == NULL) { PyErr_NoMemory(); @@ -307,7 +306,7 @@ from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { int typenum = -1; const uint8_t bits = managed->dl_tensor.dtype.bits; const npy_intp itemsize = bits / 8; - switch(managed->dl_tensor.dtype.code) { + switch (managed->dl_tensor.dtype.code) { case kDLInt: switch (bits) { @@ -356,8 +355,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. - if (managed->dl_tensor.strides != NULL) - { + if (managed->dl_tensor.strides != NULL) { strides[i] = managed->dl_tensor.strides[i] * itemsize; } } From ba8fcbe2ba0a26cd52dfa9bf40dd2e945e5b298f Mon Sep 17 00:00:00 2001 From: mattip Date: Tue, 2 Nov 2021 11:30:32 +0200 Subject: [PATCH 15/18] change from_dlpack to _dlpack, remove unused header --- doc/neps/nep-0047-array-api-standard.rst | 7 +++--- numpy/__init__.pyi | 2 +- numpy/array_api/__init__.py | 4 ++-- numpy/array_api/_creation_functions.py | 2 +- numpy/core/_add_newdocs.py | 4 ++-- numpy/core/multiarray.py | 17 +++++++------- numpy/core/numeric.py | 4 ++-- numpy/core/src/common/npy_dlpack.h | 2 +- numpy/core/src/multiarray/dlpack.c | 2 +- numpy/core/src/multiarray/dlpack.h | 18 --------------- numpy/core/src/multiarray/multiarraymodule.c | 2 +- numpy/core/tests/test_dlpack.py | 24 ++++++++++---------- 12 files changed, 36 insertions(+), 52 deletions(-) delete mode 100644 numpy/core/src/multiarray/dlpack.h diff --git a/doc/neps/nep-0047-array-api-standard.rst b/doc/neps/nep-0047-array-api-standard.rst index 3e63602ccab0..53b8e35b001f 100644 --- a/doc/neps/nep-0047-array-api-standard.rst +++ b/doc/neps/nep-0047-array-api-standard.rst @@ -338,9 +338,10 @@ the options already present in NumPy are: Adding support for DLPack to NumPy entails: -- Adding a ``ndarray.__dlpack__`` method. -- Adding a ``from_dlpack`` function, which takes as input an object - supporting ``__dlpack__``, and returns an ``ndarray``. +- Adding a ``ndarray.__dlpack__()`` method which returns a ``dlpack`` C + structure wrapped in a ``PyCapsule``. +- Adding a ``np._from_dlpack(obj)`` function, where ``obj`` supports + ``__dlpack__()``, and returns an ``ndarray``. DLPack is currently a ~200 LoC header, and is meant to be included directly, so no external dependency is needed. Implementation should be straightforward. diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 1562ce89ed98..b2e9eec778b6 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -4335,5 +4335,5 @@ class chararray(ndarray[_ShapeType, _CharDType]): class _SupportsDLPack(Protocol[_T_contra]): def __dlpack__(self, *, stream: None | _T_contra = ...) -> _PyCapsule: ... -def from_dlpack(__obj: _SupportsDLPack[None]) -> NDArray[Any]: ... +def _from_dlpack(__obj: _SupportsDLPack[None]) -> NDArray[Any]: ... diff --git a/numpy/array_api/__init__.py b/numpy/array_api/__init__.py index d8b29057ecef..89f5e9cbad49 100644 --- a/numpy/array_api/__init__.py +++ b/numpy/array_api/__init__.py @@ -136,7 +136,7 @@ empty, empty_like, eye, - from_dlpack, + _from_dlpack, full, full_like, linspace, @@ -155,7 +155,7 @@ "empty", "empty_like", "eye", - "from_dlpack", + "_from_dlpack", "full", "full_like", "linspace", diff --git a/numpy/array_api/_creation_functions.py b/numpy/array_api/_creation_functions.py index e36807468e21..c3644ac2c6ec 100644 --- a/numpy/array_api/_creation_functions.py +++ b/numpy/array_api/_creation_functions.py @@ -151,7 +151,7 @@ def eye( return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype)) -def from_dlpack(x: object, /) -> Array: +def _from_dlpack(x: object, /) -> Array: # Note: dlpack support is not yet implemented on Array raise NotImplementedError("DLPack support is not yet implemented") diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py index f1a42dffeb54..cae5bc281789 100644 --- a/numpy/core/_add_newdocs.py +++ b/numpy/core/_add_newdocs.py @@ -1573,9 +1573,9 @@ array_function_like_doc, )) -add_newdoc('numpy.core.multiarray', 'from_dlpack', +add_newdoc('numpy.core.multiarray', '_from_dlpack', """ - from_dlpack(x, /) + _from_dlpack(x, /) Create a NumPy array from an object implementing the ``__dlpack__`` protocol. diff --git a/numpy/core/multiarray.py b/numpy/core/multiarray.py index 9f431f01b7c6..f96274263448 100644 --- a/numpy/core/multiarray.py +++ b/numpy/core/multiarray.py @@ -14,8 +14,9 @@ # do not change them. issue gh-15518 # _get_ndarray_c_version is semi-public, on purpose not added to __all__ from ._multiarray_umath import ( - _fastCopyAndTranspose, _flagdict, _insert, _reconstruct, _vec_string, - _ARRAY_API, _monotonicity, _get_ndarray_c_version, _set_madvise_hugepage, + _fastCopyAndTranspose, _flagdict, _from_dlpack, _insert, _reconstruct, + _vec_string, _ARRAY_API, _monotonicity, _get_ndarray_c_version, + _set_madvise_hugepage, ) __all__ = [ @@ -23,15 +24,15 @@ 'ITEM_HASOBJECT', 'ITEM_IS_POINTER', 'LIST_PICKLE', 'MAXDIMS', 'MAY_SHARE_BOUNDS', 'MAY_SHARE_EXACT', 'NEEDS_INIT', 'NEEDS_PYAPI', 'RAISE', 'USE_GETITEM', 'USE_SETITEM', 'WRAP', '_fastCopyAndTranspose', - '_flagdict', '_insert', '_reconstruct', '_vec_string', '_monotonicity', - 'add_docstring', 'arange', 'array', 'asarray', 'asanyarray', - 'ascontiguousarray', 'asfortranarray', 'bincount', 'broadcast', - 'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast', + '_flagdict', '_from_dlpack', '_insert', '_reconstruct', '_vec_string', + '_monotonicity', 'add_docstring', 'arange', 'array', 'asarray', + 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'bincount', + 'broadcast', 'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast', 'compare_chararrays', 'concatenate', 'copyto', 'correlate', 'correlate2', 'count_nonzero', 'c_einsum', 'datetime_as_string', 'datetime_data', 'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype', 'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat', - 'frombuffer', 'from_dlpack', 'fromfile', 'fromiter', 'fromstring', + 'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'get_handler_name', 'inner', 'interp', 'interp_complex', 'is_busday', 'lexsort', 'matmul', 'may_share_memory', 'min_scalar_type', 'ndarray', 'nditer', 'nested_iters', 'normalize_axis_index', 'packbits', @@ -46,6 +47,7 @@ scalar.__module__ = 'numpy.core.multiarray' +_from_dlpack.__module__ = 'numpy' arange.__module__ = 'numpy' array.__module__ = 'numpy' asarray.__module__ = 'numpy' @@ -55,7 +57,6 @@ datetime_data.__module__ = 'numpy' empty.__module__ = 'numpy' frombuffer.__module__ = 'numpy' -from_dlpack.__module__ = 'numpy' fromfile.__module__ = 'numpy' fromiter.__module__ = 'numpy' frompyfunc.__module__ = 'numpy' diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index 987470f9239e..344d40d934cf 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -13,7 +13,7 @@ WRAP, arange, array, asarray, asanyarray, ascontiguousarray, asfortranarray, broadcast, can_cast, compare_chararrays, concatenate, copyto, dot, dtype, empty, - empty_like, flatiter, frombuffer, from_dlpack, fromfile, fromiter, + empty_like, flatiter, frombuffer, _from_dlpack, fromfile, fromiter, fromstring, inner, lexsort, matmul, may_share_memory, min_scalar_type, ndarray, nditer, nested_iters, promote_types, putmask, result_type, set_numeric_ops, shares_memory, vdot, where, @@ -41,7 +41,7 @@ 'newaxis', 'ndarray', 'flatiter', 'nditer', 'nested_iters', 'ufunc', 'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'zeros', 'count_nonzero', 'empty', 'broadcast', 'dtype', - 'fromstring', 'fromfile', 'frombuffer', 'from_dlpack', 'where', + 'fromstring', 'fromfile', 'frombuffer', '_from_dlpack', 'where', 'argwhere', 'copyto', 'concatenate', 'fastCopyAndTranspose', 'lexsort', 'set_numeric_ops', 'can_cast', 'promote_types', 'min_scalar_type', 'result_type', 'isfortran', 'empty_like', 'zeros_like', 'ones_like', diff --git a/numpy/core/src/common/npy_dlpack.h b/numpy/core/src/common/npy_dlpack.h index cb926a26271d..14ca352c01a7 100644 --- a/numpy/core/src/common/npy_dlpack.h +++ b/numpy/core/src/common/npy_dlpack.h @@ -23,6 +23,6 @@ array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)); NPY_NO_EXPORT PyObject * -from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj); +_from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj); #endif diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c index 9de3043795f7..21930b0efa1b 100644 --- a/numpy/core/src/multiarray/dlpack.c +++ b/numpy/core/src/multiarray/dlpack.c @@ -260,7 +260,7 @@ array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)) } NPY_NO_EXPORT PyObject * -from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { +_from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { PyObject *capsule = PyObject_CallMethod((PyObject *)obj->ob_type, "__dlpack__", "O", obj); if (capsule == NULL) { diff --git a/numpy/core/src/multiarray/dlpack.h b/numpy/core/src/multiarray/dlpack.h deleted file mode 100644 index 7d4a289d9747..000000000000 --- a/numpy/core/src/multiarray/dlpack.h +++ /dev/null @@ -1,18 +0,0 @@ -#define NPY_NO_DEPRECATED_API NPY_API_VERSION -#define _MULTIARRAYMODULE - -#define PY_SSIZE_T_CLEAN -#include -#include "numpy/arrayobject.h" - -static PyObject * -array_dlpack(PyArrayObject *self, PyObject *const *args, Py_ssize_t len_args, - PyObject *kwnames); - - -static PyObject * -array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)); - - -NPY_NO_EXPORT PyObject * -from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj); diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index ac4c1ae0b99d..3cb4e83f195f 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -4446,7 +4446,7 @@ static struct PyMethodDef array_module_methods[] = { {"_reload_guard", (PyCFunction)_reload_guard, METH_NOARGS, "Give a warning on reload and big warning in sub-interpreters."}, - {"from_dlpack", (PyCFunction)from_dlpack, + {"_from_dlpack", (PyCFunction)_from_dlpack, METH_O, NULL}, {NULL, NULL, 0, NULL} /* sentinel */ }; diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py index 2561991e85ec..06fc042ec805 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -27,12 +27,12 @@ def test_strides_not_multiple_of_itemsize(self): z = y['int'] with pytest.raises(RuntimeError): - np.from_dlpack(z) + 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) + y = np._from_dlpack(x) assert sys.getrefcount(x) == 3 del y assert sys.getrefcount(x) == 2 @@ -45,7 +45,7 @@ def test_from_dlpack_refcount(self): ]) def test_dtype_passthrough(self, dtype): x = np.arange(5, dtype=dtype) - y = np.from_dlpack(x) + y = np._from_dlpack(x) assert y.dtype == x.dtype assert_array_equal(x, y) @@ -54,44 +54,44 @@ def test_invalid_dtype(self): x = np.asarray(np.datetime64('2021-05-27')) with pytest.raises(TypeError): - np.from_dlpack(x) + np._from_dlpack(x) def test_invalid_byte_swapping(self): dt = np.dtype('=i8').newbyteorder() x = np.arange(5, dtype=dt) with pytest.raises(TypeError): - np.from_dlpack(x) + np._from_dlpack(x) def test_non_contiguous(self): x = np.arange(25).reshape((5, 5)) y1 = x[0] - assert_array_equal(y1, np.from_dlpack(y1)) + assert_array_equal(y1, np._from_dlpack(y1)) y2 = x[:, 0] - assert_array_equal(y2, np.from_dlpack(y2)) + assert_array_equal(y2, np._from_dlpack(y2)) y3 = x[1, :] - assert_array_equal(y3, np.from_dlpack(y3)) + assert_array_equal(y3, np._from_dlpack(y3)) y4 = x[1] - assert_array_equal(y4, np.from_dlpack(y4)) + assert_array_equal(y4, np._from_dlpack(y4)) y5 = np.diagonal(x) - assert_array_equal(y5, np.from_dlpack(y5)) + assert_array_equal(y5, np._from_dlpack(y5)) @pytest.mark.parametrize("ndim", range(33)) def test_higher_dims(self, ndim): shape = (1,) * ndim x = np.zeros(shape, dtype=np.float64) - assert shape == np.from_dlpack(x).shape + assert shape == np._from_dlpack(x).shape def test_dlpack_device(self): x = np.arange(5) assert x.__dlpack_device__() == (1, 0) - assert np.from_dlpack(x).__dlpack_device__() == (1, 0) + assert np._from_dlpack(x).__dlpack_device__() == (1, 0) def dlpack_deleter_exception(self): x = np.arange(5) From c38b7f328b87d626be544c5ffb542f2b6a97adee Mon Sep 17 00:00:00 2001 From: mattip Date: Thu, 4 Nov 2021 00:11:04 +0200 Subject: [PATCH 16/18] make a.__dlpack__() fail if a is readonly --- numpy/core/src/multiarray/dlpack.c | 6 ++++++ numpy/core/tests/test_dlpack.py | 8 +++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c index 21930b0efa1b..f061a6bf904c 100644 --- a/numpy/core/src/multiarray/dlpack.c +++ b/numpy/core/src/multiarray/dlpack.c @@ -131,6 +131,12 @@ array_dlpack(PyArrayObject *self, return NULL; } + if ( !(PyArray_FLAGS(self) & NPY_ARRAY_WRITEABLE)) { + PyErr_SetString(PyExc_TypeError, "NumPy currently only supports " + "dlpack for writeable arrays"); + return NULL; + } + npy_intp itemsize = PyArray_ITEMSIZE(self); int ndim = PyArray_NDIM(self); npy_intp *strides = PyArray_STRIDES(self); diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py index 06fc042ec805..f848b2008cf9 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -78,7 +78,7 @@ def test_non_contiguous(self): y4 = x[1] assert_array_equal(y4, np._from_dlpack(y4)) - y5 = np.diagonal(x) + y5 = np.diagonal(x).copy() assert_array_equal(y5, np._from_dlpack(y5)) @pytest.mark.parametrize("ndim", range(33)) @@ -101,3 +101,9 @@ def dlpack_deleter_exception(self): def test_dlpack_destructor_exception(self): with pytest.raises(RuntimeError): self.dlpack_deleter_exception() + + def test_readonly(self): + x = np.arange(5) + x.flags.writeable = False + with pytest.raises(TypeError): + x.__dlpack__() From f6076ee90985e7839d7ec6ed83ea1f85c0ece44d Mon Sep 17 00:00:00 2001 From: mattip Date: Tue, 9 Nov 2021 11:16:33 +0200 Subject: [PATCH 17/18] add release note, error out if offset is used --- doc/release/upcoming_changes/19083.new_feature.rst | 6 ++++++ numpy/core/src/multiarray/dlpack.c | 8 +++++++- 2 files changed, 13 insertions(+), 1 deletion(-) create mode 100644 doc/release/upcoming_changes/19083.new_feature.rst diff --git a/doc/release/upcoming_changes/19083.new_feature.rst b/doc/release/upcoming_changes/19083.new_feature.rst new file mode 100644 index 000000000000..92f00c0d6111 --- /dev/null +++ b/doc/release/upcoming_changes/19083.new_feature.rst @@ -0,0 +1,6 @@ +Add NEP 47-compatible dlpack support +------------------------------------ + +Add a ``ndarray.__dlpack__()`` method which returns a ``dlpack`` C structure +wrapped in a ``PyCapsule``. Also add a ``np._from_dlpack(obj)`` function, where +``obj`` supports ``__dlpack__()``, and returns an ``ndarray``. diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c index f061a6bf904c..b0eaa7786b67 100644 --- a/numpy/core/src/multiarray/dlpack.c +++ b/numpy/core/src/multiarray/dlpack.c @@ -206,6 +206,12 @@ array_dlpack(PyArrayObject *self, if (data == NULL) { return NULL; } + if ((char *)PyArray_DATA(self) - data != 0) { + PyErr_SetString(PyExc_TypeError, + "Offsets not clearly supported by this " + "version of DLPack."); + return NULL; + } DLManagedTensor *managed = PyMem_Malloc(sizeof(DLManagedTensor)); if (managed == NULL) { @@ -238,7 +244,7 @@ array_dlpack(PyArrayObject *self, if (PyArray_SIZE(self) != 1 && !PyArray_IS_C_CONTIGUOUS(self)) { managed->dl_tensor.strides = managed_strides; } - managed->dl_tensor.byte_offset = (char *)PyArray_DATA(self) - data; + managed->dl_tensor.byte_offset = 0; managed->manager_ctx = self; managed->deleter = array_dlpack_deleter; From 5b94a03b93d171232e0a64dc2160e4f2139e9b1a Mon Sep 17 00:00:00 2001 From: Sebastian Berg Date: Tue, 9 Nov 2021 12:14:13 -0600 Subject: [PATCH 18/18] MAINT: Simplify `byte_offset` handling in dlpack.h and add comment If `byte_offset = 0` is forced anyway, there is no point in trying to preserve a previous `data` information from the capsule. (And probably it should have used a base array also, and not just a base DLPack capsule, anyway.) --- numpy/core/src/multiarray/dlpack.c | 40 ++++++++++++------------------ 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c index b0eaa7786b67..291e60a226a7 100644 --- a/numpy/core/src/multiarray/dlpack.c +++ b/numpy/core/src/multiarray/dlpack.c @@ -3,6 +3,7 @@ #define PY_SSIZE_T_CLEAN #include +#include #include "numpy/arrayobject.h" #include "common/npy_argparse.h" @@ -100,19 +101,6 @@ array_get_dl_device(PyArrayObject *self) { return ret; } -static char * -array_get_dl_data(PyArrayObject *self) { - PyObject *base = PyArray_BASE(self); - if (PyCapsule_IsValid(base, NPY_DLPACK_INTERNAL_CAPSULE_NAME)) { - DLManagedTensor *managed = PyCapsule_GetPointer( - base, NPY_DLPACK_INTERNAL_CAPSULE_NAME); - if (managed == NULL) { - return NULL; - } - return managed->dl_tensor.data; - } - return PyArray_DATA(self); -} PyObject * array_dlpack(PyArrayObject *self, @@ -202,16 +190,6 @@ array_dlpack(PyArrayObject *self, if (PyErr_Occurred()) { return NULL; } - char *data = array_get_dl_data(self); - if (data == NULL) { - return NULL; - } - if ((char *)PyArray_DATA(self) - data != 0) { - PyErr_SetString(PyExc_TypeError, - "Offsets not clearly supported by this " - "version of DLPack."); - return NULL; - } DLManagedTensor *managed = PyMem_Malloc(sizeof(DLManagedTensor)); if (managed == NULL) { @@ -219,7 +197,21 @@ array_dlpack(PyArrayObject *self, return NULL; } - managed->dl_tensor.data = data; + /* + * Note: the `dlpack.h` header suggests/standardizes that `data` must be + * 256-byte aligned. We ignore this intentionally, because `__dlpack__` + * standardizes that `byte_offset` must be 0 (for now) to not break pytorch: + * https://github.com/data-apis/array-api/issues/293#issuecomment-964111413 + * + * We further assume that exporting fully unaligned data is OK even without + * `byte_offset` since the standard does not reject it. + * Presumably, pytorch will support importing `byte_offset != 0` and NumPy + * can choose to use it starting about 2023. At that point, it may be + * that NumPy MUST use `byte_offset` to adhere to the standard (as + * specified in the header)! + */ + managed->dl_tensor.data = PyArray_DATA(self); + managed->dl_tensor.byte_offset = 0; managed->dl_tensor.device = device; managed->dl_tensor.dtype = managed_dtype;