[NFC][SYCL][Graph] Add iterator_range::to<Container>() helper (#19431)

Author: aelovikov-intel

sycl/source/detail/device_impl.hpp

@@ -2295,7 +2295,7 @@ struct devices_deref_impl {
   }
 };
 using devices_iterator =
-    variadic_iterator<devices_deref_impl,
+    variadic_iterator<devices_deref_impl, device,
                       std::vector<std::shared_ptr<device_impl>>::const_iterator,
                       std::vector<device>::const_iterator, device_impl *>;
 

sycl/source/detail/graph/graph_impl.cpp

@@ -1970,15 +1970,11 @@ void modifiable_command_graph::print_graph(sycl::detail::string_view pathstr,
 
 std::vector<node> modifiable_command_graph::get_nodes() const {
   graph_impl::ReadLock Lock(impl->MMutex);
-  return createNodesFromImpls(impl->MNodeStorage);
+  return impl->nodes().to<std::vector<node>>();
 }
 std::vector<node> modifiable_command_graph::get_root_nodes() const {
   graph_impl::ReadLock Lock(impl->MMutex);
-  auto &Roots = impl->MRoots;
-  std::vector<node_impl *> Impls{};
-
-  std::copy(Roots.begin(), Roots.end(), std::back_inserter(Impls));
-  return createNodesFromImpls(Impls);
+  return impl->roots().to<std::vector<node>>();
 }
 
 void modifiable_command_graph::checkNodePropertiesAndThrow(

sycl/source/detail/graph/graph_impl.hpp

@@ -293,6 +293,7 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   std::vector<std::shared_ptr<node_impl>> MNodeStorage;
 
   nodes_range roots() const { return MRoots; }
+  nodes_range nodes() const { return MNodeStorage; }
 
   /// Find the last node added to this graph from an in-order queue.
   /// @param Queue In-order queue to find the last node added to the graph from.
@@ -665,6 +666,8 @@ class exec_graph_impl {
     return MPartitions;
   }
 
+  nodes_range nodes() const { return MNodeStorage; }
+
   /// Query whether the graph contains any host-task nodes.
   /// @return True if the graph contains any host-task nodes. False otherwise.
   bool containsHostTask() const { return MContainsHostTask; }

sycl/source/detail/graph/memory_pool.cpp

@@ -114,12 +114,8 @@ graph_mem_pool::tryReuseExistingAllocation(size_t Size, usm::alloc AllocType,
   // free nodes. We do this in a breadth-first approach because we want to find
   // the shortest path to a reusable allocation.
 
-  std::queue<node_impl *> NodesToCheck;
-
   // Add all the dependent nodes to the queue, they will be popped first
-  for (node_impl &Dep : DepNodes) {
-    NodesToCheck.push(&Dep);
-  }
+  auto NodesToCheck = DepNodes.to<std::queue<node_impl *>>();
 
   // Called when traversing over nodes to check if the current node is a free
   // node for one of the available allocations. If it is we populate AllocInfo

sycl/source/detail/graph/node_impl.cpp

@@ -15,42 +15,14 @@ inline namespace _V1 {
 namespace ext {
 namespace oneapi {
 namespace experimental {
-namespace detail {
-
-/// Takes a vector of weak_ptrs to node_impls and returns a vector of node
-/// objects created from those impls, in the same order.
-std::vector<node> createNodesFromImpls(
-    const std::vector<std::weak_ptr<detail::node_impl>> &Impls) {
-  std::vector<node> Nodes{};
-  Nodes.reserve(Impls.size());
-
-  for (std::weak_ptr<detail::node_impl> Impl : Impls) {
-    Nodes.push_back(sycl::detail::createSyclObjFromImpl<node>(Impl.lock()));
-  }
-
-  return Nodes;
-}
-
-std::vector<node> createNodesFromImpls(nodes_range Impls) {
-  std::vector<node> Nodes{};
-  Nodes.reserve(Impls.size());
-
-  for (detail::node_impl &Impl : Impls) {
-    Nodes.push_back(sycl::detail::createSyclObjFromImpl<node>(Impl));
-  }
-
-  return Nodes;
-}
-} // namespace detail
-
 node_type node::get_type() const { return impl->MNodeType; }
 
 std::vector<node> node::get_predecessors() const {
-  return detail::createNodesFromImpls(impl->MPredecessors);
+  return impl->predecessors().to<std::vector<node>>();
 }
 
 std::vector<node> node::get_successors() const {
-  return detail::createNodesFromImpls(impl->MSuccessors);
+  return impl->successors().to<std::vector<node>>();
 }
 
 node node::get_node_from_event(event nodeEvent) {

sycl/source/detail/graph/node_impl.hpp

@@ -36,13 +36,6 @@ class node_impl;
 class nodes_range;
 class exec_graph_impl;
 
-/// Takes a vector of weak_ptrs to node_impls and returns a vector of node
-/// objects created from those impls, in the same order.
-std::vector<node>
-createNodesFromImpls(const std::vector<std::weak_ptr<node_impl>> &Impls);
-
-std::vector<node> createNodesFromImpls(nodes_range Impls);
-
 inline node_type getNodeTypeFromCG(sycl::detail::CGType CGType) {
   using sycl::detail::CG;
 
@@ -780,7 +773,7 @@ struct nodes_deref_impl {
 
 template <typename... ContainerTy>
 using nodes_iterator_impl =
-    variadic_iterator<nodes_deref_impl,
+    variadic_iterator<nodes_deref_impl, node,
                       typename ContainerTy::const_iterator...>;
 
 using nodes_iterator = nodes_iterator_impl<

sycl/source/detail/helpers.hpp

@@ -8,11 +8,16 @@
 
 #pragma once
 
+#include <sycl/detail/impl_utils.hpp>
 #include <sycl/detail/kernel_name_str_t.hpp>
+#include <sycl/detail/type_traits.hpp>
 
 #include <ur_api.h>
 
+#include <algorithm>
+#include <iterator>
 #include <memory>
+#include <queue>
 #include <tuple>
 #include <variant>
 #include <vector>
@@ -30,13 +35,26 @@ const RTDeviceBinaryImage *
 retrieveKernelBinary(queue_impl &Queue, KernelNameStrRefT KernelName,
                      CGExecKernel *CGKernel = nullptr);
 
-template <typename DereferenceImpl, typename... Iterators>
+template <typename DereferenceImpl, typename SyclTy, typename... Iterators>
 class variadic_iterator {
   using storage_iter = std::variant<Iterators...>;
 
   storage_iter It;
 
 public:
+  using iterator_category = std::forward_iterator_tag;
+  using difference_type = std::ptrdiff_t;
+  using reference = decltype(DereferenceImpl::dereference(
+      *std::declval<nth_type_t<0, Iterators...>>()));
+  using value_type = std::remove_reference_t<reference>;
+  using sycl_type = SyclTy;
+  using pointer = value_type *;
+  static_assert(std::is_same_v<reference, value_type &>);
+
+  variadic_iterator(const variadic_iterator &) = default;
+  variadic_iterator(variadic_iterator &&) = default;
+  variadic_iterator(variadic_iterator &) = default;
+
   template <typename IterTy>
   variadic_iterator(IterTy &&It) : It(std::forward<IterTy>(It)) {}
 
@@ -52,6 +70,9 @@ class variadic_iterator {
   bool operator!=(const variadic_iterator &Other) const {
     return It != Other.It;
   }
+  bool operator==(const variadic_iterator &Other) const {
+    return It == Other.It;
+  }
 
   decltype(auto) operator*() {
     return std::visit(
@@ -64,6 +85,16 @@ class variadic_iterator {
 
 // Non-owning!
 template <typename iterator> class iterator_range {
+  using value_type = typename iterator::value_type;
+  using sycl_type = typename iterator::sycl_type;
+
+  template <typename Container, typename = void>
+  struct has_reserve : public std::false_type {};
+  template <typename Container>
+  struct has_reserve<
+      Container, std::void_t<decltype(std::declval<Container>().reserve(1))>>
+      : public std::true_type {};
+
 public:
   iterator_range(const iterator_range &Other) = default;
 
@@ -81,6 +112,40 @@ template <typename iterator> class iterator_range {
   bool empty() const { return Size == 0; }
   decltype(auto) front() const { return *begin(); }
 
+  template <typename Container>
+  std::enable_if_t<
+      check_type_in_v<Container, std::vector<sycl_type>,
+                      std::queue<value_type *>, std::vector<value_type *>,
+                      std::vector<std::shared_ptr<value_type>>>,
+      Container>
+  to() const {
+    std::conditional_t<std::is_same_v<Container, std::queue<value_type *>>,
+                       typename std::queue<value_type *>::container_type,
+                       Container>
+        Result;
+    if constexpr (has_reserve<decltype(Result)>::value)
+      Result.reserve(size());
+    std::transform(
+        begin(), end(), std::back_inserter(Result), [](value_type &E) {
+          if constexpr (std::is_same_v<Container, std::vector<sycl_type>>)
+            return createSyclObjFromImpl<sycl_type>(E);
+          else if constexpr (std::is_same_v<
+                                 Container,
+                                 std::vector<std::shared_ptr<value_type>>>)
+            return E.shared_from_this();
+          else
+            return &E;
+        });
+    if constexpr (std::is_same_v<Container, decltype(Result)>)
+      return Result;
+    else
+      return Container{std::move(Result)};
+  }
+
+protected:
+  template <typename Container>
+  static constexpr bool has_reserve_v = has_reserve<Container>::value;
+
 private:
   iterator Begin;
   iterator End;

sycl/source/detail/scheduler/commands.cpp

@@ -3802,10 +3802,10 @@ bool ExecCGCommand::readyForCleanup() const {
 UpdateCommandBufferCommand::UpdateCommandBufferCommand(
     queue_impl *Queue,
     ext::oneapi::experimental::detail::exec_graph_impl *Graph,
-    std::vector<std::shared_ptr<ext::oneapi::experimental::detail::node_impl>>
-        Nodes)
+    ext::oneapi::experimental::detail::nodes_range Nodes)
     : Command(CommandType::UPDATE_CMD_BUFFER, Queue), MGraph(Graph),
-      MNodes(std::move(Nodes)) {}
+      MNodes(Nodes.to<std::vector<std::shared_ptr<
+                 ext::oneapi::experimental::detail::node_impl>>>()) {}
 
 ur_result_t UpdateCommandBufferCommand::enqueueImp() {
   waitForPreparedHostEvents();

sycl/source/detail/scheduler/commands.hpp

@@ -33,6 +33,7 @@ inline namespace _V1 {
 namespace ext::oneapi::experimental::detail {
 class exec_graph_impl;
 class node_impl;
+class nodes_range;
 } // namespace ext::oneapi::experimental::detail
 namespace detail {
 
@@ -723,8 +724,7 @@ class UpdateCommandBufferCommand : public Command {
   explicit UpdateCommandBufferCommand(
       queue_impl *Queue,
       ext::oneapi::experimental::detail::exec_graph_impl *Graph,
-      std::vector<std::shared_ptr<ext::oneapi::experimental::detail::node_impl>>
-          Nodes);
+      ext::oneapi::experimental::detail::nodes_range Nodes);
 
   void printDot(std::ostream &Stream) const final;
   void emitInstrumentationData() final;