[C++] Add device-specific synchronization API to Buffer

[pitrou]

Describe the enhancement requested

For the C device data interface, and other applications, we'll need to add optional synchronization information to buffers.

Here is for example a possible API. It tries to avoid or minimize additional footprint, especially for the CPU case:

// in device.h
class DeviceSyncEvent {
 public:
  virtual ~DeviceSyncEvent() = default;
  /// \brief Block until synchronization event is ready
  virtual Status Wait() = 0;
};

// in buffer.h
class Buffer {
 public:
  ...
  /// \brief Device-specific synchronization event
  ///
  /// If nullptr is returned (which is always the case for CPU buffers),
  /// no synchronization is required.
  virtual DeviceSyncEvent* sync_event() { return nullptr; }
  // XXX or perhaps:
  virtual std::shared_ptr<DeviceSyncEvent> sync_event() { return nullptr; }
  ...
};

// in arrow/gpu
class CudaSyncEvent : public DeviceSyncEvent {
 public:
  void* cuda_event();  // actually a CUevent
  Status Wait() override;
};

class CudaBuffer : public Buffer {
 public:
  DeviceSyncEvent* sync_event() override { return &sync_event_; }
  ...
 protected:
  CudaSyncEvent sync_event_;
  ...
};

Or, as an alternative, define a DeviceSyncStream instead of a DeviceSyncEvent.

TODO: define lifetime semantics.

Component(s)

C++

[pitrou]

@kkraus14 @zeroshade

[zeroshade]

@pitrou Is the idea that the Wait() is actually having the CPU wait for synchronization? Or that it puts a wait on the GPU stream?

[pitrou]

Wait is a CPU wait call for convenience. You don't have to use it, you can instead use the CUDA event directly. Perhaps it's not a very useful method, in which case it can be removed.

[pitrou]

Also cc @felipecrv , if you know a bit about GPU APIs

[zeroshade]

We should probably put a void* raw_event() in the base interface, while the derived impls would have properly typed methods. i.e.: CudaSyncEvent::cuda_event would return a cudaEvent_t* and CudaSyncEvent::event would return a void*. etc. But I guess nitpicks like this can be discussed on the eventual PR for this rather than hashing it out here. I think we can agree that this is definitely needed.

[pitrou]

AFAICT, we don't expose cuda.h in our own headers. We might want to change this if really necessary, otherwise I think it's better to keep this policy.

[pitrou]

In any case, the right lifetime semantics will have to be decided, which may entail putting a shared_ptr somewhere in the API.

[kkraus14]

I'm not sure if an event is the right thing for Buffers. In addition to being able to synchronize, device buffers are often managed in a stream ordered fashion. I.E. allocated and more importantly freed asynchronously from a CPU perspective with the ordering being handled via CUDA streams.

I.E. RMM (libcudf's memory manager) has a couple of stream APIs associated to its device_buffer class as well as a private stream member which gets used in its destructor for deallocating. https://github.com/rapidsai/rmm/blob/0c08dd585031f58e2a9dcfbba5608cce10c423b2/include/rmm/device_buffer.hpp#L374-L400

[pitrou]

Ok, so instead of exposing an event-like API, we could expose a stream-like API.