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>
2025-01-13 22:06:51 +08:00
parent f7d2df7486
commit 5db21f53da
2 changed files with 11 additions and 0 deletions
@@ -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();