From 039e26ee0ff9683913535a0f8fce5bfe470db615 Mon Sep 17 00:00:00 2001
From: Anusha GodavarthySurya <anusha.godavarthysurya@amd.com>
Date: Mon, 10 Oct 2022 18:18:03 +0000
Subject: [PATCH] SWDEV-357759, SWDEV-360041, SWDEV-361145 Fix Stream end
 capture on forked streams

Change-Id: If0dc6242d2d3ca680e37e14a5dea5cf68dc295df
---
 hipamd/src/hip_graph.cpp          | 17 ++++++++++++++---
 hipamd/src/hip_graph_internal.cpp |  1 +
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp
index 311b0cadb1..e419ca79f6 100644
--- a/hipamd/src/hip_graph.cpp
+++ b/hipamd/src/hip_graph.cpp
@@ -151,6 +151,7 @@ hipError_t capturehipLaunchKernel(hipStream_t& stream, const void*& hostFunction
                                   dim3& blockDim, void**& args, size_t& sharedMemBytes) {
   ClPrint(amd::LOG_INFO, amd::LOG_API,
           "[hipGraph] current capture node kernel launch on stream : %p", stream);
+
   if (!hip::isValid(stream)) {
     return hipErrorInvalidValue;
   }
@@ -775,7 +776,7 @@ hipError_t capturehipStreamWaitEvent(hipEvent_t& event, hipStream_t& stream, uns
     s->SetCaptureGraph(reinterpret_cast<hip::Stream*>(e->GetCaptureStream())->GetCaptureGraph());
     s->SetCaptureMode(reinterpret_cast<hip::Stream*>(e->GetCaptureStream())->GetCaptureMode());
     s->SetParentStream(e->GetCaptureStream());
-    s->SetParallelCaptureStream(stream);
+    reinterpret_cast<hip::Stream*>(s->GetParentStream())->SetParallelCaptureStream(stream);
   }
   s->AddCrossCapturedNode(e->GetNodesPrevToRecorded());
   return hipSuccess;
@@ -923,6 +924,13 @@ hipError_t hipStreamEndCapture_common(hipStream_t stream, hipGraph_t* pGraph) {
   // check if all parallel streams have joined
   // Nodes that are removed from the dependency set via API hipStreamUpdateCaptureDependencies do
   // not result in hipErrorStreamCaptureUnjoined
+  // add temporary node to check if all parallel streams have joined
+  hipGraphNode_t pGraphNode;
+  pGraphNode = new hipGraphEmptyNode();
+  hipError_t status =
+      ihipGraphAddNode(pGraphNode, s->GetCaptureGraph(), s->GetLastCapturedNodes().data(),
+                       s->GetLastCapturedNodes().size());
+
   if (s->GetCaptureGraph()->GetLeafNodeCount() > 1) {
     std::vector<hipGraphNode_t> leafNodes = s->GetCaptureGraph()->GetLeafNodes();
     const std::vector<hipGraphNode_t>& removedDepNodes = s->GetRemovedDependencies();
@@ -934,9 +942,14 @@ hipError_t hipStreamEndCapture_common(hipStream_t stream, hipGraph_t* pGraph) {
         }
       }
     }
+    // remove temporary node
+    s->GetCaptureGraph()->RemoveNode(pGraphNode);
     if (foundInRemovedDep == false) {
       return hipErrorStreamCaptureUnjoined;
     }
+  } else {
+    // remove temporary node
+    s->GetCaptureGraph()->RemoveNode(pGraphNode);
   }
   *pGraph = s->GetCaptureGraph();
   // end capture on all streams/events part of graph capture
@@ -1660,8 +1673,6 @@ hipError_t hipGraphDestroyNode(hipGraphNode_t node) {
     HIP_RETURN(hipErrorInvalidValue);
   }
   node->GetParentGraph()->RemoveNode(node);
-  // Takescare of removing its dependencies and dependent nodes
-  delete node;
   HIP_RETURN(hipSuccess);
 }
 
diff --git a/hipamd/src/hip_graph_internal.cpp b/hipamd/src/hip_graph_internal.cpp
index 8c48a19858..8a93f1c1f5 100644
--- a/hipamd/src/hip_graph_internal.cpp
+++ b/hipamd/src/hip_graph_internal.cpp
@@ -493,6 +493,7 @@ void ihipGraph::AddNode(const Node& node) {
 
 void ihipGraph::RemoveNode(const Node& node) {
   vertices_.erase(std::remove(vertices_.begin(), vertices_.end(), node), vertices_.end());
+  delete node;
 }
 
 // root nodes are all vertices with 0 in-degrees