wsl/libhsakmt: add same process check for ipc buffer

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/85>

memory.cpp

@@ -304,6 +304,14 @@ HSAKMT_STATUS hsaKmtFreeMemoryInternal(void *MemoryAddress,
     if (gpu_mem->IsQueueReferenced())
       return HSAKMT_STATUS_ERROR;
 
+    wsl::thunk::GpuMemoryDescFlags flags;
+    flags.reserved = it->second.mem_flags_value;
+    if (flags.is_imported_vram_ipc &&
+      gpu_mem->DecSharedReference()) {
+      pr_info("memory is still referenced\n");
+      return HSAKMT_STATUS_SUCCESS;
+    }
+
     if (it->second.dmabuf_fd >= 0) {
       close(it->second.dmabuf_fd);
       it->second.dmabuf_fd = -1;
@@ -546,30 +554,45 @@ HSAKMT_STATUS import_dmabuf_fd(int DMABufFd,
     }
   }
 
-  auto code = dev->CreateGpuMemory(create_info, &gpu_mem);
-  if (code == ErrorCode::Success) {
-    void *MemoryAddress;
-    if (alloc_va)
-      MemoryAddress = reinterpret_cast<void *>(gpu_mem->GpuAddress());
-    else
-      MemoryAddress = reinterpret_cast<void*>(gpu_mem->HandleApeAddress());
-
-    *GpuMemHandle = gpu_mem->GetGpuMemoryHandle();
-
+  gpusize gpu_va = 0;
+  auto code = dev->CreateGpuMemory(create_info, &gpu_mem, &gpu_va);
+  if (code == ErrorCode::SameProcessSameDevice) {
+    /* Unit_hipMemPoolExportToShareableHandle_SameProc */
+    pr_info("imported from same process, use the old one\n");
     std::lock_guard<std::mutex> gard(*allocation_map_lock_);
-    /*
-     * the gpu_mem->Flags() need convert back from GpuMemoryCreateFlags to
-     * HsaMemFlags, reference hsaKmtAllocMemoryAlign
-     * */
-   (*allocation_map_)[MemoryAddress] = Allocation(
-      *GpuMemHandle, MemoryAddress, (uint64_t)MemoryAddress,
-      gpu_mem->Size(), false, nullptr, gpu_mem->ClientSize(),
-      NodeId, gpu_mem->Flags());
-
+    auto it = allocation_map_->find((void*)gpu_va);
+    if (it == allocation_map_->end()) {
+      pr_err("where's the conflict buffer? va %#lx\n", create_info.va_hint);
+      return HSAKMT_STATUS_ERROR;
+    }
+    wsl::thunk::GpuMemory *conflict_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
+    conflict_mem->IncSharedReference();
+    *GpuMemHandle = it->second.handle;
     return HSAKMT_STATUS_SUCCESS;
+  } else if (code != ErrorCode::Success) {
+    pr_err("fail to import fd, ret %d\n", (int)code);
+    return HSAKMT_STATUS_ERROR;
   }
 
-  return HSAKMT_STATUS_ERROR;
+  void *MemoryAddress;
+  if (alloc_va)
+    MemoryAddress = reinterpret_cast<void *>(gpu_mem->GpuAddress());
+  else
+    MemoryAddress = reinterpret_cast<void*>(gpu_mem->HandleApeAddress());
+
+  *GpuMemHandle = gpu_mem->GetGpuMemoryHandle();
+
+  std::lock_guard<std::mutex> gard(*allocation_map_lock_);
+  /*
+   * the gpu_mem->Flags() need convert back from GpuMemoryCreateFlags to
+   * HsaMemFlags, reference hsaKmtAllocMemoryAlign
+   * */
+  (*allocation_map_)[MemoryAddress] = Allocation(
+    *GpuMemHandle, MemoryAddress, (uint64_t)MemoryAddress,
+    gpu_mem->Size(), false, nullptr, gpu_mem->ClientSize(),
+    NodeId, gpu_mem->Flags());
+
+  return HSAKMT_STATUS_SUCCESS;
 
 }
 
@@ -645,7 +668,8 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtDeregisterMemory(void *MemoryAddress) {
     wsl::thunk::GpuMemoryDescFlags flags;
     flags.reserved = it->second.mem_flags_value;
     // IPC mem(vram)
-    if (flags.is_imported_vram_ipc) {
+    if (flags.is_imported_vram_ipc &&
+      gpu_mem->DecSharedReference() == 0) {
       allocation_map_->erase(it);
       delete gpu_mem;
       return HSAKMT_STATUS_SUCCESS;
@@ -811,7 +835,8 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtUnmapMemoryToGPU(void *MemoryAddress) {
     // IPC mem
     wsl::thunk::GpuMemoryDescFlags flags;
     flags.reserved = it->second.mem_flags_value;
-    if (flags.is_imported_vram_ipc) {
+    if (flags.is_imported_vram_ipc &&
+        !gpu_mem->IsSharedFromSameProcess()) {
       auto code = gpu_mem->UnmapGpuVirtualAddress(gpu_mem->GpuAddress(), gpu_mem->Size());
       if (code != ErrorCode::Success)
         return HSAKMT_STATUS_ERROR;

wddm/device.cpp

@@ -569,13 +569,14 @@ void WDDMDevice::UpdatePageFence(uint64_t fence_value) {
   } while (!page_fence_value_.compare_exchange_weak(current, fence_value));
 }
 
-ErrorCode WDDMDevice::CreateGpuMemory(const GpuMemoryCreateInfo &create_info, GpuMemory **gpu_mem) {
+ErrorCode WDDMDevice::CreateGpuMemory(const GpuMemoryCreateInfo &create_info,
+                                        GpuMemory **gpu_mem, gpusize *gpu_va) {
   ErrorCode ret;
 
   *gpu_mem = nullptr;
   auto mem = new GpuMemory(this);
   if (create_info.dmabuf_fd > 0)
-    ret = mem->ImportPhysicalHandle(create_info);
+    ret = mem->ImportPhysicalHandle(create_info, gpu_va);
   else 
     ret = mem->Init(create_info);
   if (ret == ErrorCode::Success)

wddm/gpu_memory.cpp

@@ -341,6 +341,8 @@ ErrorCode GpuMemory::CreatePhysicalMemory() {
     shared_info.adapter_luid = desc_.adapter_luid;
     shared_info.flags = reinterpret_cast<uint32_t>(desc_.flags.reserved);
     shared_info.mem_flags = desc_.mem_flags;
+    shared_info.pid = dxg_runtime->parent_pid;
+    shared_info.gpu_addr = desc_.gpu_addr;
     args.pPrivateRuntimeData = &shared_info;
     args.PrivateRuntimeDataSize = sizeof(shared_info);
     args.Flags.NtSecuritySharing = 1;
@@ -416,7 +418,7 @@ ErrorCode GpuMemory::ExportPhysicalHandle(int* dmabuf_fd, uint32_t flags) {
 }
 
 
-ErrorCode GpuMemory::ImportPhysicalHandle(const GpuMemoryCreateInfo &create_info) {
+ErrorCode GpuMemory::ImportPhysicalHandle(const GpuMemoryCreateInfo &create_info, gpusize *gpu_addr) {
   D3DKMT_QUERYRESOURCEINFOFROMNTHANDLE query_args;
   int dmabuf_fd = create_info.dmabuf_fd;
 
@@ -542,6 +544,16 @@ ErrorCode GpuMemory::ImportPhysicalHandle(const GpuMemoryCreateInfo &create_info
       pr_err("open resource failed %d\n", static_cast<int>(ret));
       return ret;
     }
+    if (shared_info.pid == dxg_runtime->parent_pid &&
+      create_info.flags.alloc_va &&
+      IsSameAdapter(shared_info.adapter_luid) &&
+      shared_info.gpu_addr) {
+      pr_info("import from same device and samve process, va is required. "
+               "a buffer can't be mapped to 2 va. delete the imported buffer, use the existing one.\n");
+      if (gpu_addr)
+        *gpu_addr = shared_info.gpu_addr;
+      return ErrorCode::SameProcessSameDevice;
+    }
 
     desc_.size = shared_info.size;
     desc_.client_size = shared_info.client_size;