From 5db21f53da8a1c7dc31887a46c7c1ab40ee071b4 Mon Sep 17 00:00:00 2001
From: Flora Cui <flora.cui@amd.com>
Date: Mon, 13 Jan 2025 22:06:51 +0800
Subject: [PATCH] wsl/hsakmt: add get_gpu_mem()

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>
---
 libhsakmt.h |  1 +
 memory.cpp  | 10 ++++++++++
 2 files changed, 11 insertions(+)

diff --git a/libhsakmt.h b/libhsakmt.h
index 9995026591..088e4baf2f 100644
--- a/libhsakmt.h
+++ b/libhsakmt.h
@@ -36,6 +36,7 @@
 #include "impl/wddm/device.h"
 
 wsl::thunk::WDDMDevice* get_wddmdev(uint32_t node_id);
+wsl::thunk::GpuMemory *get_gpu_mem(void *MemoryAddress);
 
 extern unsigned long dxg_open_count;
 extern bool hsakmt_forked;
diff --git a/memory.cpp b/memory.cpp
index 950a5dedba..f044ef6aac 100644
--- a/memory.cpp
+++ b/memory.cpp
@@ -355,6 +355,16 @@ bool queue_release_buffer(void *MemoryAddress) {
   return true;
 }
 
+wsl::thunk::GpuMemory *get_gpu_mem(void *MemoryAddress) {
+  std::lock_guard<std::mutex> gard(*allocation_map_lock_);
+  auto it = allocation_map_.find(MemoryAddress);
+  if (it == allocation_map_.end()) {
+    return nullptr;
+  }
+
+  return wsl::thunk::GpuMemory::Convert(it->second.handle);
+}
+
 HSAKMT_STATUS HSAKMTAPI hsaKmtAvailableMemory(HSAuint32 Node,
                                               HSAuint64 *AvailableBytes) {
   CHECK_DXG_OPEN();