From 0e37a53fae633a213eae120d9b942bc40db6fbe6 Mon Sep 17 00:00:00 2001
From: Sourabh Betigeri <sourabh.betigeri@amd.com>
Date: Wed, 8 Mar 2023 17:04:58 -0800
Subject: [PATCH] SWDEV-385089 - Fixes hipGraphLaunch functionality accounting
 for the existence of MemFree node or already freed memory when the same graph
 is launched multiple times

Change-Id: I49beb49ad4e6db4a2dd5b8c8cc8ed11ff0e4e132


[ROCm/clr commit: ec8ab9b29c1d69344c05c6699eeed9c0b9ce4821]
---
 projects/clr/hipamd/src/hip_graph.cpp          |  6 +-----
 projects/clr/hipamd/src/hip_graph_internal.cpp | 14 ++++++++++++++
 projects/clr/hipamd/src/hip_graph_internal.hpp |  7 +++++++
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_graph.cpp b/projects/clr/hipamd/src/hip_graph.cpp
index 2a499a02b1..f2d7b38c6c 100644
--- a/projects/clr/hipamd/src/hip_graph.cpp
+++ b/projects/clr/hipamd/src/hip_graph.cpp
@@ -1260,11 +1260,7 @@ hipError_t hipGraphInstantiateWithFlags(hipGraphExec_t* pGraphExec, hipGraph_t g
     HIP_RETURN(hipErrorInvalidValue);
   }
 
-  // enable when change is merged to hip
-  // if (flags == hipGraphInstantiateFlagAutoFreeOnLaunch) {
-  // Free any unfreed memory allocations before the graph is relaunched
-  //}
-  HIP_RETURN_DURATION(ihipGraphInstantiate(pGraphExec, graph));
+  HIP_RETURN_DURATION(ihipGraphInstantiate(pGraphExec, graph, flags));
 }
 
 hipError_t hipGraphExecDestroy(hipGraphExec_t pGraphExec) {
diff --git a/projects/clr/hipamd/src/hip_graph_internal.cpp b/projects/clr/hipamd/src/hip_graph_internal.cpp
index 2d1e74c728..3c118bc198 100644
--- a/projects/clr/hipamd/src/hip_graph_internal.cpp
+++ b/projects/clr/hipamd/src/hip_graph_internal.cpp
@@ -579,6 +579,20 @@ hipError_t hipGraphExec::Run(hipStream_t stream) {
       levelOrder_[0]->GetParentGraph()->FreeAllMemory();
     }
   }
+
+  // If this is a repeat launch, make sure corresponding MemFreeNode exists for a MemAlloc node
+  if (repeatLaunch_ == true) {
+    for (auto& node : levelOrder_) {
+      if (node->GetType() == hipGraphNodeTypeMemAlloc &&
+          static_cast<hipGraphMemAllocNode*>(node)->IsActiveMem() == true) {
+          return hipErrorInvalidValue;
+      }
+    }
+  }
+  else {
+    repeatLaunch_ = true;
+  }
+
   auto hip_stream = (stream == nullptr) ? hip::getCurrentDevice()->NullStream()
                                         : reinterpret_cast<hip::Stream*>(stream);
   UpdateStream(parallelLists_, hip_stream, this);
diff --git a/projects/clr/hipamd/src/hip_graph_internal.hpp b/projects/clr/hipamd/src/hip_graph_internal.hpp
index 1dcc0e8f39..76e9adc7df 100644
--- a/projects/clr/hipamd/src/hip_graph_internal.hpp
+++ b/projects/clr/hipamd/src/hip_graph_internal.hpp
@@ -544,6 +544,7 @@ struct hipGraphExec {
   std::unordered_set<hipUserObject*> graphExeUserObj_;
   static amd::Monitor graphExecSetLock_;
   uint64_t flags_ = 0;
+  bool repeatLaunch_ = false;
  public:
   hipGraphExec(std::vector<Node>& levelOrder, std::vector<std::vector<Node>>& lists,
                std::unordered_map<Node, std::vector<Node>>& nodeWaitLists,
@@ -1879,6 +1880,12 @@ class hipGraphMemAllocNode : public hipGraphNode {
     return node_params_.dptr;
   }
 
+  bool IsActiveMem() {
+    auto graph = GetParentGraph();
+    return graph->ProbeMemory(node_params_.dptr);
+  }
+
+
   void GetParams(hipMemAllocNodeParams* params) const {
     std::memcpy(params, &node_params_, sizeof(hipMemAllocNodeParams));
   }