[NFC][SYCL][Graph] Introduce nodes_range utility

sycl/source/detail/graph/graph_impl.cpp

@@ -452,8 +452,7 @@ void graph_impl::markCGMemObjs(
   }
 }
 
-std::shared_ptr<node_impl>
-graph_impl::add(std::vector<std::shared_ptr<node_impl>> &Deps) {
+std::shared_ptr<node_impl> graph_impl::add(nodes_range Deps) {
   const std::shared_ptr<node_impl> &NodeImpl = std::make_shared<node_impl>();
 
   MNodeStorage.push_back(NodeImpl);
@@ -542,12 +541,12 @@ graph_impl::add(std::function<void(handler &)> CGF,
 std::shared_ptr<node_impl>
 graph_impl::add(const std::vector<sycl::detail::EventImplPtr> Events) {
 
-  std::vector<std::shared_ptr<node_impl>> Deps;
+  std::vector<node_impl *> Deps;
 
   // Add any nodes specified by event dependencies into the dependency list
   for (const auto &Dep : Events) {
     if (auto NodeImpl = MEventsMap.find(Dep); NodeImpl != MEventsMap.end()) {
-      Deps.push_back(NodeImpl->second);
+      Deps.push_back(NodeImpl->second.get());
     } else {
       throw sycl::exception(sycl::make_error_code(errc::invalid),
                             "Event dependency from handler::depends_on does "
@@ -561,23 +560,22 @@ graph_impl::add(const std::vector<sycl::detail::EventImplPtr> Events) {
 std::shared_ptr<node_impl>
 graph_impl::add(node_type NodeType,
                 std::shared_ptr<sycl::detail::CG> CommandGroup,
-                std::vector<std::shared_ptr<node_impl>> &Deps) {
+                nodes_range Deps) {
 
   // A unique set of dependencies obtained by checking requirements and events
   std::set<std::shared_ptr<node_impl>> UniqueDeps = getCGEdges(CommandGroup);
 
   // Track and mark the memory objects being used by the graph.
   markCGMemObjs(CommandGroup);
 
-  // Add any deps determined from requirements and events into the dependency
-  // list
-  Deps.insert(Deps.end(), UniqueDeps.begin(), UniqueDeps.end());
-
   const std::shared_ptr<node_impl> &NodeImpl =
       std::make_shared<node_impl>(NodeType, std::move(CommandGroup));
   MNodeStorage.push_back(NodeImpl);
 
+  // Add any deps determined from requirements and events into the dependency
+  // list
   addDepsToNode(NodeImpl, Deps);
+  addDepsToNode(NodeImpl, UniqueDeps);
 
   if (NodeType == node_type::async_free) {
     auto AsyncFreeCG =
@@ -592,7 +590,7 @@ graph_impl::add(node_type NodeType,
 
 std::shared_ptr<node_impl>
 graph_impl::add(std::shared_ptr<dynamic_command_group_impl> &DynCGImpl,
-                std::vector<std::shared_ptr<detail::node_impl>> &Deps) {
+                nodes_range Deps) {
   // Set of Dependent nodes based on CG event and accessor dependencies.
   std::set<std::shared_ptr<node_impl>> DynCGDeps =
       getCGEdges(DynCGImpl->MCommandGroups[0]);

sycl/source/detail/graph/graph_impl.hpp

@@ -147,7 +147,7 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// @return Created node in the graph.
   std::shared_ptr<node_impl> add(node_type NodeType,
                                  std::shared_ptr<sycl::detail::CG> CommandGroup,
-                                 std::vector<std::shared_ptr<node_impl>> &Deps);
+                                 nodes_range Deps);
 
   /// Create a CGF node in the graph.
   /// @param CGF Command-group function to create node with.
@@ -161,7 +161,7 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// Create an empty node in the graph.
   /// @param Deps List of predecessor nodes.
   /// @return Created node in the graph.
-  std::shared_ptr<node_impl> add(std::vector<std::shared_ptr<node_impl>> &Deps);
+  std::shared_ptr<node_impl> add(nodes_range Deps);
 
   /// Create an empty node in the graph.
   /// @param Events List of events associated to this node.
@@ -174,8 +174,7 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// @param Deps List of predecessor nodes.
   /// @return Created node in the graph.
   std::shared_ptr<node_impl>
-  add(std::shared_ptr<dynamic_command_group_impl> &DynCGImpl,
-      std::vector<std::shared_ptr<node_impl>> &Deps);
+  add(std::shared_ptr<dynamic_command_group_impl> &DynCGImpl, nodes_range Deps);
 
   /// Add a queue to the set of queues which are currently recording to this
   /// graph.
@@ -542,14 +541,12 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// added as a root node.
   /// @param Node The node to add deps for
   /// @param Deps List of dependent nodes
-  void addDepsToNode(const std::shared_ptr<node_impl> &Node,
-                     std::vector<std::shared_ptr<node_impl>> &Deps) {
-    if (!Deps.empty()) {
-      for (auto &N : Deps) {
-        N->registerSuccessor(Node);
-        this->removeRoot(Node);
-      }
-    } else {
+  void addDepsToNode(const std::shared_ptr<node_impl> &Node, nodes_range Deps) {
+    for (node_impl &N : Deps) {
+      N.registerSuccessor(Node);
+      this->removeRoot(Node);
+    }
+    if (Node->MPredecessors.empty()) {
       this->addRoot(Node);
     }
   }

sycl/source/detail/graph/node_impl.hpp

@@ -14,10 +14,11 @@
 #include <sycl/detail/cg_types.hpp> // for CGType
 #include <sycl/detail/kernel_desc.hpp> // for kernel_param_kind_t
 
-#include <cstring> // for memcpy
-#include <fstream> // for fstream, ostream
-#include <iomanip> // for setw, setfill
-#include <vector>  // for vector
+#include <cstring>
+#include <fstream>
+#include <iomanip>
+#include <set>
+#include <vector>
 
 namespace sycl {
 inline namespace _V1 {
@@ -753,6 +754,82 @@ class node_impl : public std::enable_shared_from_this<node_impl> {
     return std::make_unique<CGT>(*static_cast<CGT *>(MCommandGroup.get()));
   }
 };
+
+// Non-owning!
+class nodes_range {
+  template <typename... Containers>
+  using storage_iter_impl =
+      std::variant<typename Containers::const_iterator...>;
+
+  using storage_iter = storage_iter_impl<
+      std::vector<std::shared_ptr<node_impl>>, std::vector<node_impl *>,
+      // Next one is temporary. It looks like `weak_ptr`s aren't
+      // used for the actual lifetime management and the objects are
+      // always guaranteed to be alive. Once the code is cleaned
+      // from `weak_ptr`s this alternative should be removed too.
+      std::vector<std::weak_ptr<node_impl>>,
+      //
+      std::set<std::shared_ptr<node_impl>>>;
+
+  storage_iter Begin;
+  storage_iter End;
+  const size_t Size;
+
+public:
+  nodes_range(const nodes_range &Other) = default;
+
+  template <
+      typename ContainerTy,
+      typename = std::enable_if_t<!std::is_same_v<nodes_range, ContainerTy>>>
+  nodes_range(ContainerTy &Container)
+      : Begin{Container.begin()}, End{Container.end()}, Size{Container.size()} {
+  }
+
+  class iterator {
+    storage_iter It;
+
+    iterator(storage_iter It) : It(It) {}
+    friend class nodes_range;
+
+  public:
+    iterator &operator++() {
+      It = std::visit(
+          [](auto &&It) {
+            ++It;
+            return storage_iter{It};
+          },
+          It);
+      return *this;
+    }
+    bool operator!=(const iterator &Other) const { return It != Other.It; }
+
+    node_impl &operator*() {
+      return std::visit(
+          [](auto &&It) -> node_impl & {
+            auto &Elem = *It;
+            if constexpr (std::is_same_v<std::decay_t<decltype(Elem)>,
+                                         std::weak_ptr<node_impl>>) {
+              // This assumes that weak_ptr doesn't actually manage lifetime and
+              // the object is guaranteed to be alive (which seems to be the
+              // assumption across all graph code).
+              return *Elem.lock();
+            } else {
+              return *Elem;
+            }
+          },
+          It);
+    }
+  };
+
+  iterator begin() const {
+    return {std::visit([](auto &&It) { return storage_iter{It}; }, Begin)};
+  }
+  iterator end() const {
+    return {std::visit([](auto &&It) { return storage_iter{It}; }, End)};
+  }
+  size_t size() const { return Size; }
+  bool empty() const { return Size == 0; }
+};
 } // namespace detail
 } // namespace experimental
 } // namespace oneapi