wsl/hsakmt: implement ipc mem of rocr non-legacy mode

The legacy mode means buffer sharing through KFD, KFD provide a buffer
id to exporter, exporter pass it to importer, importer pass buffer id
to KFD to query and import this buffer.

The non-legcay mode relys on socket to pass dmabuf fd between processes.

In hsa-runtime, the legcay mode is the default mode, setting environment
variable HSA_ENABLE_IPC_MODE_LEGACY to 0 can force hsa-runtime to new
mode code path.

Reviewed-by: Flora Cui <flora.cui@amd.com>
Reviewed-by: Longlong Yao <Longlong.Yao@amd.com>
Signed-off-by: tiancyin <tianci.yin@amd.com>

libhsakmt.h

@@ -200,6 +200,8 @@ bool queue_release_buffer(void *MemoryAddress);
 uint32_t get_vgpr_size_per_cu(HSA_ENGINE_ID id);
 #define SGPR_SIZE_PER_CU 0x4000
 
-HSAKMT_STATUS hsaKmtImportDMABufHandle(int DMABufFd, HsaGraphicsResourceInfo *GraphicsResourceInfo);
+HSAKMT_STATUS hsaKmtImportDMABufHandle(int DMABufFd,
+				       HsaGraphicsResourceInfo *GraphicsResourceInfo,
+				       bool requiresVAddr = false);
 
 #endif

memory.cpp

@@ -38,7 +38,8 @@
 struct Allocation {
   Allocation()
       : handle(0), cpu_addr(0), gpu_addr(0), size(0), userptr(false),
-        user_data(nullptr), size_requested(0), node_id(0), mem_flags_value(0) {}
+        user_data(nullptr), size_requested(0), node_id(0), mem_flags_value(0),
+        dmabuf_fd(-1) {}
   Allocation(wsl::thunk::GpuMemoryHandle handle_arg, void *cpu_addr_arg,
              uint64_t gpu_addr_arg, size_t size_arg, bool userptr_arg = false,
              void *user_data_arg = nullptr, size_t user_size_arg = 0,
@@ -46,7 +47,7 @@ struct Allocation {
       : handle(handle_arg), cpu_addr(cpu_addr_arg), gpu_addr(gpu_addr_arg),
         size(size_arg), userptr(userptr_arg), user_data(user_data_arg),
         size_requested(user_size_arg), node_id(node_id_arg),
-        mem_flags_value(mem_flags_value_arg) {}
+        mem_flags_value(mem_flags_value_arg), dmabuf_fd(-1) {}
 
   wsl::thunk::GpuMemoryHandle handle;
   void *cpu_addr;
@@ -57,6 +58,7 @@ struct Allocation {
   size_t size_requested; /* size requested by user */
   HSAuint32 node_id;
   HSAuint32 mem_flags_value;
+  int dmabuf_fd;
 };
 
 static std::map<const void *, Allocation> allocation_map_;
@@ -298,6 +300,10 @@ HSAKMT_STATUS hsaKmtFreeMemoryInternal(void *MemoryAddress,
     }
 
     gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
+    if (it->second.dmabuf_fd >= 0) {
+      close(it->second.dmabuf_fd);
+      it->second.dmabuf_fd = -1;
+    }
     allocation_map_.erase(it);
   }
 
@@ -447,10 +453,11 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtRegisterGraphicsHandleToNodesExt(HSAuint64 Graphic
   pr_debug("number of nodes %lu\n", NumberOfNodes);
 
   GraphicsResourceInfo->NodeId = 1;
-  return hsaKmtImportDMABufHandle(GraphicsResourceHandle, GraphicsResourceInfo);
+  return hsaKmtImportDMABufHandle(GraphicsResourceHandle,
+				  GraphicsResourceInfo,
+				  !!RegisterFlags.ui32.requiresVAddr);
 }
 
-
 HSAKMT_STATUS HSAKMTAPI hsaKmtExportDMABufHandle(void *MemoryAddress,
                                                  HSAuint64 MemorySizeInBytes,
                                                  int *DMABufFd,
@@ -458,21 +465,31 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtExportDMABufHandle(void *MemoryAddress,
   CHECK_DXG_OPEN();
 
   std::lock_guard<std::mutex> gard(*allocation_map_lock_);
-  auto it = allocation_map_.find(MemoryAddress);
-  if (it == allocation_map_.end())
-    return HSAKMT_STATUS_ERROR;
 
-  auto gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
-  auto code = gpu_mem->ExportPhysicalHandle(DMABufFd);
-  if (code != ErrorCode::Success)
-    return HSAKMT_STATUS_ERROR;
+  auto it = allocation_map_.upper_bound(MemoryAddress);
+  if (it != allocation_map_.begin()) {
+    --it;
+    if (it->second.dmabuf_fd == -1) {
+      auto gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle);
+      auto code = gpu_mem->ExportPhysicalHandle(DMABufFd);
+      if (code != ErrorCode::Success)
+        return HSAKMT_STATUS_ERROR;
+      it->second.dmabuf_fd = *DMABufFd;
+    }
+    *DMABufFd = dup(it->second.dmabuf_fd);
 
-  return HSAKMT_STATUS_SUCCESS;
+    *Offset = reinterpret_cast<uint64_t>(MemoryAddress) - it->second.gpu_addr;
+    return HSAKMT_STATUS_SUCCESS;
+  }
+
+  return HSAKMT_STATUS_ERROR;
 }
 
 
 HSAKMT_STATUS hsaKmtImportDMABufHandle(int DMABufFd,
-                                       HsaGraphicsResourceInfo *GraphicsResourceInfo) {
+                                       HsaGraphicsResourceInfo *GraphicsResourceInfo,
+                                       bool requiresVAddr) {
+
 
   CHECK_DXG_OPEN();
 
@@ -480,10 +497,16 @@ HSAKMT_STATUS hsaKmtImportDMABufHandle(int DMABufFd,
   wsl::thunk::GpuMemory *gpu_mem = nullptr;
   wsl::thunk::GpuMemoryCreateInfo create_info{};
   create_info.dmabuf_fd = DMABufFd;
+  create_info.flags.imported_vram_alloc_va = requiresVAddr;
 
   auto code = dev->CreateGpuMemory(create_info, &gpu_mem);
   if (code == ErrorCode::Success) {
-    void *MemoryAddress = reinterpret_cast<void *>(gpu_mem->HandleApeAddress());
+    void *MemoryAddress;
+    if (requiresVAddr)
+      MemoryAddress = reinterpret_cast<void *>(gpu_mem->GpuAddress());
+    else
+      MemoryAddress = reinterpret_cast<void*>(gpu_mem->HandleApeAddress());
+
     std::lock_guard<std::mutex> gard(*allocation_map_lock_);
     /*
      * the gpu_mem->Flags() need convert back from GpuMemoryCreateFlags to
@@ -563,9 +586,27 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtDeregisterMemory(void *MemoryAddress) {
 
   pr_debug("address %p\n", MemoryAddress);
 
+  {
+    std::lock_guard<std::mutex> gard(*allocation_map_lock_);
+
+    // IPC mem
+    auto it_ipc = allocation_map_.find(MemoryAddress);
+    if (it_ipc != allocation_map_.end()) {
+      wsl::thunk::GpuMemoryDescFlags flags;
+      flags.reserved = it_ipc->second.mem_flags_value;
+      if (flags.is_imported_vram_alloc_va) {
+        wsl::thunk::GpuMemory *gpu_mem;
+        gpu_mem = wsl::thunk::GpuMemory::Convert(it_ipc->second.handle);
+        allocation_map_.erase(it_ipc);
+        delete gpu_mem;
+        return HSAKMT_STATUS_SUCCESS;
+      }
+    }
+  }
   return HSAKMT_STATUS_SUCCESS;
 }
 
+
 HSAKMT_STATUS HSAKMTAPI hsaKmtMapMemoryToGPU(void *MemoryAddress,
                                              HSAuint64 MemorySizeInBytes,
                                              HSAuint64 *AlternateVAGPU) {
@@ -591,6 +632,31 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtMapMemoryToGPU(void *MemoryAddress,
 
   {
     std::lock_guard<std::mutex> gard(*allocation_map_lock_);
+    // IPC mem
+    auto it_ipc = allocation_map_.find(aligned_ptr);
+    if (it_ipc != allocation_map_.end()) {
+      wsl::thunk::GpuMemoryDescFlags flags;
+      flags.reserved = it_ipc->second.mem_flags_value;
+      if (flags.is_imported_vram_alloc_va) {
+        wsl::thunk::GpuMemory *gpu_mem;
+        gpu_mem = wsl::thunk::GpuMemory::Convert(it_ipc->second.handle);
+
+        auto code = gpu_mem->MapGpuVirtualAddress(gpu_mem->GpuAddress(), gpu_mem->Size());
+        if (code != ErrorCode::Success)
+          return HSAKMT_STATUS_ERROR;
+
+        code = gpu_mem->MakeResident();
+        if (code != ErrorCode::Success)
+          return HSAKMT_STATUS_ERROR;
+
+        wsl::thunk::WDDMDevice *dev = get_wddmdev(1);
+        if (!dev->WaitOnPagingFenceFromCpu())
+          return HSAKMT_STATUS_ERROR;
+
+        return HSAKMT_STATUS_SUCCESS;
+      }
+    }
+
     // GTT mem
     auto it_gtt = allocation_map_.find(aligned_ptr);
     if (it_gtt != allocation_map_.end()) {
@@ -677,6 +743,25 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtUnmapMemoryToGPU(void *MemoryAddress) {
   wsl::thunk::GpuMemoryHandle handle = nullptr;
   {
     std::lock_guard<std::mutex> gard(*allocation_map_lock_);
+
+    // IPC mem
+    auto it_ipc = allocation_map_.find(MemoryAddress);
+    if (it_ipc != allocation_map_.end()) {
+      wsl::thunk::GpuMemoryDescFlags flags;
+      flags.reserved = it_ipc->second.mem_flags_value;
+      if (flags.is_imported_vram_alloc_va) {
+        wsl::thunk::GpuMemory *gpu_mem;
+        gpu_mem = wsl::thunk::GpuMemory::Convert(it_ipc->second.handle);
+
+        auto code = gpu_mem->UnmapGpuVirtualAddress(gpu_mem->GpuAddress(), gpu_mem->Size());
+        if (code != ErrorCode::Success)
+          return HSAKMT_STATUS_ERROR;
+        gpu_mem->Evict();
+
+        return HSAKMT_STATUS_SUCCESS;
+      }
+    }
+
     auto it = allocation_map_.find(MemoryAddress);
     if (it == allocation_map_.end()) {
       return HSAKMT_STATUS_ERROR;

wddm/device.cpp

@@ -491,7 +491,7 @@ ErrorCode WDDMDevice::CreateGpuMemory(const GpuMemoryCreateInfo &create_info, Gp
   *gpu_mem = nullptr;
   auto mem = new GpuMemory(this);
   if (create_info.dmabuf_fd > 0)
-    ret = mem->ImportPhysicalHandle(create_info.dmabuf_fd);
+    ret = mem->ImportPhysicalHandle(create_info);
   else 
     ret = mem->Init(create_info);
   if (ret == ErrorCode::Success)

wddm/gpu_memory.cpp

@@ -60,7 +60,14 @@ ErrorCode GpuMemory::Init(const GpuMemoryCreateInfo &create_info) {
   desc_.flags.is_virtual = create_info.flags.virtual_alloc;
   desc_.flags.is_physical_only = create_info.flags.physical_only;
   desc_.flags.is_physical_contiguous = create_info.flags.physical_contiguous;
-  desc_.flags.is_shared = create_info.flags.interprocess;
+
+  /* we can't tell the allocation is regular vmm or ipc mem at creation stage,
+     they share same creation parameters, so forcing all vram allocations to
+     sharable to support IPC mem */
+  if (create_info.flags.interprocess ||
+      desc_.domain == thunk_proxy::AllocDomain::kLocal)
+    desc_.flags.is_shared = true;
+
   desc_.flags.is_locked = create_info.flags.locked;
 
   desc_.size = AdjustSize(desc_.client_size);
@@ -386,8 +393,9 @@ ErrorCode GpuMemory::ExportPhysicalHandle(int* dmabuf_fd, uint32_t flags) {
 }
 
 
-ErrorCode GpuMemory::ImportPhysicalHandle(int dmabuf_fd) {
+ErrorCode GpuMemory::ImportPhysicalHandle(const GpuMemoryCreateInfo &create_info) {
   D3DKMT_QUERYRESOURCEINFOFROMNTHANDLE query_args;
+  int dmabuf_fd = create_info.dmabuf_fd;
 
   if (dmabuf_fd <= 0)
     return ErrorCode::InvalidateParams;
@@ -426,6 +434,8 @@ ErrorCode GpuMemory::ImportPhysicalHandle(int dmabuf_fd) {
     return ErrorCode::OutOfMemory;
   }
 
+  auto guard = MakeScopeGuard([&open_info]() { free(open_info); });
+
   alloc_handles_ptr_ = new WinAllocationHandle[query_args.NumAllocations];
 
   D3DKMT_OPENRESOURCEFROMNTHANDLE open_args;
@@ -448,7 +458,7 @@ ErrorCode GpuMemory::ImportPhysicalHandle(int dmabuf_fd) {
   if (ret != ErrorCode::Success) {
     ret = ErrorCode::InvalidateParams;
     pr_err("open resource failed %d\n", static_cast<int>(ret));
-    goto err_out;
+    return ret;
   }
 
   desc_.size = shared_info.size;
@@ -462,14 +472,18 @@ ErrorCode GpuMemory::ImportPhysicalHandle(int dmabuf_fd) {
   for (int i = 0; i < num_allocations_; i++)
     alloc_handles_ptr_[i] = open_info[i].hAllocation;
 
-  free(open_info);
-  return device_->HandleApertureAlloc(desc_.size, &desc_.handle_ape_addr);
 
-err_out:
-  delete[] alloc_handles_ptr_;
-  alloc_handles_ptr_ = nullptr;
-  free(open_info);
-  return ret;
+  if (create_info.flags.imported_vram_alloc_va) {
+    desc_.flags.is_imported_vram_alloc_va = true;
+
+    ret = ReserveGpuVirtualAddress(create_info.va_hint, desc_.size, create_info.alignment);
+    if (ret != ErrorCode::Success)
+      pr_err("failed to allocate svm range, error:%d\n", static_cast<int>(ret));
+
+    return ret;
+  } else {
+    return device_->HandleApertureAlloc(desc_.size, &desc_.handle_ape_addr);
+  }
 }
 
 } // namespace thunk