Add retain handle and get allocation properties

Support function to retain allocation handle for memory mappings. The get allocation properties function will return the current allocation properties for existing memory mappings. This is part of patch series for Virtual Memory API. Change-Id: I0a53a11b6efc2b5bf9d463512a489a2abd812551 [ROCm/ROCR-Runtime commit: 687eb043d4]
2022-12-19 00:09:08 +00:00
parent 0bcc573ed7
commit bf41567189
9 changed files with 124 additions and 0 deletions
@@ -1281,6 +1281,17 @@ hsa_status_t HSA_API hsa_amd_vmem_import_shareable_handle(int dmabuf_fd,
  return amdExtTable->hsa_amd_vmem_import_shareable_handle_fn(dmabuf_fd, handle);
 }

+hsa_status_t HSA_API hsa_amd_vmem_retain_alloc_handle(hsa_amd_vmem_alloc_handle_t* handle,
+                                                      void* addr) {
+  return amdExtTable->hsa_amd_vmem_retain_alloc_handle_fn(handle, addr);
+}
+
+hsa_status_t HSA_API hsa_amd_vmem_get_alloc_properties_from_handle(
+    hsa_amd_vmem_alloc_handle_t alloc_handle, hsa_amd_memory_pool_t* pool,
+    hsa_amd_memory_type_t* type) {
+  return amdExtTable->hsa_amd_vmem_get_alloc_properties_from_handle_fn(alloc_handle, pool, type);
+}
+
 // Tools only table interfaces.
 namespace rocr {

@@ -337,6 +337,15 @@ hsa_status_t hsa_amd_vmem_export_shareable_handle(int* dmabuf_fd,
 // Mirrors Amd Extension Apis
 hsa_status_t hsa_amd_vmem_import_shareable_handle(int dmabuf_fd,
                                                  hsa_amd_vmem_alloc_handle_t* handle);
+
+// Mirrors Amd Extension Apis
+hsa_status_t hsa_amd_vmem_retain_alloc_handle(hsa_amd_vmem_alloc_handle_t* allocHandle, void* addr);
+
+// Mirrors Amd Extension Apis
+hsa_status_t hsa_amd_vmem_get_alloc_properties_from_handle(hsa_amd_vmem_alloc_handle_t allocHandle,
+                                                           hsa_amd_memory_pool_t* pool,
+                                                           hsa_amd_memory_type_t* type);
+
 }  // namespace amd
 }  // namespace rocr

@@ -382,6 +382,12 @@ class Runtime {
  hsa_status_t VMemoryImportShareableHandle(const int dmabuf_fd,
                                            hsa_amd_vmem_alloc_handle_t* handle);

+  hsa_status_t VMemoryRetainAllocHandle(hsa_amd_vmem_alloc_handle_t* memoryHandle, void* addr);
+
+  hsa_status_t VMemoryGetAllocPropertiesFromHandle(const hsa_amd_vmem_alloc_handle_t memoryHandle,
+                                                   const core::MemoryRegion** mem_region,
+                                                   hsa_amd_memory_type_t* type);
+
  const std::vector<Agent*>& cpu_agents() { return cpu_agents_; }

  const std::vector<Agent*>& gpu_agents() { return gpu_agents_; }
@@ -412,6 +412,9 @@ void HsaApiTable::UpdateAmdExts() {
  amd_ext_api.hsa_amd_vmem_get_access_fn = AMD::hsa_amd_vmem_get_access;
  amd_ext_api.hsa_amd_vmem_export_shareable_handle_fn = AMD::hsa_amd_vmem_export_shareable_handle;
  amd_ext_api.hsa_amd_vmem_import_shareable_handle_fn = AMD::hsa_amd_vmem_import_shareable_handle;
+  amd_ext_api.hsa_amd_vmem_retain_alloc_handle_fn = AMD::hsa_amd_vmem_retain_alloc_handle;
+  amd_ext_api.hsa_amd_vmem_get_alloc_properties_from_handle_fn =
+      AMD::hsa_amd_vmem_get_alloc_properties_from_handle;
 }

 void LoadInitialHsaApiTable() {
@@ -1326,5 +1326,35 @@ hsa_status_t hsa_amd_vmem_import_shareable_handle(int dmabuf_fd,
  CATCH;
 }

+hsa_status_t hsa_amd_vmem_retain_alloc_handle(hsa_amd_vmem_alloc_handle_t* allocHandle,
+                                              void* addr) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(addr);
+
+  return core::Runtime::runtime_singleton_->VMemoryRetainAllocHandle(allocHandle, addr);
+  CATCH;
+}
+
+hsa_status_t hsa_amd_vmem_get_alloc_properties_from_handle(hsa_amd_vmem_alloc_handle_t allocHandle,
+                                                           hsa_amd_memory_pool_t* pool,
+                                                           hsa_amd_memory_type_t* type) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(pool);
+  IS_BAD_PTR(type);
+
+  const core::MemoryRegion* mem_region = NULL;
+  hsa_status_t ret = core::Runtime::runtime_singleton_->VMemoryGetAllocPropertiesFromHandle(
+      allocHandle, &mem_region, type);
+  if (ret == HSA_STATUS_SUCCESS) {
+    hsa_region_t region = core::MemoryRegion::Convert(mem_region);
+    pool->handle = region.handle;
+  }
+
+  return ret;
+  CATCH;
+}
+
 }   //  namespace amd
 }   //  namespace rocr
@@ -2869,5 +2869,31 @@ hsa_status_t Runtime::VMemoryImportShareableHandle(int dmabuf_fd,
  return HSA_STATUS_SUCCESS;
 }

+hsa_status_t Runtime::VMemoryRetainAllocHandle(hsa_amd_vmem_alloc_handle_t* mapped_handle,
+                                               void* va) {
+  auto mappedHandleIt = mapped_handle_map_.find(va);
+  if (mappedHandleIt == mapped_handle_map_.end()) return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+  MemoryHandle* memoryHandle = mappedHandleIt->second.mem_handle;
+  memoryHandle->ref_count++;
+  *mapped_handle = MemoryHandle::Convert(memoryHandle->thunk_handle);
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t Runtime::VMemoryGetAllocPropertiesFromHandle(hsa_amd_vmem_alloc_handle_t allocHandle,
+                                                          const core::MemoryRegion** mem_region,
+                                                          hsa_amd_memory_type_t* type) {
+  auto memoryHandleIt = memory_handle_map_.find(reinterpret_cast<void*>(allocHandle.handle));
+  if (memoryHandleIt == memory_handle_map_.end()) return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+  *mem_region = memoryHandleIt->second.region;
+  *type = (memoryHandleIt->second.alloc_flag & core::MemoryRegion::AllocatePinned)
+      ? MEMORY_TYPE_PINNED
+      : MEMORY_TYPE_NONE;
+
+  return HSA_STATUS_SUCCESS;
+}
+
 }  // namespace core
 }  // namespace rocr