From 5a89405bf87be24019d3e6d45995213e137a8e2a Mon Sep 17 00:00:00 2001
From: Flora Cui <flora.cui@amd.com>
Date: Sat, 5 Jul 2025 09:02:53 +0800
Subject: [PATCH] wsl/libhsakmt: fix check for queue buffer

Signed-off-by: Flora Cui <flora.cui@amd.com>
Reviewed-by: Tianci Yin <tianci.yin@amd.com>
Part-of: <http://10.67.69.192/wsl/rocr-runtime/-/merge_requests/83>
---
 memory.cpp | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/memory.cpp b/memory.cpp
index 75bc01f331..ea863996a0 100644
--- a/memory.cpp
+++ b/memory.cpp
@@ -300,6 +300,9 @@ HSAKMT_STATUS hsaKmtFreeMemoryInternal(void *MemoryAddress,
     }
 
     gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
+    if (gpu_mem->IsQueueReferenced())
+      return HSAKMT_STATUS_ERROR;
+
     if (it->second.dmabuf_fd >= 0) {
       close(it->second.dmabuf_fd);
       it->second.dmabuf_fd = -1;
@@ -307,9 +310,6 @@ HSAKMT_STATUS hsaKmtFreeMemoryInternal(void *MemoryAddress,
     allocation_map_->erase(it);
   }
 
-  if (gpu_mem->IsQueueReferenced())
-    return HSAKMT_STATUS_ERROR;
-
   delete gpu_mem;
   return HSAKMT_STATUS_SUCCESS;
 }
@@ -763,7 +763,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtUnmapMemoryToGPU(void *MemoryAddress) {
       return HSAKMT_STATUS_SUCCESS;
   }
 
-  wsl::thunk::GpuMemoryHandle handle = nullptr;
+  wsl::thunk::GpuMemory *gpu_mem = nullptr;
   {
     std::lock_guard<std::mutex> gard(*allocation_map_lock_);
 
@@ -794,14 +794,13 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtUnmapMemoryToGPU(void *MemoryAddress) {
       return HSAKMT_STATUS_SUCCESS;
     }
 
-    handle = it->second.handle;
+    gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
+    if (gpu_mem->IsQueueReferenced())
+      return HSAKMT_STATUS_ERROR;
 
     allocation_map_->erase((void *)it->second.gpu_addr);
     allocation_map_->erase(it);
   }
-  auto gpu_mem = wsl::thunk::GpuMemory::Convert(handle);
-  if (gpu_mem->IsQueueReferenced())
-    return HSAKMT_STATUS_ERROR;
 
   delete gpu_mem;
   return HSAKMT_STATUS_SUCCESS;