Move callback queue lifetime handling to Rust

Author: szymon-zadworny

oneapi-rs-sys/include/event.hpp

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

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

@@ -22,8 +22,11 @@ pub mod ffi {
         #[namespace = "sycl_shims"]
         type Event = crate::types::ffi::Event;
 
+        #[namespace = "sycl_shims"]
+        type Queue = crate::types::ffi::Queue;
+
         fn wait(event: &mut UniquePtr<Event>);
-        fn register_callback(event: &Event, waker: Box<Waker>);
+        fn register_callback(queue: &mut UniquePtr<Queue>, event: &Event, waker: Box<Waker>);
         fn get_command_execution_status(event: &Event) -> EventCommandStatus;
         fn clone(event: &Event) -> UniquePtr<Event>;
     }

oneapi-rs-sys/src/event.cpp

@@ -10,6 +10,8 @@
 
 using sycl::info::event_command_status;
 
+namespace syclintel = sycl::ext::intel;
+
 namespace sycl_shims::event {
 void wait(std::unique_ptr<Event> & event) {
   event->wait();
@@ -30,10 +32,10 @@ EventCommandStatus get_command_execution_status(Event const & event) {
       return EventCommandStatus::Unknown;
   }
 }
-void register_callback(Event const & event, rust::Box<Waker> waker) {
-  sycl::queue().submit([waker = std::move(waker), event](sycl::handler& cgh) {
+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) {
     cgh.depends_on(event);
-    cgh.host_task([&]() mutable { wake(waker); });
+    cgh.host_task([&]() { wake(waker); });
   });
 }
 } // namespace sycl_shims::event

oneapi-rs/src/event.rs

@@ -12,7 +12,7 @@ use oneapi_rs_sys::event::ffi;
 
 use pin_project::pin_project;
 
-use crate::info::{EventCommandStatus, event::{CommandExecutionStatus, EventInfo}};
+use crate::{info::{EventCommandStatus, event::{CommandExecutionStatus, EventInfo}}, queue::Queue};
 
 pub struct Event(pub(crate) cxx::UniquePtr<ffi::Event>);
 
@@ -41,7 +41,8 @@ impl Clone for Event {
 #[pin_project]
 pub struct EventFuture {
     event: Event,
-    set_callback: bool
+    set_callback: bool,
+    queue: Queue,
 }
 
 impl Future for EventFuture {
@@ -56,7 +57,7 @@ impl Future for EventFuture {
                 let this = self.project();
                 *this.set_callback = true;
                 let waker = Box::new(cx.waker().clone().into());
-                ffi::register_callback(&this.event.0, waker);
+                ffi::register_callback(&mut this.queue.0, &this.event.0, waker);
             }
             Poll::Pending
         }
@@ -69,6 +70,7 @@ impl IntoFuture for Event {
 
     fn into_future(self) -> Self::IntoFuture {
         EventFuture {
+            queue: Queue::new_immediate(),
             event: self,
             set_callback: false
         }