From 4b9ef909b44843dda4c49f4f9118081bcc2ca1dc Mon Sep 17 00:00:00 2001
From: Flora Cui <flora.cui@amd.com>
Date: Wed, 6 Nov 2024 11:02:47 +0800
Subject: [PATCH] wsl/hsakmt: relax the check conditions for the pointer

Signed-off-by: Flora Cui <flora.cui@amd.com>
Reviewed-by: Longlong Yao <Longlong.Yao@amd.com>
Reviewed-by: tiancyin <tianci.yin@amd.com>
Part-of: <http://10.67.69.192/wsl/rocr-runtime/-/merge_requests/9>
---
 memory.cpp | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/memory.cpp b/memory.cpp
index 92047a5eba..e0a8f7548c 100644
--- a/memory.cpp
+++ b/memory.cpp
@@ -725,26 +725,36 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtQueryPointerInfo(const void *Pointer,
     return HSAKMT_STATUS_INVALID_PARAMETER;
 
   pr_debug("pointer %p\n", Pointer);
-  void *pointer = const_cast<void*>(Pointer);
+  void *ptr = const_cast<void*>(Pointer);
 
   memset(PointerInfo, 0, sizeof(HsaPointerInfo));
   void* block_base = nullptr;
   {
     std::lock_guard<std::mutex> gard(*fragment_allocator_lock_);
-      block_base = fragment_allocator_.block_base(pointer);
+      block_base = fragment_allocator_.block_base(ptr);
     if (block_base != nullptr)
-      pointer = block_base;
+      ptr = block_base;
   }
 
   Allocation allocation_info;
+  bool found = false;
   {
     std::lock_guard<std::mutex> gard(*allocation_map_lock_);
-    auto it = allocation_map_.find(pointer);
-    if (it == allocation_map_.end()) {
-      PointerInfo->Type = HSA_POINTER_UNKNOWN;
-      return HSAKMT_STATUS_ERROR;
+    auto it = allocation_map_.upper_bound(ptr);
+    if (it != allocation_map_.begin()) {
+      --it;
+      if (ptr >= it->first &&
+        (ptr < reinterpret_cast<const uint8_t*>(it->first) + it->second.size_requested)) {
+        allocation_info = it->second;
+        found = true;
+      }
     }
-    allocation_info = it->second;
+  }
+
+  if (!found) {
+    pr_debug("can't found allocation for %p\n", Pointer);
+    PointerInfo->Type = HSA_POINTER_UNKNOWN;
+    return HSAKMT_STATUS_ERROR;
   }
 
   if (allocation_info.userptr) {