[NFC][SYCL][Graph] Use raw node_impl ptr/ref in node<->event mapping (#19366)

Continuation of
intel/llvm#19295
intel/llvm#19332
intel/llvm#19334
intel/llvm#19350
intel/llvm#19352

Author: aelovikov-intel

sycl/source/detail/async_alloc.cpp

@@ -33,7 +33,7 @@ getUrEvents(const std::vector<std::shared_ptr<detail::event_impl>> &DepEvents) {
   return RetUrEvents;
 }
 
-std::vector<std::shared_ptr<detail::node_impl>> getDepGraphNodes(
+std::vector<detail::node_impl *> getDepGraphNodes(
     sycl::handler &Handler, detail::queue_impl *Queue,
     const std::shared_ptr<detail::graph_impl> &Graph,
     const std::vector<std::shared_ptr<detail::event_impl>> &DepEvents) {
@@ -42,14 +42,14 @@ std::vector<std::shared_ptr<detail::node_impl>> getDepGraphNodes(
   auto DepNodes = Graph->getNodesForEvents(DepEvents);
   // If this node was added explicitly we may have node deps in the handler as
   // well, so add them to the list
-  DepNodes.insert(DepNodes.end(), HandlerImpl.MNodeDeps.begin(),
-                  HandlerImpl.MNodeDeps.end());
+  for (auto &N : HandlerImpl.MNodeDeps)
+    DepNodes.push_back(N.get());
   // If this is being recorded from an in-order queue we need to get the last
   // in-order node if any, since this will later become a dependency of the
   // node being processed here.
   if (detail::node_impl *LastInOrderNode = Graph->getLastInorderNode(Queue);
       LastInOrderNode) {
-    DepNodes.push_back(LastInOrderNode->shared_from_this());
+    DepNodes.push_back(LastInOrderNode);
   }
   return DepNodes;
 }

sycl/source/detail/graph/graph_impl.cpp

@@ -369,7 +369,7 @@ std::set<std::shared_ptr<node_impl>> graph_impl::getCGEdges(
                             "Event dependency from handler::depends_on does "
                             "not correspond to a node within the graph");
     } else {
-      UniqueDeps.insert(NodeImpl->second);
+      UniqueDeps.insert(NodeImpl->second->shared_from_this());
     }
   }
 
@@ -417,7 +417,7 @@ std::shared_ptr<node_impl> graph_impl::add(nodes_range Deps) {
   addDepsToNode(NodeImpl, Deps);
   // Add an event associated with this explicit node for mixed usage
   addEventForNode(sycl::detail::event_impl::create_completed_host_event(),
-                  NodeImpl);
+                  *NodeImpl);
   return NodeImpl;
 }
 
@@ -476,7 +476,7 @@ graph_impl::add(std::function<void(handler &)> CGF,
 
   // Add an event associated with this explicit node for mixed usage
   addEventForNode(sycl::detail::event_impl::create_completed_host_event(),
-                  NodeImpl);
+                  *NodeImpl);
 
   // Retrieve any dynamic parameters which have been registered in the CGF and
   // register the actual nodes with them.
@@ -556,7 +556,7 @@ graph_impl::add(std::shared_ptr<dynamic_command_group_impl> &DynCGImpl,
 
   // Add an event associated with this explicit node for mixed usage
   addEventForNode(sycl::detail::event_impl::create_completed_host_event(),
-                  NodeImpl);
+                  *NodeImpl);
 
   // Track the dynamic command-group used inside the node object
   DynCGImpl->MNodes.push_back(NodeImpl);
@@ -689,9 +689,9 @@ std::vector<sycl::detail::EventImplPtr> graph_impl::getExitNodesEvents(
   auto RecordedQueueSP = RecordedQueue.lock();
   for (auto &Node : MNodeStorage) {
     if (Node->MSuccessors.empty()) {
-      auto EventForNode = getEventForNode(Node);
+      auto EventForNode = getEventForNode(*Node);
       if (EventForNode->getSubmittedQueue() == RecordedQueueSP) {
-        Events.push_back(getEventForNode(Node));
+        Events.push_back(getEventForNode(*Node));
       }
     }
   }

sycl/source/detail/graph/graph_impl.hpp

@@ -192,21 +192,21 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// @param EventImpl Event to associate with a node in map.
   /// @param NodeImpl Node to associate with event in map.
   void addEventForNode(std::shared_ptr<sycl::detail::event_impl> EventImpl,
-                       const std::shared_ptr<node_impl> &NodeImpl) {
+                       node_impl &NodeImpl) {
     if (!(EventImpl->hasCommandGraph()))
       EventImpl->setCommandGraph(shared_from_this());
-    MEventsMap[EventImpl] = NodeImpl;
+    MEventsMap[EventImpl] = &NodeImpl;
   }
 
   /// Find the sycl event associated with a node.
   /// @param NodeImpl Node to find event for.
   /// @return Event associated with node.
   std::shared_ptr<sycl::detail::event_impl>
-  getEventForNode(std::shared_ptr<node_impl> NodeImpl) const {
+  getEventForNode(node_impl &NodeImpl) const {
     ReadLock Lock(MMutex);
     if (auto EventImpl = std::find_if(
             MEventsMap.begin(), MEventsMap.end(),
-            [NodeImpl](auto &it) { return it.second == NodeImpl; });
+            [&NodeImpl](auto &it) { return it.second == &NodeImpl; });
         EventImpl != MEventsMap.end()) {
       return EventImpl->first;
     }
@@ -220,13 +220,14 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// the given event.
   /// @param EventImpl Event to find the node for.
   /// @return Node associated with the event.
-  std::shared_ptr<node_impl>
+  node_impl &
   getNodeForEvent(std::shared_ptr<sycl::detail::event_impl> EventImpl) {
     ReadLock Lock(MMutex);
 
     if (auto NodeFound = MEventsMap.find(EventImpl);
         NodeFound != std::end(MEventsMap)) {
-      return NodeFound->second;
+      // TODO: Is it guaranteed to be non-null?
+      return *NodeFound->second;
     }
 
     throw sycl::exception(
@@ -238,9 +239,9 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// found for a given event.
   /// @param Events Events to find nodes for.
   /// @return A list of node counterparts for each event, in the same order.
-  std::vector<std::shared_ptr<node_impl>> getNodesForEvents(
+  std::vector<node_impl *> getNodesForEvents(
       const std::vector<std::shared_ptr<sycl::detail::event_impl>> &Events) {
-    std::vector<std::shared_ptr<node_impl>> NodeList{};
+    std::vector<node_impl *> NodeList{};
     NodeList.reserve(Events.size());
 
     ReadLock Lock(MMutex);
@@ -544,8 +545,7 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
            std::owner_less<std::weak_ptr<sycl::detail::queue_impl>>>
       MRecordingQueues;
   /// Map of events to their associated recorded nodes.
-  std::unordered_map<std::shared_ptr<sycl::detail::event_impl>,
-                     std::shared_ptr<node_impl>>
+  std::unordered_map<std::shared_ptr<sycl::detail::event_impl>, node_impl *>
       MEventsMap;
   /// Map for every in-order queue thats recorded a node to the graph, what
   /// the last node added was. We can use this to create new edges on the last

sycl/source/handler.cpp

@@ -916,7 +916,7 @@ event handler::finalize() {
     }
 
     // Associate an event with this new node and return the event.
-    GraphImpl->addEventForNode(EventImpl, std::move(NodeImpl));
+    GraphImpl->addEventForNode(EventImpl, *NodeImpl);
 
 #ifdef __INTEL_PREVIEW_BREAKING_CHANGES
     return EventImpl;