From 8b6d919b4bc105908afb505012b7676c7b49f2e1 Mon Sep 17 00:00:00 2001 From: Flora Cui Date: Sat, 5 Jul 2025 12:38:07 +0800 Subject: [PATCH] wsl/libhsakmt: add same process check for ipc buffer Signed-off-by: Flora Cui Reviewed-by: Tianci Yin Part-of: --- memory.cpp | 69 ++++++++++++++++++++++++++++++--------------- wddm/device.cpp | 5 ++-- wddm/gpu_memory.cpp | 14 ++++++++- 3 files changed, 63 insertions(+), 25 deletions(-) diff --git a/memory.cpp b/memory.cpp index 4f516ad3a5..7085c034e2 100644 --- a/memory.cpp +++ b/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(gpu_mem->GpuAddress()); - else - MemoryAddress = reinterpret_cast(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 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(gpu_mem->GpuAddress()); + else + MemoryAddress = reinterpret_cast(gpu_mem->HandleApeAddress()); + + *GpuMemHandle = gpu_mem->GetGpuMemoryHandle(); + + std::lock_guard 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; diff --git a/wddm/device.cpp b/wddm/device.cpp index 8d5c18d162..6024f21fa2 100644 --- a/wddm/device.cpp +++ b/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) diff --git a/wddm/gpu_memory.cpp b/wddm/gpu_memory.cpp index 83ba1cffc5..582310bb84 100644 --- a/wddm/gpu_memory.cpp +++ b/wddm/gpu_memory.cpp @@ -341,6 +341,8 @@ ErrorCode GpuMemory::CreatePhysicalMemory() { shared_info.adapter_luid = desc_.adapter_luid; shared_info.flags = reinterpret_cast(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(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;