8000 ENH: Support dlpack version 1 ABI · numpy/numpy@fb60521 · GitHub
[go: up one dir, main page]

Skip to content

Commit fb60521

Browse files
committed
ENH: Support dlpack version 1 ABI
1 parent 2a9b913 commit fb60521

File tree

6 files changed

+362
-111
lines changed

6 files changed

+362
-111
lines changed

numpy/_core/src/common/dlpack/dlpack.h

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// Taken from:
2-
// https://github.com/dmlc/dlpack/blob/ca4d00ad3e2e0f410eeab3264d21b8a39397f362/include/dlpack/dlpack.h
2+
// https://github.com/dmlc/dlpack/blob/bbd2f4d32427e548797929af08cfe2a9cbb3cf12/include/dlpack/dlpack.h
3+
// but added typedef to DLManagedTensorVersioned
34
/*!
45
* Copyright (c) 2017 by Contributors
56
* \file dlpack.h
@@ -118,6 +119,8 @@ typedef enum {
118119
kDLWebGPU = 15,
119120
/*! \brief Qualcomm Hexagon DSP */
120121
kDLHexagon = 16,
122+
/*! \brief Microsoft MAIA devices */
123+
kDLMAIA = 17,
121124
} DLDeviceType;
122125

123126
/*!
@@ -197,7 +200,7 @@ typedef struct {
197200
* `byte_offset` field should be used to point to the beginning of the data.
198201
*
199202
* Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow,
200-
* TVM, perhaps others) do not adhere to this 256 byte alignment requirement
203+
* TVM, perhaps others) do not adhere to this 256 byte aligment requirement
201204
* on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed
202205
* (after which this note will be updated); at the moment it is recommended
203206
* to not rely on the data pointer being correctly aligned.
@@ -215,6 +218,9 @@ typedef struct {
215218
* return size;
216219
* }
217220
* \endcode
221+
*
222+
* Note that if the tensor is of size zero, then the data pointer should be
223+
* set to `NULL`.
218224
*/
219225
void* data;
220226
/*! \brief The device of the tensor */
@@ -259,7 +265,7 @@ typedef struct DLManagedTensor {
2592 10000 65
* \brief Destructor - this should be called
260266
* to destruct the manager_ctx which backs the DLManagedTensor. It can be
261267
* NULL if there is no way for the caller to provide a reasonable destructor.
262-
* The destructors deletes the argument self as well.
268+
* The destructor deletes the argument self as well.
263269
*/
264270
void (*deleter)(struct DLManagedTensor * self);
265271
} DLManagedTensor;
@@ -269,6 +275,14 @@ typedef struct DLManagedTensor {
269275
/*! \brief bit mask to indicate that the tensor is read only. */
270276
#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL)
271277

278+
/*!
279+
* \brief bit mask to indicate that the tensor is a copy made by the producer.
280+
*
281+
* If set, the tensor is considered solely owned throughout its lifetime by the
282+
* consumer, until the producer-provided deleter is invoked.
283+
*/
284+
#define DLPACK_FLAG_BITMASK_IS_COPIED (1UL << 1UL)
285+
272286
/*!
273287
* \brief A versioned and managed C Tensor object, manage memory of DLTensor.
274288
*
@@ -279,7 +293,7 @@ typedef struct DLManagedTensor {
279293
*
280294
* \note This is the current standard DLPack exchange data structure.
281295
*/
282-
struct DLManagedTensorVersioned {
296+
typedef struct DLManagedTensorVersioned {
283297
/*!
284298
* \brief The API and ABI version of the current managed Tensor
285299
*/
@@ -296,7 +310,7 @@ struct DLManagedTensorVersioned {
296310
*
297311
* This should be called to destruct manager_ctx which holds the DLManagedTensorVersioned.
298312
* It can be NULL if there is no way for the caller to provide a reasonable
299-
* destructor. The destructors deletes the argument self as well.
313+
* destructor. The destructor deletes the argument self as well.
300314
*/
301315
void (*deleter)(struct DLManagedTensorVersioned *self);
302316
/*!
@@ -308,11 +322,12 @@ struct DLManagedTensorVersioned {
308322
* stable, to ensure that deleter can be correctly called.
309323
*
310324
* \sa DLPACK_FLAG_BITMASK_READ_ONLY
325+
* \sa DLPACK_FLAG_BITMASK_IS_COPIED
311326
*/
312327
uint64_t flags;
313328
/*! \brief DLTensor which is being memory managed */
314329
DLTensor dl_tensor;
315-
};
330+
} DLManagedTensorVersioned;
316331

317332
#ifdef __cplusplus
318333
} // DLPACK_EXTERN_C

numpy/_core/src/common/npy_dlpack.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,15 @@
66

77
// Part of the Array API specification.
88
#define NPY_DLPACK_CAPSULE_NAME "dltensor"
9+
#define NPY_DLPACK_VERSIONED_CAPSULE_NAME "dltensor_versioned"
910
#define NPY_DLPACK_USED_CAPSULE_NAME "used_dltensor"
11+
#define NPY_DLPACK_VERSIONED_USED_CAPSULE_NAME "used_dltensor_versioned"
1012

1113
// Used internally by NumPy to store a base object
1214
// as it has to release a reference to the original
1315
// capsule.
1416
#define NPY_DLPACK_INTERNAL_CAPSULE_NAME "numpy_dltensor"
17+
#define NPY_DLPACK_VERSIONED_INTERNAL_CAPSULE_NAME "numpy_dltensor_versioned"
1518

1619
PyObject *
1720
array_dlpack(PyArrayObject *self, PyObject *const *args, Py_ssize_t len_args,

0 commit comments

Comments
 (0)
0