kkraus14 commented on code in PR #37040:
URL: https://github.com/apache/arrow/pull/37040#discussion_r1286477832
##########
cpp/src/arrow/buffer.h:
##########
@@ -346,6 +346,8 @@ class ARROW_EXPORT Buffer {
static Result<std::shared_ptr<Buffer>> ViewOrCopy(
std::shared_ptr<Buffer> source, const std::shared_ptr<MemoryManager>&
to);
+ virtual std::shared_ptr<DeviceSync> get_device_sync() { return nullptr; }
Review Comment:
This name is somewhat ambiguous. Is this a sync event? Maybe
`get_device_sync_event`?
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -714,11 +713,11 @@ Result<std::pair<std::optional<DeviceAllocationType>,
int64_t>> ValidateDeviceIn
return std::make_pair(device_type, device_id);
}
-Status ExportDeviceArray(const Array& array, RawSyncEvent sync_event,
+Status ExportDeviceArray(const Array& array, std::shared_ptr<DeviceSync>& sync,
Review Comment:
nitpick: since we're copying the `shared_ptr` at line 741 anyway, I think
it's technically more efficient to take `sync` by value here and then it can be
moved at line 741.
##########
cpp/src/arrow/c/bridge_test.cc:
##########
@@ -3552,6 +3580,190 @@ TEST_F(TestArrayRoundtrip, RecordBatch) {
}
}
+class TestDeviceArrayRoundtrip : public ::testing::Test {
+ public:
+ using ArrayFactory = std::function<Result<std::shared_ptr<Array>>()>;
+
+ void SetUp() override { pool_ = default_memory_pool(); }
+
+ static Result<std::shared_ptr<MemoryManager>> DeviceMapper(ArrowDeviceType
type,
+ int64_t id) {
+ if (type != kMyDeviceType) {
+ return Status::NotImplemented("should only be MyDevice");
+ }
+
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(id);
+ return device->default_memory_manager();
+ }
+
+ static Result<std::shared_ptr<ArrayData>> ToDeviceData(
+ const std::shared_ptr<MemoryManager>& mm, const ArrayData& data) {
+ arrow::BufferVector buffers;
+ for (const auto& buf : data.buffers) {
+ if (buf) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, mm->CopyBuffer(buf, mm));
+ buffers.push_back(dest);
+ } else {
+ buffers.push_back(nullptr);
+ }
+ }
+
+ arrow::ArrayDataVector children;
+ for (const auto& child : data.child_data) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, ToDeviceData(mm, *child));
+ children.push_back(dest);
+ }
+
+ return ArrayData::Make(data.type, data.length, buffers, children,
data.null_count,
+ data.offset);
+ }
+
+ static Result<std::shared_ptr<Array>> ToDevice(const
std::shared_ptr<MemoryManager>& mm,
+ const ArrayData& data) {
+ ARROW_ASSIGN_OR_RAISE(auto result, ToDeviceData(mm, data));
+ return MakeArray(result);
+ }
+
+ static ArrayFactory ToDeviceFactory(const std::shared_ptr<MemoryManager>& mm,
+ ArrayFactory&& factory) {
+ return [&]() -> Result<std::shared_ptr<Array>> {
+ ARROW_ASSIGN_OR_RAISE(auto arr, factory());
+ return ToDevice(mm, *arr->data());
+ };
+ }
+
+ static ArrayFactory JSONArrayFactory(const std::shared_ptr<MemoryManager>&
mm,
+ std::shared_ptr<DataType> type, const
char* json) {
Review Comment:
Instead of a `const char*` here can we take a `string_view`? Not sure how we
know what length the json is currently (unless it's assumed to be null
terminated?)
##########
cpp/src/arrow/gpu/cuda_context.cc:
##########
@@ -293,6 +293,12 @@ std::shared_ptr<CudaDevice>
CudaMemoryManager::cuda_device() const {
return checked_pointer_cast<CudaDevice>(device_);
}
+Result<std::shared_ptr<DeviceSync>> CudaMemoryManager::MakeDeviceSync(void*
sync_event) {
+ return nullptr;
+ // auto ev = reinterpret_cast<CUstream*>(sync_event);
+ // return std::make_shared<CudaDeviceSync>(ev);
Review Comment:
Is this TODO?
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -739,19 +738,20 @@ Status ExportDeviceArray(const Array& array, RawSyncEvent
sync_event,
exporter.Finish(&out->array);
auto* pdata =
reinterpret_cast<ExportedArrayPrivateData*>(out->array.private_data);
- pdata->sync_event_ = sync_event;
- out->sync_event = sync_event.sync_event;
+ pdata->sync_ = sync;
+ out->sync_event = sync_event;
guard.Detach();
return Status::OK();
}
-Status ExportDeviceRecordBatch(const RecordBatch& batch, RawSyncEvent
sync_event,
+Status ExportDeviceRecordBatch(const RecordBatch& batch,
+ std::shared_ptr<DeviceSync>& sync,
Review Comment:
nitpick: Same here
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
