From b4b0b54c796b41a7f6f8ef89ffa48a76611a0f0f Mon Sep 17 00:00:00 2001
From: lyndonli <Lyndon.Li@amd.com>
Date: Mon, 9 Sep 2024 17:05:08 +0800
Subject: [PATCH] wsl/hsakmt: Add clear_allocation_map

If fork() is called, clear allocation map
that is invalid in the child process.

Signed-off-by: lyndonli <Lyndon.Li@amd.com>
Reviewed-by: Shi.Leslie Yuliang.Shi@amd.com
Part-of: <http://10.67.69.192/wsl/libhsakmt/-/merge_requests/3>
---
 libhsakmt.h   |  2 ++
 memory.cpp    | 21 ++++++++++++++-------
 openclose.cpp |  2 +-
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/libhsakmt.h b/libhsakmt.h
index 096f9d8829..96b59a462c 100644
--- a/libhsakmt.h
+++ b/libhsakmt.h
@@ -151,6 +151,8 @@ uint32_t get_num_sysfs_nodes(void);
 
 bool is_forked_child(void);
 
+void clear_allocation_map(void);
+
 /* Calculate VGPR and SGPR register file size per CU */
 uint32_t get_vgpr_size_per_cu(HSA_ENGINE_ID id);
 #define SGPR_SIZE_PER_CU 0x4000
diff --git a/memory.cpp b/memory.cpp
index c69f1eef76..02e28b4266 100644
--- a/memory.cpp
+++ b/memory.cpp
@@ -59,7 +59,14 @@ struct Allocation {
 };
 
 static std::map<const void *, Allocation> allocation_map_;
-static std::mutex allocation_map_lock_;
+static std::unique_ptr<std::mutex> allocation_map_lock_ = std::make_unique<std::mutex>();
+
+void clear_allocation_map(void)
+{
+  allocation_map_lock_ = std::make_unique<std::mutex>();
+  std::lock_guard<std::mutex> lock(*allocation_map_lock_);
+  allocation_map_.clear();
+}
 
 HSAKMT_STATUS HSAKMTAPI hsaKmtSetMemoryPolicy(HSAuint32 Node,
                                               HSAuint32 DefaultPolicy,
@@ -159,7 +166,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtAllocMemoryAlign(HSAuint32 PreferredNode,
   auto code = dev->CreateGpuMemory(create_info, &gpu_mem);
   if (code == ErrorCode::Success) {
     *MemoryAddress = reinterpret_cast<void *>(gpu_mem->GpuAddress());
-    std::lock_guard<std::mutex> gard(allocation_map_lock_);
+    std::lock_guard<std::mutex> gard(*allocation_map_lock_);
     allocation_map_[*MemoryAddress] = Allocation(
         gpu_mem->GetGpuMemoryHandle(), *MemoryAddress, (uint64_t)*MemoryAddress,
         create_info.size, false, nullptr, SizeInBytes,
@@ -179,7 +186,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtFreeMemory(void *MemoryAddress,
 
   wsl::thunk::GpuMemory *gpu_mem = nullptr;
   {
-    std::lock_guard<std::mutex> gard(allocation_map_lock_);
+    std::lock_guard<std::mutex> gard(*allocation_map_lock_);
     auto it = allocation_map_.find(MemoryAddress);
     if (it == allocation_map_.end()) {
       return HSAKMT_STATUS_ERROR;
@@ -350,7 +357,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtMapMemoryToGPU(void *MemoryAddress,
   size_t aligned_size = end - start;
 
   {
-    std::lock_guard<std::mutex> gard(allocation_map_lock_);
+    std::lock_guard<std::mutex> gard(*allocation_map_lock_);
     // GTT mem
     auto it_gtt = allocation_map_.find(aligned_ptr);
     if (it_gtt != allocation_map_.end()) {
@@ -397,7 +404,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtMapMemoryToGPU(void *MemoryAddress,
   }
 
   {
-    std::lock_guard<std::mutex> guard(allocation_map_lock_);
+    std::lock_guard<std::mutex> guard(*allocation_map_lock_);
     allocation_map_[MemoryAddress] =
         Allocation(handle, aligned_ptr, addr, aligned_size, true, MemoryAddress,
                    MemorySizeInBytes);
@@ -430,7 +437,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtUnmapMemoryToGPU(void *MemoryAddress) {
 
   wsl::thunk::GpuMemoryHandle handle = nullptr;
   {
-    std::lock_guard<std::mutex> gard(allocation_map_lock_);
+    std::lock_guard<std::mutex> gard(*allocation_map_lock_);
     auto it = allocation_map_.find(MemoryAddress);
     if (it == allocation_map_.end()) {
       return HSAKMT_STATUS_ERROR;
@@ -495,7 +502,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtQueryPointerInfo(const void *Pointer,
 
   Allocation allocation_info;
   {
-    std::lock_guard<std::mutex> gard(allocation_map_lock_);
+    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;
diff --git a/openclose.cpp b/openclose.cpp
index b2b3c04460..9c66c1a9c7 100644
--- a/openclose.cpp
+++ b/openclose.cpp
@@ -96,7 +96,7 @@ static void child_fork_handler(void) {
  * in the child process so it is not cleared
  */
 static void clear_after_fork(void) {
-  // TODO: fmm_clear_all_mem();
+  clear_allocation_map();
   if (dxg_fd) {
     close(dxg_fd);
     dxg_fd = -1;