From 36c16db5b4c77ba67be4da3a464f2b89e3c5283a Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 6 Aug 2026 15:23:31 +0200 Subject: [PATCH 1/5] Bump DLPack 1.3.0 --- cpp/src/arrow/c/dlpack_abi.h | 356 +++++++++++++++++++++++++++++++++-- 1 file changed, 343 insertions(+), 13 deletions(-) diff --git a/cpp/src/arrow/c/dlpack_abi.h b/cpp/src/arrow/c/dlpack_abi.h index fbe2a56a344b..f705f75c157b 100644 --- a/cpp/src/arrow/c/dlpack_abi.h +++ b/cpp/src/arrow/c/dlpack_abi.h @@ -1,7 +1,8 @@ // Taken from: -// https://github.com/dmlc/dlpack/blob/ca4d00ad3e2e0f410eeab3264d21b8a39397f362/include/dlpack/dlpack.h +// https://github.com/dmlc/dlpack/blob/84d107bf416c6bab9ae68ad285876600d230490d/include/dlpack/dlpack.h + /*! - * Copyright (c) 2017 by Contributors + * Copyright (c) 2017 - by Contributors * \file dlpack.h * \brief The common header of DLPack. */ @@ -21,7 +22,7 @@ #define DLPACK_MAJOR_VERSION 1 /*! \brief The current minor version of dlpack */ -#define DLPACK_MINOR_VERSION 0 +#define DLPACK_MINOR_VERSION 3 /*! \brief DLPACK_DLL prefix for windows */ #ifdef _WIN32 @@ -118,6 +119,10 @@ typedef enum { kDLWebGPU = 15, /*! \brief Qualcomm Hexagon DSP */ kDLHexagon = 16, + /*! \brief Microsoft MAIA devices */ + kDLMAIA = 17, + /*! \brief AWS Trainium */ + kDLTrn = 18, } DLDeviceType; /*! @@ -157,6 +162,26 @@ typedef enum { kDLComplex = 5U, /*! \brief boolean */ kDLBool = 6U, + /*! \brief FP8 data types */ + kDLFloat8_e3m4 = 7U, + kDLFloat8_e4m3 = 8U, + kDLFloat8_e4m3b11fnuz = 9U, + kDLFloat8_e4m3fn = 10U, + kDLFloat8_e4m3fnuz = 11U, + kDLFloat8_e5m2 = 12U, + kDLFloat8_e5m2fnuz = 13U, + kDLFloat8_e8m0fnu = 14U, + /*! \brief FP6 data types + * Setting bits != 6 is currently unspecified, and the producer must ensure it is set + * while the consumer must stop importing if the value is unexpected. + */ + kDLFloat6_e2m3fn = 15U, + kDLFloat6_e3m2fn = 16U, + /*! \brief FP4 data types + * Setting bits != 4 is currently unspecified, and the producer must ensure it is set + * while the consumer must stop importing if the value is unexpected. + */ + kDLFloat4_e2m1fn = 17U, } DLDataTypeCode; /*! @@ -171,6 +196,12 @@ typedef enum { * - std::complex: type_code = 5, bits = 64, lanes = 1 * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, * the underlying storage size of bool is 8 bits) + * - float8_e4m3: type_code = 8, bits = 8, lanes = 1 (packed in memory) + * - float6_e3m2fn: type_code = 16, bits = 6, lanes = 1 (packed in memory) + * - float4_e2m1fn: type_code = 17, bits = 4, lanes = 1 (packed in memory) + * + * When a sub-byte type is packed, DLPack requires the data to be in little bit-endian, + * i.e., for a packed data set D ((D >> (i * bits)) && bit_mask) stores the i-th element. */ typedef struct { /*! @@ -197,8 +228,8 @@ typedef struct { * types. This pointer is always aligned to 256 bytes as in CUDA. The * `byte_offset` field should be used to point to the beginning of the data. * - * Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow, - * TVM, perhaps others) do not adhere to this 256 byte aligment requirement + * Note that as of Nov 2021, multiple libraries (CuPy, PyTorch, TensorFlow, + * TVM, perhaps others) do not adhere to this 256 byte alignment requirement * on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed * (after which this note will be updated); at the moment it is recommended * to not rely on the data pointer being correctly aligned. @@ -216,6 +247,9 @@ typedef struct { * return size; * } * \endcode + * + * Note that if the tensor is of size zero, then the data pointer should be + * set to `NULL`. */ void* data; /*! \brief The device of the tensor */ @@ -224,11 +258,23 @@ typedef struct { int32_t ndim; /*! \brief The data type of the pointer*/ DLDataType dtype; - /*! \brief The shape of the tensor */ + /*! + * \brief The shape of the tensor + * + * When ndim == 0, shape can be set to NULL. + */ int64_t* shape; /*! - * \brief strides of the tensor (in number of elements, not bytes) - * can be NULL, indicating tensor is compact and row-majored. + * \brief strides of the tensor (in number of elements, not bytes), + * can not be NULL if ndim != 0, must points to + * an array of ndim elements that specifies the strides, + * so consumer can always rely on strides[dim] being valid for 0 <= dim < ndim. + * + * When ndim == 0, strides can be set to NULL. + * + * \note Before DLPack v1.2, strides can be NULL to indicate contiguous data. + * This is not allowed in DLPack v1.2 and later. The rationale + * is to simplify the consumer handling. */ int64_t* strides; /*! \brief The offset in bytes to the beginning pointer to data */ @@ -260,16 +306,32 @@ typedef struct DLManagedTensor { * \brief Destructor - this should be called * to destruct the manager_ctx which backs 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. + * The destructor deletes the argument self as well. */ void (*deleter)(struct DLManagedTensor* self); } DLManagedTensor; -// bit masks used in in the DLManagedTensorVersioned +// bit masks used in the DLManagedTensorVersioned /*! \brief bit mask to indicate that the tensor is read only. */ #define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL) +/*! + * \brief bit mask to indicate that the tensor is a copy made by the producer. + * + * If set, the tensor is considered solely owned throughout its lifetime by the + * consumer, until the producer-provided deleter is invoked. + */ +#define DLPACK_FLAG_BITMASK_IS_COPIED (1UL << 1UL) + +/*! + * \brief bit mask to indicate that whether a sub-byte type is packed or padded. + * + * The default for sub-byte types (ex: fp4/fp6) is assumed packed. This flag can + * be set by the producer to signal that a tensor of sub-byte type is padded. + */ +#define DLPACK_FLAG_BITMASK_IS_SUBBYTE_TYPE_PADDED (1UL << 2UL) + /*! * \brief A versioned and managed C Tensor object, manage memory of DLTensor. * @@ -280,7 +342,7 @@ typedef struct DLManagedTensor { * * \note This is the current standard DLPack exchange data structure. */ -struct DLManagedTensorVersioned { +typedef struct DLManagedTensorVersioned { /*! * \brief The API and ABI version of the current managed Tensor */ @@ -297,7 +359,7 @@ struct DLManagedTensorVersioned { * * This should be called to destruct manager_ctx which holds the * DLManagedTensorVersioned. 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. + * a reasonable destructor. The destructor deletes the argument self as well. */ void (*deleter)(struct DLManagedTensorVersioned* self); /*! @@ -309,11 +371,279 @@ struct DLManagedTensorVersioned { * stable, to ensure that deleter can be correctly called. * * \sa DLPACK_FLAG_BITMASK_READ_ONLY + * \sa DLPACK_FLAG_BITMASK_IS_COPIED */ uint64_t flags; /*! \brief DLTensor which is being memory managed */ DLTensor dl_tensor; -}; +} DLManagedTensorVersioned; + +//---------------------------------------------------------------------- +// DLPack `__dlpack_c_exchange_api__` fast exchange protocol definitions +//---------------------------------------------------------------------- +/*! + * \brief Request a producer library to create a new tensor. + * + * Create a new `DLManagedTensorVersioned` within the context of the producer + * library. The allocation is defined via the prototype DLTensor. + * + * This function is exposed by the framework through the DLPackExchangeAPI. + * + * \param prototype The prototype DLTensor. Only the dtype, ndim, shape, + * and device fields are used. + * \param out The output DLManagedTensorVersioned. + * \param error_ctx Context for `SetError`. + * \param SetError The function to set the error. + * \return The owning DLManagedTensorVersioned* or NULL on failure. + * SetError is called exactly when NULL is returned (the implementer + * must ensure this). + * \note - As a C function, must not thrown C++ exceptions. + * - Error propagation via SetError to avoid any direct need + * of Python API. Due to this `SetError` may have to ensure the GIL is + * held since it will presumably set a Python error. + * + * \sa DLPackExchangeAPI + */ +typedef int (*DLPackManagedTensorAllocator)( // + DLTensor* prototype, DLManagedTensorVersioned** out, void* error_ctx, // + void (*SetError)(void* error_ctx, const char* kind, const char* message) // +); + +/*! + * \brief Exports a PyObject* Tensor/NDArray to a DLManagedTensorVersioned. + * + * This function does not perform any stream synchronization. The consumer should query + * DLPackCurrentWorkStream to get the current work stream and launch kernels on it. + * + * This function is exposed by the framework through the DLPackExchangeAPI. + * + * \param py_object The Python object to convert. Must have the same type + * as the one the `DLPackExchangeAPI` was discovered from. + * \param out The output DLManagedTensorVersioned. + * \return The owning DLManagedTensorVersioned* or NULL on failure with a + * Python exception set. If the data cannot be described using DLPack + * this should be a BufferError if possible. + * \note - As a C function, must not thrown C++ exceptions. + * + * \sa DLPackExchangeAPI, DLPackCurrentWorkStream + */ +typedef int (*DLPackManagedTensorFromPyObjectNoSync)( // + void* py_object, // + DLManagedTensorVersioned** out // +); + +/*! + * \brief Exports a PyObject* Tensor/NDArray to a provided DLTensor. + * + * This function provides a faster interface for temporary, non-owning, exchange. + * The producer (implementer) still owns the memory of data, strides, shape. + * The liveness of the DLTensor and the data it views is only guaranteed until + * control is returned. + * + * This function currently assumes that the producer (implementer) can fill + * in the DLTensor shape and strides without the need for temporary allocations. + * + * This function does not perform any stream synchronization. The consumer should query + * DLPackCurrentWorkStream to get the current work stream and launch kernels on it. + * + * This function is exposed by the framework through the DLPackExchangeAPI. + * + * \param py_object The Python object to convert. Must have the same type + * as the one the `DLPackExchangeAPI` was discovered from. + * \param out The output DLTensor, whose space is pre-allocated on stack. + * \return 0 on success, -1 on failure with a Python exception set. + * \note - As a C function, must not thrown C++ exceptions. + * + * \sa DLPackExchangeAPI, DLPackCurrentWorkStream + */ +typedef int (*DLPackDLTensorFromPyObjectNoSync)( // + void* py_object, // + DLTensor* out // +); + +/*! + * \brief Obtain the current work stream of a device. + * + * Obtain the current work stream of a device from the producer framework. + * For example, it should map to torch.cuda.current_stream in PyTorch. + * + * When device_type is kDLCPU, the consumer do not have to query the stream + * and the producer can simply return NULL when queried. + * The consumer do not have to do anything on stream sync or setting. + * So CPU only framework can just provide a dummy implementation that + * always set out_current_stream[0] to NULL. + * + * \param device_type The device type. + * \param device_id The device id. + * \param out_current_stream The output current work stream. + * + * \return 0 on success, -1 on failure with a Python exception set. + * \note - As a C function, must not thrown C++ exceptions. + * + * \sa DLPackExchangeAPI + */ +typedef int (*DLPackCurrentWorkStream)( // + DLDeviceType device_type, // + int32_t device_id, // + void** out_current_stream // +); + +/*! + * \brief Imports a DLManagedTensorVersioned to a PyObject* Tensor/NDArray. + * + * Convert an owning DLManagedTensorVersioned* to the Python tensor of the + * producer (implementer) library with the correct type. + * + * This function does not perform any stream synchronization. + * + * This function is exposed by the framework through the DLPackExchangeAPI. + * + * \param tensor The DLManagedTensorVersioned to convert the ownership of the + * tensor is stolen. + * \param out_py_object The output Python object. + * \return 0 on success, -1 on failure with a Python exception set. + * + * \sa DLPackExchangeAPI + */ +typedef int (*DLPackManagedTensorToPyObjectNoSync)( // + DLManagedTensorVersioned* tensor, // + void** out_py_object // +); + +/*! + * \brief DLPackExchangeAPI stable header. + * \sa DLPackExchangeAPI + */ +typedef struct DLPackExchangeAPIHeader { + /*! + * \brief The provided DLPack version the consumer must check major version + * compatibility before using this struct. + */ + DLPackVersion version; + /*! + * \brief Optional pointer to an older DLPackExchangeAPI in the chain. + * + * It must be NULL if the framework does not support older versions. + * If the current major version is larger than the one supported by the + * consumer, the consumer may walk this to find an earlier supported version. + * + * \sa DLPackExchangeAPI + */ + struct DLPackExchangeAPIHeader* prev_api; +} DLPackExchangeAPIHeader; + +/*! + * \brief Framework-specific function pointers table for DLPack exchange. + * + * Additionally to `__dlpack__()` we define a C function table sharable by + * + * Python implementations via `__dlpack_c_exchange_api__`. + * This attribute must be set on the type as a Python PyCapsule + * with name "dlpack_exchange_api". + * + * A consumer library may use a pattern such as: + * + * \code + * + * PyObject *api_capsule = PyObject_GetAttrString( + * (PyObject *)Py_TYPE(tensor_obj), "__dlpack_c_exchange_api__") + * ); + * if (api_capsule == NULL) { goto handle_error; } + * MyDLPackExchangeAPI *api = (MyDLPackExchangeAPI *)PyCapsule_GetPointer( + * api_capsule, "dlpack_exchange_api" + * ); + * Py_DECREF(api_capsule); + * if (api == NULL) { goto handle_error; } + * + * \endcode + * + * Note that this must be defined on the type. The consumer should look up the + * attribute on the type and may cache the result for each unique type. + * + * The precise API table is given by: + * \code + * struct MyDLPackExchangeAPI : public DLPackExchangeAPI { + * MyDLPackExchangeAPI() { + * header.version.major = DLPACK_MAJOR_VERSION; + * header.version.minor = DLPACK_MINOR_VERSION; + * header.prev_version_api = nullptr; + * + * managed_tensor_allocator = MyDLPackManagedTensorAllocator; + * managed_tensor_from_py_object_no_sync = MyDLPackManagedTensorFromPyObjectNoSync; + * managed_tensor_to_py_object_no_sync = MyDLPackManagedTensorToPyObjectNoSync; + * dltensor_from_py_object_no_sync = MyDLPackDLTensorFromPyObjectNoSync; + * current_work_stream = MyDLPackCurrentWorkStream; + * } + * + * static const DLPackExchangeAPI* Global() { + * static MyDLPackExchangeAPI inst; + * return &inst; + * } + * }; + * \endcode + * + * Guidelines for leveraging DLPackExchangeAPI: + * + * There are generally two kinds of consumer needs for DLPack exchange: + * - N0: library support, where consumer.kernel(x, y, z) would like to run a kernel + * with the data from x, y, z. The consumer is also expected to run the kernel with + * the same stream context as the producer. For example, when x, y, z is torch.Tensor, + * consumer should query exchange_api->current_work_stream to get the + * current stream and launch the kernel with the same stream. + * This setup is necessary for no synchronization in kernel launch and maximum + * compatibility with CUDA graph capture in the producer. This is the desirable behavior + * for library extension support for frameworks like PyTorch. + * - N1: data ingestion and retention + * + * Note that obj.__dlpack__() API should provide useful ways for N1. + * The primary focus of the current DLPackExchangeAPI is to enable faster exchange N0 + * with the support of the function pointer current_work_stream. + * + * Array/Tensor libraries should statically create and initialize this structure + * then return a pointer to DLPackExchangeAPI as an int value in Tensor/Array. + * The DLPackExchangeAPI* must stay alive throughout the lifetime of the process. + * + * One simple way to do so is to create a static instance of DLPackExchangeAPI + * within the framework and return a pointer to it. The following code + * shows an example to do so in C++. It should also be reasonably easy + * to do so in other languages. + */ +typedef struct DLPackExchangeAPI { + /*! + * \brief The header that remains stable across versions. + */ + DLPackExchangeAPIHeader header; + /*! + * \brief Producer function pointer for DLPackManagedTensorAllocator + * This function must not be NULL. + * \sa DLPackManagedTensorAllocator + */ + DLPackManagedTensorAllocator managed_tensor_allocator; + /*! + * \brief Producer function pointer for DLPackManagedTensorFromPyObject + * This function must be not NULL. + * \sa DLPackManagedTensorFromPyObject + */ + DLPackManagedTensorFromPyObjectNoSync managed_tensor_from_py_object_no_sync; + /*! + * \brief Producer function pointer for DLPackManagedTensorToPyObject + * This function must be not NULL. + * \sa DLPackManagedTensorToPyObjectNoSync + */ + DLPackManagedTensorToPyObjectNoSync managed_tensor_to_py_object_no_sync; + /*! + * \brief Producer function pointer for DLPackDLTensorFromPyObject + * This function can be NULL when the producer does not support this function. + * \sa DLPackDLTensorFromPyObjectNoSync + */ + DLPackDLTensorFromPyObjectNoSync dltensor_from_py_object_no_sync; + /*! + * \brief Producer function pointer for DLPackCurrentWorkStream + * This function must be not NULL. + * \sa DLPackCurrentWorkStream + */ + DLPackCurrentWorkStream current_work_stream; +} DLPackExchangeAPI; #ifdef __cplusplus } // DLPACK_EXTERN_C From 4c0ed7a170377331255df5a03c8bc6687b5f5988 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 6 Aug 2026 15:49:57 +0200 Subject: [PATCH 2/5] Unsure non-null strides --- cpp/src/arrow/c/dlpack.cc | 4 +++- cpp/src/arrow/c/dlpack_test.cc | 6 +++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc index 94d7816d78e6..0c1cc2bed5dc 100644 --- a/cpp/src/arrow/c/dlpack.cc +++ b/cpp/src/arrow/c/dlpack.cc @@ -62,6 +62,7 @@ Result GetDLDataType(const DataType& type) { struct ManagerCtx { std::shared_ptr array; DLManagedTensor tensor; + int64_t strides = 1; }; } // namespace @@ -94,8 +95,9 @@ Result ExportArray(const std::shared_ptr& arr) { ctx->tensor.dl_tensor.ndim = 1; ctx->tensor.dl_tensor.dtype = dlpack_type; ctx->tensor.dl_tensor.shape = const_cast(&data->length); - ctx->tensor.dl_tensor.strides = NULL; ctx->tensor.dl_tensor.byte_offset = 0; + // Strides must be non-null when ndim > 0 + ctx->tensor.dl_tensor.strides = &ctx->strides; ctx->array = std::move(data); ctx->tensor.manager_ctx = ctx.get(); diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc index f0119e8aef71..9e35a771d295 100644 --- a/cpp/src/arrow/c/dlpack_test.cc +++ b/cpp/src/arrow/c/dlpack_test.cc @@ -28,7 +28,7 @@ namespace arrow::dlpack { class TestExportArray : public ::testing::Test { public: - void SetUp() {} + void SetUp() override {} }; void CheckDLTensor(const std::shared_ptr& arr, @@ -44,9 +44,9 @@ void CheckDLTensor(const std::shared_ptr& arr, ASSERT_EQ(sliced_buffer->data(), dltensor.data); ASSERT_EQ(0, dltensor.byte_offset); - ASSERT_EQ(NULL, dltensor.strides); ASSERT_EQ(length, dltensor.shape[0]); ASSERT_EQ(1, dltensor.ndim); + ASSERT_EQ(1, *dltensor.strides); // Must be non-null with ndim>0 since 1.2 ASSERT_EQ(dlpack_type, dltensor.dtype.code); ASSERT_EQ(arrow_type->bit_width(), dltensor.dtype.bits); @@ -128,7 +128,7 @@ TEST_F(TestExportArray, TestErrors) { class TestExportTensor : public ::testing::Test { public: - void SetUp() {} + void SetUp() override {} }; void CheckDLTensor(const std::shared_ptr& t, From 4a95262f98fa9ca74d8ea23b53ae6a7b8099fc4c Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Thu, 6 Aug 2026 16:18:47 +0200 Subject: [PATCH 3/5] Add DLManagedTensorVersioned --- cpp/src/arrow/c/dlpack.cc | 81 +++++++++++++++++------ cpp/src/arrow/c/dlpack.h | 8 +++ cpp/src/arrow/c/dlpack_test.cc | 116 ++++++++++++++++++++++++++------- 3 files changed, 162 insertions(+), 43 deletions(-) diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc index 0c1cc2bed5dc..05ec4f0d83b0 100644 --- a/cpp/src/arrow/c/dlpack.cc +++ b/cpp/src/arrow/c/dlpack.cc @@ -17,6 +17,10 @@ #include "arrow/c/dlpack.h" +#include +#include +#include + #include "arrow/array/array_base.h" #include "arrow/c/dlpack_abi.h" #include "arrow/device.h" @@ -59,32 +63,35 @@ Result GetDLDataType(const DataType& type) { } } +template struct ManagerCtx { std::shared_ptr array; - DLManagedTensor tensor; + DT tensor; int64_t strides = 1; }; -} // namespace - -Result ExportArray(const std::shared_ptr& arr) { +template +Result ExportArrayImpl(const std::shared_ptr& arr) { // Define DLDevice struct and check if array type is supported // by the DLPack protocol at the same time. Raise TypeError if not. // Supported data types: int, uint, float with no validity buffer. - ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(arr)) + ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(arr)); // Define the DLDataType struct const DataType& type = *arr->type(); - std::shared_ptr data = arr->data(); ARROW_ASSIGN_OR_RAISE(auto dlpack_type, GetDLDataType(type)); // Create ManagerCtx that will serve as the owner of the DLManagedTensor - auto ctx = std::make_unique(); + auto ctx = std::make_unique>(); + + // Assign the Array data into the context + ctx->array = arr->data(); + auto& data = ctx->array; // Define the data pointer to the DLTensor // If array is of length 0, data pointer should be NULL if (arr->length() == 0) { - ctx->tensor.dl_tensor.data = NULL; + ctx->tensor.dl_tensor.data = nullptr; } else { const auto data_offset = data->offset * type.byte_width(); ctx->tensor.dl_tensor.data = @@ -98,15 +105,30 @@ Result ExportArray(const std::shared_ptr& arr) { ctx->tensor.dl_tensor.byte_offset = 0; // Strides must be non-null when ndim > 0 ctx->tensor.dl_tensor.strides = &ctx->strides; + if constexpr (std::is_same_v) { + ctx->tensor.version = {.major = DLPACK_MAJOR_VERSION, .minor = DLPACK_MINOR_VERSION}; + // Arrow contract is that array data is immutable once constructed + ctx->tensor.flags = DLPACK_FLAG_BITMASK_READ_ONLY; + } - ctx->array = std::move(data); ctx->tensor.manager_ctx = ctx.get(); - ctx->tensor.deleter = [](struct DLManagedTensor* self) { - delete reinterpret_cast(self->manager_ctx); + ctx->tensor.deleter = [](DT* self) { + delete reinterpret_cast*>(self->manager_ctx); }; return &ctx.release()->tensor; } +} // namespace + +Result ExportArray(const std::shared_ptr& arr) { + return ExportArrayImpl(arr); +} + +Result ExportArrayVersioned( + const std::shared_ptr& arr) { + return ExportArrayImpl(arr); +} + Result ExportDevice(const std::shared_ptr& arr) { // Check if array is supported by the DLPack protocol. if (arr->null_count() > 0) { @@ -133,30 +155,34 @@ Result ExportDevice(const std::shared_ptr& arr) { } } +namespace { + +template struct TensorManagerCtx { std::shared_ptr t; std::vector strides; std::vector shape; - DLManagedTensor tensor; + DT tensor; }; -Result ExportTensor(const std::shared_ptr& t) { +template +Result ExportTensorImpl(const std::shared_ptr& t) { // Define the DLDataType struct const DataType& type = *t->type(); ARROW_ASSIGN_OR_RAISE(auto dlpack_type, GetDLDataType(type)); // Define DLDevice struct - ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(t)) + ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(t)); // Create TensorManagerCtx that will serve as the owner of the DLManagedTensor - auto ctx = std::make_unique(); + auto ctx = std::make_unique>(); // Define the data pointer to the DLTensor // If tensor is of length 0, data pointer should be NULL if (t->size() == 0) { - ctx->tensor.dl_tensor.data = NULL; + ctx->tensor.dl_tensor.data = nullptr; } else { - ctx->tensor.dl_tensor.data = t->raw_mutable_data(); + ctx->tensor.dl_tensor.data = const_cast(t->raw_data()); } ctx->tensor.dl_tensor.device = device; @@ -173,20 +199,35 @@ Result ExportTensor(const std::shared_ptr& t) { std::vector& strides_arr = ctx->strides; strides_arr.reserve(t->ndim()); - auto byte_width = t->type()->byte_width(); + const auto byte_width = t->type()->byte_width(); for (auto i : t->strides()) { strides_arr.emplace_back(i / byte_width); } ctx->tensor.dl_tensor.strides = strides_arr.data(); + if constexpr (std::is_same_v) { + ctx->tensor.version = {.major = DLPACK_MAJOR_VERSION, .minor = DLPACK_MINOR_VERSION}; + ctx->tensor.flags = t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY; + } ctx->t = std::move(t); ctx->tensor.manager_ctx = ctx.get(); - ctx->tensor.deleter = [](struct DLManagedTensor* self) { - delete reinterpret_cast(self->manager_ctx); + ctx->tensor.deleter = [](DT* self) { + delete reinterpret_cast*>(self->manager_ctx); }; return &ctx.release()->tensor; } +} // namespace + +Result ExportTensor(const std::shared_ptr& t) { + return ExportTensorImpl(t); +} + +Result ExportTensorVersioned( + const std::shared_ptr& t) { + return ExportTensorImpl(t); +} + Result ExportDevice(const std::shared_ptr& t) { // Define DLDevice struct DLDevice device; diff --git a/cpp/src/arrow/c/dlpack.h b/cpp/src/arrow/c/dlpack.h index 65da38423c2a..776f0be15800 100644 --- a/cpp/src/arrow/c/dlpack.h +++ b/cpp/src/arrow/c/dlpack.h @@ -34,14 +34,22 @@ namespace arrow::dlpack { /// memory region which means Arrow Arrays with validity buffers /// are not supported. /// +/// \note Deprecated in DLPack 1.0. Use ExportArrayVersioned instead. +/// /// \param[in] arr Arrow array /// \return DLManagedTensor struct ARROW_EXPORT Result ExportArray(const std::shared_ptr& arr); +ARROW_EXPORT +Result ExportArrayVersioned(const std::shared_ptr& arr); + ARROW_EXPORT Result ExportTensor(const std::shared_ptr& t); +ARROW_EXPORT +Result ExportTensorVersioned(const std::shared_ptr& t); + /// \brief Get DLDevice with enumerator specifying the /// type of the device data is stored on and index of the /// device which is 0 by default for CPU. diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc index 9e35a771d295..892f7b5c7e7d 100644 --- a/cpp/src/arrow/c/dlpack_test.cc +++ b/cpp/src/arrow/c/dlpack_test.cc @@ -17,7 +17,11 @@ #include +#include +#include + #include "arrow/array/array_base.h" +#include "arrow/buffer.h" #include "arrow/c/dlpack.h" #include "arrow/c/dlpack_abi.h" #include "arrow/memory_pool.h" @@ -26,15 +30,49 @@ namespace arrow::dlpack { -class TestExportArray : public ::testing::Test { - public: - void SetUp() override {} +struct LegacyProducer { + using ManagedTensor = DLManagedTensor; + static constexpr const char* name = "Legacy"; + + static Result Export(const std::shared_ptr& arr) { + return ExportArray(arr); + } + static Result Export(const std::shared_ptr& t) { + return ExportTensor(t); + } +}; + +struct VersionedProducer { + using ManagedTensor = DLManagedTensorVersioned; + static constexpr const char* name = "Versioned"; + + static Result Export(const std::shared_ptr& arr) { + return ExportArrayVersioned(arr); + } + static Result Export(const std::shared_ptr& t) { + return ExportTensorVersioned(t); + } +}; + +using ProducerTypes = ::testing::Types; + +struct ProducerNames { + template + static std::string GetName(int) { + return Producer::name; + } }; +template +class TestExportArray : public ::testing::Test {}; + +TYPED_TEST_SUITE(TestExportArray, ProducerTypes, ProducerNames); + +template void CheckDLTensor(const std::shared_ptr& arr, const std::shared_ptr& arrow_type, DLDataTypeCode dlpack_type, int64_t length) { - ASSERT_OK_AND_ASSIGN(auto dlmtensor, arrow::dlpack::ExportArray(arr)); + ASSERT_OK_AND_ASSIGN(auto dlmtensor, Producer::Export(arr)); auto dltensor = dlmtensor->dl_tensor; const auto byte_width = arr->type()->byte_width(); @@ -58,10 +96,17 @@ void CheckDLTensor(const std::shared_ptr& arr, ASSERT_EQ(DLDeviceType::kDLCPU, device.device_type); ASSERT_EQ(0, device.device_id); + if constexpr (std::is_same_v) { + ASSERT_EQ(DLPACK_MAJOR_VERSION, dlmtensor->version.major); + ASSERT_EQ(DLPACK_MINOR_VERSION, dlmtensor->version.minor); + // Arrow array data is immutable once constructed + ASSERT_EQ(DLPACK_FLAG_BITMASK_READ_ONLY, dlmtensor->flags); + } + dlmtensor->deleter(dlmtensor); } -TEST_F(TestExportArray, TestSupportedArray) { +TYPED_TEST(TestExportArray, TestSupportedArray) { const std::vector, DLDataTypeCode>> cases = { {int8(), DLDataTypeCode::kDLInt}, {uint8(), DLDataTypeCode::kDLUInt}, @@ -89,36 +134,36 @@ TEST_F(TestExportArray, TestSupportedArray) { for (auto [arrow_type, dlpack_type] : cases) { const std::shared_ptr array = ArrayFromJSON(arrow_type, "[1, 0, 10, 0, 2, 1, 3, 5, 1, 0]"); - CheckDLTensor(array, arrow_type, dlpack_type, 10); + CheckDLTensor(array, arrow_type, dlpack_type, 10); ASSERT_OK_AND_ASSIGN(auto sliced_1, array->SliceSafe(1, 5)); - CheckDLTensor(sliced_1, arrow_type, dlpack_type, 5); + CheckDLTensor(sliced_1, arrow_type, dlpack_type, 5); ASSERT_OK_AND_ASSIGN(auto sliced_2, array->SliceSafe(0, 5)); - CheckDLTensor(sliced_2, arrow_type, dlpack_type, 5); + CheckDLTensor(sliced_2, arrow_type, dlpack_type, 5); ASSERT_OK_AND_ASSIGN(auto sliced_3, array->SliceSafe(3)); - CheckDLTensor(sliced_3, arrow_type, dlpack_type, 7); + CheckDLTensor(sliced_3, arrow_type, dlpack_type, 7); } ASSERT_EQ(allocated_bytes, arrow::default_memory_pool()->bytes_allocated()); } -TEST_F(TestExportArray, TestErrors) { +TYPED_TEST(TestExportArray, TestErrors) { const std::shared_ptr array_null = ArrayFromJSON(null(), "[]"); ASSERT_RAISES_WITH_MESSAGE(TypeError, "Type error: DataType is not compatible with DLPack spec: " + array_null->type()->ToString(), - arrow::dlpack::ExportArray(array_null)); + TypeParam::Export(array_null)); const std::shared_ptr array_with_null = ArrayFromJSON(int8(), "[1, 100, null]"); ASSERT_RAISES_WITH_MESSAGE(TypeError, "Type error: Can only use DLPack on arrays with no nulls.", - arrow::dlpack::ExportArray(array_with_null)); + TypeParam::Export(array_with_null)); const std::shared_ptr array_string = ArrayFromJSON(utf8(), R"(["itsy", "bitsy", "spider"])"); ASSERT_RAISES_WITH_MESSAGE(TypeError, "Type error: DataType is not compatible with DLPack spec: " + array_string->type()->ToString(), - arrow::dlpack::ExportArray(array_string)); + TypeParam::Export(array_string)); const std::shared_ptr array_boolean = ArrayFromJSON(boolean(), "[true, false]"); ASSERT_RAISES_WITH_MESSAGE( @@ -126,16 +171,17 @@ TEST_F(TestExportArray, TestErrors) { arrow::dlpack::ExportDevice(array_boolean)); } -class TestExportTensor : public ::testing::Test { - public: - void SetUp() override {} -}; +template +class TestExportTensor : public ::testing::Test {}; + +TYPED_TEST_SUITE(TestExportTensor, ProducerTypes, ProducerNames); +template void CheckDLTensor(const std::shared_ptr& t, const std::shared_ptr& tensor_type, DLDataTypeCode dlpack_type, std::vector shape, std::vector strides) { - ASSERT_OK_AND_ASSIGN(auto dlmtensor, arrow::dlpack::ExportTensor(t)); + ASSERT_OK_AND_ASSIGN(auto dlmtensor, Producer::Export(t)); auto dltensor = dlmtensor->dl_tensor; ASSERT_EQ(t->data()->data(), dltensor.data); @@ -156,10 +202,16 @@ void CheckDLTensor(const std::shared_ptr& t, ASSERT_EQ(DLDeviceType::kDLCPU, device.device_type); ASSERT_EQ(0, device.device_id); + if constexpr (std::is_same_v) { + ASSERT_EQ(DLPACK_MAJOR_VERSION, dlmtensor->version.major); + ASSERT_EQ(DLPACK_MINOR_VERSION, dlmtensor->version.minor); + ASSERT_EQ(t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY, dlmtensor->flags); + } + dlmtensor->deleter(dlmtensor); } -TEST_F(TestExportTensor, TestTensor) { +TYPED_TEST(TestExportTensor, TestTensor) { const std::vector, DLDataTypeCode>> cases = { {int8(), DLDataTypeCode::kDLInt}, {uint8(), DLDataTypeCode::kDLUInt}, @@ -190,13 +242,29 @@ TEST_F(TestExportTensor, TestTensor) { std::shared_ptr tensor = TensorFromJSON( arrow_type, "[1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9]", shape); - CheckDLTensor(tensor, arrow_type, dlpack_type, shape, dlpack_strides); + CheckDLTensor(tensor, arrow_type, dlpack_type, shape, dlpack_strides); } ASSERT_EQ(allocated_bytes, arrow::default_memory_pool()->bytes_allocated()); } -TEST_F(TestExportTensor, TestTensorStrided) { +TYPED_TEST(TestExportTensor, TestTensorReadOnly) { + const std::vector shape = {2, 2}; + const std::vector dlpack_strides = {2, 1}; + std::shared_ptr tensor = TensorFromJSON(float32(), "[1, 2, 3, 4]", shape); + ASSERT_TRUE(tensor->is_mutable()); + + // Slicing yields an immutable view of the same data + ASSERT_OK_AND_ASSIGN(auto read_only_buffer, SliceBufferSafe(tensor->data(), 0)); + ASSERT_OK_AND_ASSIGN(auto read_only_tensor, + Tensor::Make(float32(), read_only_buffer, shape)); + ASSERT_FALSE(read_only_tensor->is_mutable()); + + CheckDLTensor(read_only_tensor, float32(), DLDataTypeCode::kDLFloat, shape, + dlpack_strides); +} + +TYPED_TEST(TestExportTensor, TestTensorStrided) { std::vector shape = {2, 2, 2}; std::vector strides = {sizeof(float) * 4, sizeof(float) * 2, sizeof(float) * 1}; @@ -204,7 +272,8 @@ TEST_F(TestExportTensor, TestTensorStrided) { std::shared_ptr tensor = TensorFromJSON(float32(), "[1, 2, 3, 4, 5, 6, 1, 1]", shape, strides); - CheckDLTensor(tensor, float32(), DLDataTypeCode::kDLFloat, shape, dlpack_strides); + CheckDLTensor(tensor, float32(), DLDataTypeCode::kDLFloat, shape, + dlpack_strides); std::vector f_strides = {sizeof(float) * 1, sizeof(float) * 2, sizeof(float) * 4}; @@ -212,7 +281,8 @@ TEST_F(TestExportTensor, TestTensorStrided) { std::shared_ptr f_tensor = TensorFromJSON(float32(), "[1, 2, 3, 4, 5, 6, 1, 1]", shape, f_strides); - CheckDLTensor(f_tensor, float32(), DLDataTypeCode::kDLFloat, shape, f_dlpack_strides); + CheckDLTensor(f_tensor, float32(), DLDataTypeCode::kDLFloat, shape, + f_dlpack_strides); } } // namespace arrow::dlpack From 61ccdd711de83ff69440c290358e446777c26376 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 7 Aug 2026 15:25:28 +0200 Subject: [PATCH 4/5] Add Copy supoprt --- cpp/src/arrow/c/dlpack.cc | 190 ++++++++++++++++++++------------- cpp/src/arrow/c/dlpack.h | 11 +- cpp/src/arrow/c/dlpack_test.cc | 49 ++++++--- 3 files changed, 158 insertions(+), 92 deletions(-) diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc index 05ec4f0d83b0..4fda7189d34e 100644 --- a/cpp/src/arrow/c/dlpack.cc +++ b/cpp/src/arrow/c/dlpack.cc @@ -17,6 +17,7 @@ #include "arrow/c/dlpack.h" +#include #include #include #include @@ -63,70 +64,119 @@ Result GetDLDataType(const DataType& type) { } } -template +template struct ManagerCtx { - std::shared_ptr array; + /// Arrow buffer into of the data + std::shared_ptr buffer; + /// DLPack managed tensor structure. + /// Legacy `DLManagedTensor` or newer `DLManagedTensorVersioned`. DT tensor; - int64_t strides = 1; + Vec strides; + Vec shape; }; -template -Result ExportArrayImpl(const std::shared_ptr& arr) { - // Define DLDevice struct and check if array type is supported - // by the DLPack protocol at the same time. Raise TypeError if not. - // Supported data types: int, uint, float with no validity buffer. - ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(arr)); - - // Define the DLDataType struct - const DataType& type = *arr->type(); - ARROW_ASSIGN_OR_RAISE(auto dlpack_type, GetDLDataType(type)); +template +struct ExportBufferParams { + std::shared_ptr buffer = nullptr; + int64_t buffer_offset = 0; + /// Total number of values, i.e. the product of the shape. + int64_t size; + int32_t ndim; + Vec strides; + Vec shape; + DLDevice device; + DLDataType dtype; + uint64_t flags = 0; +}; +template +DT* ExportBuffer(ExportBufferParams&& p) { // Create ManagerCtx that will serve as the owner of the DLManagedTensor - auto ctx = std::make_unique>(); + using Ctx = ManagerCtx; + auto ctx = std::make_unique(); - // Assign the Array data into the context - ctx->array = arr->data(); - auto& data = ctx->array; + // Assign the Array data, shape, and strides into the context. + ctx->buffer = std::move(p.buffer); + ctx->shape = std::move(p.shape); + ctx->strides = std::move(p.strides); // Define the data pointer to the DLTensor // If array is of length 0, data pointer should be NULL - if (arr->length() == 0) { + if (p.size == 0) { ctx->tensor.dl_tensor.data = nullptr; } else { - const auto data_offset = data->offset * type.byte_width(); ctx->tensor.dl_tensor.data = - const_cast(data->buffers[1]->data() + data_offset); + const_cast(ctx->buffer->data() + p.buffer_offset); } - ctx->tensor.dl_tensor.device = device; - ctx->tensor.dl_tensor.ndim = 1; - ctx->tensor.dl_tensor.dtype = dlpack_type; - ctx->tensor.dl_tensor.shape = const_cast(&data->length); + ctx->tensor.dl_tensor.device = p.device; + ctx->tensor.dl_tensor.dtype = p.dtype; + ctx->tensor.dl_tensor.ndim = p.ndim; + ctx->tensor.dl_tensor.shape = ctx->shape.data(); ctx->tensor.dl_tensor.byte_offset = 0; // Strides must be non-null when ndim > 0 - ctx->tensor.dl_tensor.strides = &ctx->strides; + ctx->tensor.dl_tensor.strides = ctx->strides.data(); if constexpr (std::is_same_v) { ctx->tensor.version = {.major = DLPACK_MAJOR_VERSION, .minor = DLPACK_MINOR_VERSION}; - // Arrow contract is that array data is immutable once constructed - ctx->tensor.flags = DLPACK_FLAG_BITMASK_READ_ONLY; + ctx->tensor.flags = p.flags; } ctx->tensor.manager_ctx = ctx.get(); ctx->tensor.deleter = [](DT* self) { - delete reinterpret_cast*>(self->manager_ctx); + delete reinterpret_cast(self->manager_ctx); }; return &ctx.release()->tensor; } +template +Result ExportArrayImpl(const std::shared_ptr& arr, bool copy) { + // Define DLDevice struct and check if array type is supported + // by the DLPack protocol at the same time. Raise TypeError if not. + // Supported data types: int, uint, float with no validity buffer. + ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(arr)); + + // Define the DLDataType struct + const auto& type = *arr->type(); + ARROW_ASSIGN_OR_RAISE(auto dtype, GetDLDataType(type)); + + auto params = ExportBufferParams>{ + .size = arr->length(), + .ndim = 1, + .strides = {1}, + .shape = {arr->length()}, + .device = device, + .dtype = dtype, + }; + + const auto& data = *arr->data(); + if (copy) { + // We manually copy a buffer instead of using Array copy functions to avoid copying + // unused values outside of offset/length (e.g. with Slice). + const int64_t nbytes = data.length * type.byte_width(); + ARROW_ASSIGN_OR_RAISE(params.buffer, AllocateBuffer(nbytes)); + std::memcpy(params.buffer->mutable_data(), + data.buffers[1]->data() + data.offset * type.byte_width(), nbytes); + // Since we make a copy only for the consumer, we do not need to mark it readonly. + params.flags = DLPACK_FLAG_BITMASK_IS_COPIED; + } else { + // Shared buffer with Arrow Array. Arrays are readonly once constructed. + params.buffer = data.buffers[1]; + params.buffer_offset = data.offset * type.byte_width(); + params.flags = DLPACK_FLAG_BITMASK_READ_ONLY; + } + + return ExportBuffer
(std::move(params)); +} + } // namespace Result ExportArray(const std::shared_ptr& arr) { - return ExportArrayImpl(arr); + return ExportArrayImpl(arr, /* copy= */ false); } -Result ExportArrayVersioned( - const std::shared_ptr& arr) { - return ExportArrayImpl(arr); +Result ExportArrayVersioned(std::shared_ptr arr, + bool copy) { + return ExportArrayImpl(arr, copy); } Result ExportDevice(const std::shared_ptr& arr) { @@ -166,66 +216,54 @@ struct TensorManagerCtx { }; template -Result ExportTensorImpl(const std::shared_ptr& t) { - // Define the DLDataType struct - const DataType& type = *t->type(); - ARROW_ASSIGN_OR_RAISE(auto dlpack_type, GetDLDataType(type)); - +Result ExportTensorImpl(const std::shared_ptr& t, bool copy) { // Define DLDevice struct ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(t)); - // Create TensorManagerCtx that will serve as the owner of the DLManagedTensor - auto ctx = std::make_unique>(); + // Define the DLDataType struct + const auto& type = *t->type(); + ARROW_ASSIGN_OR_RAISE(auto dtype, GetDLDataType(type)); - // Define the data pointer to the DLTensor - // If tensor is of length 0, data pointer should be NULL - if (t->size() == 0) { - ctx->tensor.dl_tensor.data = nullptr; - } else { - ctx->tensor.dl_tensor.data = const_cast(t->raw_data()); + // Compute strides + std::vector strides = {}; + strides.reserve(t->ndim()); + const auto byte_width = type.byte_width(); + for (auto i : t->strides()) { + strides.emplace_back(i / byte_width); } - ctx->tensor.dl_tensor.device = device; - ctx->tensor.dl_tensor.ndim = t->ndim(); - ctx->tensor.dl_tensor.dtype = dlpack_type; - ctx->tensor.dl_tensor.byte_offset = 0; - - std::vector& shape_arr = ctx->shape; - shape_arr.reserve(t->ndim()); - for (auto i : t->shape()) { - shape_arr.emplace_back(i); - } - ctx->tensor.dl_tensor.shape = shape_arr.data(); + auto params = ExportBufferParams>{ + .size = t->size(), + .ndim = t->ndim(), + .strides = std::move(strides), + .shape = t->shape(), + .device = device, + .dtype = dtype, + }; - std::vector& strides_arr = ctx->strides; - strides_arr.reserve(t->ndim()); - const auto byte_width = t->type()->byte_width(); - for (auto i : t->strides()) { - strides_arr.emplace_back(i / byte_width); - } - ctx->tensor.dl_tensor.strides = strides_arr.data(); - if constexpr (std::is_same_v) { - ctx->tensor.version = {.major = DLPACK_MAJOR_VERSION, .minor = DLPACK_MINOR_VERSION}; - ctx->tensor.flags = t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY; + if (copy) { + ARROW_ASSIGN_OR_RAISE(params.buffer, MemoryManager::CopyBuffer( + t->data(), default_cpu_memory_manager())); + // Since we make a copy only for the consumer, we do not need to mark it readonly. + params.flags = DLPACK_FLAG_BITMASK_IS_COPIED; + } else { + // Shared buffer with Arrow Tensor. + params.buffer = t->data(); + params.flags |= t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY; } - ctx->t = std::move(t); - ctx->tensor.manager_ctx = ctx.get(); - ctx->tensor.deleter = [](DT* self) { - delete reinterpret_cast*>(self->manager_ctx); - }; - return &ctx.release()->tensor; + return ExportBuffer
(std::move(params)); } } // namespace Result ExportTensor(const std::shared_ptr& t) { - return ExportTensorImpl(t); + return ExportTensorImpl(t, /* copy= */ false); } -Result ExportTensorVersioned( - const std::shared_ptr& t) { - return ExportTensorImpl(t); +Result ExportTensorVersioned(std::shared_ptr t, + bool copy) { + return ExportTensorImpl(t, copy); } Result ExportDevice(const std::shared_ptr& t) { diff --git a/cpp/src/arrow/c/dlpack.h b/cpp/src/arrow/c/dlpack.h index 776f0be15800..71236d9b7f0d 100644 --- a/cpp/src/arrow/c/dlpack.h +++ b/cpp/src/arrow/c/dlpack.h @@ -17,9 +17,12 @@ #pragma once -#include "arrow/array/array_base.h" #include "arrow/c/dlpack_abi.h" +#include + +#include "arrow/array/array_base.h" + namespace arrow::dlpack { /// \brief Export Arrow array as DLPack tensor. @@ -42,13 +45,15 @@ ARROW_EXPORT Result ExportArray(const std::shared_ptr& arr); ARROW_EXPORT -Result ExportArrayVersioned(const std::shared_ptr& arr); +Result ExportArrayVersioned(std::shared_ptr arr, + bool copy); ARROW_EXPORT Result ExportTensor(const std::shared_ptr& t); ARROW_EXPORT -Result ExportTensorVersioned(const std::shared_ptr& t); +Result ExportTensorVersioned(std::shared_ptr, + bool copy); /// \brief Get DLDevice with enumerator specifying the /// type of the device data is stored on and index of the diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc index 892f7b5c7e7d..7d3141aebe21 100644 --- a/cpp/src/arrow/c/dlpack_test.cc +++ b/cpp/src/arrow/c/dlpack_test.cc @@ -17,6 +17,7 @@ #include +#include #include #include @@ -32,6 +33,7 @@ namespace arrow::dlpack { struct LegacyProducer { using ManagedTensor = DLManagedTensor; + static constexpr bool copy = false; // Unsupported static constexpr const char* name = "Legacy"; static Result Export(const std::shared_ptr& arr) { @@ -42,19 +44,22 @@ struct LegacyProducer { } }; +template struct VersionedProducer { using ManagedTensor = DLManagedTensorVersioned; - static constexpr const char* name = "Versioned"; + static constexpr bool copy = kCopy; + static constexpr const char* name = copy ? "VersionedCopied" : "Versioned"; static Result Export(const std::shared_ptr& arr) { - return ExportArrayVersioned(arr); + return ExportArrayVersioned(arr, copy); } static Result Export(const std::shared_ptr& t) { - return ExportTensorVersioned(t); + return ExportTensorVersioned(t, copy); } }; -using ProducerTypes = ::testing::Types; +using ProducerTypes = + ::testing::Types, VersionedProducer>; struct ProducerNames { template @@ -72,14 +77,19 @@ template void CheckDLTensor(const std::shared_ptr& arr, const std::shared_ptr& arrow_type, DLDataTypeCode dlpack_type, int64_t length) { - ASSERT_OK_AND_ASSIGN(auto dlmtensor, Producer::Export(arr)); + ASSERT_OK_AND_ASSIGN(auto* dlmtensor, Producer::Export(arr)); auto dltensor = dlmtensor->dl_tensor; const auto byte_width = arr->type()->byte_width(); const auto start = arr->offset() * byte_width; ASSERT_OK_AND_ASSIGN(auto sliced_buffer, SliceBufferSafe(arr->data()->buffers[1], start)); - ASSERT_EQ(sliced_buffer->data(), dltensor.data); + if constexpr (Producer::copy) { + ASSERT_NE(sliced_buffer->data(), dltensor.data); + ASSERT_EQ(0, std::memcmp(sliced_buffer->data(), dltensor.data, length * byte_width)); + } else { + ASSERT_EQ(sliced_buffer->data(), dltensor.data); + } ASSERT_EQ(0, dltensor.byte_offset); ASSERT_EQ(length, dltensor.shape[0]); @@ -96,11 +106,15 @@ void CheckDLTensor(const std::shared_ptr& arr, ASSERT_EQ(DLDeviceType::kDLCPU, device.device_type); ASSERT_EQ(0, device.device_id); - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { ASSERT_EQ(DLPACK_MAJOR_VERSION, dlmtensor->version.major); ASSERT_EQ(DLPACK_MINOR_VERSION, dlmtensor->version.minor); - // Arrow array data is immutable once constructed - ASSERT_EQ(DLPACK_FLAG_BITMASK_READ_ONLY, dlmtensor->flags); + if constexpr (Producer::copy) { + // Arrow array data is immutable once constructed, but a copy is ours to hand out + ASSERT_EQ(dlmtensor->flags, DLPACK_FLAG_BITMASK_IS_COPIED); + } else { + ASSERT_EQ(dlmtensor->flags, DLPACK_FLAG_BITMASK_READ_ONLY); + } } dlmtensor->deleter(dlmtensor); @@ -181,10 +195,15 @@ void CheckDLTensor(const std::shared_ptr& t, const std::shared_ptr& tensor_type, DLDataTypeCode dlpack_type, std::vector shape, std::vector strides) { - ASSERT_OK_AND_ASSIGN(auto dlmtensor, Producer::Export(t)); + ASSERT_OK_AND_ASSIGN(auto* dlmtensor, Producer::Export(t)); auto dltensor = dlmtensor->dl_tensor; - ASSERT_EQ(t->data()->data(), dltensor.data); + if constexpr (Producer::copy) { + ASSERT_NE(t->data()->data(), dltensor.data); + ASSERT_EQ(0, std::memcmp(t->data()->data(), dltensor.data, t->data()->size())); + } else { + ASSERT_EQ(t->data()->data(), dltensor.data); + } ASSERT_EQ(t->ndim(), dltensor.ndim); ASSERT_EQ(0, dltensor.byte_offset); for (int i = 0; i < t->ndim(); i++) { @@ -202,10 +221,14 @@ void CheckDLTensor(const std::shared_ptr& t, ASSERT_EQ(DLDeviceType::kDLCPU, device.device_type); ASSERT_EQ(0, device.device_id); - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { ASSERT_EQ(DLPACK_MAJOR_VERSION, dlmtensor->version.major); ASSERT_EQ(DLPACK_MINOR_VERSION, dlmtensor->version.minor); - ASSERT_EQ(t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY, dlmtensor->flags); + if constexpr (Producer::copy) { + ASSERT_EQ(dlmtensor->flags, DLPACK_FLAG_BITMASK_IS_COPIED); + } else { + ASSERT_EQ(dlmtensor->flags, (t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY)); + } } dlmtensor->deleter(dlmtensor); From 97d75260224fb5d106301251b9f281856a75dbc8 Mon Sep 17 00:00:00 2001 From: AntoinePrv Date: Fri, 7 Aug 2026 16:00:50 +0200 Subject: [PATCH 5/5] Add max_version and copy to __dlpack__ --- cpp/src/arrow/c/dlpack.cc | 24 +++------ cpp/src/arrow/c/dlpack.h | 42 ++++++++++++++- python/pyarrow/_dlpack.pxi | 28 ++++++++++ python/pyarrow/array.pxi | 47 +++++++++++++---- python/pyarrow/includes/libarrow.pxd | 10 ++++ python/pyarrow/tensor.pxi | 47 +++++++++++++---- python/pyarrow/tests/test_dlpack.py | 79 ++++++++++++++++++++++++++++ 7 files changed, 237 insertions(+), 40 deletions(-) diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc index 4fda7189d34e..8b1be56bf097 100644 --- a/cpp/src/arrow/c/dlpack.cc +++ b/cpp/src/arrow/c/dlpack.cc @@ -23,6 +23,7 @@ #include #include "arrow/array/array_base.h" +#include "arrow/buffer.h" #include "arrow/c/dlpack_abi.h" #include "arrow/device.h" #include "arrow/tensor.h" @@ -150,12 +151,11 @@ Result ExportArrayImpl(const std::shared_ptr& arr, bool copy) { const auto& data = *arr->data(); if (copy) { - // We manually copy a buffer instead of using Array copy functions to avoid copying + // We copy the buffer slice instead of using Array copy functions to avoid copying // unused values outside of offset/length (e.g. with Slice). - const int64_t nbytes = data.length * type.byte_width(); - ARROW_ASSIGN_OR_RAISE(params.buffer, AllocateBuffer(nbytes)); - std::memcpy(params.buffer->mutable_data(), - data.buffers[1]->data() + data.offset * type.byte_width(), nbytes); + const auto start = data.offset * type.byte_width(); + const auto nbytes = data.length * type.byte_width(); + ARROW_ASSIGN_OR_RAISE(params.buffer, data.buffers[1]->CopySlice(start, nbytes)); // Since we make a copy only for the consumer, we do not need to mark it readonly. params.flags = DLPACK_FLAG_BITMASK_IS_COPIED; } else { @@ -174,7 +174,7 @@ Result ExportArray(const std::shared_ptr& arr) { return ExportArrayImpl(arr, /* copy= */ false); } -Result ExportArrayVersioned(std::shared_ptr arr, +Result ExportArrayVersioned(const std::shared_ptr& arr, bool copy) { return ExportArrayImpl(arr, copy); } @@ -207,14 +207,6 @@ Result ExportDevice(const std::shared_ptr& arr) { namespace { -template -struct TensorManagerCtx { - std::shared_ptr t; - std::vector strides; - std::vector shape; - DT tensor; -}; - template Result ExportTensorImpl(const std::shared_ptr& t, bool copy) { // Define DLDevice struct @@ -249,7 +241,7 @@ Result ExportTensorImpl(const std::shared_ptr& t, bool copy) { } else { // Shared buffer with Arrow Tensor. params.buffer = t->data(); - params.flags |= t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY; + params.flags = t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY; } return ExportBuffer
(std::move(params)); @@ -261,7 +253,7 @@ Result ExportTensor(const std::shared_ptr& t) { return ExportTensorImpl(t, /* copy= */ false); } -Result ExportTensorVersioned(std::shared_ptr t, +Result ExportTensorVersioned(const std::shared_ptr& t, bool copy) { return ExportTensorImpl(t, copy); } diff --git a/cpp/src/arrow/c/dlpack.h b/cpp/src/arrow/c/dlpack.h index 71236d9b7f0d..8a9084f36c73 100644 --- a/cpp/src/arrow/c/dlpack.h +++ b/cpp/src/arrow/c/dlpack.h @@ -44,15 +44,53 @@ namespace arrow::dlpack { ARROW_EXPORT Result ExportArray(const std::shared_ptr& arr); +/// \brief Export Arrow array as a versioned DLPack tensor. +/// +/// Same restrictions on data types as ExportArray, but produces the +/// DLManagedTensorVersioned structure introduced in DLPack 1.0. +/// +/// The returned tensor is owned by the caller, who must release it by +/// calling its ``deleter``. +/// +/// Arrow arrays are immutable, so the exported tensor is flagged with +/// DLPACK_FLAG_BITMASK_READ_ONLY unless a copy is made, in which case it is +/// flagged with DLPACK_FLAG_BITMASK_IS_COPIED and the consumer is free to +/// mutate it. +/// +/// \param[in] arr Arrow array +/// \param[in] copy Whether to copy the data instead of sharing it with the array +/// \return DLManagedTensorVersioned struct ARROW_EXPORT -Result ExportArrayVersioned(std::shared_ptr arr, +Result ExportArrayVersioned(const std::shared_ptr& arr, bool copy); +/// \brief Export Arrow tensor as DLPack tensor. +/// +/// \note Deprecated in DLPack 1.0. Use ExportTensorVersioned instead. +/// +/// \param[in] t Arrow tensor +/// \return DLManagedTensor struct ARROW_EXPORT Result ExportTensor(const std::shared_ptr& t); +/// \brief Export Arrow tensor as a versioned DLPack tensor. +/// +/// Same as ExportTensor, but produces the DLManagedTensorVersioned structure +/// introduced in DLPack 1.0. +/// +/// The returned tensor is owned by the caller, who must release it by +/// calling its ``deleter``. +/// +/// When the data is shared with the Arrow tensor, the exported tensor is +/// flagged with DLPACK_FLAG_BITMASK_READ_ONLY if the Arrow tensor is not +/// mutable. When a copy is made, it is flagged with +/// DLPACK_FLAG_BITMASK_IS_COPIED and the consumer is free to mutate it. +/// +/// \param[in] t Arrow tensor +/// \param[in] copy Whether to copy the data instead of sharing it with the tensor +/// \return DLManagedTensorVersioned struct ARROW_EXPORT -Result ExportTensorVersioned(std::shared_ptr, +Result ExportTensorVersioned(const std::shared_ptr& t, bool copy); /// \brief Get DLDevice with enumerator specifying the diff --git a/python/pyarrow/_dlpack.pxi b/python/pyarrow/_dlpack.pxi index c2f4cff64069..03c4a8c64dbd 100644 --- a/python/pyarrow/_dlpack.pxi +++ b/python/pyarrow/_dlpack.pxi @@ -44,3 +44,31 @@ cdef void dlpack_pycapsule_deleter(object dltensor) noexcept: # Set the error indicator from err_type, err_value, err_traceback cpython.PyErr_Restore(err_type, err_value, err_traceback) + + +cdef void dlpack_versioned_pycapsule_deleter(object dltensor) noexcept: + cdef DLManagedTensorVersioned* dlm_tensor + cdef PyObject* err_type + cdef PyObject* err_value + cdef PyObject* err_traceback + + # Do nothing if the capsule has been consumed + if cpython.PyCapsule_IsValid(dltensor, "used_dltensor_versioned"): + return + + # An exception may be in-flight, we must save it in case + # we create another one + cpython.PyErr_Fetch(&err_type, &err_value, &err_traceback) + + dlm_tensor = cpython.PyCapsule_GetPointer( + dltensor, 'dltensor_versioned') + if dlm_tensor == NULL: + cpython.PyErr_WriteUnraisable(dltensor) + # The deleter can be NULL if there is no way for the caller + # to provide a reasonable destructor + elif dlm_tensor.deleter: + dlm_tensor.deleter(dlm_tensor) + assert (not cpython.PyErr_Occurred()) + + # Set the error indicator from err_type, err_value, err_traceback + cpython.PyErr_Restore(err_type, err_value, err_traceback) diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index 97cd9a8ebf82..3b732b58fef1 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -2246,7 +2246,7 @@ cdef class Array(_PandasConvertible): return pyarrow_wrap_array(array) - def __dlpack__(self, stream=None): + def __dlpack__(self, stream=None, max_version=None, dl_device=None, copy=None): """ Export a primitive array as a DLPack capsule. @@ -2256,20 +2256,45 @@ cdef class Array(_PandasConvertible): A Python integer representing a pointer to a stream. Currently not supported. Stream is provided by the consumer to the producer to instruct the producer to ensure that operations can safely be performed on the array. + max_version : tuple[int, int], optional + The maximum DLPack version the consumer supports, as (major, minor). + A capsule of a different version may be returned, so the consumer must + check it. Default is None, exporting the legacy unversioned capsule. + dl_device : tuple[enum.Enum, int], optional + The device of the exported capsule, in the format returned by + :meth:`__dlpack_device__`. Default is None, meaning the device of the + array itself. Since only CPU arrays are supported, any other device + raises ``BufferError``. + copy : bool, optional + If True, the data is always copied. If False, it is never copied and + ``BufferError`` is raised if a copy is required. If None (default), the + data is copied only if needed, which for CPU arrays is never. + A copy is reported to the consumer with ``DLPACK_FLAG_BITMASK_IS_COPIED``. Returns ------- capsule : PyCapsule - A DLPack capsule for the array, pointing to a DLManagedTensor. - """ - if stream is None: - dlm_tensor = GetResultValue(ExportArrayToDLPack(self.sp_array)) - - return PyCapsule_New(dlm_tensor, 'dltensor', dlpack_pycapsule_deleter) - else: - raise NotImplementedError( - "Only stream=None is supported." - ) + A DLPack capsule for the array, pointing to a DLManagedTensorVersioned, + or to a DLManagedTensor if ``max_version`` is below (1, 0). + """ + if stream is not None: + raise NotImplementedError("Only stream=None is supported.") + if dl_device is not None: + device = GetResultValue(ExportDevice(self.sp_array)) + if dl_device != (device.device_type, device.device_id): + raise BufferError( + f"Cannot export to device {dl_device}, " + f"array is on {(device.device_type, device.device_id)}." + ) + if max_version is None or max_version < (1, 0): + # Note: from March 2025 onwards, it's okay to raise BufferError here. + # Still we keep the V0 version that was added in August 2026. + legacy_tensor = GetResultValue(ExportArrayToDLPack(self.sp_array)) + return PyCapsule_New(legacy_tensor, 'dltensor', dlpack_pycapsule_deleter) + + # Currently no major version other than legacy 0 and current 1.3 + dlm_tensor = GetResultValue(ExportArrayVersionedToDLPack(self.sp_array, copy)) + return PyCapsule_New(dlm_tensor, 'dltensor_versioned', dlpack_versioned_pycapsule_deleter) def __dlpack_device__(self): """ diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index efc9602a3a81..3e6a19ffe9e5 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -1468,6 +1468,9 @@ cdef extern from "arrow/c/dlpack_abi.h" nogil: ctypedef struct DLManagedTensor: void (*deleter)(DLManagedTensor*) + ctypedef struct DLManagedTensorVersioned: + void (*deleter)(DLManagedTensorVersioned*) + cdef extern from "arrow/c/dlpack.h" namespace "arrow::dlpack" nogil: CResult[DLManagedTensor*] ExportArrayToDLPack" arrow::dlpack::ExportArray"( @@ -1475,6 +1478,13 @@ cdef extern from "arrow/c/dlpack.h" namespace "arrow::dlpack" nogil: CResult[DLManagedTensor*] ExportTensorToDLPack" arrow::dlpack::ExportTensor"( const shared_ptr[CTensor]& tensor) + CResult[DLManagedTensorVersioned*] \ + ExportArrayVersionedToDLPack" arrow::dlpack::ExportArrayVersioned"( + const shared_ptr[CArray]& arr, c_bool copy) + CResult[DLManagedTensorVersioned*] \ + ExportTensorVersionedToDLPack" arrow::dlpack::ExportTensorVersioned"( + const shared_ptr[CTensor]& tensor, c_bool copy) + CResult[DLDevice] ExportDevice(const shared_ptr[CArray]& arr) CResult[DLDevice] ExportDevice(const shared_ptr[CTensor]& tensor) diff --git a/python/pyarrow/tensor.pxi b/python/pyarrow/tensor.pxi index edb136b410c3..4a820bd3f4e3 100644 --- a/python/pyarrow/tensor.pxi +++ b/python/pyarrow/tensor.pxi @@ -300,7 +300,7 @@ strides: {self.strides}""" buffer.strides = cp.PyBytes_AsString(self._ssize_t_strides) buffer.suboffsets = NULL - def __dlpack__(self, stream=None): + def __dlpack__(self, stream=None, max_version=None, dl_device=None, copy=None): """ Export a Tensor as a DLPack capsule. @@ -310,20 +310,45 @@ strides: {self.strides}""" A Python integer representing a pointer to a stream. Currently not supported. Stream is provided by the consumer to the producer to instruct the producer to ensure that operations can safely be performed on the array. + max_version : tuple[int, int], optional + The maximum DLPack version the consumer supports, as (major, minor). + A capsule of a different version may be returned, so the consumer must + check it. Default is None, exporting the unversioned capsule. + dl_device : tuple[enum.Enum, int], optional + The device of the exported capsule, in the format returned by + :meth:`__dlpack_device__`. Default is None, meaning the device of the + tensor itself. Since only CPU tensors are supported, any other device + raises ``BufferError``. + copy : bool, optional + If True, the data is always copied. If False, it is never copied and + ``BufferError`` is raised if a copy is required. If None (default), the + data is copied only if needed, which for CPU tensors is never. + A copy is reported to the consumer with ``DLPACK_FLAG_BITMASK_IS_COPIED``. Returns ------- capsule : PyCapsule - A DLPack capsule for the tensor, pointing to a DLManagedTensor. - """ - if stream is None: - dlm_tensor = GetResultValue(ExportTensorToDLPack(self.sp_tensor)) - - return PyCapsule_New(dlm_tensor, 'dltensor', dlpack_pycapsule_deleter) - else: - raise NotImplementedError( - "Only stream=None is supported." - ) + A DLPack capsule for the tensor, pointing to a DLManagedTensorVersioned, + or to a DLManagedTensor if ``max_version`` is below (1, 0). + """ + if stream is not None: + raise NotImplementedError("Only stream=None is supported.") + if dl_device is not None: + device = GetResultValue(ExportDevice(self.sp_tensor)) + if dl_device != (device.device_type, device.device_id): + raise BufferError( + f"Cannot export to device {dl_device}, " + f"tensor is on {(device.device_type, device.device_id)}." + ) + if max_version is None or max_version < (1, 0): + # Note: from March 2025 onwards, it's okay to raise BufferError here. + # Still we keep the V0 version that was added in August 2026. + legacy_tensor = GetResultValue(ExportTensorToDLPack(self.sp_tensor)) + return PyCapsule_New(legacy_tensor, 'dltensor', dlpack_pycapsule_deleter) + + # Currently no major version other than legacy 0 and current 1.3 + dlm_tensor = GetResultValue(ExportTensorVersionedToDLPack(self.sp_tensor, copy)) + return PyCapsule_New(dlm_tensor, 'dltensor_versioned', dlpack_versioned_pycapsule_deleter) def __dlpack_device__(self): """ diff --git a/python/pyarrow/tests/test_dlpack.py b/python/pyarrow/tests/test_dlpack.py index 51eaa9036980..0e96394729c4 100644 --- a/python/pyarrow/tests/test_dlpack.py +++ b/python/pyarrow/tests/test_dlpack.py @@ -43,6 +43,24 @@ def check_dlpack_export(arr, expected_arr): assert arr.__dlpack_device__() == (1, 0) +class DLPackForwarder: + """Forward ``__dlpack__`` to a wrapped object with forced keyword arguments. + + Consumers such as ``np.from_dlpack`` do not expose every ``__dlpack__`` + keyword, so this makes them reachable from a consumer's point of view. + """ + + def __init__(self, obj, **forced): + self._obj = obj + self._forced = forced + + def __dlpack__(self, **kwargs): + return self._obj.__dlpack__(**{**kwargs, **self._forced}) + + def __dlpack_device__(self): + return self._obj.__dlpack_device__() + + def check_bytes_allocated(f): @wraps(f) def wrapper(*args, **kwargs): @@ -126,6 +144,67 @@ def test_tensor_dlpack(np_type): check_dlpack_export(t, expected) +def dlpack_objects(): + arr = pa.array([1, 2, 3], type=pa.int32()) + return [ + pytest.param(arr, id="array"), + pytest.param(arr.slice(1), id="sliced_array"), + pytest.param(pa.array([], type=pa.int32()), id="empty_array"), + pytest.param( + pa.Tensor.from_numpy(np.array([[1, 2], [3, 4]], dtype=np.int32)), + id="tensor", + ), + ] + + +@check_bytes_allocated +@pytest.mark.parametrize('obj', dlpack_objects()) +@pytest.mark.parametrize('max_version', [None, (0, 8)]) +def test_dlpack_legacy_capsule(obj, max_version): + capsule = obj.__dlpack__(max_version=max_version) + assert PyCapsule_IsValid(capsule, b"dltensor") is True + + +@check_bytes_allocated +@pytest.mark.parametrize('obj', dlpack_objects()) +@pytest.mark.parametrize('max_version', [(1, 0), (1, 3), (2, 0)]) +@pytest.mark.parametrize('copy', [None, False, True]) +def test_dlpack_versioned_capsule(obj, max_version, copy): + capsule = obj.__dlpack__(max_version=max_version, copy=copy) + assert PyCapsule_IsValid(capsule, b"dltensor_versioned") is True + + +@check_bytes_allocated +@pytest.mark.parametrize('obj', dlpack_objects()) +def test_dlpack_versioned_roundtrip(obj): + if Version(np.__version__) < Version("2.1.0"): + pytest.skip("Versioned DLPack capsules require numpy 2.1.0 or later") + + expected = np.from_dlpack(DLPackForwarder(obj, max_version=None)) + for copy in [None, False, True]: + result = np.from_dlpack( + DLPackForwarder(obj, max_version=(1, 0), copy=copy)) + np.testing.assert_array_equal(result, expected, strict=True) + + +@check_bytes_allocated +def test_dlpack_copy_is_writeable(): + if Version(np.__version__) < Version("2.1.0"): + pytest.skip("Read-only DLPack flag requires numpy 2.1.0 or later") + + arr = pa.array([1, 2, 3], type=pa.int32()) + + # Arrow arrays are immutable, so a shared export is read-only + shared = np.from_dlpack(DLPackForwarder(arr, max_version=(1, 3))) + assert not shared.flags.writeable + + # A copy is solely owned by the consumer, who may mutate it + copied = np.from_dlpack(DLPackForwarder(arr, max_version=(1, 3), copy=True)) + assert copied.flags.writeable + copied[0] = 100 + assert arr.to_pylist() == [1, 2, 3] + + def test_dlpack_not_supported(): if Version(np.__version__) < Version("1.22.0"): pytest.skip("No dlpack support in numpy versions older than 1.22.0.")