From 21c223f8df87be24079b1e06fed2f6d0924980ad Mon Sep 17 00:00:00 2001
From: Ioannis Assiouras <Ioannis.Assiouras@amd.com>
Date: Wed, 22 Jan 2025 02:40:10 +0000
Subject: [PATCH] SWDEV-510319 - Fixed random segfaults in graph tests

This change fixes random segfaults in graph tests that
are seen after the change make internal callbacks non-blocking.
The callback thread that decreases the GraphExec ref count
may now run after the runtime shutdown. This can cause a segfault
because the hip::device that is accessed in GraphExec destructor
is already destroyed during runtime shutdown. This patch ensures
that the hip::device object  stays alive until after the
callback thread completes.

Change-Id: I75a6ac01f27a0b2250bbd10ed389ebfb322927af
---
 hipamd/src/hip_graph_internal.cpp | 1 +
 hipamd/src/hip_graph_internal.hpp | 7 +++----
 rocclr/device/rocm/rocvirtual.cpp | 4 +++-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/hipamd/src/hip_graph_internal.cpp b/hipamd/src/hip_graph_internal.cpp
index f6c7221724..1205f44470 100644
--- a/hipamd/src/hip_graph_internal.cpp
+++ b/hipamd/src/hip_graph_internal.cpp
@@ -363,6 +363,7 @@ hipError_t GraphExec::Init() {
     status = CaptureAQLPackets();
   }
   instantiateDeviceId_ = hip::getCurrentDevice()->deviceId();
+  static_cast<ReferenceCountedObject*>( hip::getCurrentDevice())->retain();
   return status;
 }
 
diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp
index 3fd3be5f4c..856aa839d4 100644
--- a/hipamd/src/hip_graph_internal.hpp
+++ b/hipamd/src/hip_graph_internal.hpp
@@ -505,7 +505,6 @@ struct Graph {
     amd::ScopedLock lock(graphSetLock_);
     graphSet_.insert(this);
     mem_pool_ = device->GetGraphMemoryPool();
-    mem_pool_->retain();
     graphInstantiated_ = false;
     roots_.resize(DEBUG_HIP_FORCE_GRAPH_QUEUES);
     leafs_.resize(DEBUG_HIP_FORCE_GRAPH_QUEUES);
@@ -543,9 +542,6 @@ struct Graph {
       }
     }
     graphUserObj_.clear();
-    if (mem_pool_ != nullptr) {
-      mem_pool_->release();
-    }
     memAllocNodePtrs_.clear();
   }
 
@@ -764,6 +760,9 @@ struct GraphExec : public amd::ReferenceCountedObject, public Graph {
         kernArgManager_->release();
       }
     }
+    if (instantiateDeviceId_ != -1) {
+      static_cast<ReferenceCountedObject*>(g_devices[instantiateDeviceId_])->release();
+    }
   }
 
   Node GetClonedNode(Node node) {
diff --git a/rocclr/device/rocm/rocvirtual.cpp b/rocclr/device/rocm/rocvirtual.cpp
index 778f7f6be1..d1367cae5f 100644
--- a/rocclr/device/rocm/rocvirtual.cpp
+++ b/rocclr/device/rocm/rocvirtual.cpp
@@ -235,11 +235,13 @@ bool HsaAmdSignalHandler(hsa_signal_value_t value, void* arg) {
   // Reset last used SDMA engine mask
   gpu->setLastUsedSdmaEngine(0);
 
+  bool isBlocking = ts->GetBlocking();
+
   // Update the batch, since signal is complete
   gpu->updateCommandsState(ts->command().GetBatchHead());
 
   // Reset API callback signal. It will release AQL queue and start commands processing
-  if (callback_signal.handle != 0 && ts->GetBlocking()) {
+  if (callback_signal.handle != 0 && isBlocking) {
     hsa_signal_subtract_relaxed(callback_signal, 1);
   }