[NFC][SYCL][Graph] Use node_impl &/nodes_range in memory_pool

sycl/source/detail/graph/graph_impl.cpp

@@ -527,7 +527,7 @@ graph_impl::add(node_type NodeType,
         static_cast<CGAsyncFree *>(NodeImpl->MCommandGroup.get());
     // If this is an async free node mark that it is now available for reuse,
     // and pass the async free node for tracking.
-    MGraphMemPool.markAllocationAsAvailable(AsyncFreeCG->getPtr(), NodeImpl);
+    MGraphMemPool.markAllocationAsAvailable(AsyncFreeCG->getPtr(), *NodeImpl);
   }
 
   return NodeImpl;

sycl/source/detail/graph/memory_pool.cpp

@@ -19,10 +19,8 @@ namespace oneapi {
 namespace experimental {
 namespace detail {
 
-void *
-graph_mem_pool::malloc(size_t Size, usm::alloc AllocType,
-                       const std::vector<std::shared_ptr<node_impl>> &DepNodes,
-                       memory_pool_impl *MemPool) {
+void *graph_mem_pool::malloc(size_t Size, usm::alloc AllocType,
+                             nodes_range DepNodes, memory_pool_impl *MemPool) {
   // We are potentially modifying contents of this memory pool and the owning
   // graph, so take a lock here.
   graph_impl::WriteLock Lock(MGraph.MMutex);
@@ -81,9 +79,9 @@ graph_mem_pool::malloc(size_t Size, usm::alloc AllocType,
 }
 
 std::optional<graph_mem_pool::alloc_info>
-graph_mem_pool::tryReuseExistingAllocation(
-    size_t Size, usm::alloc AllocType, bool ReadOnly,
-    const std::vector<std::shared_ptr<node_impl>> &DepNodes) {
+graph_mem_pool::tryReuseExistingAllocation(size_t Size, usm::alloc AllocType,
+                                           bool ReadOnly,
+                                           nodes_range DepNodes) {
   // If we have no dependencies this is a no-op because allocations must connect
   // to a free node for reuse to be possible.
   if (DepNodes.empty()) {
@@ -119,8 +117,8 @@ graph_mem_pool::tryReuseExistingAllocation(
   std::queue<node_impl *> NodesToCheck;
 
   // Add all the dependent nodes to the queue, they will be popped first
-  for (auto &Dep : DepNodes) {
-    NodesToCheck.push(&*Dep);
+  for (node_impl &Dep : DepNodes) {
+    NodesToCheck.push(&Dep);
   }
 
   // Called when traversing over nodes to check if the current node is a free
@@ -175,10 +173,9 @@ graph_mem_pool::tryReuseExistingAllocation(
   return std::nullopt;
 }
 
-void graph_mem_pool::markAllocationAsAvailable(
-    void *Ptr, const std::shared_ptr<node_impl> &FreeNode) {
+void graph_mem_pool::markAllocationAsAvailable(void *Ptr, node_impl &FreeNode) {
   MFreeAllocations.push_back(Ptr);
-  MAllocations.at(Ptr).LastFreeNode = FreeNode.get();
+  MAllocations.at(Ptr).LastFreeNode = &FreeNode;
 }
 
 } // namespace detail

sycl/source/detail/graph/memory_pool.hpp

@@ -23,6 +23,7 @@ namespace detail {
 
 // Forward declarations
 class node_impl;
+class nodes_range;
 
 /// Class handling graph-owned memory allocations. Device allocations are
 /// managed using virtual memory.
@@ -82,8 +83,7 @@ class graph_mem_pool {
   /// @param MemPool Optional memory pool from which allocations will not be
   /// made directly but properties may be respected.
   /// @return A pointer to the start of the allocation
-  void *malloc(size_t Size, usm::alloc AllocType,
-               const std::vector<std::shared_ptr<node_impl>> &DepNodes,
+  void *malloc(size_t Size, usm::alloc AllocType, nodes_range DepNodes,
                memory_pool_impl *MemPool = nullptr);
 
   /// Return the total amount of memory being used by this pool
@@ -160,8 +160,7 @@ class graph_mem_pool {
   /// @param Ptr The pointer to the allocation.
   /// @param FreeNode The graph node of node_type::async_free which is freeing
   /// the allocation.
-  void markAllocationAsAvailable(void *Ptr,
-                                 const std::shared_ptr<node_impl> &FreeNode);
+  void markAllocationAsAvailable(void *Ptr, node_impl &FreeNode);
 
 private:
   /// Tries to reuse an existing allocation which has been marked free in the
@@ -173,9 +172,10 @@ class graph_mem_pool {
   /// reusable allocations.
   /// @returns An optional allocation info value, where a null value indicates
   /// that no allocation could be reused.
-  std::optional<alloc_info> tryReuseExistingAllocation(
-      size_t Size, usm::alloc AllocType, bool ReadOnly,
-      const std::vector<std::shared_ptr<node_impl>> &DepNodes);
+  std::optional<alloc_info> tryReuseExistingAllocation(size_t Size,
+                                                       usm::alloc AllocType,
+                                                       bool ReadOnly,
+                                                       nodes_range DepNodes);
 
   /// Returns an aligned byte size given a required granularity
   /// @param UnalignedByteSize The original requested allocation size