[SYCL] Fix memory dependency leaks caused by failed kernel enqueue

sycl/source/detail/scheduler/commands.hpp

@@ -136,6 +136,10 @@ class Command {
     return MEnqueueStatus == EnqueueResultT::SyclEnqueueBlocked;
   }
 
+  bool isEnqueueFailed() const {
+    return MEnqueueStatus == EnqueueResultT::SyclEnqueueFailed;
+  }
+
   const QueueImplPtr &getQueue() const { return MQueue; }
 
   const QueueImplPtr &getSubmittedQueue() const { return MSubmittedQueue; }

sycl/source/detail/scheduler/graph_builder.cpp

@@ -1125,6 +1125,15 @@ void Scheduler::GraphBuilder::cleanupFinishedCommands(
         MCmdsToVisit.push(Dep.MDepCommand);
     }
 
+    // If the command has failed to enqueue it must be removed from its leaves.
+    if (Cmd->isEnqueueFailed()) {
+      for (const DepDesc &Dep : Cmd->MDeps) {
+        const Requirement *Req = Dep.MDepRequirement;
+        MemObjRecord *Record = getMemObjRecord(Req->MSYCLMemObj);
+        updateLeaves({Cmd}, Record, Req->MAccessMode);
+      }
+    }
+
     // Do not clean up the node if it is a leaf for any memory object
     if (Cmd->MLeafCounter > 0)
       continue;

sycl/source/detail/scheduler/scheduler.cpp

@@ -108,7 +108,7 @@ EventImplPtr Scheduler::addCG(std::unique_ptr<detail::CG> CommandGroup,
     NewEvent = NewCmd->getEvent();
   }
 
-  {
+  try {
     ReadLockT Lock(MGraphLock);
 
     Command *NewCmd = static_cast<Command *>(NewEvent->getCommand());
@@ -128,38 +128,29 @@ EventImplPtr Scheduler::addCG(std::unique_ptr<detail::CG> CommandGroup,
 
     for (Command *Cmd : AuxiliaryCmds) {
       Enqueued = GraphProcessor::enqueueCommand(Cmd, Res);
-      try {
-        if (!Enqueued && EnqueueResultT::SyclEnqueueFailed == Res.MResult)
-          throw runtime_error("Auxiliary enqueue process failed.",
-                              PI_INVALID_OPERATION);
-      } catch (...) {
-        // enqueueCommand() func and if statement above may throw an exception,
-        // so destroy required resources to avoid memory leak
-        CleanUp();
-        std::rethrow_exception(std::current_exception());
-      }
+      if (!Enqueued && EnqueueResultT::SyclEnqueueFailed == Res.MResult)
+        throw runtime_error("Auxiliary enqueue process failed.",
+                            PI_INVALID_OPERATION);
     }
 
     if (NewCmd) {
       // TODO: Check if lazy mode.
       EnqueueResultT Res;
-      try {
-        bool Enqueued = GraphProcessor::enqueueCommand(NewCmd, Res);
-        if (!Enqueued && EnqueueResultT::SyclEnqueueFailed == Res.MResult)
-          throw runtime_error("Enqueue process failed.", PI_INVALID_OPERATION);
-      } catch (...) {
-        // enqueueCommand() func and if statement above may throw an exception,
-        // so destroy required resources to avoid memory leak
-        CleanUp();
-        std::rethrow_exception(std::current_exception());
-      }
+      bool Enqueued = GraphProcessor::enqueueCommand(NewCmd, Res);
+      if (!Enqueued && EnqueueResultT::SyclEnqueueFailed == Res.MResult)
+        throw runtime_error("Enqueue process failed.", PI_INVALID_OPERATION);
 
       // 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.
       CleanUp();
     }
+  } catch (...) {
+    // If enqueuing has failed we need to clean up the command to remove it
+    // from the graph so it does not cause issues for other related commands.
+    cleanupFinishedCommands(NewEvent);
+    std::rethrow_exception(std::current_exception());
   }
 
   for (auto StreamImplPtr : Streams) {

sycl/unittests/scheduler/FailedCommands.cpp

@@ -9,8 +9,58 @@
 #include "SchedulerTest.hpp"
 #include "SchedulerTestUtils.hpp"
 
+#include <helpers/CommonRedefinitions.hpp>
+#include <helpers/PiImage.hpp>
+#include <helpers/PiMock.hpp>
+
 using namespace cl::sycl;
 
+class TestKernel;
+
+__SYCL_INLINE_NAMESPACE(cl) {
+namespace sycl {
+namespace detail {
+template <> struct KernelInfo<TestKernel> {
+  static constexpr unsigned getNumParams() { return 0; }
+  static const kernel_param_desc_t &getParamDesc(int) {
+    static kernel_param_desc_t Dummy;
+    return Dummy;
+  }
+  static constexpr const char *getName() { return "TestKernel"; }
+  static constexpr bool isESIMD() { return false; }
+  static constexpr bool callsThisItem() { return false; }
+  static constexpr bool callsAnyThisFreeFunction() { return false; }
+};
+
+} // namespace detail
+} // namespace sycl
+} // __SYCL_INLINE_NAMESPACE(cl)
+
+static sycl::unittest::PiImage generateDefaultImage() {
+  using namespace sycl::unittest;
+
+  PiPropertySet PropSet;
+
+  std::vector<unsigned char> Bin{0, 1, 2, 3, 4, 5}; // Random data
+
+  PiArray<PiOffloadEntry> Entries = makeEmptyKernels({"TestKernel"});
+
+  PiImage Img{PI_DEVICE_BINARY_TYPE_SPIRV,            // Format
+              __SYCL_PI_DEVICE_BINARY_TARGET_SPIRV64, // DeviceTargetSpec
+              "",                                     // Compile options
+              "",                                     // Link options
+              std::move(Bin),
+              std::move(Entries),
+              std::move(PropSet)};
+
+  return Img;
+}
+
+static sycl::unittest::PiImage Img = generateDefaultImage();
+static sycl::unittest::PiImageArray<1> ImgArray{&Img};
+
+using namespace sycl;
+
 TEST_F(SchedulerTest, FailedDependency) {
   detail::Requirement MockReq = getMockRequirement();
   MockCommand MDep(detail::getSyclObjImpl(MQueue));
@@ -35,3 +85,132 @@ TEST_F(SchedulerTest, FailedDependency) {
   ASSERT_EQ(MDep.MEnqueueStatus, detail::EnqueueResultT::SyclEnqueueFailed)
       << "MDep should be marked as failed\n";
 }
+
+pi_result redefinedFailingEnqueueKernelLaunch(pi_queue, pi_kernel, pi_uint32,
+                                              const size_t *, const size_t *,
+                                              const size_t *, pi_uint32,
+                                              const pi_event *, pi_event *) {
+  throw sycl::runtime_error(
+      "Exception from redefinedFailingEnqueueKernelLaunch.",
+      PI_INVALID_OPERATION);
+}
+
+size_t MemBufRefCount = 0u;
+
+pi_result redefinedMemBufferCreate(pi_context, pi_mem_flags, size_t, void *,
+                                   pi_mem *ret_mem, const pi_mem_properties *) {
+  *ret_mem = reinterpret_cast<pi_mem>(0x1);
+  ++MemBufRefCount;
+  return PI_SUCCESS;
+}
+
+pi_result redefinedMemBufferPartition(pi_mem, pi_mem_flags,
+                                      pi_buffer_create_type, void *,
+                                      pi_mem *ret_mem) {
+  *ret_mem = reinterpret_cast<pi_mem>(0x1);
+  ++MemBufRefCount;
+  return PI_SUCCESS;
+}
+
+pi_result redefinedMemRetain(pi_mem) {
+  ++MemBufRefCount;
+  return PI_SUCCESS;
+}
+
+pi_result redefinedMemRelease(pi_mem) {
+  --MemBufRefCount;
+  return PI_SUCCESS;
+}
+
+TEST_F(SchedulerTest, FailedCommandAccessorCleanup) {
+  default_selector Selector;
+  platform Plt{default_selector()};
+  if (Plt.is_host()) {
+    std::cout << "Not run due to host-only environment\n";
+    return;
+  }
+  if (Plt.get_backend() == sycl::backend::ext_oneapi_cuda ||
+      Plt.get_backend() == sycl::backend::ext_oneapi_hip) {
+    std::cout << "CUDA and HIP backends do not currently support this test\n";
+    return;
+  }
+
+  unittest::PiMock Mock{Plt};
+  setupDefaultMockAPIs(Mock);
+  MemBufRefCount = 0u;
+  Mock.redefine<detail::PiApiKind::piEnqueueKernelLaunch>(
+      redefinedFailingEnqueueKernelLaunch);
+  Mock.redefine<detail::PiApiKind::piMemBufferCreate>(redefinedMemBufferCreate);
+  Mock.redefine<detail::PiApiKind::piMemRetain>(redefinedMemRetain);
+  Mock.redefine<detail::PiApiKind::piMemRelease>(redefinedMemRelease);
+
+  {
+    context Ctx{Plt};
+    queue Q{Ctx, Selector};
+
+    kernel_bundle KernelBundle =
+        sycl::get_kernel_bundle<sycl::bundle_state::input>(Ctx);
+    auto ExecBundle = sycl::build(KernelBundle);
+
+    buffer<int, 1> Buff{cl::sycl::range<1>(1)};
+
+    try {
+      Q.submit([&](sycl::handler &CGH) {
+        auto Acc = Buff.get_access<cl::sycl::access::mode::read_write>(CGH);
+        CGH.use_kernel_bundle(ExecBundle);
+        CGH.single_task<TestKernel>([=] {});
+      });
+      FAIL() << "No exception was thrown.";
+    } catch (...) {
+    }
+  }
+
+  ASSERT_EQ(MemBufRefCount, 0u) << "Memory leak detected.";
+}
+
+TEST_F(SchedulerTest, FailedCommandStreamCleanup) {
+  default_selector Selector;
+  platform Plt{default_selector()};
+  if (Plt.is_host()) {
+    std::cout << "Not run due to host-only environment\n";
+    return;
+  }
+  if (Plt.get_backend() == sycl::backend::ext_oneapi_cuda ||
+      Plt.get_backend() == sycl::backend::ext_oneapi_hip) {
+    std::cout << "CUDA and HIP backends do not currently support this test\n";
+    return;
+  }
+
+  unittest::PiMock Mock{Plt};
+  setupDefaultMockAPIs(Mock);
+  MemBufRefCount = 0u;
+  Mock.redefine<detail::PiApiKind::piEnqueueKernelLaunch>(
+      redefinedFailingEnqueueKernelLaunch);
+  Mock.redefine<detail::PiApiKind::piMemBufferCreate>(redefinedMemBufferCreate);
+  Mock.redefine<detail::PiApiKind::piMemBufferPartition>(
+      redefinedMemBufferPartition);
+  Mock.redefine<detail::PiApiKind::piMemRetain>(redefinedMemRetain);
+  Mock.redefine<detail::PiApiKind::piMemRelease>(redefinedMemRelease);
+
+  {
+    context Ctx{Plt};
+    queue Q{Ctx, Selector};
+
+    kernel_bundle KernelBundle =
+        sycl::get_kernel_bundle<sycl::bundle_state::input>(Ctx);
+    auto ExecBundle = sycl::build(KernelBundle);
+
+    try {
+      Q.submit([&](sycl::handler &CGH) {
+        sycl::stream KernelStream(108 * 64 + 128, 64, CGH);
+        CGH.use_kernel_bundle(ExecBundle);
+        CGH.single_task<TestKernel>([=] {});
+      });
+      FAIL() << "No exception was thrown.";
+    } catch (...) {
+    }
+    Q.wait();
+  }
+
+  ASSERT_EQ(MemBufRefCount, 0u) << "Memory leak detected.";
+}