[SYCL] Fix waiting for events completion by commands on host-queue (#1737)

Signed-off-by: Sergey Kanaev <[email protected]>

Author: s-kanaev

sycl/source/detail/scheduler/commands.cpp

@@ -237,11 +237,44 @@ void Command::waitForEvents(QueueImplPtr Queue,
                             RT::PiEvent &Event) {
 
   if (!EventImpls.empty()) {
-    std::vector<RT::PiEvent> RawEvents = getPiEvents(EventImpls);
     if (Queue->is_host()) {
-      const detail::plugin &Plugin = EventImpls[0]->getPlugin();
-      Plugin.call<PiApiKind::piEventsWait>(RawEvents.size(), &RawEvents[0]);
+      // Host queue can wait for events from different contexts, i.e. it may
+      // contain events with different contexts in its MPreparedDepsEvents.
+      // OpenCL 2.1 spec says that clWaitForEvents will return
+      // CL_INVALID_CONTEXT if events specified in the list do not belong to
+      // the same context. Thus we split all the events into per-context map.
+      // An example. We have two queues for the same CPU device: Q1, Q2. Thus
+      // we will have two different contexts for the same CPU device: C1, C2.
+      // Also we have default host queue. This queue is accessible via
+      // Scheduler. Now, let's assume we have three different events: E1(C1),
+      // E2(C1), E3(C2). Also, we have an EmptyCommand which is to be executed
+      // on host queue. The command's MPreparedDepsEvents will contain all three
+      // events (E1, E2, E3). Now, if piEventsWait is called for all three
+      // events we'll experience failure with CL_INVALID_CONTEXT 'cause these
+      // events refer to different contexts.
+      std::map<context_impl *, std::vector<EventImplPtr>>
+          RequiredEventsPerContext;
+
+      for (const EventImplPtr &Event : EventImpls) {
+        ContextImplPtr Context = Event->getContextImpl();
+        assert(Context.get() &&
+               "Only non-host events are expected to be waited for here");
+        RequiredEventsPerContext[Context.get()].push_back(Event);
+      }
+
+      for (auto &CtxWithEvents : RequiredEventsPerContext) {
+        std::vector<RT::PiEvent> RawEvents = getPiEvents(CtxWithEvents.second);
+        CtxWithEvents.first->getPlugin().call<PiApiKind::piEventsWait>(
+            RawEvents.size(), RawEvents.data());
+      }
     } else {
+#ifndef NDEBUG
+      for (const EventImplPtr &Event : EventImpls)
+        assert(Event->getContextImpl().get() &&
+               "Only non-host events are expected to be waited for here");
+#endif
+
+      std::vector<RT::PiEvent> RawEvents = getPiEvents(EventImpls);
       const detail::plugin &Plugin = Queue->getPlugin();
       Plugin.call<PiApiKind::piEnqueueEventsWait>(
           Queue->getHandleRef(), RawEvents.size(), &RawEvents[0], &Event);

sycl/unittests/scheduler/CMakeLists.txt

@@ -4,5 +4,6 @@ add_sycl_unittest(SchedulerTests OBJECT
     FinishedCmdCleanup.cpp
     LeafLimit.cpp
     MemObjCommandCleanup.cpp
+    CommandsWaitForEvents.cpp
     utils.cpp
-)
+)

sycl/unittests/scheduler/CommandsWaitForEvents.cpp

@@ -0,0 +1,115 @@
+//==-------- CommandsWaitForEvents.cpp --- Scheduler unit tests ------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "SchedulerTest.hpp"
+#include "SchedulerTestUtils.hpp"
+#include <helpers/PiMock.hpp>
+
+using namespace cl::sycl;
+
+struct TestCtx {
+  queue &Q1;
+  queue &Q2;
+
+  std::shared_ptr<detail::context_impl> Ctx1;
+  std::shared_ptr<detail::context_impl> Ctx2;
+
+  pi_event EventCtx1 = reinterpret_cast<pi_event>(0x01);
+  pi_event EventCtx2 = reinterpret_cast<pi_event>(0x02);
+
+  bool EventCtx1WasWaited = false;
+  bool EventCtx2WasWaited = false;
+
+  TestCtx(queue &Queue1, queue &Queue2)
+      : Q1(Queue1), Q2(Queue2), Ctx1{detail::getSyclObjImpl(Q1.get_context())},
+        Ctx2{detail::getSyclObjImpl(Q2.get_context())} {}
+};
+
+std::unique_ptr<TestCtx> TestContext;
+
+pi_result waitFunc(pi_uint32 N, const pi_event *List) {
+  EXPECT_EQ(N, 1u) << "piEventsWait called for different contexts\n";
+
+  EXPECT_TRUE((TestContext->EventCtx1 == *List) ||
+              (TestContext->EventCtx2 == *List))
+      << "piEventsWait called for unknown event";
+
+  if (TestContext->EventCtx1 == *List)
+    TestContext->EventCtx1WasWaited = true;
+
+  if (TestContext->EventCtx2 == *List)
+    TestContext->EventCtx2WasWaited = true;
+
+  return PI_SUCCESS;
+}
+
+pi_result retainReleaseFunc(pi_event) { return PI_SUCCESS; }
+
+pi_result getEventInfoFunc(pi_event Event, pi_event_info PName, size_t PVSize,
+                           void *PV, size_t *PVSizeRet) {
+  EXPECT_EQ(PName, PI_EVENT_INFO_CONTEXT) << "Unknown param name";
+
+  if (Event == TestContext->EventCtx1)
+    *reinterpret_cast<pi_context *>(PV) =
+        reinterpret_cast<pi_context>(TestContext->Ctx1->get());
+  else if (Event == TestContext->EventCtx2)
+    *reinterpret_cast<pi_context *>(PV) =
+        reinterpret_cast<pi_context>(TestContext->Ctx2->get());
+
+  return PI_SUCCESS;
+}
+
+TEST_F(SchedulerTest, CommandsWaitForEvents) {
+  default_selector Selector{};
+  if (Selector.select_device().is_host()) {
+    std::cerr << "Not run due to host-only environment\n";
+    return;
+  }
+
+  queue Q1;
+  queue Q2;
+
+  unittest::PiMock Mock1(Q1);
+  unittest::PiMock Mock2(Q2);
+
+  Mock1.redefine<detail::PiApiKind::piEventsWait>(waitFunc);
+  Mock1.redefine<detail::PiApiKind::piEventRetain>(retainReleaseFunc);
+  Mock1.redefine<detail::PiApiKind::piEventRelease>(retainReleaseFunc);
+  Mock1.redefine<detail::PiApiKind::piEventGetInfo>(getEventInfoFunc);
+
+  Mock2.redefine<detail::PiApiKind::piEventsWait>(waitFunc);
+  Mock2.redefine<detail::PiApiKind::piEventRetain>(retainReleaseFunc);
+  Mock2.redefine<detail::PiApiKind::piEventRelease>(retainReleaseFunc);
+  Mock2.redefine<detail::PiApiKind::piEventGetInfo>(getEventInfoFunc);
+
+  TestContext.reset(new TestCtx(Q1, Q2));
+
+  std::shared_ptr<detail::event_impl> E1(
+      new detail::event_impl(TestContext->EventCtx1, Q1.get_context()));
+  std::shared_ptr<detail::event_impl> E2(
+      new detail::event_impl(TestContext->EventCtx2, Q2.get_context()));
+
+  sycl::device HostDevice;
+  std::shared_ptr<detail::queue_impl> DefaultHostQueue(new detail::queue_impl(
+      detail::getSyclObjImpl(HostDevice), /*AsyncHandler=*/{},
+      detail::QueueOrder::Ordered, /*PropList=*/{}));
+
+  MockCommand Cmd(DefaultHostQueue);
+
+  std::vector<std::shared_ptr<detail::event_impl>> Events;
+  Events.push_back(E1);
+  Events.push_back(E2);
+
+  pi_event EventResult = nullptr;
+
+  Cmd.waitForEventsCall(DefaultHostQueue, Events, EventResult);
+
+  ASSERT_TRUE(TestContext->EventCtx1WasWaited &&
+              TestContext->EventCtx2WasWaited)
+      << "Not all events were waited for";
+}

sycl/unittests/scheduler/SchedulerTestUtils.hpp

@@ -40,6 +40,13 @@ class MockCommand : public cl::sycl::detail::Command {
 
   cl_int MRetVal = CL_SUCCESS;
 
+  void waitForEventsCall(
+      std::shared_ptr<cl::sycl::detail::queue_impl> Queue,
+      std::vector<std::shared_ptr<cl::sycl::detail::event_impl>> &RawEvents,
+      pi_event &Event) {
+    Command::waitForEvents(Queue, RawEvents, Event);
+  }
+
 protected:
   cl::sycl::detail::Requirement MRequirement;
 };