From 23bc53e9a8f739f8f1ff9d287a7a8c5c92fae254 Mon Sep 17 00:00:00 2001
From: Flora Cui <flora.cui@amd.com>
Date: Thu, 17 Jul 2025 13:31:40 +0800
Subject: [PATCH] wsl/libhsakmt: fix abused memory flag in allocation_map_

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/89>
---
 memory.cpp | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/memory.cpp b/memory.cpp
index 52843ff6d3..21f4e6248e 100644
--- a/memory.cpp
+++ b/memory.cpp
@@ -305,7 +305,7 @@ HSAKMT_STATUS hsaKmtFreeMemoryInternal(void *MemoryAddress,
       return HSAKMT_STATUS_ERROR;
 
     wsl::thunk::GpuMemoryDescFlags flags;
-    flags.reserved = it->second.mem_flags_value;
+    flags.reserved = gpu_mem->Flags();
     if (flags.is_imported_vram_ipc &&
       gpu_mem->DecSharedReference()) {
       pr_info("memory is still referenced\n");
@@ -666,7 +666,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtDeregisterMemory(void *MemoryAddress) {
 
     auto *gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
     wsl::thunk::GpuMemoryDescFlags flags;
-    flags.reserved = it->second.mem_flags_value;
+    flags.reserved = gpu_mem->Flags();
     // IPC mem(vram)
     if (flags.is_imported_vram_ipc &&
       gpu_mem->DecSharedReference() == 0) {
@@ -722,12 +722,11 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtMapMemoryToGPUNodes(
     std::lock_guard<std::mutex> gard(*allocation_map_lock_);
     auto it = allocation_map_->find(aligned_ptr);
     if (it != allocation_map_->end()) {
+      wsl::thunk::GpuMemory *gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
       wsl::thunk::GpuMemoryDescFlags flags;
-      flags.reserved = it->second.mem_flags_value;
+      flags.reserved = gpu_mem->Flags();
       // IPC mem
       if (flags.is_imported_vram_ipc) {
-        wsl::thunk::GpuMemory *gpu_mem;
-        gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
 
         auto code = gpu_mem->MapGpuVirtualAddress(gpu_mem->GpuAddress(), gpu_mem->Size());
         if (code != ErrorCode::Success)
@@ -834,7 +833,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtUnmapMemoryToGPU(void *MemoryAddress) {
 
     // IPC mem
     wsl::thunk::GpuMemoryDescFlags flags;
-    flags.reserved = it->second.mem_flags_value;
+    flags.reserved = gpu_mem->Flags();
     if (flags.is_imported_vram_ipc &&
         !gpu_mem->IsSharedFromSameProcess()) {
       auto code = gpu_mem->UnmapGpuVirtualAddress(gpu_mem->GpuAddress(), gpu_mem->Size());