[SYCL][USM] Fix queue::wait() behavior after USM operations

sycl/include/CL/sycl/detail/queue_impl.hpp

@@ -352,13 +352,15 @@ class queue_impl {
     handler Handler(std::move(Self), MHostQueue);
     CGF(Handler);
     event Event = Handler.finalize();
-    {
-      std::lock_guard<mutex_class> Guard(MMutex);
-      MEvents.push_back(Event);
-    }
+    addEvent(Event);
     return Event;
   }
 
+  /// Stores an event that should be associated with the queue
+  ///
+  /// @param Event is the event to be stored
+  void addEvent(event Event);
+
   /// Protects all the fields that can be changed by class' methods.
   mutex_class MMutex;
 

sycl/source/detail/queue_impl.cpp

@@ -45,7 +45,9 @@ event queue_impl::memset(shared_ptr_class<detail::queue_impl> Impl, void *Ptr,
   if (Context.is_host())
     return event();
 
-  return event(pi::cast<cl_event>(Event), Context);
+  event ResEvent{pi::cast<cl_event>(Event), Context};
+  addEvent(ResEvent);
+  return ResEvent;
 }
 
 event queue_impl::memcpy(shared_ptr_class<detail::queue_impl> Impl, void *Dest,
@@ -57,7 +59,9 @@ event queue_impl::memcpy(shared_ptr_class<detail::queue_impl> Impl, void *Dest,
   if (Context.is_host())
     return event();
 
-  return event(pi::cast<cl_event>(Event), Context);
+  event ResEvent{pi::cast<cl_event>(Event), Context};
+  addEvent(ResEvent);
+  return ResEvent;
 }
 
 event queue_impl::mem_advise(const void *Ptr, size_t Length, int Advice) {
@@ -72,8 +76,16 @@ event queue_impl::mem_advise(const void *Ptr, size_t Length, int Advice) {
   Plugin.call<PiApiKind::piextUSMEnqueueMemAdvise>(getHandleRef(), Ptr, Length,
                                                    Advice, &Event);
 
-  return event(pi::cast<cl_event>(Event), Context);
+  event ResEvent{pi::cast<cl_event>(Event), Context};
+  addEvent(ResEvent);
+  return ResEvent;
 }
+
+void queue_impl::addEvent(event Event) {
+  std::lock_guard<mutex_class> Guard(MMutex);
+  MEvents.push_back(std::move(Event));
+}
+
 } // namespace detail
 } // namespace sycl
 } // __SYCL_INLINE_NAMESPACE(cl)

sycl/test/usm/queue_wait.cpp

@@ -0,0 +1,43 @@
+// RUN: %clangxx -fsycl %s -o %t.out
+// RUN: env SYCL_DEVICE_TYPE=HOST %t.out
+// RUN: %CPU_RUN_PLACEHOLDER %t.out
+// RUN: %GPU_RUN_PLACEHOLDER %t.out
+
+#include <CL/sycl.hpp>
+
+#include <cassert>
+#include <cstddef>
+
+using namespace cl::sycl;
+
+// This test checks that queue USM functions are properly waited for during
+// calls to queue::wait().
+
+int main() {
+  const std::size_t Size = 32;
+  queue Q;
+  std::cout << Q.is_host() << std::endl;
+  device Dev = Q.get_device();
+  context Ctx = Q.get_context();
+  if (!(Dev.get_info<info::device::usm_device_allocations>() &&
+        Dev.get_info<info::device::usm_host_allocations>()))
+    return 0;
+
+  unsigned char *DevArr = (unsigned char *)malloc_device(Size, Dev, Ctx);
+  assert(DevArr);
+  unsigned char *HostArr = (unsigned char *)malloc_host(Size, Ctx);
+  assert(HostArr);
+
+  Q.memset(DevArr, 42, Size);
+  Q.wait();
+  Q.memcpy(HostArr, DevArr, Size);
+  Q.wait();
+
+  for (std::size_t i = 0; i < Size; ++i)
+    assert(HostArr[i] == 42);
+
+  free(DevArr, Ctx);
+  free(HostArr, Ctx);
+
+  return 0;
+}