Support mapping and unmapping memory handles

Add support for mapping and unmapping memory handles to virtual address ranges. This is part of patch series for Virtual Memory API. Change-Id: If512d49ff4211e68f2064249add607a3200e458a [ROCm/ROCR-Runtime commit: 179dcf1c77]
2022-12-18 23:57:44 +00:00
@@ -1250,6 +1250,15 @@ hsa_status_t HSA_API hsa_amd_vmem_handle_release(hsa_amd_vmem_alloc_handle_t mem
  return amdExtTable->hsa_amd_vmem_handle_release_fn(memory_handle);
 }

+hsa_status_t HSA_API hsa_amd_vmem_map(void* va, size_t size, size_t in_offset,
+                                      hsa_amd_vmem_alloc_handle_t memory_handle, uint64_t flags) {
+  return amdExtTable->hsa_amd_vmem_map_fn(va, size, in_offset, memory_handle, flags);
+}
+
+hsa_status_t HSA_API hsa_amd_vmem_unmap(void* va, size_t size) {
+  return amdExtTable->hsa_amd_vmem_unmap_fn(va, size);
+}
+
 // Tools only table interfaces.
 namespace rocr {

@@ -313,6 +313,13 @@ hsa_status_t hsa_amd_vmem_handle_create(hsa_amd_memory_pool_t pool, size_t size,
 // Mirrors Amd Extension Apis
 hsa_status_t hsa_amd_vmem_handle_release(hsa_amd_vmem_alloc_handle_t memory_handle);

+// Mirrors Amd Extension Apis
+hsa_status_t hsa_amd_vmem_map(void* va, size_t size, size_t in_offset,
+                              hsa_amd_vmem_alloc_handle_t memory_handle, uint64_t flags);
+
+// Mirrors Amd Extension Apis
+hsa_status_t hsa_amd_vmem_unmap(void* va, size_t size);
+
 }  // namespace amd
 }  // namespace rocr

@@ -364,6 +364,11 @@ class Runtime {

  hsa_status_t VMemoryHandleRelease(hsa_amd_vmem_alloc_handle_t memoryHandle);

+  hsa_status_t VMemoryHandleMap(void* va, size_t size, size_t in_offset,
+                                hsa_amd_vmem_alloc_handle_t memoryHandle, uint64_t flags);
+
+  hsa_status_t VMemoryHandleUnmap(void* va, size_t size);
+
  const std::vector<Agent*>& cpu_agents() { return cpu_agents_; }

  const std::vector<Agent*>& gpu_agents() { return gpu_agents_; }
@@ -681,6 +686,7 @@ class Runtime {

 private:
  void CheckVirtualMemApiSupport();
+  int GetAmdgpuDeviceArgs(Agent* agent, amdgpu_bo_handle bo, int* drm_fd, uint64_t* cpu_addr);

  bool virtual_mem_api_supported_;

@@ -721,9 +727,44 @@ class Runtime {
    ThunkHandle thunk_handle;  // handle returned by hsaKmtAllocMemory(NoAddress = 1)
    MemoryRegion::AllocateFlags alloc_flag;
  };
-
  std::map<ThunkHandle, MemoryHandle> memory_handle_map_;

+  struct MappedHandle {
+    MappedHandle()
+        : mem_handle(NULL),
+          address_handle(NULL),
+          offset(0),
+          mmap_offset(0),
+          size(0),
+          drm_fd(-1),
+          drm_cpu_addr(NULL),
+          ldrm_bo(0) {}
+
+    MappedHandle(MemoryHandle* mem_handle, AddressHandle* address_handle, uint64_t offset,
+                 size_t size, int drm_fd, void* drm_cpu_addr, hsa_access_permission_t perm,
+                 amdgpu_bo_handle bo)
+        : mem_handle(mem_handle),
+          address_handle(address_handle),
+          offset(offset),
+          mmap_offset(0),
+          size(size),
+          drm_fd(drm_fd),
+          drm_cpu_addr(drm_cpu_addr),
+          ldrm_bo(bo) {}
+
+    __forceinline core::Agent* agentOwner() const { return mem_handle->region->owner(); }
+
+    MemoryHandle* mem_handle;
+    AddressHandle* address_handle;
+    uint64_t offset;
+    uint64_t mmap_offset;
+    size_t size;
+    int drm_fd;
+    void* drm_cpu_addr;  // CPU Buffer address
+    amdgpu_bo_handle ldrm_bo;
+  };
+  std::map<const void*, MappedHandle> mapped_handle_map_;  // Indexed by VA
+
  // Frees runtime memory when the runtime library is unloaded if safe to do so.
  // Failure to release the runtime indicates an incorrect application but is
  // common (example: calls library routines at process exit).
@@ -406,6 +406,8 @@ void HsaApiTable::UpdateAmdExts() {
  amd_ext_api.hsa_amd_vmem_address_free_fn = AMD::hsa_amd_vmem_address_free;
  amd_ext_api.hsa_amd_vmem_handle_create_fn = AMD::hsa_amd_vmem_handle_create;
  amd_ext_api.hsa_amd_vmem_handle_release_fn = AMD::hsa_amd_vmem_handle_release;
+  amd_ext_api.hsa_amd_vmem_map_fn = AMD::hsa_amd_vmem_map;
+  amd_ext_api.hsa_amd_vmem_unmap_fn = AMD::hsa_amd_vmem_unmap;
 }

 void LoadInitialHsaApiTable() {
@@ -1254,5 +1254,27 @@ hsa_status_t hsa_amd_vmem_handle_release(hsa_amd_vmem_alloc_handle_t memory_hand
  CATCH;
 }

+hsa_status_t hsa_amd_vmem_map(void* va, size_t size, size_t in_offset,
+                              hsa_amd_vmem_alloc_handle_t memory_handle, uint64_t flags) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(va);
+  IS_ZERO(size);
+
+  return core::Runtime::runtime_singleton_->VMemoryHandleMap(va, size, in_offset, memory_handle,
+                                                             flags);
+  CATCH;
+}
+
+hsa_status_t hsa_amd_vmem_unmap(void* va, size_t size) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(va);
+  IS_ZERO(size);
+
+  return core::Runtime::runtime_singleton_->VMemoryHandleUnmap(va, size);
+  CATCH;
+}
+
 }   //  namespace amd
 }   //  namespace rocr
@@ -49,6 +49,8 @@
 #include <vector>
 #include <list>
 #include <dlfcn.h>
+#include <amdgpu_drm.h>
+#include <sys/mman.h>

 #include "core/common/shared.h"
 #include "core/inc/hsa_ext_interface.h"
@@ -1594,6 +1596,27 @@ int fn_amdgpu_device_get_fd_nosupport(HsaAMDGPUDeviceHandle device_handle) {
  return -1;
 }

+int Runtime::GetAmdgpuDeviceArgs(Agent* agent, amdgpu_bo_handle bo, int* drm_fd,
+                                 uint64_t* cpu_addr) {
+  int renderFd = fn_amdgpu_device_get_fd(static_cast<AMD::GpuAgent*>(agent)->libDrmDev());
+  if (renderFd < 0) return HSA_STATUS_ERROR;
+
+  uint32_t gem_handle = 0;
+  if (amdgpu_bo_export(bo, amdgpu_bo_handle_type_kms, &gem_handle)) return HSA_STATUS_ERROR;
+
+  union drm_amdgpu_gem_mmap args;
+  memset(&args, 0, sizeof(args));
+  /* Query the buffer address (args.addr_ptr).
+   * The kernel driver ignores the offset and size parameters. */
+  args.in.handle = gem_handle;
+  if (drmCommandWriteRead(renderFd, DRM_AMDGPU_GEM_MMAP, &args, sizeof(args)))
+    return HSA_STATUS_ERROR;
+
+  *drm_fd = renderFd;
+  *cpu_addr = args.out.addr_ptr;
+  return HSA_STATUS_SUCCESS;
+}
+
 void Runtime::CheckVirtualMemApiSupport() {
  virtual_mem_api_supported_ = false;
  // TODO: May have to change the minor version required once Thunk merges changes into amd-staging
@@ -2492,5 +2515,131 @@ hsa_status_t Runtime::VMemoryHandleRelease(hsa_amd_vmem_alloc_handle_t memoryOnl
  return HSA_STATUS_SUCCESS;
 }

+__forceinline uint64_t drm_perm(hsa_access_permission_t perm) {
+  switch (perm) {
+    case HSA_ACCESS_PERMISSION_RO:
+      return AMDGPU_VM_PAGE_READABLE;
+    case HSA_ACCESS_PERMISSION_WO:
+      return AMDGPU_VM_PAGE_WRITEABLE;
+    case HSA_ACCESS_PERMISSION_RW:
+      return AMDGPU_VM_PAGE_READABLE | AMDGPU_VM_PAGE_WRITEABLE;
+    case HSA_ACCESS_PERMISSION_NONE:
+      return 0;
+    default:
+      break;
+  }
+
+  return 0;
+}
+
+__forceinline int mmap_perm(hsa_access_permission_t perms) {
+  switch (perms) {
+    case HSA_ACCESS_PERMISSION_RO:
+      return PROT_READ;
+    case HSA_ACCESS_PERMISSION_WO:
+      return PROT_WRITE;
+    case HSA_ACCESS_PERMISSION_RW:
+      return PROT_READ | PROT_WRITE;
+    case HSA_ACCESS_PERMISSION_NONE:
+      return PROT_NONE;
+    default:
+      break;
+  }
+
+  return 0;
+}
+
+hsa_status_t Runtime::VMemoryHandleMap(void* va, size_t size, size_t in_offset,
+                                       hsa_amd_vmem_alloc_handle_t memoryOnlyHandle,
+                                       uint64_t flags) {
+  int drm_fd, dmabuf_fd = 0;
+  uint64_t offset = 0, ret;
+  uint64_t drm_cpu_addr = 0;
+  amdgpu_bo_handle ldrm_bo = 0;
+  bool reservedAddressFound = false;
+
+  ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
+  auto reservedAddressIt = reserved_address_map_.upper_bound(va);
+  if (reservedAddressIt != reserved_address_map_.begin()) {
+    reservedAddressIt--;
+    if ((reservedAddressIt->first <= va) &&
+        ((va + size) <= (reservedAddressIt->first + reservedAddressIt->second.size))) {
+      reservedAddressFound = true;
+    }
+  }
+  if (!reservedAddressFound) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+
+  /* Confirm that this VA range has not been mapped yet */
+  auto upperMappedHandleIt = mapped_handle_map_.upper_bound(va);
+  if (upperMappedHandleIt != mapped_handle_map_.begin()) {
+    upperMappedHandleIt--;
+    if (upperMappedHandleIt->first + upperMappedHandleIt->second.size > va)
+      return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+  auto lowerMappedHandleIt = mapped_handle_map_.lower_bound(va);
+  if (lowerMappedHandleIt != mapped_handle_map_.end()) {
+    if (va + size > lowerMappedHandleIt->first) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  auto memoryHandleIt = memory_handle_map_.find(reinterpret_cast<void*>(memoryOnlyHandle.handle));
+  if (memoryHandleIt == memory_handle_map_.end()) {
+    debug_warning(false && "Can't find memory handle");
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  ret = hsaKmtExportDMABufHandle(memoryHandleIt->first, size, &dmabuf_fd, &offset);
+  if (ret != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  assert(offset == 0);
+
+  AMD::GpuAgent* agent = static_cast<AMD::GpuAgent*>(memoryHandleIt->second.agentOwner());
+  amdgpu_bo_import_result res;
+  ret = amdgpu_bo_import(agent->libDrmDev(), amdgpu_bo_handle_type_dma_buf_fd, dmabuf_fd, &res);
+  if (ret) return HSA_STATUS_ERROR;
+
+  close(dmabuf_fd);
+
+  ldrm_bo = res.buf_handle;
+  ret = GetAmdgpuDeviceArgs(agent, ldrm_bo, &drm_fd, &drm_cpu_addr);
+  if (ret) return HSA_STATUS_ERROR;
+  mapped_handle_map_[va] =
+      MappedHandle(&memoryHandleIt->second, &reservedAddressIt->second, offset, size, drm_fd,
+                   reinterpret_cast<void*>(drm_cpu_addr), HSA_ACCESS_PERMISSION_NONE, ldrm_bo);
+
+  reservedAddressIt->second.use_count++;
+  memoryHandleIt->second.use_count++;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t Runtime::VMemoryHandleUnmap(void* va, size_t size) {
+  int ret;
+  ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
+
+  auto mappedHandleIt = mapped_handle_map_.find(va);
+  if (mappedHandleIt == mapped_handle_map_.end()) return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+  if (mappedHandleIt->second.size != size) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+
+  if (mappedHandleIt->second.ldrm_bo) ret = amdgpu_bo_free(mappedHandleIt->second.ldrm_bo);
+
+  if (ret) return HSA_STATUS_ERROR;
+
+  assert(mappedHandleIt->second.address_handle->use_count >= 1);
+  mappedHandleIt->second.address_handle->use_count--;
+  assert(mappedHandleIt->second.mem_handle->use_count >= 1);
+  mappedHandleIt->second.mem_handle->use_count--;
+
+  if (!mappedHandleIt->second.mem_handle->use_count &&
+      !mappedHandleIt->second.mem_handle->ref_count) {
+    // User called VMemoryHandleRelease while this mapping was still outstanding. We need to delete
+    // the MemoryHandle as is the last MappedHandle that was using it
+    mappedHandleIt->second.mem_handle->region->Free(mappedHandleIt->second.mem_handle->thunk_handle,
+                                                    mappedHandleIt->second.mem_handle->size);
+    memory_handle_map_.erase(mappedHandleIt->second.mem_handle->thunk_handle);
+  }
+
+  mapped_handle_map_.erase(mappedHandleIt);
+  return HSA_STATUS_SUCCESS;
+}
 }  // namespace core
 }  // namespace rocr
@@ -237,6 +237,8 @@ global:
 	hsa_amd_vmem_address_free;
 	hsa_amd_vmem_handle_create;
 	hsa_amd_vmem_handle_release;
+	hsa_amd_vmem_map;
+	hsa_amd_vmem_unmap;

 local:
    *;
@@ -313,6 +313,10 @@ typedef struct hsa_dim3_s {
 * @brief Access permissions.
 */
 typedef enum {
+  /**
+   * Used to remove existing access
+   */
+  HSA_ACCESS_PERMISSION_NONE = 0,
  /**
   * Read-only access.
   */
@@ -198,6 +198,8 @@ struct AmdExtTable {
  decltype(hsa_amd_vmem_address_free)* hsa_amd_vmem_address_free_fn;
  decltype(hsa_amd_vmem_handle_create)* hsa_amd_vmem_handle_create_fn;
  decltype(hsa_amd_vmem_handle_release)* hsa_amd_vmem_handle_release_fn;
+  decltype(hsa_amd_vmem_map)* hsa_amd_vmem_map_fn;
+  decltype(hsa_amd_vmem_unmap)* hsa_amd_vmem_unmap_fn;
 };

 // Table to export HSA Core Runtime Apis
@@ -2800,6 +2800,50 @@ hsa_status_t hsa_amd_vmem_handle_create(hsa_amd_memory_pool_t pool, size_t size,
 */
 hsa_status_t hsa_amd_vmem_handle_release(hsa_amd_vmem_alloc_handle_t memory_handle);

+/*
+ * @brief Map a virtual memory handle
+ *
+ * Map a virtual memory handle to a reserved address range. The virtual address requested must be
+ * within a previously reserved address range. @p va and (@p va + size) must be must be within
+ * (va + size) of the previous allocated address range.
+ * @p size must be equal to size of the @p memory_handle
+ * hsa_amd_vmem_set_access needs to be called to make the memory accessible to specific agents
+ *
+ * @param[in] va virtual address range where memory will be mapped
+ * @param[in] size of memory mapping
+ * @param[in] in_offset offset into memory. Currently unsupported
+ * @param[in] memory_handle virtual memory handle to be mapped
+ * @param[in] flags. Currently unsupported
+ *
+ * @retval ::HSA_STATUS_SUCCESS Memory mapped successfully
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT va, size or memory_handle are invalid
+ *
+ * @retval ::HSA_STATUS_ERROR_OUT_OF_RESOURCES Insufficient resources
+ *
+ * @retval ::HSA_STATUS_ERROR Unexpected internal error
+ */
+hsa_status_t hsa_amd_vmem_map(void* va, size_t size, size_t in_offset,
+                              hsa_amd_vmem_alloc_handle_t memory_handle, uint64_t flags);
+
+/*
+ * @brief Unmap a virtual memory handle
+ *
+ * Unmap previously mapped virtual address range
+ *
+ * @param[in] va virtual address range where memory will be mapped
+ * @param[in] size of memory mapping
+ *
+ * @retval ::HSA_STATUS_SUCCESS Memory backing unmapped successfully
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ALLOCATION memory_handle is invalid
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT size is invalid
+ *
+ * @retval ::HSA_STATUS_ERROR Unexpected internal error
+ */
+hsa_status_t hsa_amd_vmem_unmap(void* va, size_t size);
+
 #ifdef __cplusplus
 }  // end extern "C" block
 #endif