You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "westonpace (via GitHub)" <gi...@apache.org> on 2023/07/19 00:04:12 UTC
[GitHub] [arrow] westonpace commented on a diff in pull request #36489: GH-36488: [C++] Import/Export ArrowDeviceArray
westonpace commented on code in PR #36489:
URL: https://github.com/apache/arrow/pull/36489#discussion_r1267394276
##########
cpp/src/arrow/buffer.h:
##########
@@ -57,18 +58,31 @@ class ARROW_EXPORT Buffer {
///
/// \note The passed memory must be kept alive through some other means
Buffer(const uint8_t* data, int64_t size)
- : is_mutable_(false), is_cpu_(true), data_(data), size_(size), capacity_(size) {
+ : is_mutable_(false),
+ is_cpu_(true),
+ data_(data),
+ size_(size),
+ capacity_(size),
+ device_type_(DeviceAllocationType::kCPU) {
SetMemoryManager(default_cpu_memory_manager());
}
Buffer(const uint8_t* data, int64_t size, std::shared_ptr<MemoryManager> mm,
- std::shared_ptr<Buffer> parent = NULLPTR)
+ std::shared_ptr<Buffer> parent = NULLPTR,
+ std::optional<DeviceAllocationType> device_type = std::nullopt)
: is_mutable_(false),
data_(data),
size_(size),
capacity_(size),
parent_(std::move(parent)) {
+ // will set device_type from the memory manager
Review Comment:
```suggestion
// SetMemoryManager will also set device_type_
```
##########
cpp/src/arrow/buffer.h:
##########
@@ -294,6 +310,7 @@ class ARROW_EXPORT Buffer {
const uint8_t* data_;
int64_t size_;
int64_t capacity_;
+ std::optional<DeviceAllocationType> device_type_;
Review Comment:
Do we still need `is_cpu_` if we have `device_type_`?
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -166,6 +167,139 @@ Result<std::shared_ptr<RecordBatch>> ImportRecordBatch(struct ArrowArray* array,
/// @}
+/// \defgroup c-data-device-interface Functions for working with the C data device
+/// interface.
+///
+/// @{
+
+/// \brief EXPERIMENTAL: Type for freeing a sync event
+///
+/// If synchronization is necessary for accessing the data on a device,
+/// a pointer to an event needs to be passed when exporting the device
+/// array. It's the responsibility of the release function for the array
+/// to release the event. Both can be null if no sync'ing is necessary.
+struct RawSyncEvent {
Review Comment:
Why "raw" and not just `SyncEvent`?
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -166,6 +167,139 @@ Result<std::shared_ptr<RecordBatch>> ImportRecordBatch(struct ArrowArray* array,
/// @}
+/// \defgroup c-data-device-interface Functions for working with the C data device
+/// interface.
+///
+/// @{
+
+/// \brief EXPERIMENTAL: Type for freeing a sync event
+///
+/// If synchronization is necessary for accessing the data on a device,
+/// a pointer to an event needs to be passed when exporting the device
+/// array. It's the responsibility of the release function for the array
+/// to release the event. Both can be null if no sync'ing is necessary.
+struct RawSyncEvent {
+ void* sync_event = NULL;
+ std::function<void(void*)> release_func;
+};
+
+/// \brief EXPERIMENTAL: Export C++ Array as an ArrowDeviceArray.
+///
+/// The resulting ArrowDeviceArray struct keeps the array data and buffers alive
+/// until its release callback is called by the consumer. All buffers in
+/// the provided array MUST have the same device_type, otherwise an error
+/// will be returned.
+///
+/// If a non-null sync_event is provided, then the sync_release func must also be
+/// non-null. If the sync_event is null, then the sync_release parameter is not called.
+///
+/// \param[in] array Array object to export
+/// \param[in] sync_event A struct containing what is needed for syncing if necessary
+/// \param[out] out C struct to export the array to
+/// \param[out] out_schema optional C struct to export the array type to
+ARROW_EXPORT
+Status ExportDeviceArray(const Array& array, RawSyncEvent sync_event,
+ struct ArrowDeviceArray* out,
+ struct ArrowSchema* out_schema = NULLPTR);
+
+/// \brief EXPERIMENTAL: Export C++ RecordBatch as an ArrowDeviceArray.
+///
+/// The record batch is exported as if it were a struct array.
+/// The resulting ArrowDeviceArray struct keeps the record batch data and buffers alive
+/// until its release callback is called by the consumer.
+///
+/// All buffers of all columns in the record batch must have the same device_type
+/// otherwise an error will be returned. If columns are on different devices,
+/// they should be exported using different ArrowDeviceArray instances.
+///
+/// If a non-null sync_event is provided, then the sync_release func must also be
+/// non-null. If the sync_event is null, then the sync_release parameter is ignored.
+///
+/// \param[in] batch Record batch to export
+/// \param[in] sync_event A struct containing what is needed for syncing if necessary
+/// \param[out] out C struct where to export the record batch
+/// \param[out] out_schema optional C struct where to export the record batch schema
+ARROW_EXPORT
+Status ExportDeviceRecordBatch(const RecordBatch& batch, RawSyncEvent sync_event,
+ struct ArrowDeviceArray* out,
+ struct ArrowSchema* out_schema = NULLPTR);
+
+using DeviceMemoryMapper =
+ std::function<Result<std::shared_ptr<MemoryManager>>(ArrowDeviceType, int64_t)>;
+
+// ARROW_EXPORT
+// Result<std::shared_ptr<MemoryManager>> DefaultMemoryMapper(ArrowDeviceType device_type,
+// int64_t device_id);
+
+/// \brief EXPERIMENTAL: Import C++ device array from the C data interface.
+///
+/// The ArrowArray struct has its contents moved (as per the C data interface
+/// specification) to a private object held alive by the resulting array. The
+/// buffers of the Array are located on the device indicated by the device_type.
+///
+/// \param[in,out] array C data interface struct holding the array data
+/// \param[in] type type of the imported array
+/// \param[in] mapper A function to map device + id to memory manager
+/// \return Imported array object
+ARROW_EXPORT
+Result<std::shared_ptr<Array>> ImportDeviceArray(struct ArrowDeviceArray* array,
+ std::shared_ptr<DataType> type,
+ const DeviceMemoryMapper& mapper);
+
+/// \brief EXPERIMENTAL: Import C++ device array and its type from the C data interface.
+///
+/// The ArrowArray struct has its contents moved (as per the C data interface
+/// specification) to a private object held alive by the resulting array.
+/// The ArrowSchema struct is released, even if this function fails. The
+/// buffers of the Array are located on the device indicated by the device_type.
+///
+/// \param[in,out] array C data interface struct holding the array data
+/// \param[in,out] type C data interface struct holding the array type
+/// \param[in] mapper A function to map device + id to memory manager
+/// \return Imported array object
+ARROW_EXPORT
+Result<std::shared_ptr<Array>> ImportDeviceArray(struct ArrowDeviceArray* array,
+ struct ArrowSchema* type,
+ const DeviceMemoryMapper& mapper);
+
+/// \brief EXPERIMENTAL: Import C++ record batch with buffers on a device from the C data
+/// interface.
+///
+/// The ArrowArray struct has its contents moved (as per the C data interface
+/// specification) to a private object held alive by the resulting record batch.
+/// The buffers of all columns of the record batch are located on the device
+/// indicated by the device type.
+///
+/// \param[in,out] array C data interface struct holding the record batch data
+/// \param[in] schema schema of the imported record batch
+/// \param[in] mapper A function to map device + id to memory manager
+/// \return Imported record batch object
+ARROW_EXPORT
+Result<std::shared_ptr<RecordBatch>> ImportDeviceRecordBatch(
+ struct ArrowArray* array, std::shared_ptr<Schema> schema,
Review Comment:
```suggestion
struct ArrowDeviceArray* array, std::shared_ptr<Schema> schema,
```
##########
cpp/src/arrow/buffer.h:
##########
@@ -57,18 +58,31 @@ class ARROW_EXPORT Buffer {
///
/// \note The passed memory must be kept alive through some other means
Buffer(const uint8_t* data, int64_t size)
- : is_mutable_(false), is_cpu_(true), data_(data), size_(size), capacity_(size) {
+ : is_mutable_(false),
+ is_cpu_(true),
+ data_(data),
+ size_(size),
+ capacity_(size),
+ device_type_(DeviceAllocationType::kCPU) {
SetMemoryManager(default_cpu_memory_manager());
}
Buffer(const uint8_t* data, int64_t size, std::shared_ptr<MemoryManager> mm,
- std::shared_ptr<Buffer> parent = NULLPTR)
+ std::shared_ptr<Buffer> parent = NULLPTR,
+ std::optional<DeviceAllocationType> device_type = std::nullopt)
: is_mutable_(false),
data_(data),
size_(size),
capacity_(size),
parent_(std::move(parent)) {
+ // will set device_type from the memory manager
SetMemoryManager(std::move(mm));
+ // if a device type is specified, use that instead. for example:
+ // CUDA_HOST. The CudaMemoryManager will set device_type_ to CUDA,
+ // but you can specify CUDA_HOST as the device type to override it.
Review Comment:
Are `CUDA` and `CUDA_HOST` two different things? Why would I want to override the device type set by the memory manager?
##########
cpp/src/arrow/device.h:
##########
@@ -29,6 +29,24 @@
namespace arrow {
+/// \brief EXPERIMENTAL: Device type enum which matches up with C Data Device types
+enum class DeviceAllocationType : char {
Review Comment:
Why `DeviceAllocationType` and not `DeviceType`?
##########
cpp/src/arrow/device.h:
##########
@@ -58,6 +76,12 @@ class ARROW_EXPORT Device : public std::enable_shared_from_this<Device>,
/// \brief Whether this instance points to the same device as another one.
virtual bool Equals(const Device&) const = 0;
+ /// \brief A device ID to identify this device if there are multiple of this type.
+ ///
+ /// If there is no "device_id" equivalent (such as for the main CPU device)
Review Comment:
Well, the main CPU on a non-numa system :)
##########
cpp/src/arrow/gpu/cuda_memory.cc:
##########
@@ -480,5 +485,21 @@ Result<uint8_t*> GetHostAddress(uintptr_t device_ptr) {
return static_cast<uint8_t*>(ptr);
}
+Result<std::shared_ptr<MemoryManager>> DefaultMemoryMapper(ArrowDeviceType device_type,
+ int64_t device_id) {
+ switch (device_type) {
+ case ARROW_DEVICE_CPU:
+ return default_cpu_memory_manager();
+ case ARROW_DEVICE_CUDA:
+ case ARROW_DEVICE_CUDA_HOST:
+ case ARROW_DEVICE_CUDA_MANAGED: {
+ ARROW_ASSIGN_OR_RAISE(auto device, arrow::cuda::CudaDevice::Make(device_id));
Review Comment:
This is a little weird but caused by this PR. My first thought, at looking at this, is that we are returning a dangling object since `device` will be destroyed almost immediately.
I think `arrow::cuda::CudaDevice::Make` should probably be `arrow::cuda::CudaDevice::Get` and return `const CudaDevice&` or `const CudaDevice*`.
##########
cpp/src/arrow/device.h:
##########
@@ -29,6 +29,24 @@
namespace arrow {
+/// \brief EXPERIMENTAL: Device type enum which matches up with C Data Device types
+enum class DeviceAllocationType : char {
+ kCPU = 1,
+ kCUDA = 2,
+ kCUDA_HOST = 3,
+ kOPENCL = 4,
+ kVULKAN = 7,
+ kMETAL = 8,
+ kVPI = 9,
+ kROCM = 10,
+ kROCM_HOST = 11,
+ kEXT_DEV = 12,
+ kCUDA_MANAGED = 13,
+ kONEAPI = 14,
+ kWEBGPU = 15,
+ kHEXAGON = 16,
Review Comment:
Do these names come from somewhere else? I'd prefer names like `kCpu` and `kExtDev` but if this list is mirroring some other list that is probably ok.
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -166,6 +167,139 @@ Result<std::shared_ptr<RecordBatch>> ImportRecordBatch(struct ArrowArray* array,
/// @}
+/// \defgroup c-data-device-interface Functions for working with the C data device
+/// interface.
+///
+/// @{
+
+/// \brief EXPERIMENTAL: Type for freeing a sync event
+///
+/// If synchronization is necessary for accessing the data on a device,
+/// a pointer to an event needs to be passed when exporting the device
+/// array. It's the responsibility of the release function for the array
+/// to release the event. Both can be null if no sync'ing is necessary.
+struct RawSyncEvent {
+ void* sync_event = NULL;
+ std::function<void(void*)> release_func;
+};
+
+/// \brief EXPERIMENTAL: Export C++ Array as an ArrowDeviceArray.
+///
+/// The resulting ArrowDeviceArray struct keeps the array data and buffers alive
+/// until its release callback is called by the consumer. All buffers in
+/// the provided array MUST have the same device_type, otherwise an error
+/// will be returned.
+///
+/// If a non-null sync_event is provided, then the sync_release func must also be
+/// non-null. If the sync_event is null, then the sync_release parameter is not called.
+///
+/// \param[in] array Array object to export
+/// \param[in] sync_event A struct containing what is needed for syncing if necessary
+/// \param[out] out C struct to export the array to
+/// \param[out] out_schema optional C struct to export the array type to
+ARROW_EXPORT
+Status ExportDeviceArray(const Array& array, RawSyncEvent sync_event,
+ struct ArrowDeviceArray* out,
+ struct ArrowSchema* out_schema = NULLPTR);
+
+/// \brief EXPERIMENTAL: Export C++ RecordBatch as an ArrowDeviceArray.
+///
+/// The record batch is exported as if it were a struct array.
+/// The resulting ArrowDeviceArray struct keeps the record batch data and buffers alive
+/// until its release callback is called by the consumer.
+///
+/// All buffers of all columns in the record batch must have the same device_type
+/// otherwise an error will be returned. If columns are on different devices,
+/// they should be exported using different ArrowDeviceArray instances.
+///
+/// If a non-null sync_event is provided, then the sync_release func must also be
+/// non-null. If the sync_event is null, then the sync_release parameter is ignored.
+///
+/// \param[in] batch Record batch to export
+/// \param[in] sync_event A struct containing what is needed for syncing if necessary
+/// \param[out] out C struct where to export the record batch
+/// \param[out] out_schema optional C struct where to export the record batch schema
+ARROW_EXPORT
+Status ExportDeviceRecordBatch(const RecordBatch& batch, RawSyncEvent sync_event,
+ struct ArrowDeviceArray* out,
+ struct ArrowSchema* out_schema = NULLPTR);
+
+using DeviceMemoryMapper =
+ std::function<Result<std::shared_ptr<MemoryManager>>(ArrowDeviceType, int64_t)>;
+
+// ARROW_EXPORT
+// Result<std::shared_ptr<MemoryManager>> DefaultMemoryMapper(ArrowDeviceType device_type,
+// int64_t device_id);
Review Comment:
?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org