[SYCL][Graph] Support handler-based restricted host task with native recording

sycl/source/handler.cpp

@@ -364,12 +364,36 @@ fill_copy_args(detail::handler_impl *impl,
                  DestOffset, DestExtent, CopyExtent);
 }
 
+// A command submission may transition a queue into a recording state. To obtain
+// the native graph for this scenario, we must check the recording status of the
+// queue on each wait event.
 static std::shared_ptr<ext::oneapi::experimental::detail::graph_impl>
-getNativeGraphImpl(queue_impl &Queue) {
+getNativeGraphImpl(queue_impl &Queue,
+                   const std::vector<EventImplPtr> &WaitEvents) {
+  std::shared_ptr<queue_impl> RecordingQueue = nullptr;
+  if (Queue.isNativeRecording()) {
+    RecordingQueue = Queue.shared_from_this();
+  } else {
+    for (const EventImplPtr &Event : WaitEvents) {
+      std::shared_ptr<queue_impl> EventQueue =
+          Event ? Event->getWorkerQueue() : nullptr;
+      if (EventQueue && EventQueue->isNativeRecording()) {
+        RecordingQueue = EventQueue;
+        break;
+      }
+    }
+  }
+  if (!RecordingQueue)
+    return nullptr;
+
   ur_exp_graph_handle_t UrGraphHandle = nullptr;
-  Queue.getAdapter().call<UrApiKind::urQueueGetGraphExp>(Queue.getHandleRef(),
-                                                         &UrGraphHandle);
-  return Queue.getContextImpl().getNativeGraph(UrGraphHandle);
+  ur_result_t Result =
+      RecordingQueue->getAdapter().call_nocheck<UrApiKind::urQueueGetGraphExp>(
+          RecordingQueue->getHandleRef(), &UrGraphHandle);
+  if (Result != UR_RESULT_SUCCESS)
+    return nullptr;
+
+  return RecordingQueue->getContextImpl().getNativeGraph(UrGraphHandle);
 }
 
 } // namespace detail
@@ -781,22 +805,39 @@ detail::EventImplPtr handler::finalize() {
   }
 
   // Host tasks in native recording mode are captured into the native graph
-  // rather than submitted to the scheduler.
-  if (type == detail::CGType::NativeHostTask && Queue->isNativeRecording()) {
-    auto GraphImpl = detail::getNativeGraphImpl(*Queue);
-    assert(GraphImpl && "Native graph handle expired while recording");
-
+  // rather than submitted to the scheduler. The submitting queue may already be
+  // recording, or it may join an in-progress capture by depending on an event
+  // from a recording queue (a fork).
+  std::shared_ptr<ext::oneapi::experimental::detail::graph_impl> HTNativeGraph =
+      nullptr;
+  if (type == detail::CGType::NativeHostTask)
+    HTNativeGraph =
+        detail::getNativeGraphImpl(*Queue, CommandGroup->getEvents());
+  if (HTNativeGraph) {
     auto *HT = static_cast<detail::CGHostTask *>(CommandGroup.get());
     // Store callback in the graph to manage its lifetime
-    auto *CallbackData = GraphImpl->addNativeHostTaskCallback(
+    auto *CallbackData = HTNativeGraph->addNativeHostTaskCallback(
         std::make_unique<detail::EnqueueHostTaskData>(
             detail::HandlerAccess::getHostTaskFunc(*HT->MHostTask)));
 
+    // All event dependencies supported with native recording go through UR
+    auto WaitEventList = detail::Command::getUrEvents(
+        HT->getEvents(), Queue, /*IsHostTaskCommand*/ true);
+    ur_event_handle_t SignalEvent = nullptr;
+
     Queue->getAdapter().call<detail::UrApiKind::urEnqueueHostTaskExp>(
         Queue->getHandleRef(), detail::NativeHostTask<false>, CallbackData,
-        nullptr, 0, nullptr, nullptr);
-
-    return detail::event_impl::create_completed_host_event();
+        nullptr, static_cast<uint32_t>(WaitEventList.size()),
+        WaitEventList.empty() ? nullptr : WaitEventList.data(), &SignalEvent);
+
+    auto EventImpl = detail::event_impl::create_device_event(*Queue);
+    EventImpl->setWorkerQueue(Queue->weak_from_this());
+    EventImpl->setSubmittedQueue(Queue);
+    EventImpl->setPotentiallyNativeRecorded(true);
+    EventImpl->setSubmissionTime();
+    EventImpl->setHandle(SignalEvent);
+    EventImpl->setEnqueued();
+    return EventImpl;
   }
   if (!CommandGroup->getRequirements().empty() && Queue->isNativeRecording()) {
     throw sycl::exception(

sycl/test-e2e/Graph/RecordReplay/Inputs/enqueue_func_host_task_event.cpp

@@ -0,0 +1,107 @@
+// Tests that a restricted host task submitted through the
+// handler path can be recorded into a SYCL Graph and
+// participate in event-based dependencies.
+
+#include "../../graph_common.hpp"
+
+#include <sycl/ext/oneapi/experimental/enqueue_functions.hpp>
+#include <sycl/properties/all_properties.hpp>
+
+constexpr size_t N = 1024;
+
+int main() {
+  device Dev;
+  context Ctx{Dev};
+
+  queue Queue1{Ctx, Dev, {property::queue::in_order{}}};
+  queue Queue2{Ctx, Dev, {property::queue::in_order{}}};
+  queue Queue3{Ctx, Dev, {property::queue::in_order{}}};
+
+#ifdef GRAPH_E2E_NATIVE_RECORDING
+  exp_ext::command_graph Graph{
+      Ctx, Dev, {exp_ext::property::graph::enable_native_recording{}}};
+#else
+  exp_ext::command_graph Graph{Ctx, Dev};
+#endif
+
+  uint32_t *Data = malloc_shared<uint32_t>(N, Queue1);
+  std::fill(Data, Data + N, 0);
+
+  QueueStateVerifier verifier(Queue1, Queue2, Queue3);
+  verifier.verify(EXECUTING, EXECUTING, EXECUTING);
+
+  Graph.begin_recording(Queue1);
+  verifier.verify(RECORDING, EXECUTING, EXECUTING);
+
+  exp_ext::submit_with_event(Queue1, [&](handler &CGH) {
+    exp_ext::host_task(CGH, [=] {
+      for (size_t i = 0; i < N; i++)
+        Data[i] += 5;
+    });
+  });
+
+  Queue1.submit([&](handler &CGH) {
+    CGH.parallel_for(range<1>{N}, [=](id<1> idx) {
+      Data[idx] += static_cast<uint32_t>(idx[0]) + 1;
+    });
+  });
+
+  // Fork: host task signals ForkEvent, consumed by Q2 and Q3.
+  auto ForkEvent = exp_ext::submit_with_event(Queue1, [&](handler &CGH) {
+    exp_ext::host_task(CGH, [=] {
+      for (size_t i = 0; i < N; i++)
+        Data[i] += 100;
+    });
+  });
+
+  // Q2 and Q3 run concurrently, so they operate on disjoint halves.
+  // Q2 fork: kernel on the first half.
+  auto Q2Event = exp_ext::submit_with_event(Queue2, [&](handler &CGH) {
+    CGH.depends_on(ForkEvent);
+    CGH.parallel_for(range<1>{N / 2}, [=](id<1> idx) { Data[idx] += 10; });
+  });
+  verifier.verify(RECORDING, RECORDING, EXECUTING);
+
+  // Q3 fork: host task on the second half.
+  auto Q3Event = exp_ext::submit_with_event(Queue3, [&](handler &CGH) {
+    CGH.depends_on(ForkEvent);
+    exp_ext::host_task(CGH, [=] {
+      for (size_t i = N / 2; i < N; i++)
+        Data[i] += 10;
+    });
+  });
+  verifier.verify(RECORDING, RECORDING, RECORDING);
+
+  // Join: host task waits on events from two distinct queues
+  Queue1.submit([&](handler &CGH) {
+    CGH.depends_on({Q2Event, Q3Event});
+    exp_ext::host_task(CGH, [=] {
+      for (size_t i = 0; i < N; i++)
+        Data[i] += 1;
+    });
+  });
+
+  Graph.end_recording();
+
+  auto ExecutableGraph = Graph.finalize();
+
+  // Each replay adds 5 + (i + 1) + 100 + 10 + 1 = (i + 117) to every element.
+  Queue1.submit([&](handler &CGH) { CGH.ext_oneapi_graph(ExecutableGraph); });
+  Queue1.wait();
+
+  for (size_t i = 0; i < N; i++) {
+    uint32_t Expected = static_cast<uint32_t>(i) + 117;
+    assert(check_value(i, Expected, Data[i], "Data"));
+  }
+
+  Queue1.submit([&](handler &CGH) { CGH.ext_oneapi_graph(ExecutableGraph); });
+  Queue1.wait();
+
+  for (size_t i = 0; i < N; i++) {
+    uint32_t Expected = 2 * (static_cast<uint32_t>(i) + 117);
+    assert(check_value(i, Expected, Data[i], "Data"));
+  }
+
+  free(Data, Queue1);
+  return 0;
+}

sycl/test-e2e/Graph/RecordReplay/NativeRecording/enqueue_func_host_task_event.cpp

@@ -0,0 +1,17 @@
+// REQUIRES: level_zero_v2_adapter && arch-intel_gpu_bmg_g21
+// REQUIRES: linux
+// REQUIRES: aspect-usm_shared_allocations
+// REQUIRES-INTEL-DRIVER: lin: 38146
+
+// TODO: Update windows driver once available with other tests. The current
+// driver restriction aligns with get_graph_multi_queue.cpp
+
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+// RUN: %if level_zero %{%{l0_leak_check} %{run} %t.out 2>&1 | FileCheck %s --implicit-check-not=LEAK %}
+
+// Tests fork/join of a restricted host_task using the handler path.
+
+#define GRAPH_E2E_NATIVE_RECORDING
+
+#include "../Inputs/enqueue_func_host_task_event.cpp"