GH-43541: [C++] Check accepted device allocation types before executing kernel #43542

felipecrv · 2024-08-02T17:55:38Z

Rationale for this change

Kernels shouldn't segfault when executing against GPU-allocated (or any other non-CPU device) Arrays.

What changes are included in this PR?

A way of declaring which device allocation types a kernel can handle per parameter
Removal of some of the checks from the Python code
Making drop_null callable when array is CUDA-allocated

Are these changes tested?

Tested from the Python tests.

Are there any user-facing changes?

New member functions to some classes.

GitHub Issue: [C++] Compute functions should fail gracefully when given non-CPU resident Arrays #43541

github-actions · 2024-08-02T17:56:04Z

⚠️ GitHub issue #43541 has been automatically assigned in GitHub to PR creator.

cpp/src/arrow/chunked_array.h

felipecrv · 2024-08-02T17:58:39Z

cpp/src/arrow/compute/kernel.cc

+    case Datum::RECORD_BATCH:
+    case Datum::TABLE:


I can probably handle these even though they are not passed here.

felipecrv · 2024-08-02T18:06:55Z

cpp/src/arrow/compute/kernel.cc

@@ -430,6 +431,25 @@ bool InputType::Matches(const Datum& value) const {
  return Matches(*value.type());
 }

+bool InputType::MatchesDeviceAllocationType(const Datum& value) const {
+  DCHECK(Matches(value));


I will probably remove this DCHECK and keep the ones in KernelSignature::MatchesDeviceAllocationTypes.

felipecrv · 2024-08-27T01:58:34Z

I know what is wrong with drop_null (failing in tests), but I'm still deciding the best way to fix it.

jorisvandenbossche · 2024-08-27T12:59:55Z

python/pyarrow/array.pxi

@@ -1646,7 +1630,6 @@ cdef class Array(_PandasConvertible):
        -------
        lst : list
        """
-        self._assert_cpu()


I would expect this one to not change (not using compute kernels under the hood)?

Handled by __iter__ already. Isn't that what's called in the for x in self sugar?

bkietz

In general this looks good but I'm not convinced about the changes to InputType; I don't know of any kernels which take non-CPU data so it seems more reasonable to leave the InputType refactor for a follow up. On the other hand, Users may have implemented their own kernels which do operate on non CPU data in which case this change to InputType is breaking (since those kernels were not defined with the new constructor and default to gracefully rejecting non-CPU data). I'll have to think on this some more

cpp/src/arrow/device_allocation_type.h

bkietz · 2024-08-27T15:32:13Z

cpp/src/arrow/device_allocation_type_set.cc

+  for (int i = 1; i <= kDeviceAllocationTypeMax; i++) {
+    if (device_type_bitset_.test(i)) {
+      // Skip all the unused values in the enum.
+      switch (i) {
+        case 0:
+        case 5:
+        case 6:
+          continue;
+      }


Suggested change

for (int i = 1; i <= kDeviceAllocationTypeMax; i++) {

if (device_type_bitset_.test(i)) {

// Skip all the unused values in the enum.

switch (i) {

case 0:

case 5:

case 6:

continue;

}

for (int i : {

DeviceAllocationType::kCPU,

DeviceAllocationType::kCUDA,

DeviceAllocationType::kCUDA_HOST,

DeviceAllocationType::kOPENCL,

DeviceAllocationType::kVULKAN,

DeviceAllocationType::kMETAL,

DeviceAllocationType::kVPI,

DeviceAllocationType::kROCM,

DeviceAllocationType::kROCM_HOST,

DeviceAllocationType::kEXT_DEV,

DeviceAllocationType::kCUDA_MANAGED,

DeviceAllocationType::kONEAPI,

DeviceAllocationType::kWEBGPU,

DeviceAllocationType::kHEXAGON,

}) {

This won't work when another enum entry is added. Mine will work and I put the kDeviceAllocationTypeMax right after the enum so people remember to not leave gaps.

cpp/src/arrow/chunked_array.h

felipecrv · 2024-08-27T17:21:30Z

In general this looks good but I'm not convinced about the changes to InputType; I don't know of any kernels which take non-CPU data so it seems more reasonable to leave the InputType refactor for a follow up.

@bkietz I can leave it to a follow up, but that would force me to put ad-hoc device type checks in the kernel dispatching code instead of doing it in a systematic way.

On the other hand, Users may have implemented their own kernels which do operate on non CPU data in which case this change to InputType is breaking (since those kernels were not defined with the new constructor and default to gracefully rejecting non-CPU data). I'll have to think on this some more

These users were already on thin ice: all it takes is a GetNullCount() call in kernel dispatching logic to crash. The bright side is that I already give them a chance to avoid the check by passing the device types in the type matcher. This is why I didn't go for definitive "is cpu" checks and instead allowed extension of the set.

felipecrv · 2024-08-27T18:52:45Z

@bkietz I extracted the basics from this PR into a new PR -> #43853

danepitkin

PyArrow bits LGTM!

danepitkin · 2024-08-27T22:21:38Z

cpp/src/arrow/chunked_array.h

+  DeviceAllocationTypeSet device_types() const;
+
+  /// \return true if all chunks are allocated on CPU-accessible memory.
+  bool is_cpu() const { return device_types().is_cpu_only(); }


Is it worth cacheing this after the first execution?

I might cache the entire set on the constructor. This check is very cheap, so I wouldn't cache in the C++ layer. The set is represented as a 64-bit word under the hood. It's very small.

felipecrv · 2024-08-27T22:45:32Z

@danepitkin after @bkietz I extracted some of the commits from here into a smaller PR that only adds the device_types() methods, doesn't constrain kernel execution, and doesn't touch the python code: #43853

Because it's already checked by to_numpy(). `self.null_count` property is also guarded.

Already checked by __iter__().

With a slice instead of just an index.

…/fill_null/index/sort_indices fill_null() seems to use the coalesce function since that's the first one to complain about the device allocation type.

github-actions bot added the Component: C++ label Aug 2, 2024

felipecrv commented Aug 2, 2024

View reviewed changes

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

github-actions bot added awaiting committer review Awaiting committer review awaiting changes Awaiting changes and removed awaiting committer review Awaiting committer review labels Aug 2, 2024

felipecrv commented Aug 2, 2024

View reviewed changes

felipecrv force-pushed the check_device_type branch from 8e28341 to 552046e Compare August 2, 2024 19:08

github-actions bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Aug 2, 2024

jorisvandenbossche mentioned this pull request Aug 26, 2024

GH-43728: [Python] ChunkedArray fails gracefully on non-cpu devices #43795

Merged

felipecrv force-pushed the check_device_type branch from 552046e to 493976c Compare August 27, 2024 00:31

github-actions bot added the Component: Python label Aug 27, 2024

felipecrv marked this pull request as ready for review August 27, 2024 00:32

github-actions bot added awaiting changes Awaiting changes and removed awaiting change review Awaiting change review labels Aug 27, 2024

felipecrv requested a review from danepitkin August 27, 2024 00:35

jorisvandenbossche reviewed Aug 27, 2024

View reviewed changes

felipecrv requested a review from bkietz August 27, 2024 13:54

bkietz requested changes Aug 27, 2024

View reviewed changes

felipecrv force-pushed the check_device_type branch from 493976c to 5137f8d Compare August 27, 2024 16:39

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 Aug 27, 2024

felipecrv force-pushed the check_device_type branch from 50cab9d to 2b1bd82 Compare August 27, 2024 20:52

github-actions bot added awaiting change review Awaiting change review Component: Python and removed Component: Python awaiting changes Awaiting changes labels Aug 27, 2024

felipecrv force-pushed the check_device_type branch from d86ef88 to 778ead6 Compare August 27, 2024 21:19

apache deleted a comment from github-actions bot Aug 27, 2024

danepitkin reviewed Aug 27, 2024

View reviewed changes

github-actions bot added awaiting changes Awaiting changes and removed awaiting change review Awaiting change review labels Aug 27, 2024

felipecrv added 11 commits August 28, 2024 15:20

cpp: Constrain InputTypes by a set of accepted allocation device types

948195e

cpp: Expose AdjustNonNullable on internal to use it in data.h

4abb198

cpp: Make DropNull[Chunked]Array CUDA-safe

ba87d48

array.pxi: Remove _assert_cpu() from __array__

3ddf706

Because it's already checked by to_numpy(). `self.null_count` property is also guarded.

array.pxi: Remove _assert_cpu() from to_pylist()

039aaa1

Already checked by __iter__().

test_array.py: Add another test to __getitem__

699860d

With a slice instead of just an index.

array.pxi: Let C++ check device type for 4 functions

d646f6e

array.pxi: Let C++ check device type for take/filter

eb1a0d8

array.pxi: Let C++ check device type for cast/is_null/is_nan/is_valid…

da36961

…/fill_null/index/sort_indices fill_null() seems to use the coalesce function since that's the first one to complain about the device allocation type.

array.pxi: Let C++ check device type for drop_null

a1c4a44

test_array.py: Add another invocation for drop_null

610750b

felipecrv force-pushed the check_device_type branch from 778ead6 to 610750b Compare August 28, 2024 18:20

github-actions bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Aug 28, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

GH-43541: [C++] Check accepted device allocation types before executing kernel #43542

GH-43541: [C++] Check accepted device allocation types before executing kernel #43542

felipecrv commented Aug 2, 2024 •

edited

Loading

github-actions bot commented Aug 2, 2024

felipecrv Aug 2, 2024 •

edited

Loading

felipecrv Aug 2, 2024

felipecrv commented Aug 27, 2024

jorisvandenbossche Aug 27, 2024

felipecrv Aug 27, 2024

bkietz left a comment

bkietz Aug 27, 2024

felipecrv Aug 27, 2024

felipecrv commented Aug 27, 2024 •

edited

Loading

felipecrv commented Aug 27, 2024

danepitkin left a comment

danepitkin Aug 27, 2024

felipecrv Aug 27, 2024

felipecrv commented Aug 27, 2024

-  for (int i = 1; i <= kDeviceAllocationTypeMax; i++) {
-    if (device_type_bitset_.test(i)) {
-      // Skip all the unused values in the enum.
-      switch (i) {
-        case 0:
-        case 5:
-        case 6:
-          continue;
-      }
+  for (int i : {
+    DeviceAllocationType::kCPU,
+    DeviceAllocationType::kCUDA,
+    DeviceAllocationType::kCUDA_HOST,
+    DeviceAllocationType::kOPENCL,
+    DeviceAllocationType::kVULKAN,
+    DeviceAllocationType::kMETAL,
+    DeviceAllocationType::kVPI,
+    DeviceAllocationType::kROCM,
+    DeviceAllocationType::kROCM_HOST,
+    DeviceAllocationType::kEXT_DEV,
+    DeviceAllocationType::kCUDA_MANAGED,
+    DeviceAllocationType::kONEAPI,
+    DeviceAllocationType::kWEBGPU,
+    DeviceAllocationType::kHEXAGON,
+  }) {

GH-43541: [C++] Check accepted device allocation types before executing kernel #43542

Are you sure you want to change the base?

GH-43541: [C++] Check accepted device allocation types before executing kernel #43542

Conversation

felipecrv commented Aug 2, 2024 • edited Loading

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

github-actions bot commented Aug 2, 2024

felipecrv Aug 2, 2024 • edited Loading

Choose a reason for hiding this comment

felipecrv Aug 2, 2024

Choose a reason for hiding this comment

felipecrv commented Aug 27, 2024

jorisvandenbossche Aug 27, 2024

Choose a reason for hiding this comment

felipecrv Aug 27, 2024

Choose a reason for hiding this comment

bkietz left a comment

Choose a reason for hiding this comment

bkietz Aug 27, 2024

Choose a reason for hiding this comment

felipecrv Aug 27, 2024

Choose a reason for hiding this comment

felipecrv commented Aug 27, 2024 • edited Loading

felipecrv commented Aug 27, 2024

danepitkin left a comment

Choose a reason for hiding this comment

danepitkin Aug 27, 2024

Choose a reason for hiding this comment

felipecrv Aug 27, 2024

Choose a reason for hiding this comment

felipecrv commented Aug 27, 2024

felipecrv commented Aug 2, 2024 •

edited

Loading

felipecrv Aug 2, 2024 •

edited

Loading

felipecrv commented Aug 27, 2024 •

edited

Loading