Skip to content

Commit

Permalink
feat(python): Clarify interaction between the CDeviceArray, the CArra…
Browse files Browse the repository at this point in the history
…yView, and the CArray (#409)

When device support was first added, the `CArrayView` was device-aware
but the `CArray` was not. This worked well until it was clear that
`__arrow_c_array__` needed to error if it did not represent a CPU array
(and the `CArray` had no way to check). Now, the `CArray` has a
`device_type` and `device_id`. A nice side-effect of this is that we get
back the `view()` method (whose removal @jorisvandenbossche had
lamented!).

This also implements the device array protocol to help test
apache/arrow#40717 . This protocol isn't
finalized yet and I could remove that part until it is (although it
doesn't seem likely to change).

The non-cpu case is still hard to test without real-world CUDA
support...this PR is just trying to get the right information in the
right place as early as possible.

```python
import nanoarrow as na

array = na.c_array([1, 2, 3], na.int32())
array.device_type, array.device_id
#> (1, 0)
```

---------

Co-authored-by: Dane Pitkin <48041712+danepitkin@users.noreply.github.com>
  • Loading branch information
paleolimbot and danepitkin authored Apr 9, 2024
1 parent 00aa9c3 commit bab66ac
Show file tree
Hide file tree
Showing 11 changed files with 281 additions and 68 deletions.
4 changes: 2 additions & 2 deletions extensions/nanoarrow_device/src/nanoarrow/nanoarrow_device.c
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ struct ArrowDevice* ArrowDeviceCpu(void) {

void ArrowDeviceInitCpu(struct ArrowDevice* device) {
device->device_type = ARROW_DEVICE_CPU;
device->device_id = 0;
device->device_id = -1;
device->array_init = NULL;
device->array_move = NULL;
device->buffer_init = &ArrowDeviceCpuBufferInit;
Expand All @@ -135,7 +135,7 @@ struct ArrowDevice* ArrowDeviceCuda(ArrowDeviceType device_type, int64_t device_
#endif

struct ArrowDevice* ArrowDeviceResolve(ArrowDeviceType device_type, int64_t device_id) {
if (device_type == ARROW_DEVICE_CPU && device_id == 0) {
if (device_type == ARROW_DEVICE_CPU) {
return ArrowDeviceCpu();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ TEST(NanoarrowDevice, CheckRuntime) {
TEST(NanoarrowDevice, CpuDevice) {
struct ArrowDevice* cpu = ArrowDeviceCpu();
EXPECT_EQ(cpu->device_type, ARROW_DEVICE_CPU);
EXPECT_EQ(cpu->device_id, 0);
EXPECT_EQ(cpu->device_id, -1);
EXPECT_EQ(cpu, ArrowDeviceCpu());

void* sync_event = nullptr;
Expand Down
Loading

0 comments on commit bab66ac

Please sign in to comment.