Release graph if hipStreamEndCapture fails (#738)

projects/clr/hipamd/src/hip_graph.cpp

@@ -1218,9 +1218,7 @@ hipError_t hipStreamEndCapture_common(hipStream_t stream, hip::Graph** pGraph) {
   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();
+    s->ReleaseCaptureGraph();
 
     return hipErrorStreamCaptureInvalidated;
   }
@@ -1257,6 +1255,8 @@ hipError_t hipStreamEndCapture_common(hipStream_t stream, hip::Graph** pGraph) {
     s->GetCaptureGraph()->RemoveNode(pGraphNode);
     s->GetCaptureGraph()->RemoveManualNodesDuringCapture();
     if (leafNodes.size() > 1 && foundInRemovedDep == false) {
+      // Release created graph as it can't be retrieved anymore
+      s->ReleaseCaptureGraph();
       return hipErrorStreamCaptureUnjoined;
     }
   } else {

projects/clr/hipamd/src/hip_internal.hpp

@@ -415,8 +415,8 @@ public:
       pCaptureGraph_ = pGraph;
       captureStatus_ = hipStreamCaptureStatusActive;
     }
-    /// Reset graph to nullptr when capture is invalidated, but keep the status
-    void ResetCaptureGraph() { pCaptureGraph_ = nullptr; }
+    /// Release graph when capture is invalidated
+    void ReleaseCaptureGraph();
     void SetCaptureId() {
       // ID is generated in Begin Capture i.e.. when capture status is active
       captureID_ = GenerateCaptureID();

projects/clr/hipamd/src/hip_stream.cpp

@@ -20,6 +20,7 @@
 
 #include <hip/hip_runtime.h>
 #include "hip_internal.hpp"
+#include "hip_graph_internal.hpp"
 #include "hip_event.hpp"
 #include "thread/monitor.hpp"
 #include "hip_prof_api.h"
@@ -104,6 +105,13 @@ bool isValid(hipStream_t& stream) {
   return false;
 }
 
+void Stream::ReleaseCaptureGraph() {
+  if (pCaptureGraph_ != nullptr) {
+    delete pCaptureGraph_;
+    pCaptureGraph_ = nullptr;
+  }
+}
+
 // ================================================================================================
 int Stream::DeviceId() const { return device_->deviceId(); }