diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 46a7aa910633d..00947c6275678 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -192,6 +192,7 @@ set(ARROW_SRCS type_traits.cc visitor.cc c/bridge.cc + c/dlpack.cc io/buffered.cc io/caching.cc io/compressed.cc diff --git a/cpp/src/arrow/c/CMakeLists.txt b/cpp/src/arrow/c/CMakeLists.txt index 3765477ba09cd..81a81cd3f1103 100644 --- a/cpp/src/arrow/c/CMakeLists.txt +++ b/cpp/src/arrow/c/CMakeLists.txt @@ -16,6 +16,7 @@ # under the License. add_arrow_test(bridge_test PREFIX "arrow-c") +add_arrow_test(dlpack_test) add_arrow_benchmark(bridge_benchmark) diff --git a/cpp/src/arrow/c/dlpack.cc b/cpp/src/arrow/c/dlpack.cc new file mode 100644 index 0000000000000..13ee2761b0c11 --- /dev/null +++ b/cpp/src/arrow/c/dlpack.cc @@ -0,0 +1,133 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/c/dlpack.h" + +#include "arrow/array/array_base.h" +#include "arrow/c/dlpack_abi.h" +#include "arrow/device.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" + +namespace arrow::dlpack { + +namespace { + +Result GetDLDataType(const DataType& type) { + DLDataType dtype; + dtype.lanes = 1; + dtype.bits = type.bit_width(); + switch (type.id()) { + case Type::INT8: + case Type::INT16: + case Type::INT32: + case Type::INT64: + dtype.code = DLDataTypeCode::kDLInt; + return dtype; + case Type::UINT8: + case Type::UINT16: + case Type::UINT32: + case Type::UINT64: + dtype.code = DLDataTypeCode::kDLUInt; + return dtype; + case Type::HALF_FLOAT: + case Type::FLOAT: + case Type::DOUBLE: + dtype.code = DLDataTypeCode::kDLFloat; + return dtype; + case Type::BOOL: + // DLPack supports byte-packed boolean values + return Status::TypeError("Bit-packed boolean data type not supported by DLPack."); + default: + return Status::TypeError("DataType is not compatible with DLPack spec: ", + type.ToString()); + } +} + +struct ManagerCtx { + std::shared_ptr array; + DLManagedTensor tensor; +}; + +} // namespace + +Result ExportArray(const std::shared_ptr& arr) { + // Define DLDevice struct nad 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)); + + // Create ManagerCtx that will serve as the owner of the DLManagedTensor + std::unique_ptr ctx(new ManagerCtx); + + // 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; + } else { + const auto data_offset = data->offset * type.byte_width(); + ctx->tensor.dl_tensor.data = + const_cast(data->buffers[1]->data() + data_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.byte_offset = 0; + + ctx->array = std::move(data); + ctx->tensor.manager_ctx = ctx.get(); + ctx->tensor.deleter = [](struct DLManagedTensor* self) { + delete reinterpret_cast(self->manager_ctx); + }; + return &ctx.release()->tensor; +} + +Result ExportDevice(const std::shared_ptr& arr) { + // Check if array is supported by the DLPack protocol. + if (arr->null_count() > 0) { + return Status::TypeError("Can only use DLPack on arrays with no nulls."); + } + const DataType& type = *arr->type(); + if (type.id() == Type::BOOL) { + return Status::TypeError("Bit-packed boolean data type not supported by DLPack."); + } + if (!is_integer(type.id()) && !is_floating(type.id())) { + return Status::TypeError("DataType is not compatible with DLPack spec: ", + type.ToString()); + } + + // Define DLDevice struct + DLDevice device; + if (arr->data()->buffers[1]->device_type() == DeviceAllocationType::kCPU) { + device.device_id = 0; + device.device_type = DLDeviceType::kDLCPU; + return device; + } else { + return Status::NotImplemented( + "DLPack support is implemented only for buffers on CPU device."); + } +} + +} // namespace arrow::dlpack diff --git a/cpp/src/arrow/c/dlpack.h b/cpp/src/arrow/c/dlpack.h new file mode 100644 index 0000000000000..d11ccfc1fd722 --- /dev/null +++ b/cpp/src/arrow/c/dlpack.h @@ -0,0 +1,51 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include "arrow/array/array_base.h" +#include "arrow/c/dlpack_abi.h" + +namespace arrow::dlpack { + +/// \brief Export Arrow array as DLPack tensor. +/// +/// DLMangedTensor is produced as defined by the DLPack protocol, +/// see https://dmlc.github.io/dlpack/latest/. +/// +/// Data types for which the protocol is supported are +/// integer and floating-point data types. +/// +/// DLPack protocol only supports arrays with one contiguous +/// memory region which means Arrow Arrays with validity buffers +/// are not supported. +/// +/// \param[in] arr Arrow array +/// \return DLManagedTensor struct +ARROW_EXPORT +Result ExportArray(const std::shared_ptr& arr); + +/// \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. +/// +/// \param[in] arr Arrow array +/// \return DLDevice struct +ARROW_EXPORT +Result ExportDevice(const std::shared_ptr& arr); + +} // namespace arrow::dlpack diff --git a/cpp/src/arrow/c/dlpack_abi.h b/cpp/src/arrow/c/dlpack_abi.h new file mode 100644 index 0000000000000..4af557a7ed5d7 --- /dev/null +++ b/cpp/src/arrow/c/dlpack_abi.h @@ -0,0 +1,321 @@ +// Taken from: +// https://github.com/dmlc/dlpack/blob/ca4d00ad3e2e0f410eeab3264d21b8a39397f362/include/dlpack/dlpack.h +/*! + * Copyright (c) 2017 by Contributors + * \file dlpack.h + * \brief The common header of DLPack. + */ +#ifndef DLPACK_DLPACK_H_ +#define DLPACK_DLPACK_H_ + +/** + * \brief Compatibility with C++ + */ +#ifdef __cplusplus +#define DLPACK_EXTERN_C extern "C" +#else +#define DLPACK_EXTERN_C +#endif + +/*! \brief The current major version of dlpack */ +#define DLPACK_MAJOR_VERSION 1 + +/*! \brief The current minor version of dlpack */ +#define DLPACK_MINOR_VERSION 0 + +/*! \brief DLPACK_DLL prefix for windows */ +#ifdef _WIN32 +#ifdef DLPACK_EXPORTS +#define DLPACK_DLL __declspec(dllexport) +#else +#define DLPACK_DLL __declspec(dllimport) +#endif +#else +#define DLPACK_DLL +#endif + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/*! + * \brief The DLPack version. + * + * A change in major version indicates that we have changed the + * data layout of the ABI - DLManagedTensorVersioned. + * + * A change in minor version indicates that we have added new + * code, such as a new device type, but the ABI is kept the same. + * + * If an obtained DLPack tensor has a major version that disagrees + * with the version number specified in this header file + * (i.e. major != DLPACK_MAJOR_VERSION), the consumer must call the deleter + * (and it is safe to do so). It is not safe to access any other fields + * as the memory layout will have changed. + * + * In the case of a minor version mismatch, the tensor can be safely used as + * long as the consumer knows how to interpret all fields. Minor version + * updates indicate the addition of enumeration values. + */ +typedef struct { + /*! \brief DLPack major version. */ + uint32_t major; + /*! \brief DLPack minor version. */ + uint32_t minor; +} DLPackVersion; + +/*! + * \brief The device type in DLDevice. + */ +#ifdef __cplusplus +typedef enum : int32_t { +#else +typedef enum { +#endif + /*! \brief CPU device */ + kDLCPU = 1, + /*! \brief CUDA GPU device */ + kDLCUDA = 2, + /*! + * \brief Pinned CUDA CPU memory by cudaMallocHost + */ + kDLCUDAHost = 3, + /*! \brief OpenCL devices. */ + kDLOpenCL = 4, + /*! \brief Vulkan buffer for next generation graphics. */ + kDLVulkan = 7, + /*! \brief Metal for Apple GPU. */ + kDLMetal = 8, + /*! \brief Verilog simulator buffer */ + kDLVPI = 9, + /*! \brief ROCm GPUs for AMD GPUs */ + kDLROCM = 10, + /*! + * \brief Pinned ROCm CPU memory allocated by hipMallocHost + */ + kDLROCMHost = 11, + /*! + * \brief Reserved extension device type, + * used for quickly test extension device + * The semantics can differ depending on the implementation. + */ + kDLExtDev = 12, + /*! + * \brief CUDA managed/unified memory allocated by cudaMallocManaged + */ + kDLCUDAManaged = 13, + /*! + * \brief Unified shared memory allocated on a oneAPI non-partititioned + * device. Call to oneAPI runtime is required to determine the device + * type, the USM allocation type and the sycl context it is bound to. + * + */ + kDLOneAPI = 14, + /*! \brief GPU support for next generation WebGPU standard. */ + kDLWebGPU = 15, + /*! \brief Qualcomm Hexagon DSP */ + kDLHexagon = 16, +} DLDeviceType; + +/*! + * \brief A Device for Tensor and operator. + */ +typedef struct { + /*! \brief The device type used in the device. */ + DLDeviceType device_type; + /*! + * \brief The device index. + * For vanilla CPU memory, pinned memory, or managed memory, this is set to 0. + */ + int32_t device_id; +} DLDevice; + +/*! + * \brief The type code options DLDataType. + */ +typedef enum { + /*! \brief signed integer */ + kDLInt = 0U, + /*! \brief unsigned integer */ + kDLUInt = 1U, + /*! \brief IEEE floating point */ + kDLFloat = 2U, + /*! + * \brief Opaque handle type, reserved for testing purposes. + * Frameworks need to agree on the handle data type for the exchange to be well-defined. + */ + kDLOpaqueHandle = 3U, + /*! \brief bfloat16 */ + kDLBfloat = 4U, + /*! + * \brief complex number + * (C/C++/Python layout: compact struct per complex number) + */ + kDLComplex = 5U, + /*! \brief boolean */ + kDLBool = 6U, +} DLDataTypeCode; + +/*! + * \brief The data type the tensor can hold. The data type is assumed to follow the + * native endian-ness. An explicit error message should be raised when attempting to + * export an array with non-native endianness + * + * Examples + * - float: type_code = 2, bits = 32, lanes = 1 + * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4 + * - int8: type_code = 0, bits = 8, lanes = 1 + * - std::complex: type_code = 5, bits = 64, lanes = 1 + * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, + * the underlying storage size of bool is 8 bits) + */ +typedef struct { + /*! + * \brief Type code of base types. + * We keep it uint8_t instead of DLDataTypeCode for minimal memory + * footprint, but the value should be one of DLDataTypeCode enum values. + * */ + uint8_t code; + /*! + * \brief Number of bits, common choices are 8, 16, 32. + */ + uint8_t bits; + /*! \brief Number of lanes in the type, used for vector types. */ + uint16_t lanes; +} DLDataType; + +/*! + * \brief Plain C Tensor object, does not manage memory. + */ +typedef struct { + /*! + * \brief The data pointer points to the allocated data. This will be CUDA + * device pointer or cl_mem handle in OpenCL. It may be opaque on some device + * 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 + * 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. + * + * For given DLTensor, the size of memory required to store the contents of + * data is calculated as follows: + * + * \code{.c} + * static inline size_t GetDataSize(const DLTensor* t) { + * size_t size = 1; + * for (tvm_index_t i = 0; i < t->ndim; ++i) { + * size *= t->shape[i]; + * } + * size *= (t->dtype.bits * t->dtype.lanes + 7) / 8; + * return size; + * } + * \endcode + */ + void* data; + /*! \brief The device of the tensor */ + DLDevice device; + /*! \brief Number of dimensions */ + int32_t ndim; + /*! \brief The data type of the pointer*/ + DLDataType dtype; + /*! \brief The shape of the tensor */ + int64_t* shape; + /*! + * \brief strides of the tensor (in number of elements, not bytes) + * can be NULL, indicating tensor is compact and row-majored. + */ + int64_t* strides; + /*! \brief The offset in bytes to the beginning pointer to data */ + uint64_t byte_offset; +} DLTensor; + +/*! + * \brief C Tensor object, manage memory of DLTensor. This data structure is + * intended to facilitate the borrowing of DLTensor by another framework. It is + * not meant to transfer the tensor. When the borrowing framework doesn't need + * the tensor, it should call the deleter to notify the host that the resource + * is no longer needed. + * + * \note This data structure is used as Legacy DLManagedTensor + * in DLPack exchange and is deprecated after DLPack v0.8 + * Use DLManagedTensorVersioned instead. + * This data structure may get renamed or deleted in future versions. + * + * \sa DLManagedTensorVersioned + */ +typedef struct DLManagedTensor { + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; + /*! \brief the context of the original host framework of DLManagedTensor in + * which DLManagedTensor is used in the framework. It can also be NULL. + */ + void* manager_ctx; + /*! + * \brief Destructor - 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. + */ + void (*deleter)(struct DLManagedTensor* self); +} DLManagedTensor; + +// bit masks used in in the DLManagedTensorVersioned + +/*! \brief bit mask to indicate that the tensor is read only. */ +#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL) + +/*! + * \brief A versioned and managed C Tensor object, manage memory of DLTensor. + * + * This data structure is intended to facilitate the borrowing of DLTensor by + * another framework. It is not meant to transfer the tensor. When the borrowing + * framework doesn't need the tensor, it should call the deleter to notify the + * host that the resource is no longer needed. + * + * \note This is the current standard DLPack exchange data structure. + */ +struct DLManagedTensorVersioned { + /*! + * \brief The API and ABI version of the current managed Tensor + */ + DLPackVersion version; + /*! + * \brief the context of the original host framework. + * + * Stores DLManagedTensorVersioned is used in the + * framework. It can also be NULL. + */ + void* manager_ctx; + /*! + * \brief Destructor. + * + * 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. + */ + void (*deleter)(struct DLManagedTensorVersioned* self); + /*! + * \brief Additional bitmask flags information about the tensor. + * + * By default the flags should be set to 0. + * + * \note Future ABI changes should keep everything until this field + * stable, to ensure that deleter can be correctly called. + * + * \sa DLPACK_FLAG_BITMASK_READ_ONLY + */ + uint64_t flags; + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; +}; + +#ifdef __cplusplus +} // DLPACK_EXTERN_C +#endif +#endif // DLPACK_DLPACK_H_ diff --git a/cpp/src/arrow/c/dlpack_test.cc b/cpp/src/arrow/c/dlpack_test.cc new file mode 100644 index 0000000000000..3136506bf39ad --- /dev/null +++ b/cpp/src/arrow/c/dlpack_test.cc @@ -0,0 +1,129 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include "arrow/array/array_base.h" +#include "arrow/c/dlpack.h" +#include "arrow/c/dlpack_abi.h" +#include "arrow/memory_pool.h" +#include "arrow/testing/gtest_util.h" + +namespace arrow::dlpack { + +class TestExportArray : public ::testing::Test { + public: + void SetUp() {} +}; + +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)); + 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); + + ASSERT_EQ(0, dltensor.byte_offset); + ASSERT_EQ(NULL, dltensor.strides); + ASSERT_EQ(length, dltensor.shape[0]); + ASSERT_EQ(1, dltensor.ndim); + + ASSERT_EQ(dlpack_type, dltensor.dtype.code); + + ASSERT_EQ(arrow_type->bit_width(), dltensor.dtype.bits); + ASSERT_EQ(1, dltensor.dtype.lanes); + ASSERT_EQ(DLDeviceType::kDLCPU, dltensor.device.device_type); + ASSERT_EQ(0, dltensor.device.device_id); + + ASSERT_OK_AND_ASSIGN(auto device, arrow::dlpack::ExportDevice(arr)); + ASSERT_EQ(DLDeviceType::kDLCPU, device.device_type); + ASSERT_EQ(0, device.device_id); + + dlmtensor->deleter(dlmtensor); +} + +TEST_F(TestExportArray, TestSupportedArray) { + const std::vector, DLDataTypeCode>> cases = { + {int8(), DLDataTypeCode::kDLInt}, + {uint8(), DLDataTypeCode::kDLUInt}, + { + int16(), + DLDataTypeCode::kDLInt, + }, + {uint16(), DLDataTypeCode::kDLUInt}, + { + int32(), + DLDataTypeCode::kDLInt, + }, + {uint32(), DLDataTypeCode::kDLUInt}, + { + int64(), + DLDataTypeCode::kDLInt, + }, + {uint64(), DLDataTypeCode::kDLUInt}, + {float16(), DLDataTypeCode::kDLFloat}, + {float32(), DLDataTypeCode::kDLFloat}, + {float64(), DLDataTypeCode::kDLFloat}}; + + const auto allocated_bytes = arrow::default_memory_pool()->bytes_allocated(); + + 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); + ASSERT_OK_AND_ASSIGN(auto sliced_1, array->SliceSafe(1, 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); + ASSERT_OK_AND_ASSIGN(auto sliced_3, array->SliceSafe(3)); + CheckDLTensor(sliced_3, arrow_type, dlpack_type, 7); + } + + ASSERT_EQ(allocated_bytes, arrow::default_memory_pool()->bytes_allocated()); +} + +TEST_F(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)); + + 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)); + + 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)); + + const std::shared_ptr array_boolean = ArrayFromJSON(boolean(), "[true, false]"); + ASSERT_RAISES_WITH_MESSAGE( + TypeError, "Type error: Bit-packed boolean data type not supported by DLPack.", + arrow::dlpack::ExportDevice(array_boolean)); +} + +} // namespace arrow::dlpack diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt index ce637bf839232..4f86a12afe4fb 100644 --- a/dev/release/rat_exclude_files.txt +++ b/dev/release/rat_exclude_files.txt @@ -12,6 +12,7 @@ ci/etc/*.patch ci/vcpkg/*.patch CHANGELOG.md cpp/CHANGELOG_PARQUET.md +cpp/src/arrow/c/dlpack_abi.h cpp/src/arrow/io/mman.h cpp/src/arrow/util/random.h cpp/src/arrow/status.cc diff --git a/docs/source/python/dlpack.rst b/docs/source/python/dlpack.rst new file mode 100644 index 0000000000000..f612ebabde5c9 --- /dev/null +++ b/docs/source/python/dlpack.rst @@ -0,0 +1,93 @@ +.. Licensed to the Apache Software Foundation (ASF) under one +.. or more contributor license agreements. See the NOTICE file +.. distributed with this work for additional information +.. regarding copyright ownership. The ASF licenses this file +.. to you under the Apache License, Version 2.0 (the +.. "License"); you may not use this file except in compliance +.. with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, +.. software distributed under the License is distributed on an +.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +.. KIND, either express or implied. See the License for the +.. specific language governing permissions and limitations +.. under the License. + +.. _pyarrow-dlpack: + +The DLPack Protocol +=================== + +`The DLPack Protocol `_ +is a stable in-memory data structure that allows exchange +between major frameworks working with multidimensional +arrays or tensors. It is designed for cross hardware +support meaning it allows exchange of data on devices other +than the CPU (e.g. GPU). + +DLPack protocol had been +`selected as the Python array API standard `_ +by the +`Consortium for Python Data API Standards `_ +in order to enable device aware data interchange between array/tensor +libraries in the Python ecosystem. See more about the standard +in the +`protocol documentation `_ +and more about DLPack in the +`Python Specification for DLPack `_. + +Implementation of DLPack in PyArrow +----------------------------------- + +The producing side of the DLPack Protocol is implemented for ``pa.Array`` +and can be used to interchange data between PyArrow and other tensor +libraries. Supported data types are integer, unsigned integer and float. The +protocol has no missing data support meaning PyArrow arrays with +missing values cannot be transferred through the DLPack +protocol. Currently, the Arrow implementation of the protocol only supports +data on a CPU device. + +Data interchange syntax of the protocol includes + +1. ``from_dlpack(x)``: consuming an array object that implements a + ``__dlpack__`` method and creating a new array while sharing the + memory. + +2. ``__dlpack__(self, stream=None)`` and ``__dlpack_device__``: + producing a PyCapsule with the DLPack struct which is called from + within ``from_dlpack(x)``. + +PyArrow implements the second part of the protocol +(``__dlpack__(self, stream=None)`` and ``__dlpack_device__``) and can +thus be consumed by libraries implementing ``from_dlpack``. + +Example +------- + +Convert a PyArrow CPU array to NumPy array: + +.. code-block:: + + >>> import pyarrow as pa + >>> array = pa.array([2, 0, 2, 4]) + + [ + 2, + 0, + 2, + 4 + ] + + >>> import numpy as np + >>> np.from_dlpack(array) + array([2, 0, 2, 4]) + +Convert a PyArrow CPU array to PyTorch tensor: + +.. code-block:: + + >>> import torch + >>> torch.from_dlpack(array) + tensor([2, 0, 2, 4]) diff --git a/docs/source/python/index.rst b/docs/source/python/index.rst index 6a3de3d42b149..08939bc760df6 100644 --- a/docs/source/python/index.rst +++ b/docs/source/python/index.rst @@ -53,6 +53,7 @@ files into Arrow structures. numpy pandas interchange_protocol + dlpack timestamps orc csv diff --git a/docs/source/python/interchange_protocol.rst b/docs/source/python/interchange_protocol.rst index c354541a6779c..2a5ec8afede7b 100644 --- a/docs/source/python/interchange_protocol.rst +++ b/docs/source/python/interchange_protocol.rst @@ -37,7 +37,7 @@ libraries in the Python ecosystem. See more about the standard in the `protocol documentation `_. -From pyarrow to other libraries: ``__dataframe__()`` method +From PyArrow to other libraries: ``__dataframe__()`` method ----------------------------------------------------------- The ``__dataframe__()`` method creates a new exchange object that @@ -54,7 +54,7 @@ This is meant to be used by the consumer library when calling the ``from_dataframe()`` function and is not meant to be used manually by the user. -From other libraries to pyarrow: ``from_dataframe()`` +From other libraries to PyArrow: ``from_dataframe()`` ----------------------------------------------------- With the ``from_dataframe()`` function, we can construct a :class:`pyarrow.Table` @@ -63,7 +63,7 @@ from any dataframe object that implements the protocol. We can for example take a pandas dataframe and construct a -pyarrow table with the use of the interchange protocol: +PyArrow table with the use of the interchange protocol: .. code-block:: diff --git a/python/pyarrow/_dlpack.pxi b/python/pyarrow/_dlpack.pxi new file mode 100644 index 0000000000000..c2f4cff640691 --- /dev/null +++ b/python/pyarrow/_dlpack.pxi @@ -0,0 +1,46 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +cimport cpython +from cpython.pycapsule cimport PyCapsule_New + + +cdef void dlpack_pycapsule_deleter(object dltensor) noexcept: + cdef DLManagedTensor* 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"): + 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') + 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 789e30d3e9b00..74a196002bfa6 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -1779,6 +1779,44 @@ cdef class Array(_PandasConvertible): return pyarrow_wrap_array(array) + def __dlpack__(self, stream=None): + """Export a primitive array as a DLPack capsule. + + Parameters + ---------- + stream : int, optional + 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. + + Returns + ------- + capsule : PyCapsule + A DLPack capsule for the array, pointing to a DLManagedTensor. + """ + if stream is None: + dlm_tensor = GetResultValue(ExportToDLPack(self.sp_array)) + + return PyCapsule_New(dlm_tensor, 'dltensor', dlpack_pycapsule_deleter) + else: + raise NotImplementedError( + "Only stream=None is supported." + ) + + def __dlpack_device__(self): + """ + Return the DLPack device tuple this arrays resides on. + + Returns + ------- + tuple : Tuple[int, int] + Tuple with index specifying the type of the device (where + CPU = 1, see cpp/src/arrow/c/dpack_abi.h) and index of the + device which is 0 by default for CPU. + """ + device = GetResultValue(ExportDevice(self.sp_array)) + return device.device_type, device.device_id + cdef _array_like_to_pandas(obj, options, types_mapper): cdef: diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 403846a38f3fd..bad5ec606c268 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -1199,6 +1199,25 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: shared_ptr[CScalar] MakeNullScalar(shared_ptr[CDataType] type) +cdef extern from "arrow/c/dlpack_abi.h" nogil: + ctypedef enum DLDeviceType: + kDLCPU = 1 + + ctypedef struct DLDevice: + DLDeviceType device_type + int32_t device_id + + ctypedef struct DLManagedTensor: + void (*deleter)(DLManagedTensor*) + + +cdef extern from "arrow/c/dlpack.h" namespace "arrow::dlpack" nogil: + CResult[DLManagedTensor*] ExportToDLPack" arrow::dlpack::ExportArray"( + const shared_ptr[CArray]& arr) + + CResult[DLDevice] ExportDevice(const shared_ptr[CArray]& arr) + + cdef extern from "arrow/builder.h" namespace "arrow" nogil: cdef cppclass CArrayBuilder" arrow::ArrayBuilder": diff --git a/python/pyarrow/lib.pyx b/python/pyarrow/lib.pyx index 57fb0f42e38bf..29a0bed55949c 100644 --- a/python/pyarrow/lib.pyx +++ b/python/pyarrow/lib.pyx @@ -176,6 +176,9 @@ include "table.pxi" # Tensors include "tensor.pxi" +# DLPack +include "_dlpack.pxi" + # File IO include "io.pxi" diff --git a/python/pyarrow/tests/test_dlpack.py b/python/pyarrow/tests/test_dlpack.py new file mode 100644 index 0000000000000..7cf3f4acdbd40 --- /dev/null +++ b/python/pyarrow/tests/test_dlpack.py @@ -0,0 +1,142 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import ctypes +from functools import wraps +import pytest + +import numpy as np + +import pyarrow as pa +from pyarrow.vendored.version import Version + + +def PyCapsule_IsValid(capsule, name): + return ctypes.pythonapi.PyCapsule_IsValid(ctypes.py_object(capsule), name) == 1 + + +def check_dlpack_export(arr, expected_arr): + DLTensor = arr.__dlpack__() + assert PyCapsule_IsValid(DLTensor, b"dltensor") is True + + result = np.from_dlpack(arr) + np.testing.assert_array_equal(result, expected_arr, strict=True) + + assert arr.__dlpack_device__() == (1, 0) + + +def check_bytes_allocated(f): + @wraps(f) + def wrapper(*args, **kwargs): + allocated_bytes = pa.total_allocated_bytes() + try: + return f(*args, **kwargs) + finally: + assert pa.total_allocated_bytes() == allocated_bytes + return wrapper + + +@check_bytes_allocated +@pytest.mark.parametrize( + ('value_type', 'np_type'), + [ + (pa.uint8(), np.uint8), + (pa.uint16(), np.uint16), + (pa.uint32(), np.uint32), + (pa.uint64(), np.uint64), + (pa.int8(), np.int8), + (pa.int16(), np.int16), + (pa.int32(), np.int32), + (pa.int64(), np.int64), + (pa.float16(), np.float16), + (pa.float32(), np.float32), + (pa.float64(), np.float64), + ] +) +def test_dlpack(value_type, np_type): + if Version(np.__version__) < Version("1.24.0"): + pytest.skip("No dlpack support in numpy versions older than 1.22.0, " + "strict keyword in assert_array_equal added in numpy version " + "1.24.0") + + expected = np.array([1, 2, 3], dtype=np_type) + arr = pa.array(expected, type=value_type) + check_dlpack_export(arr, expected) + + arr_sliced = arr.slice(1, 1) + expected = np.array([2], dtype=np_type) + check_dlpack_export(arr_sliced, expected) + + arr_sliced = arr.slice(0, 1) + expected = np.array([1], dtype=np_type) + check_dlpack_export(arr_sliced, expected) + + arr_sliced = arr.slice(1) + expected = np.array([2, 3], dtype=np_type) + check_dlpack_export(arr_sliced, expected) + + arr_zero = pa.array([], type=value_type) + expected = np.array([], dtype=np_type) + check_dlpack_export(arr_zero, expected) + + +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.") + + arr = pa.array([1, None, 3]) + with pytest.raises(TypeError, match="Can only use DLPack " + "on arrays with no nulls."): + np.from_dlpack(arr) + + arr = pa.array( + [[0, 1], [3, 4]], + type=pa.list_(pa.int32()) + ) + with pytest.raises(TypeError, match="DataType is not compatible with DLPack spec"): + np.from_dlpack(arr) + + arr = pa.array([]) + with pytest.raises(TypeError, match="DataType is not compatible with DLPack spec"): + np.from_dlpack(arr) + + # DLPack doesn't support bit-packed boolean values + arr = pa.array([True, False, True]) + with pytest.raises(TypeError, match="Bit-packed boolean data type " + "not supported by DLPack."): + np.from_dlpack(arr) + + +def test_dlpack_cuda_not_supported(): + cuda = pytest.importorskip("pyarrow.cuda") + + schema = pa.schema([pa.field('f0', pa.int16())]) + a0 = pa.array([1, 2, 3], type=pa.int16()) + batch = pa.record_batch([a0], schema=schema) + + cbuf = cuda.serialize_record_batch(batch, cuda.Context(0)) + cbatch = cuda.read_record_batch(cbuf, batch.schema) + carr = cbatch["f0"] + + # CudaBuffers not yet supported + with pytest.raises(NotImplementedError, match="DLPack support is implemented " + "only for buffers on CPU device."): + np.from_dlpack(carr) + + with pytest.raises(NotImplementedError, match="DLPack support is implemented " + "only for buffers on CPU device."): + carr.__dlpack_device__()