wsl/libhsakmt: move local heap and va_Mgr from device to thunk runtime

In multi-GPU, local heap space is shared between all GPUs, not belongs to
specific one GPU, so move it from wddm device (which presents a specific GPU)
to thunk runtime which has gloable view, can manage local heap for all GPUs.

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

libhsakmt.h

@@ -32,6 +32,7 @@
 #include "hsakmt/hsakmt.h"
 #include "hsakmt/hsakmt_drm.h"
 
+#include "impl/wddm/va_mgr.h"
 #include "impl/wddm/types.h"
 #include "impl/wddm/device.h"
 
@@ -60,8 +61,16 @@ struct hsakmtRuntime {
     check_avail_sysram(false),
     max_single_alloc_size(0),
     enable_thunk_sub_allocator(0),
+    local_heap_space_start_(0),
+    local_heap_space_size_(0),
     default_node(1) {}
 
+  void HeapInit();
+  void HeapFini();
+  bool ReserveLocalHeapSpace();
+  bool FreeLocalHeapSpace();
+  void InitLocalHeapMgr();
+
   pthread_mutex_t hsakmt_mutex;
   const char *dxg_device_name = "/dev/dxg";
   long page_size;
@@ -79,6 +88,13 @@ struct hsakmtRuntime {
   size_t max_single_alloc_size;
   int enable_thunk_sub_allocator;
   uint32_t default_node;
+
+  /* local heap means bo's backend is vram of all GPUs */
+  uint64_t local_heap_space_start_;
+  uint64_t local_heap_space_size_;
+
+  /* manage the reserved local heap space which shared by CPU and GPUs */
+  std::unique_ptr<wsl::thunk::VaMgr> local_heap_mgr_;
 };
 
 extern hsakmtRuntime *dxg_runtime;

openclose.cpp

@@ -26,6 +26,7 @@
 #include <cstring>
 #include <sys/types.h>
 #include <sys/stat.h>
+#include <sys/mman.h>
 #include <fcntl.h>
 #include <unistd.h>
 #include <cstdio>
@@ -33,7 +34,130 @@
 #include <cassert>
 #include "libhsakmt.h"
 
+
 hsakmtRuntime *dxg_runtime = new hsakmtRuntime();
+
+void hsakmtRuntime::HeapInit() {
+    ReserveLocalHeapSpace();
+    InitLocalHeapMgr();
+}
+
+void hsakmtRuntime::HeapFini() {
+    FreeLocalHeapSpace();
+}
+
+/*
+ * To find the avaliable same range for cpu
+ * virtual space and gpu virtual space.
+ * sys_va_size of cpu va range is larger 1G
+ * than gpu va range, otherwise ReserveGPUVirtualAddress
+ * will return error.
+ */
+bool hsakmtRuntime::ReserveLocalHeapSpace() {
+    uint64_t sys_va[16] = {0};
+    uint64_t local_va;
+    uint64_t sys_va_size;
+    int match_index = -1;
+    uint64_t align = 0x40000000; /* 1G */
+    void* ptr = NULL;
+
+    wsl::thunk::WDDMDevice* device;
+    uint64_t total_local_size = 0;
+    size_t num_adapters = get_num_wddmdev();
+    for (uint32_t j = 0; j < num_adapters; j++) {
+        device = get_wddmdev(j+1);
+        if (device == nullptr)
+            return -1;
+        total_local_size += wsl::AlignUp(device->LocalHeapSize(), align) * 4;
+    }
+
+    local_heap_space_start_ = 0;
+    local_heap_space_size_ = total_local_size;
+    sys_va_size = local_heap_space_size_ + align;
+
+    /* it will retry 16 times to find the avaliable range. */
+    for (int i = 0; i < 16; i++) {
+        local_va = 0;
+        ptr = mmap(NULL, sys_va_size , PROT_NONE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
+        if (ptr == MAP_FAILED) {
+            pr_err("fail to reserve cpu va in %d time!\n", i);
+            break;
+        }
+
+        sys_va[i] = (uint64_t)ptr;
+
+        int match_cnt = 0;
+        for (uint32_t j = 0; j < num_adapters; j++) {
+            device = get_wddmdev(j+1);
+            uint64_t start = (local_heap_space_start_ == 0) ? (uint64_t)ptr : local_heap_space_start_;
+            uint64_t end = start + ((local_heap_space_start_ == 0) ? sys_va_size : local_heap_space_size_) + 1;
+
+            if (wsl::thunk::d3dthunk::ReserveGpuVirtualAddress(
+                        device->GetAdapter(), local_heap_space_size_,
+                        start,
+                        end, &local_va) == ErrorCode::Success) {
+
+                match_cnt++;
+                local_heap_space_start_ = local_va;
+                pr_debug("success to reserve gpu va %lx and va cpu %p in %d time\n",
+                        local_va, ptr, i);
+            } else {
+                pr_err("%s fail to reserve gpu va for cpu va %p in %d time!\n",
+                        __FUNCTION__, ptr, i);
+            }
+        }
+
+        if (match_cnt == num_adapters) {
+                match_index = i;
+                break;
+        }
+    }
+
+    if (match_index >= 0) {
+        /* release cpu unused ranges*/
+        uint64_t left_size = local_va - sys_va[match_index];
+        uint64_t right_size = align - left_size;
+        if ((left_size > 0) && munmap((void*)sys_va[match_index], left_size))
+            pr_err("fail to unmap left %lx with size %lx\n", sys_va[match_index], left_size);
+        if ((right_size > 0) && munmap((void*)(local_va + local_heap_space_size_), right_size))
+            pr_err("fail to unmap right %lx with size %lx\n", (local_va + local_heap_space_size_), right_size);
+    } else {
+        pr_err("fail to reserve Local Heap Space!\n");
+        local_heap_space_start_ = 0;
+        local_heap_space_size_ = 0;
+    }
+
+    /* free match fail address for cpu va */
+    int free = match_index >= 0 ? match_index : 16;
+    for (int j = 0; j < free; j++) {
+        if (sys_va[j] != 0 && munmap((void*)sys_va[j], sys_va_size)) {
+            pr_err("fail to unmap %d %lx\n", j, sys_va[j]);
+        }
+    }
+
+    return match_index >= 0;
+}
+
+bool hsakmtRuntime::FreeLocalHeapSpace() {
+    wsl::thunk::WDDMDevice* device;
+    size_t num_adapters = get_num_wddmdev();
+    for (uint32_t j = 0; j < num_adapters; j++) {
+        device = get_wddmdev(j+1);
+        if (device == nullptr)
+            return -1;
+        wsl::thunk::d3dthunk::FreeGpuVirtualAddress(device->GetAdapter(), local_heap_space_start_, local_heap_space_size_);
+    }
+
+    void *cpu = (void *)local_heap_space_start_;
+    return munmap(cpu, local_heap_space_size_) == 0;
+}
+
+void hsakmtRuntime::InitLocalHeapMgr() {
+  local_heap_mgr_ = std::make_unique<wsl::thunk::VaMgr>(local_heap_space_start_,
+                                          local_heap_space_size_,
+                                          DEFAULT_GPU_PAGE_SIZE);
+}
+
 /* is_forked_child detects when the process has forked since the last
  * time this function was called. We cannot rely on pthread_atfork
  * because the process can fork without calling the fork function in

topology.cpp

@@ -373,6 +373,7 @@ HSAKMT_STATUS topology_sysfs_get_system_props(HsaSystemProperties& props) {
 
   dxg_topology->num_sysfs_nodes = num_adapters + 1;
 
+  dxg_runtime->HeapFini();
   for (auto device : dxg_topology->wdevices_)
     delete device;
   dxg_topology->wdevices_.clear();
@@ -383,6 +384,7 @@ HSAKMT_STATUS topology_sysfs_get_system_props(HsaSystemProperties& props) {
     assert(device && "Create WDDM Device fail");
     dxg_topology->wdevices_.push_back(device);
   }
+  dxg_runtime->HeapInit();
   props.NumNodes = dxg_topology->num_sysfs_nodes;
   if (dxg_runtime->default_node > num_adapters)
     dxg_runtime->default_node = num_adapters;

wddm/device.cpp

@@ -67,16 +67,13 @@ WDDMDevice::WDDMDevice(D3DKMT_HANDLE adapter, LUID adapter_luid, uint32_t node_i
   CreateDevice();
   SetPowerOptimization(false);
   CreatePagingQueue();
-  ReserveLocalHeapSpace();
   ReserveSystemHeapSpace();
   InitHandleApertureSpace();
-  InitVaMgr();
   InitHandleApertureMgr();
   InitCmdbufInfo();
 }
 
 WDDMDevice::~WDDMDevice() {
-  FreeLocalHeapSpace();
   FreeSystemHeapSpace();
   DestroyPagingQueue();
   SetPowerOptimization(true);
@@ -308,87 +305,6 @@ bool WDDMDevice::FreeSystemHeapSpace(void) {
   return true;
 }
 
-/*
- * To find the avaliable same range for cpu
- * virtual space and gpu virtual space.
- * sys_va_size of cpu va range is larger 1G
- * than gpu va range, otherwise ReserveGPUVirtualAddress
- * will return error.
- */
-bool WDDMDevice::ReserveLocalHeapSpace(void) {
-  uint64_t sys_va[16] = {0};
-  uint64_t local_va;
-  uint64_t sys_va_size;
-  int match_index = -1;
-  uint64_t align = 0x40000000; /* 1G */
-  void* ptr = NULL;
-
-  local_heap_space_start_ = 0;
-  local_heap_space_size_ = AlignUp(LocalHeapSize(), align) * 4;
-  sys_va_size = local_heap_space_size_ + align;
-
-  /* it will retry 16 times to find the avaliable range. */
-  for (int i = 0; i < 16; i++) {
-    local_va = 0;
-    ptr = mmap(NULL, sys_va_size , PROT_NONE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
-    if (ptr == MAP_FAILED) {
-      pr_err("fail to reserve cpu va in %d time!\n", i);
-      break;
-    }
-
-    sys_va[i] = (uint64_t)ptr;
-
-    if (d3dthunk::ReserveGpuVirtualAddress(
-          adapter_, local_heap_space_size_,
-          (uint64_t)ptr,
-          (uint64_t)ptr + sys_va_size, &local_va) == ErrorCode::Success) {
-
-      match_index = i;
-      local_heap_space_start_ = local_va;
-      pr_debug("success to reserve gpu va %lx and va cpu %p in %d time\n",
-               local_va, ptr, i);
-      break;
-    } else {
-      pr_err("%s fail to reserve gpu va for cpu va %p in %d time!\n",
-              __FUNCTION__, ptr, i);
-    }
-  }
-
-  if (match_index >= 0) {
-    /* release cpu unused ranges*/
-    uint64_t left_size = local_va - sys_va[match_index];
-    uint64_t right_size = align - left_size;
-    if ((left_size > 0) && munmap((void*)sys_va[match_index], left_size))
-      pr_err("fail to unmap left %lx with size %lx\n", sys_va[match_index], left_size);
-    if ((right_size > 0) && munmap((void*)(local_va + local_heap_space_size_), right_size))
-      pr_err("fail to unmap right %lx with size %lx\n", (local_va + local_heap_space_size_), right_size);
-  } else {
-      pr_err("fail to reserve Local Heap Space!\n");
-  }
-
-  /* free match fail address for cpu va */
-  int free = match_index >= 0 ? match_index : 16;
-  for (int j = 0; j < free; j++) {
-    if (sys_va[j] != 0 && munmap((void*)sys_va[j], sys_va_size)) {
-      pr_err("fail to unmap %d %lx\n", j, sys_va[j]);
-    }
-  }
-
-  return match_index >= 0;
-}
-
-bool WDDMDevice::FreeLocalHeapSpace(void) {
-  d3dthunk::FreeGpuVirtualAddress(adapter_, local_heap_space_start_, local_heap_space_size_);
-  void *cpu = (void *)local_heap_space_start_;
-  return munmap(cpu, local_heap_space_size_) == 0;
-}
-
-void WDDMDevice::InitVaMgr() {
-  local_va_mgr_ = std::make_unique<VaMgr>(local_heap_space_start_,
-                                          local_heap_space_size_,
-                                          DEFAULT_GPU_PAGE_SIZE);
-}
-
 void WDDMDevice::InitHandleApertureMgr() {
   handle_aperture_mgr_ = std::make_unique<VaMgr>(handle_aperture_start_,
                                                  handle_aperture_size_,
@@ -474,7 +390,7 @@ ErrorCode WDDMDevice::ReserveGpuVirtualAddress(const thunk_proxy::AllocDomain do
     if (domain == thunk_proxy::kLocal && size >= GPU_HUGE_PAGE_SIZE)
       align = GPU_HUGE_PAGE_SIZE;
 
-    gpu_addr = local_va_mgr_->Alloc(size, align, hit_base_addr);
+    gpu_addr = dxg_runtime->local_heap_mgr_->Alloc(size, align, hit_base_addr);
     if (gpu_addr == 0)
       code = ErrorCode::OutOfGpuMemory;
 
@@ -499,7 +415,7 @@ ErrorCode WDDMDevice::FreeGpuVirtualAddress(const thunk_proxy::AllocDomain domai
 
       code = d3dthunk::FreeGpuVirtualAddress(&free_args);
   } else {
-    local_va_mgr_->Free(gpu_addr);
+      dxg_runtime->local_heap_mgr_->Free(gpu_addr);
   }
 
   return code;