From f342af52c575e8d88decfb94e001bf4d7df9e7f5 Mon Sep 17 00:00:00 2001
From: Flora Cui <flora.cui@amd.com>
Date: Mon, 13 Jan 2025 22:07:37 +0800
Subject: [PATCH] wsl/hsakmt: fix gpu_mem handle

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/28>
---
 libdrm.cpp | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/libdrm.cpp b/libdrm.cpp
index d85ca9aa1e..7d366916bd 100644
--- a/libdrm.cpp
+++ b/libdrm.cpp
@@ -110,7 +110,6 @@ HSAKMTAPI int amdgpu_bo_import(amdgpu_device_handle dev,
   }
 }
 
-
 HSAKMTAPI int amdgpu_bo_va_op(amdgpu_bo_handle bo,
                               uint64_t offset,
                               uint64_t size,
@@ -120,7 +119,7 @@ HSAKMTAPI int amdgpu_bo_va_op(amdgpu_bo_handle bo,
   switch(ops) {
     case AMDGPU_VA_OP_MAP:
       {
-        wsl::thunk::GpuMemory *gpu_mem = reinterpret_cast<wsl::thunk::GpuMemory *>(bo);
+        wsl::thunk::GpuMemory *gpu_mem = get_gpu_mem(bo);
         assert(gpu_mem != nullptr);
         auto code = gpu_mem->MapGpuVirtualAddress(reinterpret_cast<gpusize>(addr), size, offset);
         if (code != ErrorCode::Success)
@@ -133,7 +132,7 @@ HSAKMTAPI int amdgpu_bo_va_op(amdgpu_bo_handle bo,
       break;
     case AMDGPU_VA_OP_UNMAP:
       {
-        wsl::thunk::GpuMemory *gpu_mem = reinterpret_cast<wsl::thunk::GpuMemory *>(bo);
+        wsl::thunk::GpuMemory *gpu_mem = get_gpu_mem(bo);
         assert(gpu_mem != nullptr);
         auto code = gpu_mem->UnmapGpuVirtualAddress(reinterpret_cast<gpusize>(addr), size, offset);
         if (code != ErrorCode::Success)