Synchronize through barrier

3 files changed, 78 insertions, 20 deletions

diff --git a/zluda/src/impl/context.rs b/zluda/src/impl/context.rs
index 9ea0874..8d7a465 100644
--- a/zluda/src/impl/context.rs
+++ b/zluda/src/impl/context.rs
@@ -102,9 +102,10 @@ impl ContextData {
         l0_dev: l0::Device,
         flags: c_uint,
         is_primary: bool,
+        host_event: (l0::Event<'static>, u64),
         dev: *mut device::Device,
     ) -> Result<Self, CUresult> {
-        let default_stream = StreamData::new_unitialized(l0_ctx, l0_dev)?;
+        let default_stream = StreamData::new_unitialized(l0_ctx, l0_dev, host_event)?;
         Ok(ContextData {
             flags: AtomicU32::new(flags),
             device: dev,
@@ -136,10 +137,11 @@ pub fn create_v2(
     let mut ctx_box = GlobalState::lock_device(dev_idx, |dev| {
         let dev_ptr = dev as *mut _;
         let mut ctx_box = Box::new(LiveCheck::new(ContextData::new(
-            &mut dev.l0_context,
+            &dev.l0_context,
             dev.base,
             flags,
             false,
+            dev.host_event_pool.get(dev.base, &dev.l0_context)?,
             dev_ptr as *mut _,
         )?));
         ctx_box.late_init();
diff --git a/zluda/src/impl/device.rs b/zluda/src/impl/device.rs
index 0594252..e686f27 100644
--- a/zluda/src/impl/device.rs
+++ b/zluda/src/impl/device.rs
@@ -21,7 +21,8 @@ pub struct Device {
     pub default_queue: l0::CommandQueue<'static>,
     pub l0_context: l0::Context,
     pub primary_context: context::Context,
-    pub event_pool: DynamicEventPool,
+    pub device_event_pool: DynamicEventPool,
+    pub host_event_pool: DynamicEventPool,
     properties: Option<Box<l0::sys::ze_device_properties_t>>,
     image_properties: Option<Box<l0::sys::ze_device_image_properties_t>>,
     memory_properties: Option<Vec<l0::sys::ze_device_memory_properties_t>>,
@@ -36,21 +37,36 @@ impl Device {
     unsafe fn new(drv: &l0::Driver, l0_dev: l0::Device, idx: usize) -> Result<Self, CUresult> {
         let ctx = l0::Context::new(*drv, Some(&[l0_dev]))?;
         let queue = l0::CommandQueue::new(mem::transmute(&ctx), l0_dev)?;
+        let mut host_event_pool = DynamicEventPool::new(
+            l0_dev,
+            transmute_lifetime(&ctx),
+            l0::sys::ze_event_pool_flags_t::ZE_EVENT_POOL_FLAG_HOST_VISIBLE,
+            l0::sys::ze_event_scope_flags_t::ZE_EVENT_SCOPE_FLAG_HOST,
+        )?;
+        let host_event =
+            transmute_lifetime_mut(&mut host_event_pool).get(l0_dev, transmute_lifetime(&ctx))?;
         let primary_context = context::Context::new(context::ContextData::new(
-            mem::transmute(&ctx),
+            transmute_lifetime(&ctx),
             l0_dev,
             0,
             true,
+            host_event,
             ptr::null_mut(),
         )?);
-        let event_pool = DynamicEventPool::new(l0_dev, transmute_lifetime(&ctx))?;
+        let device_event_pool = DynamicEventPool::new(
+            l0_dev,
+            transmute_lifetime(&ctx),
+            l0::sys::ze_event_pool_flags_t(0),
+            l0::sys::ze_event_scope_flags_t(0),
+        )?;
         Ok(Self {
             index: Index(idx as c_int),
             base: l0_dev,
             default_queue: queue,
             l0_context: ctx,
             primary_context: primary_context,
-            event_pool,
+            device_event_pool,
+            host_event_pool,
             properties: None,
             image_properties: None,
             memory_properties: None,
@@ -400,14 +416,23 @@ pub(crate) fn primary_ctx_release_v2(_dev_idx: Index) -> CUresult {
 
 pub struct DynamicEventPool {
     count: usize,
+    pool_flags: l0::sys::ze_event_pool_flags_t,
+    signal_flags: l0::sys::ze_event_scope_flags_t,
     events: Vec<DynamicEventPoolEntry>,
 }
 
 impl DynamicEventPool {
-    fn new(dev: l0::Device, ctx: &'static l0::Context) -> l0::Result<Self> {
+    fn new(
+        dev: l0::Device,
+        ctx: &'static l0::Context,
+        pool_flags: l0::sys::ze_event_pool_flags_t,
+        signal_flags: l0::sys::ze_event_scope_flags_t,
+    ) -> l0::Result<Self> {
         Ok(DynamicEventPool {
             count: 0,
-            events: vec![DynamicEventPoolEntry::new(dev, ctx)?],
+            pool_flags,
+            signal_flags,
+            events: vec![DynamicEventPoolEntry::new(dev, ctx, pool_flags)?],
         })
     }
 
@@ -420,14 +445,17 @@ impl DynamicEventPool {
         let events = unsafe { transmute_lifetime_mut(&mut self.events) };
         let (global_idx, (ev, local_idx)) = {
             for (idx, entry) in self.events.iter_mut().enumerate() {
-                if let Some((ev, local_idx)) = entry.get()? {
+                if let Some((ev, local_idx)) = entry.get(self.signal_flags)? {
                     let marker = (idx << 32) as u64 | local_idx as u64;
                     return Ok((ev, marker));
                 }
             }
-            events.push(DynamicEventPoolEntry::new(dev, ctx)?);
+            events.push(DynamicEventPoolEntry::new(dev, ctx, self.pool_flags)?);
             let global_idx = (events.len() - 1) as u64;
-            (global_idx, events.last_mut().unwrap().get()?.unwrap())
+            (
+                global_idx,
+                events.last_mut().unwrap().get(self.signal_flags)?.unwrap(),
+            )
         };
         let marker = (global_idx << 32) | local_idx as u64;
         Ok((ev, marker))
@@ -452,10 +480,15 @@ struct DynamicEventPoolEntry {
 }
 
 impl DynamicEventPoolEntry {
-    fn new(dev: l0::Device, ctx: &'static l0::Context) -> l0::Result<Self> {
+    fn new(
+        dev: l0::Device,
+        ctx: &'static l0::Context,
+        flags: l0::sys::ze_event_pool_flags_t,
+    ) -> l0::Result<Self> {
         Ok(DynamicEventPoolEntry {
             event_pool: l0::EventPool::new(
                 ctx,
+                flags,
                 DYNAMIC_EVENT_POOL_ENTRY_SIZE as u32,
                 Some(&[dev]),
             )?,
@@ -463,7 +496,10 @@ impl DynamicEventPoolEntry {
         })
     }
 
-    fn get(&'static mut self) -> l0::Result<Option<(l0::Event<'static>, u32)>> {
+    fn get(
+        &'static mut self,
+        signal: l0::sys::ze_event_scope_flags_t,
+    ) -> l0::Result<Option<(l0::Event<'static>, u32)>> {
         for (idx, value) in self.bit_map.iter_mut().enumerate() {
             let shift = first_index_of_zero_u64(*value);
             if shift == 64 {
@@ -471,7 +507,12 @@ impl DynamicEventPoolEntry {
             }
             *value = *value | (1u64 << shift);
             let entry_index = (idx as u32 * 64u32) + shift;
-            let event = l0::Event::new(&self.event_pool, entry_index)?;
+            let event = l0::Event::new(
+                &self.event_pool,
+                entry_index,
+                signal,
+                l0::sys::ze_event_scope_flags_t(0),
+            )?;
             return Ok(Some((event, entry_index)));
         }
         Ok(None)
diff --git a/zluda/src/impl/stream.rs b/zluda/src/impl/stream.rs
index 59f8778..1de422b 100644
--- a/zluda/src/impl/stream.rs
+++ b/zluda/src/impl/stream.rs
@@ -38,28 +38,35 @@ pub struct StreamData {
     pub busy_events: VecDeque<(l0::Event<'static>, u64)>,
     // This could be a Vec, but I'd rather reuse earliest enqueued event not the one recently enqueued
     pub free_events: VecDeque<(l0::Event<'static>, u64)>,
+    pub synchronization_event: (l0::Event<'static>, u64),
 }
 
 impl StreamData {
     pub fn new_unitialized(
         ctx: &'static l0::Context,
         device: l0::Device,
+        host_event: (l0::Event<'static>, u64),
     ) -> Result<Self, CUresult> {
         Ok(StreamData {
             context: ptr::null_mut(),
             cmd_list: l0::CommandList::new_immediate(ctx, device)?,
             busy_events: VecDeque::new(),
             free_events: VecDeque::new(),
+            synchronization_event: host_event,
         })
     }
     pub fn new(ctx: &mut ContextData) -> Result<Self, CUresult> {
         let l0_ctx = &mut unsafe { &mut *ctx.device }.l0_context;
         let device = unsafe { &*ctx.device }.base;
+        let synchronization_event = unsafe { &mut *ctx.device }
+            .host_event_pool
+            .get(device, l0_ctx)?;
         Ok(StreamData {
             context: ctx as *mut _,
             cmd_list: l0::CommandList::new_immediate(l0_ctx, device)?,
             busy_events: VecDeque::new(),
             free_events: VecDeque::new(),
+            synchronization_event,
         })
     }
 
@@ -98,9 +105,17 @@ impl StreamData {
     }
 
     pub fn synchronize(&mut self) -> l0::Result<()> {
-        if let Some((ev, _)) = self.busy_events.back() {
-            ev.host_synchronize(u64::MAX)?;
-        }
+        let empty = [];
+        let busy_event_arr = self.busy_events.back().map(|(ev, _)| [ev]);
+        let wait_events = busy_event_arr.as_ref().map_or(&empty[..], |arr| &arr[..]);
+        unsafe {
+            self.cmd_list
+                .append_barrier(Some(&self.synchronization_event.0), wait_events)?
+        };
+        self.synchronization_event
+            .0
+            .host_synchronize(u64::max_value())?;
+        self.synchronization_event.0.host_reset()?;
         self.reuse_all_finished_events()?;
         Ok(())
     }
@@ -114,7 +129,7 @@ impl StreamData {
             .pop_front()
             .map(|x| Ok(x))
             .unwrap_or_else(|| {
-                let event_pool = unsafe { &mut (*(*self.context).device).event_pool };
+                let event_pool = unsafe { &mut (*(*self.context).device).device_event_pool };
                 event_pool.get(l0_dev, l0_ctx)
             })
     }
@@ -126,8 +141,8 @@ impl Drop for StreamData {
             return;
         }
         for (_, marker) in self.busy_events.iter().chain(self.free_events.iter()) {
-            let event_pool = unsafe { &mut (*(*self.context).device).event_pool };
-            event_pool.mark_as_free(*marker);
+            let device_event_pool = unsafe { &mut (*(*self.context).device).device_event_pool };
+            device_event_pool.mark_as_free(*marker);
         }
         unsafe { (&mut *self.context).streams.remove(&(&mut *self as *mut _)) };
     }