feat(extensions/nanoarrow_device): Draft DeviceArray interface

[paleolimbot]

After:

• https://github.com/zeroshade/arrow-non-cpu/tree/main
• https://lists.apache.org/thread/o2hsw7o1gm3qgw5z51rmz6zqxh0p7bvk
• GH-34971: [Format] Add non-CPU version of C Data Interface arrow#34972

Still in very much draft form; however, it does implement arbitrary ArrowArray copy to/from ARROW_DEVICE_METAL, ARROW_DEVICE_CUDA, ARROW_DEVICE_CUDA_HOST, and ARROW_DEVICE_CPU.

The nanoarrow_device extension as drafted here serves a similar purpose to nanoarrow: a means by which to create and consume the C ABI with the intention of shipping those structures to other libraries to do transformations, and potentially retrieving them again after the computation is complete. Perhaps another way to put it is that nanoarrow is designed to help at the edges: it can create and consume. Similarly, the nanoarrow_device extension is designed to help at the edges: it can copy/move arrays to and from CPU-land.

With this PR, you can currently do something like:

struct ArrowDevice* gpu = ArrowDeviceMetalDefaultDevice();
// Alternatively, ArrowDeviceCuda(ARROW_DEVICE_CUDA, 0)
// or  ArrowDeviceCuda(ARROW_DEVICE_CUDA_HOST, 0)
struct ArrowDevice* cpu = ArrowDeviceCpu();
struct ArrowArray array;
struct ArrowDeviceArray device_array;
struct ArrowDeviceArrayView device_array_view;

// Build a CPU array
ASSERT_EQ(ArrowArrayInitFromType(&array, NANOARROW_TYPE_STRING), NANOARROW_OK);
ASSERT_EQ(ArrowArrayStartAppending(&array), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(&array, ArrowCharView("abc")), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(&array, ArrowCharView("defg")), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendNull(&array, 1), NANOARROW_OK);
ASSERT_EQ(ArrowArrayFinishBuildingDefault(&array, nullptr), NANOARROW_OK);

// Convert to a DeviceArray, still on the CPU
ArrowDeviceArrayInit(&device_array, cpu);
ArrowArrayMove(&array, &device_array.array);

// Parse contents into a view that can be copied to another device
ArrowDeviceArrayViewInit(&device_array_view);
ArrowArrayViewInitFromType(&device_array_view.array_view, string_type);
ASSERT_EQ(ArrowDeviceArrayViewSetArray(&device_array_view, &device_array, nullptr),
          NANOARROW_OK);

// Try to zero-copy move to another device or copy if that is not possible. Zero-copy move
// is implemented for ARROW_DEVICE_METAL and ARROW_DEVICE_CUDA_HOST for the
// gpu -> cpu case.
struct ArrowDeviceArray device_array2;
device_array2.array.release = nullptr;
ASSERT_EQ(
    ArrowDeviceArrayTryMove(&device_array, &device_array_view, gpu, &device_array2),
    NANOARROW_OK);

In concrete terms, that means we to know enough about a device to (1) copy and/or move an arbitrary ArrowArray/ArrowSchema pair to a device from the CPU and (2) copy/move an arbitrary ArrowDeviceArray/ArrowSchema pair back to the CPU. The three types of copying I support (and maybe there could be fewer/need to be more) are:

• ArrowDeviceBufferInit(): Make a non-owning buffer into an owning buffer on a device. The entry point if you want to take a slice of an ArrowArrayView and ship it to a device.
• ArrowDeviceBufferMove(): Move an existing (owning) buffer to a device. For devices like the CPU, this is a true zero-copy move; for shared memory this can also sometimes be zero copy (e.g., Apple Metal -> CPU) but might also involve a copy.
• ArrowDeviceBufferCopy(): Copy a section of a buffer into a preallocated section of another buffer. I'm envisioning this to be necessary when copying a String, Binary, List...we need the first and last values of the offsets buffer in order to know what portion of the data buffer to copy. It seems unnecessary to copy 4 bytes of a buffer into an owning variant covered by the first bullet but 🤷 .

This PR currently provides support for the CPU device, Apple Metal, CUDA, and CUDA_HOST (i.e., CPU memory that has been registered with CUDA which CUDA copies under the hood).

[codecov-commenter]

Codecov Report

Merging #205 (3b5ee6c) into main (e811cfe) will decrease coverage by 0.48%.
The diff coverage is 71.26%.

@@            Coverage Diff             @@
##             main     #205      +/-   ##
==========================================
- Coverage   87.64%   87.17%   -0.48%     
==========================================
  Files          63       66       +3     
  Lines        9789    10061     +272     
==========================================
+ Hits         8580     8771     +191     
- Misses       1209     1290      +81     

Impacted Files	Coverage Δ
src/nanoarrow/array.c	91.42% <0.00%> (-0.60%)	⬇️
src/nanoarrow/array_inline.h	89.72% <ø> (ø)
src/nanoarrow/nanoarrow.h	100.00% <ø> (ø)
.../nanoarrow_device/src/nanoarrow/nanoarrow_device.c	66.35% <66.35%> (ø)
.../nanoarrow_device/src/nanoarrow/nanoarrow_device.h	100.00% <100.00%> (ø)
...anoarrow_device/src/nanoarrow/nanoarrow_device.hpp	100.00% <100.00%> (ø)

... and 2 files with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[kkraus14]

Most GPU libraries / frameworks have their own memory pool / memory management implementations that are often asynchronous (and are ordered by CUDA streams) where they won't be able to benefit from this implementation. This is generally true for most operations: free, alloc, realloc, memset, memcpy, etc.

I'm not sure if we need an actual implementation to live within nanoarrow or if we can just define an interface for downstream libraries to implement.

[paleolimbot]

An early version of ArrowDeviceBufferXXX() functions had a sync_event* argument, which I removed before I saw cudaMemcpyAsync and friends in the documentation. I don't know if that's the perfect interface, but the part of nanoarrow's generic "copy this array to the device" implementation would benefit a lot since essentially all of those buffers can be copied in parallel.

[kkraus14]

In most situations you'd want to use cudaStreamWaitEvent as opposed to this API as its much more efficient and doesn't unnecessarily block the CPU until things are done.

[paleolimbot]

I assume that is what a library doing gpu--gpu calculations would do...here I think it does need to be the blocking version (this is the method that is called before an arbitrary ArrowDeviceArray or a slice of it is copied back to the CPU).

[zeroshade]

To call cudaStreamWaitEvent you'd need to know what stream to wait on.

[paleolimbot]

Yes, this function is intended to do the bad thing and block until it's safe to do CPU things. If there's a way to avoid the sync before copying back to the CPU it could be removed.

[kkraus14]

We'd likely be better off using the CUDA driver API here instead of the runtime API as there's much stronger forward compatibility guarantees as well as easier deployment (someone can have the driver installed but not the CUDA runtime, but not the reverse).

[paleolimbot]

That does seem like a better fit (although may require implementing some reference counting of our own).

[kkraus14]

For what it's worth: this will likely be a blocker for most libraries / frameworks to be able to utilize things.

[paleolimbot]

I don't think there's a technical limitation for making ArrowDeviceArrayViewCopy() return after having kicked off all the buffer copies and populating the ArrowDeviceArray's sync_event...this is mostly a personal limitation (steep learning curve for me).

[kkraus14]

Should we handle the situations where the src is ARROW_DEVICE_CUDA_HOST and dst is ARROW_DEVICE_CUDA and vice versa?

[paleolimbot]

I did see cuMemHostGetDevicePointer(), so I assume this is possible. I think it would require that ArrowDeviceArrayViewCopy() has a device-specific implementation (probably for the best anyway).

[kkraus14]

cuMemHostGetDevicePointer() gets a pointer to that pinned host memory that can be used from device code but doesn't actually copy any memory to device memory. As far as I know it can be used anywhere that device memory can be used, but obviously has different performance characteristics where that would likely be very unexpected.

[zeroshade]

First pass

[zeroshade]

Should we specify semantics about the src and dst devices such as the src being CPU memory? etc.

[paleolimbot]

The idea was that an implementation might be able to handle a few directions, although it does result in verbose method implementations. For the CUDA case the generality is somewhat useful...it can also theoretically move a buffer from CUDA_HOST to CUDA and it might be difficult to construct a method signature that captures that. That generality might also not be useful 🤷

[zeroshade]

same question, should we put specifics as to the semantics of the device types? Should implementations have to check the device type every time for both source and destination or have to implement multiple devices?

[zeroshade]

Should there be an equivalent that doesn't do the copy?

[paleolimbot]

I'm not sure I understand where that would be used? It's definitely suboptimal to issue copies in this way (but your suggestion of skipping validation and Keith's suggestion of leveraging async memcpy may be a workaround).

[zeroshade]

It would probably be useful for device code to be able to access the value of a specific index of the array without having to perform the copy. It also allows CPU code to find the address of the specific index (pointer into non-cpu memory) that can then be used for whatever is necessary on the device side without needing to copy the value.

[kkraus14]

Definitely don't want to issue copies one by one like this even asynchronously, that would be really really bad performance wise and put significant pressure on the system via the GPU driver.

I would +1 @zeroshade's suggestion of skipping validation and generally anything that needs to introspect the data.

[paleolimbot]

Ok, I updated the function name to ArrowDeviceArrayViewResolveBufferSizes() since that's what it's actually doing and added ArrowDeviceArrayViewSetArrayMinimal() that just sets array_view->buffers[i].size_bytes to -1 if it would require a copy to calculate.

For the case of "just get me the pointer value", I don't think there needs to a be a function (array_view->buffers[i].data.as_int32 + some_logical_offset would do it).

For the case where we copy back to the CPU, I don't see a way around copying the last int32/int64 from the offsets buffer (or else there is no way to know how many bytes of the next buffer to copy). We can possibly mitigate the impact of that by asynchronously kicking off all the tiny copies at once?

[kkraus14]

Yea, issuing two single element copies for getting the starting and ending offset into the CPU is necessary, but once we have those offsets, we should just do pointer arithmetic to get a pointer and a size to feed into a copy call.

Ideally you want to issue the two individual element copies asynchronously (potentially on different streams so they can be overlapped even though they're tiny), synchronize the stream(s) since you need those values to use in host code, and then issue the actual data copy.

[kkraus14]

If we're going to continue using the runtime, may want to use the static library instead: CUDA::cudart_static.

Would still recommend using the driver library though.

[paleolimbot]

Driver library is a definite yes (just haven't gotten there yet).

[kkraus14]

Definitely don't want to issue copies one by one like this even asynchronously, that would be really really bad performance wise and put significant pressure on the system via the GPU driver.

I would +1 @zeroshade's suggestion of skipping validation and generally anything that needs to introspect the data.

[kkraus14]

+1 to the int32 function here

[paleolimbot]

I think this is handled by array_view->buffer_views[i].data.as_int32 + some_index (which would get you the pointer to an element of a buffer).

[kkraus14]

Why do we need to synchronize on the event here? This function ultimately is just responsible for setting the points in the array view from the passed in array, correct?

If so, synchronize guarantees that the data underneath the pointer is synchronized, but doesn't impact the pointers themselves at all.

[paleolimbot]

I think I correctly separated the case where it's needed (it is needed before copy to CPU, correct? Or is that synchronizatio handled by cudaMemcpy()?).

[kkraus14]

all variants of cudaMemcpy are stream ordered so you can safely use them without synchronizing the stream, but the destination data is stream ordered as well, so if you're going to operate on it from a different stream or from host code then you need to synchronize the stream in some kind of way

[paleolimbot]

It seems like we do need to do the bad thing here and wait for a sync before calling cudaMemcpy() here for the GPU -> CPU direction (although hopefully this is now isolated such that it won't get accidentally called by somebody who does not need this).

[kkraus14]

We should sync after calling the cudaMemcpy() as opposed to before it. Otherwise, you could in theory get into a situation where that device to host copy is asynchronous (if you have pinned host memory for example) and accessing it from the CPU without synchronization is a race condition.

[kkraus14]

If this buffer initialization is asynchronous then we need to set a synchronization event somewhere I think?

[paleolimbot]

Ok! There are definitely some holes in this implementation (notably around properly synchronizing memory copies). I'd propose that this PR get merged (and clearly marked as under development/experimental in the README) with some related changes grouped into some follow-ups:

• Switch to using the driver library: Device extension should use the CUDA driver library (not the CUDA runtime library) #246
• Properly order copy operations when copying to/from device: Device extension should implement async buffer copying for CUDA #245
• Expand test coverage to more than just string arrays: Expand test coverage of device extension #247

I think it's still an open question as to whether or not this particular extension will be used/will be useful...if there is no interest in using it before the next release it can always be excluded from the source release (like the Python bindings currently are) or moved back to a PR state.