bkietz commented on code in PR #37040:
URL: https://github.com/apache/arrow/pull/37040#discussion_r1296024262
##########
cpp/src/arrow/device.h:
##########
@@ -98,6 +101,71 @@ class ARROW_EXPORT Device : public
std::enable_shared_from_this<Device>,
/// \brief Return the DeviceAllocationType of this device
virtual DeviceAllocationType device_type() const = 0;
+ /// \brief EXPERIMENTAL: An object that provides event/stream sync primitives
+ class ARROW_EXPORT SyncEvent {
+ public:
+ virtual ~SyncEvent() = default;
+
+ /// @brief Block until sync event is completed.
+ virtual Status wait() = 0;
+
+ /// @brief Make the provided stream wait on the sync event.
+ ///
+ /// Tells the provided stream that it should wait until the
+ /// synchronization event is completed without blocking the CPU.
+ /// @param stream Should be appropriate for the underlying device
+ virtual Status stream_wait(void* stream) = 0;
+
+ void set_stream(void* stream) {
+ stream_ = stream;
+ }
+
+ /// @brief Returns the stored raw event or creates a new one to return.
+ ///
+ /// clear_event should always be called to cleanup afterwards. If this
+ /// creates the event, then clear_event will call release_event
+ /// internally. If this doesn't own the event, then the lifetime should
+ /// be controlled externally to this class.
+ Result<void*> get_event() {
+ if (!sync_event_) {
+ ARROW_ASSIGN_OR_RAISE(sync_event_, create_event());
+ owns_event_ = true;
+ }
Review Comment:
Using virtual methods seems to be painting us into a corner where we can't
plug API holes. Maybe device sync should use a different mode of polymorphism,
for example type erasure:
```c++
class SyncEvent final {
public:
// SyncEvent is not a base class and we don't define a derived class for
// each `Device`. Instead, SyncEvent is (exclusively) opaquely constructed
// by a memory manager. Its device affiliation (if desired)
// can be queried through its memory manager, which is an explicit
// data member
const std::shared_ptr<Device>& device() const { return
memory_manager_->device(); }
const std::shared_ptr<MemoryManager>& memory_manager() const { return
memory_manager_; }
// Since this pushes responsibility for construction and destruction out of
// SyncEvent's virtual methods (and into MemoryManager::MakeDeviceSync),
// the constructor can just assume it will receive the fully constructed
event
// and the destructor may be implicit and safe.
SyncEvent(void* stream,
std::shared_ptr<void> event,
std::function<Status(void*, void*)> stream_wait,
std::function<Status(void*, void*)> record_event_on_stream,
std::shared_ptr<MemoryManager> memory_manager);
// public wrappers for function pointers...
Status record_event() { return record_event_on_stream_(stream_,
event_.get()); }
Status wait() { return stream_wait_(stream_, event_.get()); }
private:
void* stream_;
std::shared_ptr<void> event_;
std::function<Status(void*, void*)> stream_wait_;
std::function<Status(void*, void*)> record_event_on_stream_;
std::shared_ptr<MemoryManager> memory_manager_;
};
```
With this general design, an implementation of MakeDeviceSync might look
like this:
```c++
Result<std::shared_ptr<Device::SyncEvent>>
CudaMemoryManager::MakeDeviceSync(void* stream) {
std::shared_ptr<CUevent> event{
new CUevent,
[](CUevent* event) { cuEventDestroy(*event); }
};
CU_RETURN_NOT_OK("cuEventCreate", cuEventCreate(event.get(),
CU_EVENT_DEFAULT));
auto stream_wait = [](void* stream, void* event) {
CU_RETURN_NOT_OK("cuEventSynchronize", cuEventSynchronize(
*static_cast<CUevent*>(event.get())));
return Status::OK();
};
auto record_event_on_stream = [](void* stream, void* event) {
CU_RETURN_NOT_OK("cuEventRecord", cuEventRecord(
*static_cast<CUevent*>(event.get()),
*static_cast<CUstream*>(stream)));
return Status::OK();
};
return std::make_shared<Device::SyncEvent>(
stream,
std::move(event),
stream_wait,
record_event_on_stream,
/*memory_manager=*/shared_from_this()
);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
