[SYCL] Clean up finished command nodes of the execution graph

Delete completed command nodes that are neither memory allocations
nor leaves for any memory object. This cleanup is performed at each
host-device synchronization point in order to avoid execution graph
bloat.

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

Author: sergey-semenov

sycl/doc/SYCLEnvironmentVariables.md

@@ -18,6 +18,7 @@ subject to change. Do not rely on these variables in production code.
 | SYCL_USE_KERNEL_SPV | Path to the SPIR-V binary | Load device image from the specified file. If runtime is unable to read the file, `cl::sycl::runtime_error` exception is thrown.|
 | SYCL_DUMP_IMAGES | Any(*) | Dump device image binaries to file. Control has no effect if SYCL_USE_KERNEL_SPV is set. |
 | SYCL_PRINT_EXECUTION_GRAPH | Described [below](#sycl_print_execution_graph-options) | Print execution graph to DOT text file. |
+| SYCL_DISABLE_EXECUTION_GRAPH_CLEANUP | Any(*) | Disable cleanup of finished command nodes at host-device synchronization points. |
 | SYCL_THROW_ON_BLOCK | Any(*) | Throw an exception on attempt to wait for a blocked command.  |
 | SYCL_DEVICELIB_INHIBIT_NATIVE | String of device library extensions (separated by a whitespace) | Do not rely on device native support for devicelib extensions listed in this option. |
 | SYCL_DEVICE_ALLOWLIST | A list of devices and their minimum driver version following the pattern: DeviceName:{{XXX}},DriverVersion:{{X.Y.Z.W}}. Also may contain PlatformName and PlatformVersion | Filter out devices that do not match the pattern specified. Regular expression can be passed and the SYCL RT will select only those devices which satisfy the regex. |

sycl/include/CL/sycl/detail/scheduler/commands.hpp

@@ -10,6 +10,7 @@
 
 #include <atomic>
 #include <memory>
+#include <unordered_set>
 #include <vector>
 
 #include <CL/sycl/access/access.hpp>
@@ -98,7 +99,7 @@ class Command {
 
   void addDep(EventImplPtr Event) { MDepsEvents.push_back(std::move(Event)); }
 
-  void addUser(Command *NewUser) { MUsers.push_back(NewUser); }
+  void addUser(Command *NewUser) { MUsers.insert(NewUser); }
 
   // Return type of the command, e.g. Allocate, MemoryCopy.
   CommandType getType() const { return MType; }
@@ -149,11 +150,13 @@ class Command {
   // Contains list of dependencies(edges)
   std::vector<DepDesc> MDeps;
   // Contains list of commands that depend on the command
-  std::vector<Command *> MUsers;
+  std::unordered_set<Command *> MUsers;
   // Indicates whether the command can be blocked from enqueueing
   bool MIsBlockable = false;
   // Indicates whether the command is blocked from enqueueing
   std::atomic<bool> MCanEnqueue;
+  // Counts the number of memory objects this command is a leaf for
+  unsigned MLeafCounter = 0;
 
   const char *MBlockReason = "Unknown";
 };

sycl/include/CL/sycl/detail/scheduler/scheduler.hpp

@@ -77,6 +77,10 @@ class Scheduler {
   // sycl::image destructors.
   void removeMemoryObject(detail::SYCLMemObjI *MemObj);
 
+  // Removes finished non-leaf non-alloca commands from the subgraph (assuming
+  // that all its commands have been waited for).
+  void cleanupFinishedCommands(Command *FinishedCmd);
+
   // Creates nodes in the graph, that update Req with the pointer to the host
   // memory which contains the latest data of the memory object. New operations
   // with the same memory object that have side effects are blocked until
@@ -125,8 +129,9 @@ class Scheduler {
     // Event passed and its dependencies.
     void optimize(EventImplPtr Event);
 
-    // Removes unneeded commands from the graph.
-    void cleanupCommands(bool CleanupReleaseCommands = false);
+    // Removes finished non-leaf non-alloca commands from the subgraph (assuming
+    // that all its commands have been waited for).
+    void cleanupFinishedCommands(Command *FinishedCmd);
 
     // Reschedules command passed using Queue provided. this can lead to
     // rescheduling of all dependent commands. This can be used when user
@@ -140,6 +145,9 @@ class Scheduler {
     MemObjRecord *getOrInsertMemObjRecord(const QueueImplPtr &Queue,
                                           Requirement *Req);
 
+    // Decrements leaf counters for all leaves of the record.
+    void updateLeafCountersForRecord(MemObjRecord *Record);
+
     // Removes commands that use given MemObjRecord from the graph.
     void cleanupCommandsForRecord(MemObjRecord *Record);
 

sycl/source/detail/config.def

@@ -11,5 +11,5 @@
 // underscore(__).
 
 CONFIG(SYCL_PRINT_EXECUTION_GRAPH, 32, __SYCL_PRINT_EXECUTION_GRAPH)
+CONFIG(SYCL_DISABLE_EXECUTION_GRAPH_CLEANUP, 1, __SYCL_DISABLE_EXECUTION_GRAPH_CLEANUP)
 CONFIG(SYCL_DEVICE_ALLOWLIST, 1024, __SYCL_DEVICE_ALLOWLIST)
-

sycl/source/detail/event_impl.cpp

@@ -13,6 +13,8 @@
 #include <CL/sycl/detail/queue_impl.hpp>
 #include <CL/sycl/detail/scheduler/scheduler.hpp>
 
+#include "detail/config.hpp"
+
 #include <chrono>
 
 __SYCL_INLINE namespace cl {
@@ -100,6 +102,9 @@ void event_impl::wait(
     waitInternal();
   else if (MCommand)
     detail::Scheduler::getInstance().waitForEvent(std::move(Self));
+  if (MCommand && !SYCLConfig<SYCL_DISABLE_EXECUTION_GRAPH_CLEANUP>::get())
+    detail::Scheduler::getInstance().cleanupFinishedCommands(
+        static_cast<Command *>(MCommand));
 }
 
 void event_impl::wait_and_throw(

sycl/source/detail/scheduler/graph_builder.cpp

@@ -138,13 +138,15 @@ void Scheduler::GraphBuilder::UpdateLeaves(const std::set<Command *> &Cmds,
   if (ReadOnlyReq)
     return;
 
-  for (const Command *Cmd : Cmds) {
+  for (Command *Cmd : Cmds) {
     auto NewEnd = std::remove(Record->MReadLeaves.begin(),
                               Record->MReadLeaves.end(), Cmd);
+    Cmd->MLeafCounter -= std::distance(NewEnd, Record->MReadLeaves.end());
     Record->MReadLeaves.erase(NewEnd, Record->MReadLeaves.end());
 
     NewEnd = std::remove(Record->MWriteLeaves.begin(),
                          Record->MWriteLeaves.end(), Cmd);
+    Cmd->MLeafCounter -= std::distance(NewEnd, Record->MWriteLeaves.end());
     Record->MWriteLeaves.erase(NewEnd, Record->MWriteLeaves.end());
   }
 }
@@ -166,8 +168,10 @@ void Scheduler::GraphBuilder::AddNodeToLeaves(MemObjRecord *Record,
     Dep.MDepCommand = OldLeaf;
     Cmd->addDep(Dep);
     OldLeaf->addUser(Cmd);
+    --(OldLeaf->MLeafCounter);
   }
   Leaves.push_back(Cmd);
+  ++(Cmd->MLeafCounter);
 }
 
 UpdateHostRequirementCommand *Scheduler::GraphBuilder::insertUpdateHostReqCmd(
@@ -560,6 +564,7 @@ AllocaCommandBase *Scheduler::GraphBuilder::getOrCreateAllocaForReq(
 
     Record->MAllocaCommands.push_back(AllocaCmd);
     Record->MWriteLeaves.push_back(AllocaCmd);
+    ++(AllocaCmd->MLeafCounter);
   }
   return AllocaCmd;
 }
@@ -633,6 +638,16 @@ Scheduler::GraphBuilder::addCG(std::unique_ptr<detail::CG> CommandGroup,
   return NewCmd.release();
 }
 
+void Scheduler::GraphBuilder::updateLeafCountersForRecord(
+    MemObjRecord *Record) {
+  for (Command *Cmd : Record->MReadLeaves) {
+    --(Cmd->MLeafCounter);
+  }
+  for (Command *Cmd : Record->MWriteLeaves) {
+    --(Cmd->MLeafCounter);
+  }
+}
+
 void Scheduler::GraphBuilder::cleanupCommandsForRecord(MemObjRecord *Record) {
   std::vector<AllocaCommandBase *> &AllocaCommands = Record->MAllocaCommands;
   if (AllocaCommands.empty())
@@ -683,9 +698,7 @@ void Scheduler::GraphBuilder::cleanupCommandsForRecord(MemObjRecord *Record) {
     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());
+      DepCmdIt.first->MUsers.erase(Cmd);
     }
 
     // If all dependencies have been removed this way, mark the command for
@@ -702,8 +715,48 @@ void Scheduler::GraphBuilder::cleanupCommandsForRecord(MemObjRecord *Record) {
   }
 }
 
-void Scheduler::GraphBuilder::cleanupCommands(bool CleanupReleaseCommands) {
-  // TODO: Implement.
+void Scheduler::GraphBuilder::cleanupFinishedCommands(Command *FinishedCmd) {
+  std::queue<Command *> CmdsToVisit({FinishedCmd});
+  std::set<Command *> Visited;
+
+  // Traverse the graph using BFS
+  while (!CmdsToVisit.empty()) {
+    Command *Cmd = CmdsToVisit.front();
+    CmdsToVisit.pop();
+
+    if (!Visited.insert(Cmd).second)
+      continue;
+
+    for (const DepDesc &Dep : Cmd->MDeps) {
+      if (Dep.MDepCommand)
+        CmdsToVisit.push(Dep.MDepCommand);
+    }
+
+    // Do not clean up the node if it is a leaf for any memory object
+    if (Cmd->MLeafCounter > 0)
+      continue;
+    // Do not clean up allocation commands
+    Command::CommandType CmdT = Cmd->getType();
+    if (CmdT == Command::ALLOCA || CmdT == Command::ALLOCA_SUB_BUF)
+      continue;
+
+    for (Command *UserCmd : Cmd->MUsers) {
+      for (DepDesc &Dep : UserCmd->MDeps) {
+        // Link the users of the command to the alloca command(s) instead
+        if (Dep.MDepCommand == Cmd) {
+          Dep.MDepCommand = Dep.MAllocaCmd;
+          Dep.MDepCommand->MUsers.insert(UserCmd);
+        }
+      }
+    }
+    // Update dependency users
+    for (DepDesc &Dep : Cmd->MDeps) {
+      Command *DepCmd = Dep.MDepCommand;
+      DepCmd->MUsers.erase(Cmd);
+    }
+    Cmd->getEvent()->setCommand(nullptr);
+    delete Cmd;
+  }
 }
 
 void Scheduler::GraphBuilder::removeRecordForMemObj(SYCLMemObjI *MemObject) {

sycl/source/detail/scheduler/scheduler.cpp

@@ -117,6 +117,11 @@ void Scheduler::waitForEvent(EventImplPtr Event) {
   GraphProcessor::waitForEvent(std::move(Event));
 }
 
+void Scheduler::cleanupFinishedCommands(Command *FinishedCmd) {
+  std::lock_guard<std::mutex> lock(MGraphLock);
+  MGraphBuilder.cleanupFinishedCommands(FinishedCmd);
+}
+
 void Scheduler::removeMemoryObject(detail::SYCLMemObjI *MemObj) {
   std::lock_guard<std::mutex> lock(MGraphLock);
 
@@ -125,6 +130,7 @@ void Scheduler::removeMemoryObject(detail::SYCLMemObjI *MemObj) {
     // No operations were performed on the mem object
     return;
   waitForRecordToFinish(Record);
+  MGraphBuilder.updateLeafCountersForRecord(Record);
   MGraphBuilder.cleanupCommandsForRecord(Record);
   MGraphBuilder.removeRecordForMemObj(MemObj);
 }

sycl/test/scheduler/LeafLimit.cpp

@@ -2,6 +2,7 @@
 // RUN: %t.out
 #include <CL/sycl.hpp>
 
+#include <algorithm>
 #include <cstddef>
 #include <memory>
 #include <vector>
@@ -70,7 +71,9 @@ int main() {
   FakeCommand *OldestLeaf = LeavesToAdd.front();
   FakeCommand *NewestLeaf = LeavesToAdd.back();
   assert(OldestLeaf->MUsers.size() == 1);
-  assert(OldestLeaf->MUsers[0] == NewestLeaf);
+  assert(OldestLeaf->MUsers.count(NewestLeaf));
   assert(NewestLeaf->MDeps.size() == 2);
-  assert(NewestLeaf->MDeps[1].MDepCommand == OldestLeaf);
+  assert(std::any_of(
+      NewestLeaf->MDeps.begin(), NewestLeaf->MDeps.end(),
+      [&](const detail::DepDesc &DD) { return DD.MDepCommand == OldestLeaf; }));
 }