[SYCL] Fix execution graph cleanup on memory object destruction (#1065)

The previous algorithm deleted a visited node if it had no dependencies
left, but it only visited the immediate users of deleted nodes while
leaving all potential indirect users unchanged. This resulted in
leftover dependencies on dead memory objects.

The new algorithm addresses this problem by traversing the entire graph
component of the memory object being deleted instead.

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

Author: sergey-semenov

sycl/source/detail/scheduler/graph_builder.cpp

@@ -17,6 +17,7 @@
 
 #include <cstdlib>
 #include <fstream>
+#include <map>
 #include <memory>
 #include <queue>
 #include <set>
@@ -633,43 +634,71 @@ Scheduler::GraphBuilder::addCG(std::unique_ptr<detail::CG> CommandGroup,
 }
 
 void Scheduler::GraphBuilder::cleanupCommandsForRecord(MemObjRecord *Record) {
-  if (Record->MAllocaCommands.empty())
+  std::vector<AllocaCommandBase *> &AllocaCommands = Record->MAllocaCommands;
+  if (AllocaCommands.empty())
     return;
 
-  std::queue<Command *> RemoveQueue;
+  std::queue<Command *> ToVisit;
   std::set<Command *> Visited;
+  std::vector<Command *> CmdsToDelete;
+  // First, mark all allocas for deletion and their direct users for traversal
+  // Dependencies of the users will be cleaned up during the traversal
+  for (Command *AllocaCmd : AllocaCommands) {
+    Visited.insert(AllocaCmd);
+    for (Command *UserCmd : AllocaCmd->MUsers)
+      ToVisit.push(UserCmd);
+    CmdsToDelete.push_back(AllocaCmd);
+    // These commands will be deleted later, clear users now to avoid
+    // updating them during edge removal
+    AllocaCmd->MUsers.clear();
+  }
 
-  // TODO: release commands need special handling here as they are not reachable
-  // from alloca commands
-
-  for (AllocaCommandBase *AllocaCmd : Record->MAllocaCommands) {
-    if (Visited.find(AllocaCmd) == Visited.end())
-      RemoveQueue.push(AllocaCmd);
-    // Use BFS to find and process all users of removal candidate
-    while (!RemoveQueue.empty()) {
-      Command *CandidateCommand = RemoveQueue.front();
-      RemoveQueue.pop();
-
-      if (Visited.insert(CandidateCommand).second) {
-        for (Command *UserCmd : CandidateCommand->MUsers) {
-          // As candidate command is about to be freed, we need
-          // to remove it from dependency list of other commands.
-          auto NewEnd =
-              std::remove_if(UserCmd->MDeps.begin(), UserCmd->MDeps.end(),
-                             [CandidateCommand](const DepDesc &Dep) {
-                               return Dep.MDepCommand == CandidateCommand;
-                             });
-          UserCmd->MDeps.erase(NewEnd, UserCmd->MDeps.end());
-
-          // Commands that have no unsatisfied dependencies can be executed
-          // and are good candidates for clean up.
-          if (UserCmd->MDeps.empty())
-            RemoveQueue.push(UserCmd);
-        }
-        CandidateCommand->getEvent()->setCommand(nullptr);
-        delete CandidateCommand;
-      }
+  // Traverse the graph using BFS
+  while (!ToVisit.empty()) {
+    Command *Cmd = ToVisit.front();
+    ToVisit.pop();
+
+    if (!Visited.insert(Cmd).second)
+      continue;
+
+    for (Command *UserCmd : Cmd->MUsers)
+      ToVisit.push(UserCmd);
+
+    // Delete all dependencies on any allocations being removed
+    // Track which commands should have their users updated
+    std::map<Command *, bool> ShouldBeUpdated;
+    auto NewEnd = std::remove_if(
+        Cmd->MDeps.begin(), Cmd->MDeps.end(), [&](const DepDesc &Dep) {
+          if (std::find(AllocaCommands.begin(), AllocaCommands.end(),
+                        Dep.MAllocaCmd) != AllocaCommands.end()) {
+            ShouldBeUpdated.insert({Dep.MDepCommand, true});
+            return true;
+          }
+          ShouldBeUpdated[Dep.MDepCommand] = false;
+          return false;
+        });
+    Cmd->MDeps.erase(NewEnd, Cmd->MDeps.end());
+
+    // Update users of removed dependencies
+    for (auto DepCmdIt : ShouldBeUpdated) {
+      if (!DepCmdIt.second)
+        continue;
+      std::vector<Command *> &DepUsers = DepCmdIt.first->MUsers;
+      DepUsers.erase(std::remove(DepUsers.begin(), DepUsers.end(), Cmd),
+                     DepUsers.end());
     }
+
+    // If all dependencies have been removed this way, mark the command for
+    // deletion
+    if (Cmd->MDeps.empty()) {
+      CmdsToDelete.push_back(Cmd);
+      Cmd->MUsers.clear();
+    }
+  }
+
+  for (Command *Cmd : CmdsToDelete) {
+    Cmd->getEvent()->setCommand(nullptr);
+    delete Cmd;
   }
 }
 

sycl/test/scheduler/FakeCommand.hpp

@@ -0,0 +1,23 @@
+#include <CL/sycl.hpp>
+
+// A fake command class used for testing
+class FakeCommand : public cl::sycl::detail::Command {
+public:
+  FakeCommand(cl::sycl::detail::QueueImplPtr Queue,
+              cl::sycl::detail::Requirement Req)
+      : Command{cl::sycl::detail::Command::ALLOCA, Queue},
+        MRequirement{std::move(Req)} {}
+
+  void printDot(std::ostream &Stream) const override {}
+
+  const cl::sycl::detail::Requirement *getRequirement() const final {
+    return &MRequirement;
+  };
+
+  cl_int enqueueImp() override { return MRetVal; }
+
+  cl_int MRetVal = CL_SUCCESS;
+
+protected:
+  cl::sycl::detail::Requirement MRequirement;
+};

sycl/test/scheduler/LeafLimit.cpp

@@ -6,29 +6,12 @@
 #include <memory>
 #include <vector>
 
+#include "FakeCommand.hpp"
+
 // This test checks the leaf limit imposed on the execution graph
 
 using namespace cl::sycl;
 
-class FakeCommand : public detail::Command {
-public:
-  FakeCommand(detail::QueueImplPtr Queue, detail::Requirement Req)
-      : Command{detail::Command::ALLOCA, Queue}, MRequirement{std::move(Req)} {}
-
-  void printDot(std::ostream &Stream) const override {}
-
-  const detail::Requirement *getRequirement() const final {
-    return &MRequirement;
-  };
-
-  cl_int enqueueImp() override { return MRetVal; }
-
-  cl_int MRetVal = CL_SUCCESS;
-
-protected:
-  detail::Requirement MRequirement;
-};
-
 class TestScheduler : public detail::Scheduler {
 public:
   void AddNodeToLeaves(detail::MemObjRecord *Rec, detail::Command *Cmd,

sycl/test/scheduler/MemObjCommandCleanup.cpp

@@ -0,0 +1,103 @@
+// RUN: %clangxx -fsycl %s -o %t.out
+// RUN: %t.out
+#include <CL/sycl.hpp>
+
+#include <functional>
+#include <memory>
+#include <utility>
+
+#include "FakeCommand.hpp"
+
+// This test checks that the execution graph cleanup on memory object
+// destruction traverses the entire graph, rather than only the immediate users
+// of deleted commands.
+
+using namespace cl::sycl;
+
+class TestScheduler : public detail::Scheduler {
+public:
+  void cleanupCommandsForRecord(detail::MemObjRecord *Rec) {
+    MGraphBuilder.cleanupCommandsForRecord(Rec);
+  }
+
+  void removeRecordForMemObj(detail::SYCLMemObjI *MemObj) {
+    MGraphBuilder.removeRecordForMemObj(MemObj);
+  }
+
+  detail::MemObjRecord *
+  getOrInsertMemObjRecord(const detail::QueueImplPtr &Queue,
+                          detail::Requirement *Req) {
+    return MGraphBuilder.getOrInsertMemObjRecord(Queue, Req);
+  }
+};
+
+class FakeCommandWithCallback : public FakeCommand {
+public:
+  FakeCommandWithCallback(detail::QueueImplPtr Queue, detail::Requirement Req,
+                          std::function<void()> Callback)
+      : FakeCommand(Queue, Req), MCallback(std::move(Callback)) {}
+
+  ~FakeCommandWithCallback() override { MCallback(); }
+
+protected:
+  std::function<void()> MCallback;
+};
+
+template <typename MemObjT>
+detail::Requirement getFakeRequirement(const MemObjT &MemObj) {
+  return {{0, 0, 0},
+          {0, 0, 0},
+          {0, 0, 0},
+          access::mode::read_write,
+          detail::getSyclObjImpl(MemObj).get(),
+          0,
+          0,
+          0};
+}
+
+void addEdge(detail::Command *User, detail::Command *Dep,
+             detail::AllocaCommandBase *Alloca) {
+  User->addDep(detail::DepDesc{Dep, User->getRequirement(), Alloca});
+  Dep->addUser(User);
+}
+
+int main() {
+  TestScheduler TS;
+  queue Queue;
+  buffer<int, 1> BufA(range<1>(1));
+  buffer<int, 1> BufB(range<1>(1));
+  detail::Requirement FakeReqA = getFakeRequirement(BufA);
+  detail::Requirement FakeReqB = getFakeRequirement(BufB);
+  detail::MemObjRecord *RecA =
+      TS.getOrInsertMemObjRecord(detail::getSyclObjImpl(Queue), &FakeReqA);
+
+  // Create 2 fake allocas, one of which will be cleaned up
+  detail::AllocaCommand *FakeAllocaA =
+      new detail::AllocaCommand(detail::getSyclObjImpl(Queue), FakeReqA);
+  std::unique_ptr<detail::AllocaCommand> FakeAllocaB{
+      new detail::AllocaCommand(detail::getSyclObjImpl(Queue), FakeReqB)};
+  RecA->MAllocaCommands.push_back(FakeAllocaA);
+
+  // Create a direct user of both allocas
+  std::unique_ptr<FakeCommand> FakeDirectUser{
+      new FakeCommand(detail::getSyclObjImpl(Queue), FakeReqA)};
+  addEdge(FakeDirectUser.get(), FakeAllocaA, FakeAllocaA);
+  addEdge(FakeDirectUser.get(), FakeAllocaB.get(), FakeAllocaB.get());
+
+  // Create an indirect user of the soon-to-be deleted alloca
+  bool IndirectUserDeleted = false;
+  std::function<void()> Callback = [&]() { IndirectUserDeleted = true; };
+  FakeCommand *FakeIndirectUser = new FakeCommandWithCallback(
+      detail::getSyclObjImpl(Queue), FakeReqA, Callback);
+  addEdge(FakeIndirectUser, FakeDirectUser.get(), FakeAllocaA);
+
+  TS.cleanupCommandsForRecord(RecA);
+  TS.removeRecordForMemObj(detail::getSyclObjImpl(BufA).get());
+
+  // Check that the direct user has been left with the second alloca
+  // as the only dependency, while the indirect user has been cleaned up.
+  assert(FakeDirectUser->MUsers.size() == 0);
+  assert(FakeDirectUser->MDeps.size() == 1);
+  assert(FakeDirectUser->MDeps[0].MDepCommand == FakeAllocaB.get());
+  assert(IndirectUserDeleted);
+}