From 6206139cb2917beab7ca710b50b98207bac461e3 Mon Sep 17 00:00:00 2001 From: Sebastian Berg Date: Mon, 7 Apr 2025 10:06:15 -0700 Subject: [PATCH] Implement __dlpack__ dunder for pylibcudf columns This implements the `__dlpack__` dunder (and `__dlpack_device__`), which is made available to cudf.Column. There is a bit of a clash with the old dlpack implementation. It is similar, but also different because the old one is table centric and always copies, while this is column centric and copies only if requested. Thus, I kept it as detailed API. One of the more complex things here is the stream synchronization unfortunately, it seems very hard to test reliably in practice. Signed-off-by: Sebastian Berg --- .pre-commit-config.yaml | 3 +- .../all_cuda-129_arch-aarch64.yaml | 1 - .../all_cuda-129_arch-x86_64.yaml | 1 - .../all_cuda-130_arch-aarch64.yaml | 1 - .../all_cuda-130_arch-x86_64.yaml | 1 - cpp/cmake/thirdparty/get_dlpack.cmake | 4 +- cpp/include/cudf/detail/dlpack_vendored.h | 374 ++++++++++++++++++ cpp/include/cudf/detail/interop.hpp | 36 ++ cpp/include/cudf/interop.hpp | 9 +- cpp/src/interop/dlpack.cpp | 104 ++++- cpp/tests/interop/dlpack_test.cpp | 5 +- dependencies.yaml | 1 - python/cudf/cudf/core/column/numerical.py | 19 + python/cudf/cudf/core/single_column_frame.py | 28 ++ python/pylibcudf/pylibcudf/column.pyi | 12 + python/pylibcudf/pylibcudf/column.pyx | 17 +- python/pylibcudf/pylibcudf/interop.pxd | 10 + python/pylibcudf/pylibcudf/interop.pyx | 112 +++++- .../pylibcudf/pylibcudf/libcudf/interop.pxd | 15 +- python/pylibcudf/tests/test_interop.py | 83 ++++ 20 files changed, 797 insertions(+), 39 deletions(-) create mode 100644 cpp/include/cudf/detail/dlpack_vendored.h diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e66e866c4f43..18b36ace93bc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -152,7 +152,8 @@ repos: (?x)^( cpp/include/cudf_test/cxxopts[.]hpp$| cpp/src/io/parquet/ipc/Message_generated[.]h$| - cpp/src/io/parquet/ipc/Schema_generated[.]h$ + cpp/src/io/parquet/ipc/Schema_generated[.]h$| + cpp/include/cudf/detail/dlpack_vendored.h$ ) - id: verify-alpha-spec - id: verify-codeowners diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 7a603a29ec9a..aff78e734f91 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -27,7 +27,6 @@ dependencies: - cxx-compiler - cython>=3.0.3 - dask-cuda==25.10.*,>=0.0.0a0 -- dlpack>=0.8,<1.0 - doxygen=1.9.1 - fastavro>=0.22.9 - flatbuffers==24.3.25 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 40ed1cbbc6b3..1c6cfc4fdacc 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -27,7 +27,6 @@ dependencies: - cxx-compiler - cython>=3.0.3 - dask-cuda==25.10.*,>=0.0.0a0 -- dlpack>=0.8,<1.0 - doxygen=1.9.1 - fastavro>=0.22.9 - flatbuffers==24.3.25 diff --git a/conda/environments/all_cuda-130_arch-aarch64.yaml b/conda/environments/all_cuda-130_arch-aarch64.yaml index e58e93e1aa63..3579b7128da1 100644 --- a/conda/environments/all_cuda-130_arch-aarch64.yaml +++ b/conda/environments/all_cuda-130_arch-aarch64.yaml @@ -27,7 +27,6 @@ dependencies: - cxx-compiler - cython>=3.0.3 - dask-cuda==25.10.*,>=0.0.0a0 -- dlpack>=0.8,<1.0 - doxygen=1.9.1 - fastavro>=0.22.9 - flatbuffers==24.3.25 diff --git a/conda/environments/all_cuda-130_arch-x86_64.yaml b/conda/environments/all_cuda-130_arch-x86_64.yaml index 0d4dc2b2c2de..4a0c9a8d0704 100644 --- a/conda/environments/all_cuda-130_arch-x86_64.yaml +++ b/conda/environments/all_cuda-130_arch-x86_64.yaml @@ -27,7 +27,6 @@ dependencies: - cxx-compiler - cython>=3.0.3 - dask-cuda==25.10.*,>=0.0.0a0 -- dlpack>=0.8,<1.0 - doxygen=1.9.1 - fastavro>=0.22.9 - flatbuffers==24.3.25 diff --git a/cpp/cmake/thirdparty/get_dlpack.cmake b/cpp/cmake/thirdparty/get_dlpack.cmake index 790d6367745d..a1c38cd5dbec 100644 --- a/cpp/cmake/thirdparty/get_dlpack.cmake +++ b/cpp/cmake/thirdparty/get_dlpack.cmake @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2020-2024, NVIDIA CORPORATION. +# Copyright (c) 2020-2025, NVIDIA CORPORATION. # # Licensed 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 @@ -36,6 +36,6 @@ function(find_and_configure_dlpack VERSION) endif() endfunction() -set(CUDF_MIN_VERSION_dlpack 0.8) +set(CUDF_MIN_VERSION_dlpack 1.0) find_and_configure_dlpack(${CUDF_MIN_VERSION_dlpack}) diff --git a/cpp/include/cudf/detail/dlpack_vendored.h b/cpp/include/cudf/detail/dlpack_vendored.h new file mode 100644 index 000000000000..a68eecaa453a --- /dev/null +++ b/cpp/include/cudf/detail/dlpack_vendored.h @@ -0,0 +1,374 @@ +/* + * This header is vendored from: + * https://github.com/dmlc/dlpack/blob/main/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 1 + +/*! \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-partitioned + * 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, + /*! \brief Microsoft MAIA devices */ + kDLMAIA = 17, + /*! \brief AWS Trainium */ + kDLTrn = 18, +} 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, + /*! \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; + +/*! + * \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) + * - 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 { + /*! + * \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, 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. + * + * 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 + * + * 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 */ + 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 destructor deletes the argument self as well. + */ + void (*deleter)(struct DLManagedTensor* self); +} DLManagedTensor; + +// 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. + * + * 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 destructor 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 + * \sa DLPACK_FLAG_BITMASK_IS_COPIED + */ + 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/include/cudf/detail/interop.hpp b/cpp/include/cudf/detail/interop.hpp index 80773ac4c9ca..7c7bddb7c14d 100644 --- a/cpp/include/cudf/detail/interop.hpp +++ b/cpp/include/cudf/detail/interop.hpp @@ -26,6 +26,9 @@ #include +struct DLManagedTensor; +struct DLManagedTensorVersioned; + namespace CUDF_EXPORT cudf { namespace detail { @@ -38,9 +41,42 @@ std::unique_ptr from_dlpack(DLManagedTensor const* managed_tensor, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/** + * @brief Export column as DLManagedTensorVersioned + * + * This function copies the data only if `to_cpu` is true or it is flagged. + * Note that the signature is currently designed to be called from pylibcudf. + * (Otherwise, it should for example signal the DLPack version.) + * + * @param col Column view to export. + * @param copy Whether to copy the data (must be set if `to_cpu` is true). + * If copy is false, a `delete_func` must be provided if the capsule may outlive + * the column view. WARNING: This function is not called on error! + * @param to_cpu Whether to copy to the CPU rather than making a device copy. + * @param sync_stream If passed, an event is put on this stream to ensure + * data is available for use on it. If `nullopt`, the stream is synchronized. + * @param stream The cudf stream to operate on (i.e. to copy the data safely). + * @param mr Memory resource to use for the new allocation (if on GPU). + * @param delete_func Custom delete function, needed to avoid Python API use. + * On error, the delete function is not called. + * @param delete_ctx Context for custom delete (i.e. Python owner of data). + */ +DLManagedTensorVersioned* to_dlpack_versioned( + column_view col, + bool copy, + bool to_cpu, + rmm::cuda_stream_view sync_stream, + void (*delete_func)(void*), + void* delete_ctx, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + /** * @copydoc cudf::to_dlpack * + * @note Prefer the column-based overload for new code. This overload should be + * deprecated eventually, but is public API. (Or updated if we need it.) + * * @param stream CUDA stream used for device memory operations and kernel launches. */ DLManagedTensor* to_dlpack(table_view const& input, diff --git a/cpp/include/cudf/interop.hpp b/cpp/include/cudf/interop.hpp index 2929b5cc6018..acb7306156f7 100644 --- a/cpp/include/cudf/interop.hpp +++ b/cpp/include/cudf/interop.hpp @@ -56,7 +56,7 @@ namespace CUDF_EXPORT cudf { */ /** - * @brief Convert a DLPack DLTensor into a cudf table + * @brief Convert a DLPack DLManagedTensor into a cudf table * * The `device_type` of the DLTensor must be `kDLCPU`, `kDLCuda`, or * `kDLCUDAHost`, and `device_id` must match the current device. The `ndim` @@ -79,11 +79,10 @@ std::unique_ptr
from_dlpack( rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); /** - * @brief Convert a cudf table into a DLPack DLTensor + * @brief Convert a cudf table into a DLPack DLManagedTensor. * - * All columns must have the same data type and this type must be numeric. The - * columns may be nullable, but the null count must be zero. If the input table - * is empty or has zero rows, the result will be nullptr. + * @note This exports DLPack <1 struct. The newer struct is currently available + * for columns via Python `__dlpack__`. * * @note The `deleter` method of the returned `DLManagedTensor` must be used to * free the memory allocated for the tensor. diff --git a/cpp/src/interop/dlpack.cpp b/cpp/src/interop/dlpack.cpp index 7b76139f7c85..d6ee3746ac34 100644 --- a/cpp/src/interop/dlpack.cpp +++ b/cpp/src/interop/dlpack.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ #include +#include #include #include #include @@ -22,10 +23,6 @@ #include #include -#include - -#include - #include namespace cudf { @@ -115,15 +112,21 @@ DLDataType data_type_to_DLDataType(data_type type) return type_dispatcher(type, data_type_to_DLDataType_impl{}); } -// Context object to own memory allocated for DLManagedTensor +// Context object to own memory allocated for DLManagedTensor and DLManagedTensorVersioned +template struct dltensor_context { int64_t shape[2]{}; // NOLINT int64_t strides[2]{}; // NOLINT - rmm::device_buffer buffer; + rmm::device_buffer gpu_buffer; + std::vector cpu_buffer; + /* custom deletion, needed for lifetime management */ + void (*delete_func)(void*); + void* delete_ctx; - static void deleter(DLManagedTensor* arg) + static void deleter(DLManagedTensorT* arg) { auto context = static_cast(arg->manager_ctx); + if (context->delete_func != nullptr) { context->delete_func(context->delete_ctx); } delete context; delete arg; } @@ -241,7 +244,10 @@ DLManagedTensor* to_dlpack(table_view const& input, "Input required to have null count zero"); auto managed_tensor = std::make_unique(); - auto context = std::make_unique(); + auto context = std::make_unique>(); + + context->delete_func = nullptr; + context->delete_ctx = nullptr; DLTensor& tensor = managed_tensor->dl_tensor; tensor.dtype = dltype; @@ -271,8 +277,8 @@ DLManagedTensor* to_dlpack(table_view const& input, size_t const stride_bytes = num_rows * size_of(type); size_t const total_bytes = stride_bytes * num_cols; - context->buffer = rmm::device_buffer(total_bytes, stream, mr); - tensor.data = context->buffer.data(); + context->gpu_buffer = rmm::device_buffer(total_bytes, stream, mr); + tensor.data = context->gpu_buffer.data(); auto tensor_data = reinterpret_cast(tensor.data); for (auto const& col : input) { @@ -285,7 +291,7 @@ DLManagedTensor* to_dlpack(table_view const& input, } // Defer ownership of managed tensor to caller - managed_tensor->deleter = dltensor_context::deleter; + managed_tensor->deleter = dltensor_context::deleter; managed_tensor->manager_ctx = context.release(); // synchronize the stream because after the return the data may be accessed from the host before @@ -296,6 +302,82 @@ DLManagedTensor* to_dlpack(table_view const& input, return managed_tensor.release(); } +DLManagedTensorVersioned* to_dlpack_versioned(column_view col, + bool copy, + bool to_cpu, + rmm::cuda_stream_view sync_stream, + void (*delete_func)(void*), + void* delete_ctx, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_EXPECTS(copy || !to_cpu, "Must copy data if exporting to CPU."); + // Ensure none of the columns have nulls + CUDF_EXPECTS(!col.has_nulls(), "Cannot export column with nulls to DLPack."); + + auto const num_rows = col.size(); + + // Ensure that type is convertible to DLDataType + data_type const type = col.type(); + DLDataType const dltype = data_type_to_DLDataType(type); + + auto managed_tensor = std::make_unique(); + auto context = std::make_unique>(); + + context->delete_func = delete_func; + context->delete_ctx = delete_ctx; + + DLTensor& tensor = managed_tensor->dl_tensor; + tensor.dtype = dltype; + + tensor.ndim = 1; + tensor.shape = context->shape; + tensor.shape[0] = num_rows; + + tensor.device.device_id = rmm::get_current_cuda_device().value(); + tensor.device.device_type = to_cpu ? kDLCPU : kDLCUDA; + + size_t const nbytes = num_rows * size_of(type); + + if (!copy) { + tensor.data = (void*)get_column_data(col); + } else { + cudaMemcpyKind memcpykind; + if (!to_cpu) { + context->gpu_buffer = rmm::device_buffer(nbytes, stream, mr); + tensor.data = context->gpu_buffer.data(); + memcpykind = cudaMemcpyDeviceToDevice; + } else { + context->cpu_buffer.reserve(nbytes); + tensor.data = context->cpu_buffer.data(); + memcpykind = cudaMemcpyDeviceToHost; + } + + CUDF_CUDA_TRY( + cudaMemcpyAsync(tensor.data, get_column_data(col), nbytes, memcpykind, stream.value())); + } + + managed_tensor->version = {1, 0}; + if (copy) { + managed_tensor->flags = DLPACK_FLAG_BITMASK_IS_COPIED; + } else { + managed_tensor->flags = DLPACK_FLAG_BITMASK_READ_ONLY; + } + managed_tensor->deleter = context->deleter; + + if (sync_stream.value() != stream.value()) { + cudaEvent_t event; + CUDF_CUDA_TRY(cudaEventCreate(&event)); + CUDF_CUDA_TRY(cudaEventRecord(event, stream.value())); + CUDF_CUDA_TRY(cudaStreamWaitEvent(sync_stream.value(), event, 0)); + CUDF_CUDA_TRY(cudaEventDestroy(event)); + } + + // Release ownership to the caller (we will not cleanup anymore) + managed_tensor->manager_ctx = context.release(); + return managed_tensor.release(); +} + } // namespace detail std::unique_ptr
from_dlpack(DLManagedTensor const* managed_tensor, diff --git a/cpp/tests/interop/dlpack_test.cpp b/cpp/tests/interop/dlpack_test.cpp index b7106e823dd7..ee8e73d3ad60 100644 --- a/cpp/tests/interop/dlpack_test.cpp +++ b/cpp/tests/interop/dlpack_test.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. + * Copyright (c) 2019-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -19,13 +19,12 @@ #include #include +#include #include #include #include -#include - namespace { struct dlpack_deleter { void operator()(DLManagedTensor* tensor) { tensor->deleter(tensor); } diff --git a/dependencies.yaml b/dependencies.yaml index 0e96a48ad86a..9ff8d65cb17c 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -453,7 +453,6 @@ dependencies: - c-compiler - cuda-nvcc - cxx-compiler - - dlpack>=0.8,<1.0 - zlib>=1.2.13 specific: - output_types: conda diff --git a/python/cudf/cudf/core/column/numerical.py b/python/cudf/cudf/core/column/numerical.py index c0ebc3343a13..4a8d49807042 100644 --- a/python/cudf/cudf/core/column/numerical.py +++ b/python/cudf/cudf/core/column/numerical.py @@ -964,6 +964,25 @@ def digitize(self, bins: np.ndarray, right: bool = False) -> Self: ) ) + def __dlpack__( + self, + /, + *, + stream: int | Any | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ): + return self.to_pylibcudf(mode="read").__dlpack__( + stream=stream, + max_version=max_version, + dl_device=dl_device, + copy=copy, + ) + + def __dlpack_device__(self): + return plc.interop._get_dlpack_device() + def _normalize_find_and_replace_input( input_column_dtype: DtypeObj, col_to_normalize: ColumnBase | list diff --git a/python/cudf/cudf/core/single_column_frame.py b/python/cudf/cudf/core/single_column_frame.py index dadd8e3a86b2..f345e9466b72 100644 --- a/python/cudf/cudf/core/single_column_frame.py +++ b/python/cudf/cudf/core/single_column_frame.py @@ -278,6 +278,34 @@ def __cuda_array_interface__(self): "'__cuda_array_interface__'" ) + def __dlpack__( + self, + /, + *, + stream: int | Any | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ): + dlpack = getattr(self._column, "__dlpack__", None) + if dlpack: + return dlpack( + stream=stream, + max_version=max_version, + dl_device=dl_device, + copy=copy, + ) + else: + raise BufferError("Column dtype is not supported by DLPack") + + @property + def __dlpack_device__(self): + dlpack_device = getattr(self._column, "__dlpack_device__", None) + if dlpack_device: + return dlpack_device + else: + raise BufferError("Column dtype is not supported by DLPack") + @_performance_tracking def factorize( self, sort: bool = False, use_na_sentinel: bool = True diff --git a/python/pylibcudf/pylibcudf/column.pyi b/python/pylibcudf/pylibcudf/column.pyi index ff4945c123ef..709198a43b6e 100644 --- a/python/pylibcudf/pylibcudf/column.pyi +++ b/python/pylibcudf/pylibcudf/column.pyi @@ -3,6 +3,8 @@ from collections.abc import Sequence from typing import Any, Protocol, TypedDict +from typing_extensions import CapsuleType + from rmm.pylibrmm.device_buffer import DeviceBuffer from rmm.pylibrmm.stream import Stream @@ -43,6 +45,16 @@ class Column: offset: int, children: list[Column], ) -> None: ... + def __dlpack__( + self, + /, + *, + stream: int | Any | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> CapsuleType: ... + def __dlpack_device__(self, /) -> tuple[int, int]: ... def type(self) -> DataType: ... def child(self, index: int) -> Column: ... def size(self) -> int: ... diff --git a/python/pylibcudf/pylibcudf/column.pyx b/python/pylibcudf/pylibcudf/column.pyx index e7e936554e32..18383e4f171f 100644 --- a/python/pylibcudf/pylibcudf/column.pyx +++ b/python/pylibcudf/pylibcudf/column.pyx @@ -33,7 +33,6 @@ from pylibcudf.libcudf.types cimport size_type, size_of as cpp_size_of, bitmask_ from pylibcudf.libcudf.utilities.traits cimport is_fixed_width from pylibcudf.libcudf.copying cimport get_element - from rmm.pylibrmm.device_buffer cimport DeviceBuffer from rmm.pylibrmm.stream cimport Stream from rmm.pylibrmm.memory_resource cimport ( @@ -43,7 +42,7 @@ from rmm.pylibrmm.memory_resource cimport ( from .gpumemoryview cimport gpumemoryview from .filling cimport sequence -from .gpumemoryview cimport gpumemoryview +from .interop cimport to_dlpack_col, _get_dlpack_device from .scalar cimport Scalar from .traits cimport ( is_fixed_width as plc_is_fixed_width, @@ -1339,6 +1338,20 @@ cdef class Column: return self._to_schema(), self._to_device_array() + def __dlpack__( + self, + /, + *, + stream: int | Any | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ): + return to_dlpack_col(self, stream, max_version, dl_device, copy) + + def __dlpack_device__(self): + return _get_dlpack_device() + cdef class ListColumnView: """Accessor for methods of a Column that are specific to lists.""" diff --git a/python/pylibcudf/pylibcudf/interop.pxd b/python/pylibcudf/pylibcudf/interop.pxd index a46d720ae911..78b139499166 100644 --- a/python/pylibcudf/pylibcudf/interop.pxd +++ b/python/pylibcudf/pylibcudf/interop.pxd @@ -1,5 +1,6 @@ # Copyright (c) 2024-2025, NVIDIA CORPORATION. +from pylibcudf.column cimport Column from pylibcudf.table cimport Table from rmm.pylibrmm.stream cimport Stream from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource @@ -9,3 +10,12 @@ cpdef Table from_dlpack( ) cpdef object to_dlpack(Table input, Stream stream=*, DeviceMemoryResource mr=*) + +cpdef _get_dlpack_device() + +cpdef object to_dlpack_col( + Column col, + stream=*, + max_version=*, + dl_device=*, + copy=*) diff --git a/python/pylibcudf/pylibcudf/interop.pyx b/python/pylibcudf/pylibcudf/interop.pyx index 86e265e898bb..3d3095f9b491 100644 --- a/python/pylibcudf/pylibcudf/interop.pyx +++ b/python/pylibcudf/pylibcudf/interop.pyx @@ -6,19 +6,28 @@ from cpython.pycapsule cimport ( PyCapsule_New, PyCapsule_SetName, ) +from cpython.ref cimport Py_DECREF, Py_INCREF + +from libc.stdint cimport uintptr_t from libcpp.memory cimport unique_ptr from libcpp.utility cimport move from functools import singledispatch +from rmm.librmm cimport cuda_stream_view + +from pylibcudf.libcudf.utilities.default_stream cimport get_default_stream from pylibcudf.libcudf.interop cimport ( DLManagedTensor, + DLManagedTensorVersioned, from_dlpack as cpp_from_dlpack, to_dlpack as cpp_to_dlpack, + to_dlpack_versioned as cpp_to_dlpack_versioned, ) from pylibcudf.libcudf.table.table cimport table from rmm.pylibrmm.stream cimport Stream +from rmm.librmm.device_buffer cimport get_current_cuda_device from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .column cimport Column @@ -42,6 +51,8 @@ __all__ = [ "from_dlpack", "to_arrow", "to_dlpack", + "get_dlpack_device", + "to_dlpack_col", ] @@ -225,12 +236,95 @@ cpdef object to_dlpack(Table input, Stream stream=None, DeviceMemoryResource mr= ) -cdef void dlmanaged_tensor_pycapsule_deleter(object pycap_obj) noexcept: - if PyCapsule_IsValid(pycap_obj, "used_dltensor"): - # we do not call a used capsule's deleter - return - cdef DLManagedTensor* dlpack_tensor = PyCapsule_GetPointer( - pycap_obj, "dltensor" - ) - if dlpack_tensor is not NULL: - dlpack_tensor.deleter(dlpack_tensor) +cdef void _dltensor_delete_owner(void *ctx) noexcept nogil: + """Helper to delete the column owning the dlpack tensor data.""" + with gil: + Py_DECREF(ctx) + + +cpdef _get_dlpack_device(): + # Indicate as CUDA memory (although it could be pinned in theory). + return (2, get_current_cuda_device().value()) + + +cpdef object to_dlpack_col( + Column self, + stream=None, + max_version=None, + dl_device=None, + copy=None, +): + cdef DLManagedTensorVersioned* dlmtensorv + cdef bint to_cpu = False + cdef cuda_stream_view.cuda_stream_view c_stream + + if max_version is None or max_version[0] == 0: + raise BufferError("Unable to export DLPack with version <1.0.") + + if dl_device is not None: + if dl_device == (1, 0): + to_cpu = True + if copy is None: + copy = True # moving to CPU is always a copy + elif not copy: + raise BufferError("copying to CPU requires a copy.") + + elif dl_device != self.__dlpack_device__: + raise BufferError("cudf only supports dl_device if identical or CPU.") + + if self.null_count() != 0: + # A BufferError is nicer, so explicitly check in Python + # (we could convert a custom exception from C->Cython also) + raise BufferError("Cannot export column with nulls.") + + if copy is None: + copy = False + + if to_cpu: + # just ignore stream, it should match the device, cpu has none + # (could guess it's a GPU one, but that isn't strictly standardized) + c_stream = get_default_stream() + elif stream == -1: + # Passing the default stream means we don't do any synchronization + # (so use our stream which enforces no order between streams). + c_stream = get_default_stream() + elif stream is None or stream == 1: + c_stream = cuda_stream_view.cuda_stream_legacy + elif stream == 2: + c_stream = cuda_stream_view.cuda_stream_default + elif not isinstance(stream, int) or stream < 3: + raise ValueError( + f'On CUDA, the valid stream for the DLPack protocol is -1,' + f' 1, 2, or any larger value, but {stream} was provided') + else: + # User provided a custom stream. + c_stream = cuda_stream_view.cuda_stream_view( + stream) + + dlmtensorv = cpp_to_dlpack_versioned( + self.view(), copy, to_cpu, c_stream, _dltensor_delete_owner, self) + Py_INCREF(self) # on success, dlmtensorv takes a reference to self + try: + capsule = PyCapsule_New( + dlmtensorv, 'dltensor_versioned', dlmanaged_tensor_pycapsule_deleter) + except BaseException: # ownership not transferred to capsule + dlmtensorv.deleter(dlmtensorv) + raise + + if to_cpu: + # Make sure data is CPU available + c_stream.synchronize() + + return capsule + + +cdef void dlmanaged_tensor_pycapsule_deleter(object capsule) noexcept: + """Delete the dlpack tensor stored inside the capsule (if needed).""" + cdef DLManagedTensorVersioned *dltensorv + + if PyCapsule_IsValid(capsule, "dltensor_versioned"): + dltensorv = PyCapsule_GetPointer( + capsule, "dltensor_versioned") + dltensorv.deleter(dltensorv) + + # If the capsule was renamed, the tensor is already free'd. diff --git a/python/pylibcudf/pylibcudf/libcudf/interop.pxd b/python/pylibcudf/pylibcudf/libcudf/interop.pxd index b80431c960da..2a1c4e5be809 100644 --- a/python/pylibcudf/pylibcudf/libcudf/interop.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/interop.pxd @@ -1,5 +1,6 @@ # Copyright (c) 2020-2025, NVIDIA CORPORATION. from libc.stdint cimport int32_t +from libcpp cimport bool from libcpp.memory cimport shared_ptr, unique_ptr from libcpp.string cimport string from libcpp.optional cimport optional @@ -17,7 +18,10 @@ from rmm.librmm.memory_resource cimport device_memory_resource cdef extern from "dlpack/dlpack.h" nogil: ctypedef struct DLManagedTensor: - void(*deleter)(DLManagedTensor*) except +libcudf_exception_handler + void(*deleter)(DLManagedTensor*) noexcept + + ctypedef struct DLManagedTensorVersioned: + void(*deleter)(DLManagedTensorVersioned*) noexcept # The Arrow structs are not namespaced. @@ -57,6 +61,15 @@ cdef extern from "cudf/interop.hpp" namespace "cudf" \ optional[int32_t] precision vector[column_metadata] children_meta +cdef extern from "cudf/detail/interop.hpp" nogil: + DLManagedTensorVersioned* to_dlpack_versioned ( + const column_view& col, + bool copy, + bool to_cpu, + cuda_stream_view sync_stream, + void (*delete_func)(void *), + void *delete_ctx + ) except +libcudf_exception_handler cdef extern from "cudf/interop.hpp" namespace "cudf::interop" \ nogil: diff --git a/python/pylibcudf/tests/test_interop.py b/python/pylibcudf/tests/test_interop.py index ce3f77875342..34bf0f318b39 100644 --- a/python/pylibcudf/tests/test_interop.py +++ b/python/pylibcudf/tests/test_interop.py @@ -1,6 +1,7 @@ # Copyright (c) 2024-2025, NVIDIA CORPORATION. import decimal +import sys import cupy as cp import nanoarrow @@ -253,3 +254,85 @@ def f(): except ValueError: # Ignore the exception. A failure in this test is a seg fault pass + + +@pytest.mark.parametrize("dtype", ["int32", "float32", "float64", "uint64"]) +def test__dlpack___simple(dtype): + # Going via cupy also sanity checks __dlpack_device__ works. + orig = cp.arange(1000, dtype=dtype) + col = plc.Column.from_cuda_array_interface(orig) + arr = cp.from_dlpack(col) + + cp.testing.assert_array_equal(orig, arr) + + +def test__dlpack__copy(): + orig = cp.arange(1000, dtype="int32") + col = plc.Column.from_cuda_array_interface(orig) + arr = cp.from_dlpack(col, copy=False) + cp.testing.assert_array_equal(orig, arr) + + orig_ptr = orig.__cuda_array_interface__["data"][0] + assert arr.__cuda_array_interface__["data"][0] == orig_ptr + arr = cp.from_dlpack(col, copy=True) + assert arr.__cuda_array_interface__["data"][0] != orig_ptr + + +def test__dlpack___cleanup(): + # Sanity check refcount of the original column with dlpack + col = plc.Column.from_arrow(pa.array([1, 2, 3])) + refcnt = sys.getrefcount(col) + + capsule = col.__dlpack__(max_version=(1, 0)) # No capsule user + assert sys.getrefcount(col) == refcnt + 1 + del capsule + assert sys.getrefcount(col) == refcnt + + arr = cp.from_dlpack(col) # Capsule is consumed + assert sys.getrefcount(col) == refcnt + 1 + del arr + assert sys.getrefcount(col) == refcnt + + +def test__dlpack__version(): + orig = cp.arange(1000, dtype="int32") + col = plc.Column.from_cuda_array_interface(orig) + + name = '"dltensor_versioned"' + assert name in str(col.__dlpack__(max_version=(2, 0))) + + # Must provide a version (error is more helpful if we give it) + with pytest.raises(BufferError): + col.__dlpack__(max_version=(0, 10)) + + with pytest.raises(BufferError): + col.__dlpack__(max_version=None) + + +@pytest.mark.parametrize("dtype", ["int32", "float32", "float64"]) +def test__dlpack___to_host(dtype): + orig = cp.arange(1000, dtype=dtype) + col = plc.Column.from_cuda_array_interface(orig) + try: + arr = np.from_dlpack(col, device="cpu") + except TypeError as e: + if "numpy.from_dlpack() takes no keyword arguments" in str(e): + pytest.skip("New numpy version does not support device kwarg") + raise + + cp.testing.assert_array_equal(orig, cp.asarray(arr)) + + +def test__dlpack__device_error(): + col = plc.Column.from_cuda_array_interface(cp.arange(1000)) + with pytest.raises( + BufferError, match="cudf only supports dl_device if identical or CPU" + ): + col.__dlpack__(max_version=(1, 0), dl_device=(10, 0)) + + +def test__dlpack___reject_mask(): + col = plc.Column.from_arrow(pa.array([1, 2, 3, None])) + + with pytest.raises(BufferError, match="Cannot export column with nulls."): + col.__dlpack__(max_version=(1, 0))