[Sycl][Graph] Reimplement topological sort algorithm (#17495)

The current implementation of the topological sort algorithm relies on a
recursive function. This causes stack overflow issues when the graphs
are very large. There is also a performance issue in the current
implementation related to the way std::find() is being used.

This commit reimplements the topological sort algorithm to fix these
issues:

- Uses an iterative approach instead of recursive.
- Removes the use of std::find() and instead relies on keeping the state
with a helper variable.
- Reimplements the cycle checking algorithm to use the same topological
sort function as the one used during scheduling.
- Fixes a bug with make_edge() not removing the dest node from the list
of root nodes.

Author: fabiomestre

sycl/doc/design/CommandGraph.md

@@ -100,11 +100,15 @@ Edges are stored in each node as lists of predecessor and successor nodes.
 
 ## Execution Order
 
-The current way graph nodes are linearized into execution order is using a
-reversed depth-first sorting algorithm. Alternative algorithms, such as
-breadth-first, are possible and may give better performance on certain
-workloads/hardware. In the future there might be options for allowing the
-user to control this implementation detail.
+Graph nodes are currently linearized into execution order using a topological
+sort algorithm. This algorithm uses a breadth-first search approach
+and is an implementation of Kahn's algorithm. Alternative algorithms, such as
+depth-first search, are possible and may give better performance on certain
+workloads/hardware. However, depth first searches are usually implemented 
+using recursion, which can lead to stack overflow issues on large graphs.
+In the future, there might be options for allowing the user to control this
+implementation detail. It might also be possible to automatically change
+which algorithm is used based on characteristics such as the graph's size.
 
 ## Scheduler Integration
 

sycl/source/detail/graph_impl.cpp

@@ -8,6 +8,7 @@
 
 #define __SYCL_GRAPH_IMPL_CPP
 
+#include <stack>
 #include <detail/graph_impl.hpp>
 #include <detail/handler_impl.hpp>
 #include <detail/kernel_arg_mask.hpp>
@@ -31,64 +32,47 @@ namespace experimental {
 namespace detail {
 
 namespace {
-/// Visits a node on the graph and it's successors recursively in a depth-first
-/// approach.
-/// @param[in] Node The current node being visited.
-/// @param[in,out] VisitedNodes A set of unique nodes which have already been
-/// visited.
-/// @param[in] NodeStack Stack of nodes which are currently being visited on the
-/// current path through the graph.
-/// @param[in] NodeFunc The function object to be run on each node. A return
-/// value of true indicates the search should be ended immediately and the
-/// function will return.
-/// @return True if the search should end immediately, false if not.
-bool visitNodeDepthFirst(
-    std::shared_ptr<node_impl> Node,
-    std::set<std::shared_ptr<node_impl>> &VisitedNodes,
-    std::deque<std::shared_ptr<node_impl>> &NodeStack,
-    std::function<bool(std::shared_ptr<node_impl> &,
-                       std::deque<std::shared_ptr<node_impl>> &)>
-        NodeFunc) {
-  auto EarlyReturn = NodeFunc(Node, NodeStack);
-  if (EarlyReturn) {
-    return true;
-  }
-  NodeStack.push_back(Node);
-  Node->MVisited = true;
-  VisitedNodes.emplace(Node);
-  for (auto &Successor : Node->MSuccessors) {
-    if (visitNodeDepthFirst(Successor.lock(), VisitedNodes, NodeStack,
-                            NodeFunc)) {
-      return true;
-    }
-  }
-  NodeStack.pop_back();
-  return false;
-}
-
-/// Recursively add nodes to execution stack.
-/// @param NodeImpl Node to schedule.
-/// @param Schedule Execution ordering to add node to.
-/// @param PartitionBounded If set to true, the topological sort is stopped at
-/// partition borders. Hence, nodes belonging to a partition different from the
-/// NodeImpl partition are not processed.
-void sortTopological(std::shared_ptr<node_impl> NodeImpl,
-                     std::list<std::shared_ptr<node_impl>> &Schedule,
-                     bool PartitionBounded = false) {
-  for (auto &Succ : NodeImpl->MSuccessors) {
-    auto NextNode = Succ.lock();
-    if (PartitionBounded &&
-        (NextNode->MPartitionNum != NodeImpl->MPartitionNum)) {
-      continue;
-    }
-    // Check if we've already scheduled this node
-    if (std::find(Schedule.begin(), Schedule.end(), NextNode) ==
-        Schedule.end()) {
-      sortTopological(NextNode, Schedule, PartitionBounded);
+/// Topologically sorts the graph in order to schedule nodes for execution.
+/// This implementation is based on Kahn's algorithm which uses a Breadth-first
+/// search approach.
+/// For performance reasons, this function uses the MTotalVisitedEdges
+/// member variable of the node_impl class. It's the caller responsibility to
+/// make sure that MTotalVisitedEdges is set to 0 for all nodes in the graph
+/// before calling this function.
+/// @param[in] Roots List of root nodes.
+/// @param[out] SortedNodes The graph nodes sorted in topological order.
+/// @param[in] PartitionBounded If set to true, the topological sort is stopped
+/// at partition borders. Hence, nodes belonging to a partition different from
+/// the NodeImpl partition are not processed.
+void sortTopological(std::set<std::weak_ptr<node_impl>,
+                              std::owner_less<std::weak_ptr<node_impl>>> &Roots,
+                     std::list<std::shared_ptr<node_impl>> &SortedNodes,
+                     bool PartitionBounded) {
+  std::stack<std::weak_ptr<node_impl>> Source;
+
+  for (auto &Node : Roots) {
+    Source.push(Node);
+  }
+
+  while (!Source.empty()) {
+    auto Node = Source.top().lock();
+    Source.pop();
+    SortedNodes.push_back(Node);
+
+    for (auto &SuccWP : Node->MSuccessors) {
+      auto Succ = SuccWP.lock();
+
+      if (PartitionBounded && (Succ->MPartitionNum != Node->MPartitionNum)) {
+        continue;
+      }
+
+      auto &TotalVisitedEdges = Succ->MTotalVisitedEdges;
+      ++TotalVisitedEdges;
+      if (TotalVisitedEdges == Succ->MPredecessors.size()) {
+        Source.push(Succ);
+      }
     }
   }
-
-  Schedule.push_front(NodeImpl);
 }
 
 /// Propagates the partition number `PartitionNum` to predecessors.
@@ -180,9 +164,9 @@ std::vector<node> createNodesFromImpls(
 
 void partition::schedule() {
   if (MSchedule.empty()) {
-    for (auto &Node : MRoots) {
-      sortTopological(Node.lock(), MSchedule, true);
-    }
+    // There is no need to reset MTotalVisitedEdges before calling
+    // sortTopological because this function is only called once per partition.
+    sortTopological(MRoots, MSchedule, true);
   }
 }
 
@@ -311,6 +295,7 @@ static void checkGraphPropertiesAndThrow(const property_list &Properties) {
 #define __SYCL_MANUALLY_DEFINED_PROP(NS_QUALIFIER, PROP_NAME)
     switch (PropertyKind) {
 #include <sycl/ext/oneapi/experimental/detail/properties/graph_properties.def>
+
     default:
       return false;
     }
@@ -627,44 +612,20 @@ bool graph_impl::clearQueues() {
   return AnyQueuesCleared;
 }
 
-void graph_impl::searchDepthFirst(
-    std::function<bool(std::shared_ptr<node_impl> &,
-                       std::deque<std::shared_ptr<node_impl>> &)>
-        NodeFunc) {
-  // Track nodes visited during the search which can be used by NodeFunc in
-  // depth first search queries. Currently unusued but is an
-  // integral part of depth first searches.
-  std::set<std::shared_ptr<node_impl>> VisitedNodes;
+bool graph_impl::checkForCycles() {
+  std::list<std::shared_ptr<node_impl>> SortedNodes;
+  sortTopological(MRoots, SortedNodes, false);
 
-  for (auto &Root : MRoots) {
-    std::deque<std::shared_ptr<node_impl>> NodeStack;
-    if (visitNodeDepthFirst(Root.lock(), VisitedNodes, NodeStack, NodeFunc)) {
-      break;
-    }
-  }
+  // If after a topological sort, not all the nodes in the graph are sorted,
+  // then there must be at least one cycle in the graph. This is guaranteed
+  // by Kahn's algorithm, which sortTopological() implements.
+  bool CycleFound = SortedNodes.size() != MNodeStorage.size();
 
-  // Reset the visited status of all nodes encountered in the search.
-  for (auto &Node : VisitedNodes) {
-    Node->MVisited = false;
+  // Reset the MTotalVisitedEdges variable to prepare for the next cycle check.
+  for (auto &Node : MNodeStorage) {
+    Node->MTotalVisitedEdges = 0;
   }
-}
 
-bool graph_impl::checkForCycles() {
-  // Using a depth-first search and checking if we vist a node more than once in
-  // the current path to identify if there are cycles.
-  bool CycleFound = false;
-  auto CheckFunc = [&](std::shared_ptr<node_impl> &Node,
-                       std::deque<std::shared_ptr<node_impl>> &NodeStack) {
-    // If the current node has previously been found in the current path through
-    // the graph then we have a cycle and we end the search early.
-    if (std::find(NodeStack.begin(), NodeStack.end(), Node) !=
-        NodeStack.end()) {
-      CycleFound = true;
-      return true;
-    }
-    return false;
-  };
-  searchDepthFirst(CheckFunc);
   return CycleFound;
 }
 
@@ -698,19 +659,31 @@ void graph_impl::makeEdge(std::shared_ptr<node_impl> Src,
                           "Dest must be a node inside the graph.");
   }
 
+  bool DestWasGraphRoot = Dest->MPredecessors.size() == 0;
+
   // We need to add the edges first before checking for cycles
   Src->registerSuccessor(Dest);
 
+  bool DestLostRootStatus = DestWasGraphRoot && Dest->MPredecessors.size() == 1;
+  if (DestLostRootStatus) {
+    // Dest is no longer a Root node, so we need to remove it from MRoots.
+    MRoots.erase(Dest);
+  }
+
   // We can skip cycle checks if either Dest has no successors (cycle not
   // possible) or cycle checks have been disabled with the no_cycle_check
   // property;
   if (Dest->MSuccessors.empty() || !MSkipCycleChecks) {
     bool CycleFound = checkForCycles();
 
     if (CycleFound) {
-      // Remove the added successor and predecessor
+      // Remove the added successor and predecessor.
       Src->MSuccessors.pop_back();
       Dest->MPredecessors.pop_back();
+      if (DestLostRootStatus) {
+        // Add Dest back into MRoots.
+        MRoots.insert(Dest);
+      }
 
       throw sycl::exception(make_error_code(sycl::errc::invalid),
                             "Command graphs cannot contain cycles.");

sycl/source/detail/graph_impl.hpp

@@ -102,8 +102,8 @@ class node_impl : public std::enable_shared_from_this<node_impl> {
   /// subgraph node.
   std::shared_ptr<exec_graph_impl> MSubGraphImpl;
 
-  /// Used for tracking visited status during cycle checks.
-  bool MVisited = false;
+  /// Used for tracking visited status during cycle checks and node scheduling.
+  size_t MTotalVisitedEdges = 0;
 
   /// Partition number needed to assign a Node to a a partition.
   /// Note : This number is only used during the partitionning process and
@@ -1130,17 +1130,6 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   unsigned long long getID() const { return MID; }
 
 private:
-  /// Iterate over the graph depth-first and run \p NodeFunc on each node.
-  /// @param NodeFunc A function which receives as input a node in the graph to
-  /// perform operations on as well as the stack of nodes encountered in the
-  /// current path. The return value of this function determines whether an
-  /// early exit is triggered, if true the depth-first search will end
-  /// immediately and no further nodes will be visited.
-  void
-  searchDepthFirst(std::function<bool(std::shared_ptr<node_impl> &,
-                                      std::deque<std::shared_ptr<node_impl>> &)>
-                       NodeFunc);
-
   /// Check the graph for cycles by performing a depth-first search of the
   /// graph. If a node is visited more than once in a given path through the
   /// graph, a cycle is present and the search ends immediately.

sycl/unittests/Extensions/CommandGraph/CMakeLists.txt

@@ -10,6 +10,7 @@ add_sycl_unittest(CommandGraphExtensionTests OBJECT
   Queries.cpp
   Regressions.cpp
   Subgraph.cpp
+  TopologicalSort.cpp
   Update.cpp
   Properties.cpp
 )