summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--numpy/core/code_generators/genapi.py1
-rw-r--r--numpy/core/setup.py1
-rw-r--r--numpy/core/src/common/npy_dlpack.h14
-rw-r--r--numpy/core/src/multiarray/dlpack.c382
-rw-r--r--numpy/core/src/multiarray/dlpack.h18
-rw-r--r--numpy/core/src/multiarray/methods.c238
-rw-r--r--numpy/core/src/multiarray/multiarraymodule.c158
7 files changed, 419 insertions, 393 deletions
diff --git a/numpy/core/code_generators/genapi.py b/numpy/core/code_generators/genapi.py
index c2458c2b5..b401ee6a5 100644
--- a/numpy/core/code_generators/genapi.py
+++ b/numpy/core/code_generators/genapi.py
@@ -41,6 +41,7 @@ API_FILES = [join('multiarray', 'alloc.c'),
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 5c6fd4dcf..2c99060ec 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -883,6 +883,7 @@ def configuration(parent_package='',top_path=None):
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 446d44afa..cb926a262 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 000000000..591eddfaf
--- /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 <Python.h>
+
+#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 000000000..7d4a289d9
--- /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 <Python.h>
+#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 d3162e549..f40db492c 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 <stdarg.h>
-#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 60e3875fc..ac4c1ae0b 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,