Replace Waker with AtomicWaker

Author: szymon-zadworny

Cargo.lock

@@ -20,7 +20,7 @@ enum class EventCommandStatus : std::uint8_t;
 
 namespace sycl_shims::event {
 void wait(std::unique_ptr<Event> &);
-void register_callback(std::unique_ptr<Queue> &, Event const &, rust::Box<Waker>);
+void register_callback(std::unique_ptr<Queue> &, Event const &, SharedWaker const *);
 EventCommandStatus get_command_execution_status(Event const &);
 std::unique_ptr<Event> clone(Event const &);
 } // namespace sycl_shims::event

oneapi-rs-sys/include/event.hpp

@@ -6,7 +6,7 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 //
 
-use crate::types::Waker;
+use crate::types::SharedWaker;
 
 #[cxx::bridge(namespace = "sycl_shims::event")]
 pub mod ffi {
@@ -26,17 +26,13 @@ pub mod ffi {
         type Queue = crate::types::ffi::Queue;
 
         fn wait(event: &mut UniquePtr<Event>);
-        fn register_callback(queue: &mut UniquePtr<Queue>, event: &Event, waker: Box<Waker>);
+        unsafe fn register_callback(queue: &mut UniquePtr<Queue>, event: &Event, waker: *const SharedWaker);
         fn get_command_execution_status(event: &Event) -> EventCommandStatus;
         fn clone(event: &Event) -> UniquePtr<Event>;
     }
 
     extern "Rust" {
-        type Waker;
-        fn wake(waker: &Box<Waker>);
+        type SharedWaker;
+        fn wake(&self);
     }
 }
-
-fn wake(waker: &Box<Waker>) {
-    waker.0.wake_by_ref();
-}

oneapi-rs-sys/src/event-sys.rs

@@ -32,10 +32,10 @@ EventCommandStatus get_command_execution_status(Event const & event) {
       return EventCommandStatus::Unknown;
   }
 }
-void register_callback(std::unique_ptr<Queue> & queue, Event const & event, rust::Box<Waker> waker) {
-  queue->submit([waker = std::move(waker), event](sycl::handler& cgh) {
+void register_callback(std::unique_ptr<Queue> & queue, Event const & event, SharedWaker const * waker) {
+  queue->submit([=](sycl::handler& cgh) {
     cgh.depends_on(event);
-    cgh.host_task([&]() { wake(waker); });
+    cgh.host_task([=]() { waker->wake(); });
   });
 }
 } // namespace sycl_shims::event

oneapi-rs-sys/src/event.cpp

@@ -6,11 +6,26 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 //
 
-// cxx requires Rust opaque types to be defined in the current crate.
-pub struct Waker(pub std::task::Waker);
-impl From<std::task::Waker> for Waker {
-    fn from(value: std::task::Waker) -> Self {
-        Waker(value)
+use std::sync::atomic::{Ordering::Relaxed, AtomicBool};
+
+use futures::task::AtomicWaker;
+
+pub struct SharedWaker {
+    pub waker: AtomicWaker,
+    pub done: AtomicBool
+}
+
+impl SharedWaker {
+    pub fn new() -> Self {
+        Self {
+            waker: AtomicWaker::new(),
+            done: AtomicBool::new(false)
+        }
+    }
+
+    pub fn wake(&self) {
+        self.done.store(true, Relaxed);
+        self.waker.wake();
     }
 }
 

oneapi-rs-sys/src/types-sys.rs

@@ -6,13 +6,13 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 //
 
-use std::{pin::Pin, task::{Context, Poll}};
+use std::{pin::Pin, task::{Context, Poll}, sync::atomic::Ordering::Relaxed};
 
-use oneapi_rs_sys::event::ffi;
+use oneapi_rs_sys::{event::ffi, types::SharedWaker};
 
 use pin_project::pin_project;
 
-use crate::{info::{EventCommandStatus, event::{CommandExecutionStatus, EventInfo}}, queue::Queue};
+use crate::{info::event::EventInfo, queue::Queue};
 
 pub struct Event(pub(crate) cxx::UniquePtr<ffi::Event>);
 
@@ -41,6 +41,7 @@ impl Clone for Event {
 #[pin_project]
 pub struct EventFuture {
     event: Event,
+    shared: SharedWaker,
     set_callback: bool,
     queue: Queue,
 }
@@ -49,17 +50,41 @@ impl Future for EventFuture {
     type Output = ();
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        if self.event.get_info::<CommandExecutionStatus>() == EventCommandStatus::Complete {
-            Poll::Ready(())
+        let this = self.project();
+
+        // Set the callback on first Future poll (Futures can't be active until polled)
+        if *this.set_callback == false {
+            *this.set_callback = true;
+            this.shared.waker.register(cx.waker());
+            unsafe { ffi::register_callback(&mut this.queue.0, &this.event.0, this.shared) };
+
+            // Check the event again to avoid a race condition
+            // https://docs.rs/futures/latest/futures/task/struct.AtomicWaker.html#examples
+            if this.shared.done.load(Relaxed) {
+                Poll::Ready(())
+            }
+            else {
+                Poll::Pending
+            }
         }
         else {
-            if self.set_callback == false {
-                let this = self.project();
-                *this.set_callback = true;
-                let waker = Box::new(cx.waker().clone().into());
-                ffi::register_callback(&mut this.queue.0, &this.event.0, waker);
+            // Quick check before registering to avoid wasting time
+            if this.shared.done.load(Relaxed) {
+                return Poll::Ready(());
             }
-            Poll::Pending
+
+            // Register the waker if result isn't ready. This is a slow atomic operation
+            this.shared.waker.register(cx.waker());
+
+            // Check the event again to avoid a race condition
+            // https://docs.rs/futures/latest/futures/task/struct.AtomicWaker.html#examples
+            if this.shared.done.load(Relaxed) {
+                Poll::Ready(())
+            }
+            else {
+                Poll::Pending
+            }
+
         }
     }
 }
@@ -70,9 +95,10 @@ impl IntoFuture for Event {
 
     fn into_future(self) -> Self::IntoFuture {
         EventFuture {
-            queue: Queue::new_immediate(),
             event: self,
-            set_callback: false
+            shared: SharedWaker::new(),
+            set_callback: false,
+            queue: Queue::new_immediate(),
         }
     }
 }