From 6e8cdb18d36513e890075edc600e52b59f8cbdee Mon Sep 17 00:00:00 2001
From: Maxime France-Pillois <maxime.francepillois@codeplay.com>
Date: Tue, 5 Mar 2024 12:28:49 +0000
Subject: [PATCH] [SYCL][Graph] Improve handling of the events returned from
 graph with multiple partitions. (#12870)

Add attached events as dependencies when the graph is resubmitted.
Change a class member to a function local variable.

---------

Co-authored-by: Steffen Larsen <steffen.larsen@intel.com>
---
 sycl/source/detail/event_impl.hpp |  4 ++++
 sycl/source/detail/graph_impl.cpp | 16 ++++++++++++----
 sycl/source/detail/graph_impl.hpp |  3 ---
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/sycl/source/detail/event_impl.hpp b/sycl/source/detail/event_impl.hpp
index db268539136f9..ed5b9ef07ef91 100644
--- a/sycl/source/detail/event_impl.hpp
+++ b/sycl/source/detail/event_impl.hpp
@@ -289,6 +289,10 @@ class event_impl {
     return MEventFromSubmittedExecCommandBuffer;
   }
 
+  const std::vector<EventImplPtr> &getPostCompleteEvents() const {
+    return MPostCompleteEvents;
+  }
+
 protected:
   // When instrumentation is enabled emits trace event for event wait begin and
   // returns the telemetry event generated for the wait
diff --git a/sycl/source/detail/graph_impl.cpp b/sycl/source/detail/graph_impl.cpp
index bdfc90537b520..5b0bc3b8324e9 100644
--- a/sycl/source/detail/graph_impl.cpp
+++ b/sycl/source/detail/graph_impl.cpp
@@ -764,7 +764,9 @@ exec_graph_impl::enqueue(const std::shared_ptr<sycl::detail::queue_impl> &Queue,
                          sycl::detail::CG::StorageInitHelper CGData) {
   WriteLock Lock(MMutex);
 
-  std::vector<sycl::detail::EventImplPtr> PartitionEvents;
+  // Map of the partitions to their execution events
+  std::unordered_map<std::shared_ptr<partition>, sycl::detail::EventImplPtr>
+      PartitionsExecutionEvents;
 
   auto CreateNewEvent([&]() {
     auto NewEvent = std::make_shared<sycl::detail::event_impl>(Queue);
@@ -787,7 +789,7 @@ exec_graph_impl::enqueue(const std::shared_ptr<sycl::detail::queue_impl> &Queue,
     }
 
     for (auto const &DepPartition : CurrentPartition->MPredecessors) {
-      CGData.MEvents.push_back(MPartitionsExecutionEvents[DepPartition]);
+      CGData.MEvents.push_back(PartitionsExecutionEvents[DepPartition]);
     }
 
     auto CommandBuffer =
@@ -819,7 +821,13 @@ exec_graph_impl::enqueue(const std::shared_ptr<sycl::detail::queue_impl> &Queue,
               sycl::backend::ext_oneapi_level_zero) {
             Event->wait(Event);
           } else {
+            auto &AttachedEventsList = Event->getPostCompleteEvents();
+            CGData.MEvents.reserve(AttachedEventsList.size() + 1);
             CGData.MEvents.push_back(Event);
+            // Add events of the previous execution of all graph partitions.
+            for (auto &AttachedEvent : AttachedEventsList) {
+              CGData.MEvents.push_back(AttachedEvent);
+            }
           }
           ++It;
         } else {
@@ -929,7 +937,7 @@ exec_graph_impl::enqueue(const std::shared_ptr<sycl::detail::queue_impl> &Queue,
       NewEvent->setStateIncomplete();
       NewEvent->getPreparedDepsEvents() = ScheduledEvents;
     }
-    MPartitionsExecutionEvents[CurrentPartition] = NewEvent;
+    PartitionsExecutionEvents[CurrentPartition] = NewEvent;
   }
 
   // Keep track of this execution event so we can make sure it's completed in
@@ -937,7 +945,7 @@ exec_graph_impl::enqueue(const std::shared_ptr<sycl::detail::queue_impl> &Queue,
   MExecutionEvents.push_back(NewEvent);
   // Attach events of previous partitions to ensure that when the returned event
   // is complete all execution associated with the graph have been completed.
-  for (auto const &Elem : MPartitionsExecutionEvents) {
+  for (auto const &Elem : PartitionsExecutionEvents) {
     if (Elem.second != NewEvent) {
       NewEvent->attachEventToComplete(Elem.second);
     }
diff --git a/sycl/source/detail/graph_impl.hpp b/sycl/source/detail/graph_impl.hpp
index eafb66b1dca9b..6793ab0b2229f 100644
--- a/sycl/source/detail/graph_impl.hpp
+++ b/sycl/source/detail/graph_impl.hpp
@@ -1190,9 +1190,6 @@ class exec_graph_impl {
   std::vector<sycl::detail::EventImplPtr> MExecutionEvents;
   /// List of the partitions that compose the exec graph.
   std::vector<std::shared_ptr<partition>> MPartitions;
-  /// Map of the partitions to their execution events
-  std::unordered_map<std::shared_ptr<partition>, sycl::detail::EventImplPtr>
-      MPartitionsExecutionEvents;
   /// Storage for copies of nodes from the original modifiable graph.
   std::vector<std::shared_ptr<node_impl>> MNodeStorage;
 };