From 6f2bad3998b87670cbd81b6f4e6df3976f64d666 Mon Sep 17 00:00:00 2001
From: Vladana Stojiljkovic <Vladana.Stojiljkovic@amd.com>
Date: Mon, 7 Oct 2024 17:22:15 +0200
Subject: [PATCH] SWDEV-489823 - Fix hipStreamEndCapture leak when capture is
 invalidated

Change-Id: If8f5163d70e04d34a75fd0a7ba6c0a15ea59bb8b
---
 hipamd/src/hip_graph.cpp    | 5 +++++
 hipamd/src/hip_internal.hpp | 2 ++
 2 files changed, 7 insertions(+)

diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp
index 9eb9864426..2e0da2f432 100644
--- a/hipamd/src/hip_graph.cpp
+++ b/hipamd/src/hip_graph.cpp
@@ -1016,6 +1016,11 @@ hipError_t hipStreamEndCapture_common(hipStream_t stream, hip::Graph** pGraph) {
   // If capture was invalidated, due to a violation of the rules of stream capture
   if (s->GetCaptureStatus() == hipStreamCaptureStatusInvalidated) {
     *pGraph = nullptr;
+    // When capture is invalidated, graph should be deleted, otherwise it leaks
+    hip::Graph* graph = s->GetCaptureGraph();
+    delete graph;
+    s->ResetCaptureGraph();
+
     return hipErrorStreamCaptureInvalidated;
   }
 
diff --git a/hipamd/src/hip_internal.hpp b/hipamd/src/hip_internal.hpp
index 3b8cedad76..8edaacd3c4 100644
--- a/hipamd/src/hip_internal.hpp
+++ b/hipamd/src/hip_internal.hpp
@@ -405,6 +405,8 @@ public:
       pCaptureGraph_ = pGraph;
       captureStatus_ = hipStreamCaptureStatusActive;
     }
+    /// Reset graph to nullptr when capture is invalidated, but keep the status
+    void ResetCaptureGraph() { pCaptureGraph_ = nullptr; }
     void SetCaptureId() {
       // ID is generated in Begin Capture i.e.. when capture status is active
       captureID_ = GenerateCaptureID();