From b4df9fb6ec96ed680fa8dee0fa95ec54c6a43f66 Mon Sep 17 00:00:00 2001
From: Jaydeep Patel <jaydeepkumar.patel@amd.com>
Date: Mon, 20 Jan 2025 06:13:40 +0000
Subject: [PATCH] SWDEV-457316 - Use phy memory obj stored in user data instead
 of querying from memObjs.

Change-Id: Id837eb00195d88b50904441f01cf8153fa752ecd
---
 hipamd/src/hip_graph_internal.hpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp
index 9019b2aa6d..3fd3be5f4c 100644
--- a/hipamd/src/hip_graph_internal.hpp
+++ b/hipamd/src/hip_graph_internal.hpp
@@ -2479,7 +2479,10 @@ class GraphMemFreeNode : public GraphNode {
       vaddr_sub_obj->release();
       vaddr_mem_obj->release();
       // Release the allocation back to graph's pool
-      graph_->FreeMemory(phys_mem_obj->getSvmPtr(), static_cast<hip::Stream*>(queue()));
+      auto device_id = phys_mem_obj->getUserData().deviceId;
+      if (!g_devices[device_id]->FreeMemory(phys_mem_obj, static_cast<hip::Stream*>(queue()))) {
+        LogError("Memory didn't belong to any pool!");
+      }
       amd::MemObjMap::AddMemObj(ptr(), vaddr_mem_obj);
       graph_->memalloc_nodes_--; // Decrement count of unreleased memalloc nodes
       ClPrint(amd::LOG_INFO, amd::LOG_MEM_POOL, "Graph MemFree execute: %p, %p",