[SYCL] Don't perform early free of native kernel command group (#2674)

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

Author: s-kanaev

sycl/source/detail/scheduler/commands.cpp

@@ -1724,12 +1724,18 @@ pi_result ExecCGCommand::SetKernelParamsAndLaunch(
 void DispatchNativeKernel(void *Blob) {
   // First value is a pointer to Corresponding CGExecKernel object.
   CGExecKernel *HostTask = *(CGExecKernel **)Blob;
+  bool ShouldDeleteCG = static_cast<void **>(Blob)[1] != nullptr;
 
   // Other value are pointer to the buffers.
-  void **NextArg = (void **)Blob + 1;
+  void **NextArg = static_cast<void **>(Blob) + 2;
   for (detail::Requirement *Req : HostTask->MRequirements)
     Req->MData = *(NextArg++);
   HostTask->MHostKernel->call(HostTask->MNDRDesc, nullptr);
+
+  // The command group will (if not already was) be released in scheduler.
+  // Hence we're free to deallocate it here.
+  if (ShouldDeleteCG)
+    delete HostTask;
 }
 
 cl_int ExecCGCommand::enqueueImp() {
@@ -1814,9 +1820,14 @@ cl_int ExecCGCommand::enqueueImp() {
     // piEnqueueNativeKernel takes arguments blob which is passes to user
     // function.
     // Reserve extra space for the pointer to CGExecKernel to restore context.
-    std::vector<void *> ArgsBlob(HostTask->MArgs.size() + 1);
+    std::vector<void *> ArgsBlob(HostTask->MArgs.size() + 2);
     ArgsBlob[0] = (void *)HostTask;
-    void **NextArg = ArgsBlob.data() + 1;
+    {
+      std::intptr_t ShouldDeleteCG =
+          static_cast<std::intptr_t>(MDeps.size() == 0 && MUsers.size() == 0);
+      ArgsBlob[1] = reinterpret_cast<void *>(ShouldDeleteCG);
+    }
+    void **NextArg = ArgsBlob.data() + 2;
 
     if (MQueue->is_host()) {
       for (ArgDesc &Arg : HostTask->MArgs) {

sycl/source/detail/scheduler/commands.hpp

@@ -505,6 +505,15 @@ class ExecCGCommand : public Command {
   // the cleanup process.
   EmptyCommand *MEmptyCmd = nullptr;
 
+  // This function is only usable for native kernel to prevent access to free'd
+  // memory in DispatchNativeKernel.
+  // TODO remove when native kernel support is terminated.
+  void releaseCG() {
+    assert(MCommandGroup->getType() == CG::RUN_ON_HOST_INTEL &&
+           "Only 'native kernel' is allowed to release command group");
+    MCommandGroup.release();
+  }
+
 private:
   cl_int enqueueImp() final override;
 

sycl/source/detail/scheduler/scheduler.cpp

@@ -70,6 +70,7 @@ EventImplPtr Scheduler::addCG(std::unique_ptr<detail::CG> CommandGroup,
                               QueueImplPtr Queue) {
   EventImplPtr NewEvent = nullptr;
   const bool IsKernel = CommandGroup->getType() == CG::KERNEL;
+  const bool IsHostKernel = CommandGroup->getType() == CG::RUN_ON_HOST_INTEL;
   vector_class<StreamImplPtr> Streams;
   {
     std::unique_lock<std::shared_timed_mutex> Lock(MGraphLock, std::defer_lock);
@@ -104,9 +105,15 @@ EventImplPtr Scheduler::addCG(std::unique_ptr<detail::CG> CommandGroup,
       if (IsKernel)
         Streams = ((ExecCGCommand *)NewCmd)->getStreams();
 
+      // If there are no memory dependencies decouple and free the command.
+      // Though, dismiss ownership of native kernel command group as it's
+      // resources may be in use by backend and synchronization point here is
+      // at native kernel execution finish.
       if (NewCmd->MDeps.size() == 0 && NewCmd->MUsers.size() == 0) {
-        NewEvent->setCommand(nullptr); // if there are no memory dependencies,
-                                       // decouple and free the command
+        if (IsHostKernel)
+          static_cast<ExecCGCommand *>(NewCmd)->releaseCG();
+
+        NewEvent->setCommand(nullptr);
         delete NewCmd;
       }
     }

sycl/test/basic_tests/handler/run_on_host_intel2.cpp

@@ -0,0 +1,23 @@
+// RUN: %clangxx -fsycl %s -o %t.out
+// RUN: env SYCL_DEVICE_TYPE=HOST %t.out
+// RUN: %CPU_RUN_PLACEHOLDER %t.out
+
+#include <CL/sycl.hpp>
+
+// This tests that early free of command (and, hence, the command group) won't
+// affect "native kernel" feature support.
+int main(void) {
+  cl::sycl::queue Q;
+
+  int *Ptr = new int;
+
+  auto E = Q.submit([&](cl::sycl::handler &CGH) {
+    CGH.run_on_host_intel([=] { *Ptr = 5; });
+  });
+
+  E.wait();
+
+  std::cout << "Finished successfully\n";
+
+  return 0;
+}