[SYCL][UR][OpenCL] allow passing events from different contexts

to urEventWait

Author: igchor

sycl/test-e2e/Basic/multi_context_wait.cpp

@@ -0,0 +1,88 @@
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+#include <iostream>
+#include <sycl/sycl.hpp>
+#include <vector>
+
+std::vector<sycl::event> submit_dependencies(sycl::queue q1, sycl::queue q2,
+                                             int *mem1, int *mem2) {
+  int delay_ops = 1024 * 1024;
+  auto delay = [=] {
+    volatile int value = delay_ops;
+    while (--value)
+      ;
+  };
+
+  auto ev1 =
+      q1.parallel_for(sycl::range<1>(1024), [=]([[maybe_unused]] auto u) {
+        delay();
+        mem1[u.get_id()] = 1;
+      });
+  auto ev2 =
+      q2.parallel_for(sycl::range<1>(1024), [=]([[maybe_unused]] auto u) {
+        delay();
+        mem2[u.get_id()] = 2;
+      });
+
+  return {ev1, ev2};
+}
+
+void test_host_task() {
+  sycl::context c1{};
+  sycl::context c2{};
+
+  sycl::queue q1(c1, sycl::default_selector_v);
+  sycl::queue q2(c2, sycl::default_selector_v);
+
+  auto mem1 = sycl::malloc_host<int>(1024, q1);
+  auto mem2 = sycl::malloc_host<int>(1024, q2);
+
+  auto events = submit_dependencies(q1, q2, mem1, mem2);
+
+  q2.submit([&](sycl::handler &cgh) {
+    cgh.depends_on(events[0]);
+    cgh.depends_on(events[1]);
+    cgh.host_task([=]() {
+      for (int i = 0; i < 1024; i++) {
+        assert(mem1[i] == 1);
+        assert(mem2[i] == 2);
+      }
+    });
+  });
+
+  sycl::free(mem1, c1);
+  sycl::free(mem2, c2);
+}
+
+void test_kernel() {
+  sycl::context c1{};
+  sycl::context c2{};
+
+  sycl::queue q1(c1, sycl::default_selector_v);
+  sycl::queue q2(c2, sycl::default_selector_v);
+
+  auto mem1 = sycl::malloc_device<int>(1024, q1);
+  auto mem2 = sycl::malloc_device<int>(1024, q2);
+
+  auto events = submit_dependencies(q1, q2, mem1, mem2);
+
+  q2.submit([&](sycl::handler &cgh) {
+    cgh.depends_on(events[0]);
+    cgh.depends_on(events[1]);
+    cgh.parallel_for(sycl::range<1>(1024), [=](auto item) {
+      assert(mem1[item.get_id()] == 1);
+      assert(mem2[item.get_id()] == 2);
+    });
+  });
+
+  sycl::free(mem1, c1);
+  sycl::free(mem2, c2);
+}
+
+int main() {
+  test_host_task();
+  test_kernel();
+
+  return 0;
+}

unified-runtime/source/adapters/opencl/event.cpp

@@ -149,11 +149,29 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventRetain(ur_event_handle_t hEvent) {
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urEventWait(uint32_t numEvents, const ur_event_handle_t *phEventWaitList) {
-  std::vector<cl_event> CLEvents(numEvents);
+  if (numEvents == 0 || !phEventWaitList) {
+    return UR_RESULT_SUCCESS;
+  }
+
+  ur_context_handle_t hContext = phEventWaitList[0]->Context;
+  std::vector<cl_event> CLEvents;
+  CLEvents.reserve(numEvents);
+
+  // clWaitForEvents can only be called on events from the same context.
+  // If the events are from different contexts, we need to wait for each
+  // set of events separately.
   for (uint32_t i = 0; i < numEvents; i++) {
-    CLEvents[i] = phEventWaitList[i]->CLEvent;
+    if (phEventWaitList[i]->Context != hContext) {
+      cl_int RetErr = clWaitForEvents(CLEvents.size(), CLEvents.data());
+      CL_RETURN_ON_FAILURE(RetErr);
+
+      CLEvents.clear();
+    }
+
+    CLEvents.push_back(phEventWaitList[i]->CLEvent);
+    hContext = phEventWaitList[i]->Context;
   }
-  cl_int RetErr = clWaitForEvents(numEvents, CLEvents.data());
+  cl_int RetErr = clWaitForEvents(CLEvents.size(), CLEvents.data());
   CL_RETURN_ON_FAILURE(RetErr);
   return UR_RESULT_SUCCESS;
 }