[SYCL] Fix execution graph cleanup on memory object destruction

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);
+}