GH-36488: [C++] Import/Export ArrowDeviceArray #36489

zeroshade · 2023-07-05T17:22:50Z

Rationale for this change

With the addition of the ArrowDeviceArray we should provide import/export functions just like we have for ArrowArray.

What changes are included in this PR?

Adding Import/Export functions to the bridge.h/.cc files for C Data Non-CPU data.

This requires adding a device type to buffers and memory managers to propagate through.

Are these changes tested?

Yes.

(Still need to add tests to bridge_test.cc but I wanted to get eyes on this first)

Are there any user-facing changes?

No.

Closes: [C++] Import/Export ArrowDeviceArray #36488

zeroshade · 2023-07-05T17:22:58Z

CC @kkraus14

felipecrv · 2023-07-05T17:41:07Z

cpp/src/arrow/buffer.h

+    // but you can specify CUDA_HOST as the device type to override it.
+    if (device_type != DeviceType::UNKNOWN) {
+      device_type_ = device_type;
+    }


Would a DCHECK be useful here to confirm that device_type_ ended-up being compatible with the memory manager?

Possibly, we'd have to add a Validate or Valid method to the memory manager to do so, to allow the memory manager to specify if the device_type is valid for it. Not sure how worthwhile that is though.

felipecrv · 2023-07-05T17:44:08Z

cpp/src/arrow/buffer.h

@@ -57,18 +57,25 @@ 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_(DeviceType::CPU) {


Do we still need the is_cpu_ boolean? Is there a situation in which is_cpu_ is false and device_type_ != CPU?

The is_cpu_ boolean is still useful because of pinned host memory. Both CUDA_HOST and ROCM_HOST are pinned CPU memory which would have the corresponding device type (since the memory is managed by the device drivers) but is accessible to the CPU.

While the pointers are accessible to the CPU, there's no guarantee that there isn't stream ordered work happening on them where it could cause a race condition to CPU code accessing the buffer. Given is_cpu_ is existing ABI, I'm wondering if it wouldn't be a safer option to have it return False in these cases where if someone wants to do CPU things against CUDA_HOST or ROCM_HOST buffers they should explicitly check the device type instead?

Hmm, I think the intent for the is_cpu_ flag was solely to know whether or not the data is accessible via CPU rather than caring about the stream ordered work synchronizations. I intend on expanding the existing classes/objects to make it easier to synchronize for stream ordered work as a separate abstraction, so I lean towards leaving it marked as true in this case, but I'm open to changing it.

@pitrou @felipeblazing @Rhonda85 what do you all think?

I had not envisioned this issue, but I think your approach is correct: is_cpu() == true means that address() is a CPU-accessible pointer to the data. Whether or not the "expected" data is already there will depend on the producer.

And, yes, a generalized version of Buffer could have an optional sync event or something...

cpp/src/arrow/c/bridge.cc

cpp/src/arrow/c/bridge.h

cpp/src/arrow/gpu/cuda_context.cc

felipecrv

LGTM, but I will let @kkraus14 do the final review.

benibus · 2023-07-11T19:20:34Z

cpp/src/arrow/c/bridge.cc

+Status ValidateDeviceInfo(const ArrayData& data, DeviceType* device_type,
+                          int64_t* device_id) {


Not a big deal ATM, but a DeviceInfo struct might be good have (not just here, but in general), e.g:

struct DeviceInfo { explicit DeviceInfo(DeviceType type = DeviceType::UNKNOWN, int64_t id = -1) : type(type), id(id) {} DeviceType type; int64_t id; };

Or at least return Result<std::pair<DeviceType, int64_t>> here...

That brings me back to my original implementation vs this tail recursion version that @felipecrv suggested haha.

This is the internal recursive call, below it is the intended function to call which does return Result<std::pair<DeviceType, int64_t>>

If you don't like passing a Result recursively, you can always build a struct with a method in it...

the Result isn't recursive, the recursion just passes the pointers to the device type and id

cpp/src/arrow/c/bridge.cc

cpp/src/arrow/device.h

cpp/src/arrow/c/bridge.h

westonpace

That's quite a comprehensive set of tests. I think it's probably more exhaustive than strictly needed but that's probably ok. I'll let others weigh in too but this looks good as far as I can tell.

It might also be good to exercise the sync event but I think, if you have plans to exercise it in practice in a follow-up PR, that is probably ok too.

cpp/src/arrow/c/bridge_test.cc

zeroshade · 2023-07-24T15:14:45Z

Does anyone have any further comments here? If not I'll potentially merge this at the end of the day

zeroshade · 2023-07-24T21:37:04Z

Looks like there's a memory leak identified by ASAN, will merge after I fix that problem.

pitrou · 2023-07-25T14:38:54Z

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_;


Why is this optional, are there any cases where the information is not available?

pitrou · 2023-07-25T14:42:24Z

cpp/src/arrow/c/bridge.cc

+
+    if (buf->device_type() != *device_type) {
+      return Status::Invalid(
+          "Exporting device array with buffers on more than one device.");


"with more than one device allocation type" perhaps?

pitrou · 2023-07-25T14:42:44Z

cpp/src/arrow/c/bridge.cc

+  for (const auto& child : data.child_data) {
+    RETURN_NOT_OK(ValidateDeviceInfo(*child, device_type, device_id));
+  }
+


Should also validate data.dictionary, if any.

pitrou · 2023-07-25T14:43:19Z

cpp/src/arrow/c/bridge.cc

+//////////////////////////////////////////////////////////////////////////
+// C device arrays
+
+Status ValidateDeviceInfo(const ArrayData& data,


Can you put helper functions inside the anonymous namespace?

pitrou · 2023-07-25T14:44:38Z

cpp/src/arrow/c/bridge.cc

+
+Status ExportDeviceArray(const Array& array, RawSyncEvent sync_event,
+                         struct ArrowDeviceArray* out, struct ArrowSchema* out_schema) {
+  if (sync_event.sync_event != nullptr && sync_event.release_func) {


Hmm, did you test this? It seems you're missing a logical negation here, i.e. !sync_event.release_func.

pitrou · 2023-07-25T15:05:24Z

cpp/src/arrow/gpu/cuda_memory.h

@@ -258,5 +260,9 @@ Result<uintptr_t> GetDeviceAddress(const uint8_t* cpu_data,
 ARROW_EXPORT
 Result<uint8_t*> GetHostAddress(uintptr_t device_ptr);

+ARROW_EXPORT
+Result<std::shared_ptr<MemoryManager>> DefaultMemoryMapper(ArrowDeviceType device_type,


It seems weird to have this in arrow/gpu, why not under arrow/c instead?

pitrou · 2023-07-25T15:08:33Z

cpp/src/arrow/c/bridge_test.cc

+
+class MyDevice : public Device {
+ public:
+  explicit MyDevice(int value) : Device(true), value_(value) {}


Would be nice to be more explicit: Device(/*xxx=*/ true)

pitrou · 2023-07-25T15:12:00Z

cpp/src/arrow/c/bridge_test.cc

+    for (const auto& child : data.child_data) {
+      ARROW_ASSIGN_OR_RAISE(auto dest, ToDeviceData(mm, *child));
+      children.push_back(dest);
+    }


Should take care of the dictionary array as well.

pitrou · 2023-07-25T15:16:15Z

cpp/src/arrow/c/bridge_test.cc

+    checker(&c_array, *batch, kMyDeviceType, 1, nullptr);
+  }
+}
+


It would be nice to have tests for importing device arrays.
Also tests that the sync event and its release func are taken into account.

pitrou · 2023-07-25T15:16:57Z

cpp/src/arrow/c/bridge_test.cc

+  TestPrimitive(mm, decimal(16, 4), R"(["1234.5670", null])");
+  TestPrimitive(mm, decimal256(16, 4), R"(["1234.5670", null])");
+
+  TestPrimitive(mm, month_day_nano_interval(), R"([[-1, 5, 20], null])");


TBH, I'm not sure it's useful to copy/paste all those tests, besides the readability burden of duplicate code.

We just want to check that primitive, nested, extension and dictionary are properly handled in general (detailed export behaviour is already tested previously).

conbench-apache-arrow · 2023-08-01T00:56:27Z

After merging your PR, Conbench analyzed the 6 benchmarking runs that have been run so far on merge-commit 969f4b4.

There were no benchmark performance regressions. 🎉

The full Conbench report has more details. It also includes information about possible false positives for unstable benchmarks that are known to sometimes produce them.

### Rationale for this change With the addition of the `ArrowDeviceArray` we should provide import/export functions just like we have for `ArrowArray`. ### What changes are included in this PR? Adding Import/Export functions to the `bridge.h/.cc` files for C Data Non-CPU data. This requires adding a device type to buffers and memory managers to propagate through. ### Are these changes tested? Yes. (Still need to add tests to `bridge_test.cc` but I wanted to get eyes on this first) ### Are there any user-facing changes? No. * Closes: apache#36488 Lead-authored-by: Matt Topol <zotthewizard@gmail.com> Co-authored-by: Weston Pace <weston.pace@gmail.com> Signed-off-by: Matt Topol <zotthewizard@gmail.com>

### Rationale for this change The original PRs for adding support for importing and exporting the new C Device interface (#36488 / #36489) only added support for the Arrays themselves, not for the stream structure. We should support both. ### What changes are included in this PR? Adding parallel functions for Import/Export of streams that accept `ArrowDeviceArrayStream`. ### Are these changes tested? Test writing in progress, wanted to get this up for review while I write tests. ### Are there any user-facing changes? No, only new functions have been added. * GitHub Issue: #40078 Lead-authored-by: Matt Topol <zotthewizard@gmail.com> Co-authored-by: Felipe Oliveira Carvalho <felipekde@gmail.com> Co-authored-by: Benjamin Kietzman <bengilgit@gmail.com> Co-authored-by: Antoine Pitrou <pitrou@free.fr> Signed-off-by: Matt Topol <zotthewizard@gmail.com>

) ### Rationale for this change The original PRs for adding support for importing and exporting the new C Device interface (apache#36488 / apache#36489) only added support for the Arrays themselves, not for the stream structure. We should support both. ### What changes are included in this PR? Adding parallel functions for Import/Export of streams that accept `ArrowDeviceArrayStream`. ### Are these changes tested? Test writing in progress, wanted to get this up for review while I write tests. ### Are there any user-facing changes? No, only new functions have been added. * GitHub Issue: apache#40078 Lead-authored-by: Matt Topol <zotthewizard@gmail.com> Co-authored-by: Felipe Oliveira Carvalho <felipekde@gmail.com> Co-authored-by: Benjamin Kietzman <bengilgit@gmail.com> Co-authored-by: Antoine Pitrou <pitrou@free.fr> Signed-off-by: Matt Topol <zotthewizard@gmail.com>

apacheGH-36488: [C++] Import/Export ArrowDeviceArray

ab98249

zeroshade requested review from felipecrv, lidavidm, pitrou and benibus July 5, 2023 17:22

github-actions bot added Component: C++ awaiting committer review Awaiting committer review labels Jul 5, 2023

linter

4d0ac93

felipecrv requested changes Jul 5, 2023

View reviewed changes

Linting

30256c1

github-actions bot added awaiting changes Awaiting changes and removed awaiting committer review Awaiting committer review labels Jul 6, 2023

changes from review comments

d591208

github-actions bot added awaiting change review Awaiting change review awaiting changes Awaiting changes and removed awaiting changes Awaiting changes awaiting change review Awaiting change review labels Jul 6, 2023

update from feedback

58c15b0

github-actions bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Jul 10, 2023

linting

2d3bcb6

kkraus14 reviewed Jul 11, 2023

View reviewed changes

cpp/src/arrow/gpu/cuda_context.cc Outdated Show resolved Hide resolved

github-actions bot added awaiting changes Awaiting changes and removed awaiting change review Awaiting change review labels Jul 11, 2023

felipecrv reviewed Jul 11, 2023

View reviewed changes

benibus reviewed Jul 11, 2023

View reviewed changes

kkraus14 approved these changes Jul 12, 2023

View reviewed changes

pitrou reviewed Jul 12, 2023

View reviewed changes

cpp/src/arrow/device.h Outdated Show resolved Hide resolved

pitrou reviewed Jul 12, 2023

View reviewed changes

cpp/src/arrow/c/bridge.h Outdated Show resolved Hide resolved

add tests

9244ea4

github-actions bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Jul 20, 2023

linting

f79d2a2

westonpace approved these changes Jul 20, 2023

View reviewed changes

cpp/src/arrow/c/bridge_test.cc Outdated Show resolved Hide resolved

github-actions bot added awaiting merge Awaiting merge awaiting changes Awaiting changes and removed awaiting change review Awaiting change review awaiting merge Awaiting merge labels Jul 20, 2023

update from feedback

5e1939f

github-actions bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Jul 21, 2023

zeroshade added 2 commits July 24, 2023 15:36

mark device_id override to fix compiler warning

1ca39b8

add schema export guard for memory leak

d90ce9e

zeroshade merged commit 969f4b4 into apache:main Jul 25, 2023

zeroshade removed the awaiting change review Awaiting change review label Jul 25, 2023

zeroshade deleted the import-export-devicearray branch July 25, 2023 14:34

pitrou reviewed Jul 25, 2023

View reviewed changes

github-actions bot added awaiting review Awaiting review awaiting committer review Awaiting committer review and removed awaiting review Awaiting review labels Jul 25, 2023

jorisvandenbossche mentioned this pull request Feb 14, 2024

[C++] Import/Export ArrowDeviceArrayStream #40078

Closed

zeroshade mentioned this pull request Mar 26, 2024

GH-40078: [C++] Import/Export ArrowDeviceArrayStream #40807

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

GH-36488: [C++] Import/Export ArrowDeviceArray #36489

GH-36488: [C++] Import/Export ArrowDeviceArray #36489

zeroshade commented Jul 5, 2023 •

edited by github-actions bot

Loading

zeroshade commented Jul 5, 2023

felipecrv Jul 5, 2023

zeroshade Jul 6, 2023

felipecrv Jul 5, 2023

zeroshade Jul 6, 2023

kkraus14 Jul 10, 2023

zeroshade Jul 17, 2023

pitrou Jul 17, 2023

felipecrv left a comment

benibus Jul 11, 2023

pitrou Jul 12, 2023

zeroshade Jul 12, 2023

pitrou Jul 12, 2023

zeroshade Jul 12, 2023

westonpace left a comment

zeroshade commented Jul 24, 2023

zeroshade commented Jul 24, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

pitrou Jul 25, 2023

conbench-apache-arrow bot commented Aug 1, 2023

		Status ValidateDeviceInfo(const ArrayData& data, DeviceType* device_type,
		int64_t* device_id) {

GH-36488: [C++] Import/Export ArrowDeviceArray #36489

GH-36488: [C++] Import/Export ArrowDeviceArray #36489

Conversation

zeroshade commented Jul 5, 2023 • edited by github-actions bot Loading

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

zeroshade commented Jul 5, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felipecrv left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

westonpace left a comment

Choose a reason for hiding this comment

zeroshade commented Jul 24, 2023

zeroshade commented Jul 24, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

conbench-apache-arrow bot commented Aug 1, 2023

zeroshade commented Jul 5, 2023 •

edited by github-actions bot

Loading