diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc index 94d7816d78e6..8b1be56bf097 100644 --- a/cpp/src/arrow/c/dlpack.cc +++ b/cpp/src/arrow/c/dlpack.cc @@ -17,7 +17,13 @@ #include "arrow/c/dlpack.h" +#include +#include +#include +#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" @@ -59,52 +65,120 @@ Result GetDLDataType(const DataType& type) { } } +template struct ManagerCtx { - std::shared_ptr array; - DLManagedTensor tensor; + /// Arrow buffer into of the data + std::shared_ptr buffer; + /// DLPack managed tensor structure. + /// Legacy `DLManagedTensor` or newer `DLManagedTensorVersioned`. + DT tensor; + Vec strides; + Vec shape; }; -} // namespace - -Result ExportArray(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(); - std::shared_ptr data = arr->data(); - 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, 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) { - ctx->tensor.dl_tensor.data = NULL; + 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.strides = NULL; + 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.data(); + if constexpr (std::is_same_v) { + ctx->tensor.version = {.major = DLPACK_MAJOR_VERSION, .minor = DLPACK_MINOR_VERSION}; + ctx->tensor.flags = p.flags; + } - 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; } +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 copy the buffer slice instead of using Array copy functions to avoid copying + // unused values outside of offset/length (e.g. with Slice). + 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 { + // 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, /* copy= */ false); +} + +Result ExportArrayVersioned(const std::shared_ptr& arr, + bool copy) { + return ExportArrayImpl(arr, copy); +} + Result ExportDevice(const std::shared_ptr& arr) { // Check if array is supported by the DLPack protocol. if (arr->null_count() > 0) { @@ -131,58 +205,57 @@ Result ExportDevice(const std::shared_ptr& arr) { } } -struct TensorManagerCtx { - std::shared_ptr t; - std::vector strides; - std::vector shape; - DLManagedTensor tensor; -}; +namespace { + +template +Result ExportTensorImpl(const std::shared_ptr& t, bool copy) { + // Define DLDevice struct + ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(t)); -Result ExportTensor(const std::shared_ptr& t) { // Define the DLDataType struct - const DataType& type = *t->type(); - ARROW_ASSIGN_OR_RAISE(auto dlpack_type, GetDLDataType(type)); + const auto& type = *t->type(); + ARROW_ASSIGN_OR_RAISE(auto dtype, GetDLDataType(type)); - // Define DLDevice struct - ARROW_ASSIGN_OR_RAISE(auto device, ExportDevice(t)) + // 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); + } - // Create TensorManagerCtx that will serve as the owner of the DLManagedTensor - auto ctx = std::make_unique(); + auto params = ExportBufferParams>{ + .size = t->size(), + .ndim = t->ndim(), + .strides = std::move(strides), + .shape = t->shape(), + .device = device, + .dtype = dtype, + }; - // 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; + 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 { - ctx->tensor.dl_tensor.data = t->raw_mutable_data(); + // Shared buffer with Arrow Tensor. + params.buffer = t->data(); + params.flags = t->is_mutable() ? 0 : DLPACK_FLAG_BITMASK_READ_ONLY; } - 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; + return ExportBuffer
(std::move(params)); +} - 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(); +} // namespace - std::vector& strides_arr = ctx->strides; - strides_arr.reserve(t->ndim()); - 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(); +Result ExportTensor(const std::shared_ptr& t) { + return ExportTensorImpl(t, /* copy= */ false); +} - ctx->t = std::move(t); - ctx->tensor.manager_ctx = ctx.get(); - ctx->tensor.deleter = [](struct DLManagedTensor* self) { - delete reinterpret_cast(self->manager_ctx); - }; - return &ctx.release()->tensor; +Result ExportTensorVersioned(const 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 65da38423c2a..8a9084f36c73 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. @@ -34,14 +37,62 @@ 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); +/// \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(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(const std::shared_ptr& t, + bool copy); + /// \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_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 diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc index f0119e8aef71..7d3141aebe21 100644 --- a/cpp/src/arrow/c/dlpack_test.cc +++ b/cpp/src/arrow/c/dlpack_test.cc @@ -17,7 +17,12 @@ #include +#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,27 +31,70 @@ namespace arrow::dlpack { -class TestExportArray : public ::testing::Test { - public: - void SetUp() {} +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) { + return ExportArray(arr); + } + static Result Export(const std::shared_ptr& t) { + return ExportTensor(t); + } +}; + +template +struct VersionedProducer { + using ManagedTensor = DLManagedTensorVersioned; + static constexpr bool copy = kCopy; + static constexpr const char* name = copy ? "VersionedCopied" : "Versioned"; + + static Result Export(const std::shared_ptr& arr) { + return ExportArrayVersioned(arr, copy); + } + static Result Export(const std::shared_ptr& t) { + return ExportTensorVersioned(t, copy); + } }; +using ProducerTypes = + ::testing::Types, VersionedProducer>; + +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(); 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(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); @@ -58,10 +106,21 @@ 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); + 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); } -TEST_F(TestExportArray, TestSupportedArray) { +TYPED_TEST(TestExportArray, TestSupportedArray) { const std::vector, DLDataTypeCode>> cases = { {int8(), DLDataTypeCode::kDLInt}, {uint8(), DLDataTypeCode::kDLUInt}, @@ -89,36 +148,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,19 +185,25 @@ TEST_F(TestExportArray, TestErrors) { arrow::dlpack::ExportDevice(array_boolean)); } -class TestExportTensor : public ::testing::Test { - public: - void SetUp() {} -}; +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); + 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++) { @@ -156,10 +221,20 @@ 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); + 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); } -TEST_F(TestExportTensor, TestTensor) { +TYPED_TEST(TestExportTensor, TestTensor) { const std::vector, DLDataTypeCode>> cases = { {int8(), DLDataTypeCode::kDLInt}, {uint8(), DLDataTypeCode::kDLUInt}, @@ -190,13 +265,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 +295,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 +304,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 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.")