Support exporting and importing memory mappings

Support exporting and importing dmabuf file descriptors for memory mappings. The exported dmabuf file descriptors are shareable posix file descriptors that can be used for cross-vendor, cross-device and cross-process memory sharing. This is part of patch series for Virtual Memory API. Change-Id: I3673fc009f7e73bc26be8349e19f66e20d0607c5 [ROCm/ROCR-Runtime commit: b03c96c264]
2022-12-19 00:06:33 +00:00
@@ -1270,6 +1270,17 @@ hsa_status_t HSA_API hsa_amd_vmem_get_access(void* va, hsa_access_permission_t*
  return amdExtTable->hsa_amd_vmem_get_access_fn(va, perms, agent_handle);
 }

+hsa_status_t HSA_API hsa_amd_vmem_export_shareable_handle(int* dmabuf_fd,
+                                                          hsa_amd_vmem_alloc_handle_t handle,
+                                                          uint64_t flags) {
+  return amdExtTable->hsa_amd_vmem_export_shareable_handle_fn(dmabuf_fd, handle, flags);
+}
+
+hsa_status_t HSA_API hsa_amd_vmem_import_shareable_handle(int dmabuf_fd,
+                                                          hsa_amd_vmem_alloc_handle_t* handle) {
+  return amdExtTable->hsa_amd_vmem_import_shareable_handle_fn(dmabuf_fd, handle);
+}
+
 // Tools only table interfaces.
 namespace rocr {

@@ -329,6 +329,14 @@ hsa_status_t hsa_amd_vmem_set_access(void* va, size_t size,
 hsa_status_t hsa_amd_vmem_get_access(void* va, hsa_access_permission_t* flags,
                                     const hsa_agent_t agent_handle);

+// Mirrors Amd Extension Apis
+hsa_status_t hsa_amd_vmem_export_shareable_handle(int* dmabuf_fd,
+                                                  hsa_amd_vmem_alloc_handle_t handle,
+                                                  uint64_t flags);
+
+// Mirrors Amd Extension Apis
+hsa_status_t hsa_amd_vmem_import_shareable_handle(int dmabuf_fd,
+                                                  hsa_amd_vmem_alloc_handle_t* handle);
 }  // namespace amd
 }  // namespace rocr

@@ -375,6 +375,13 @@ class Runtime {
  hsa_status_t VMemoryGetAccess(const void* va, hsa_access_permission_t* perms,
                                hsa_agent_t agent_handle);

+  hsa_status_t VMemoryExportShareableHandle(int* dmabuf_fd,
+                                            const hsa_amd_vmem_alloc_handle_t handle,
+                                            const uint64_t flags);
+
+  hsa_status_t VMemoryImportShareableHandle(const int dmabuf_fd,
+                                            hsa_amd_vmem_alloc_handle_t* handle);
+
  const std::vector<Agent*>& cpu_agents() { return cpu_agents_; }

  const std::vector<Agent*>& gpu_agents() { return gpu_agents_; }
@@ -410,6 +410,8 @@ void HsaApiTable::UpdateAmdExts() {
  amd_ext_api.hsa_amd_vmem_unmap_fn = AMD::hsa_amd_vmem_unmap;
  amd_ext_api.hsa_amd_vmem_set_access_fn = AMD::hsa_amd_vmem_set_access;
  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;
 }

 void LoadInitialHsaApiTable() {
@@ -110,6 +110,11 @@ struct ValidityError<const T*> {
    if ((arg) == 0) return HSA_STATUS_ERROR_INVALID_ARGUMENT;                                      \
  } while (false)

+#define IS_VALID_FD(fd)                                                                            \
+  do {                                                                                             \
+    if ((fd) < 0) return HSA_STATUS_ERROR_INVALID_ARGUMENT;                                        \
+  } while (false)
+
 #define IS_VALID(ptr)                                           \
  do {                                                          \
    if ((ptr) == NULL || !(ptr)->IsValid())                     \
@@ -1300,5 +1305,26 @@ hsa_status_t hsa_amd_vmem_get_access(void* va, hsa_access_permission_t* perms,
  CATCH;
 }

+hsa_status_t hsa_amd_vmem_export_shareable_handle(int* dmabuf_fd,
+                                                  hsa_amd_vmem_alloc_handle_t handle,
+                                                  uint64_t flags) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(dmabuf_fd);
+
+  return core::Runtime::runtime_singleton_->VMemoryExportShareableHandle(dmabuf_fd, handle, flags);
+  CATCH;
+}
+
+hsa_status_t hsa_amd_vmem_import_shareable_handle(int dmabuf_fd,
+                                                  hsa_amd_vmem_alloc_handle_t* handle) {
+  TRY;
+  IS_BAD_PTR(handle);
+  IS_VALID_FD(dmabuf_fd);
+
+  return core::Runtime::runtime_singleton_->VMemoryImportShareableHandle(dmabuf_fd, handle);
+  CATCH;
+}
+
 }   //  namespace amd
 }   //  namespace rocr
@@ -2786,5 +2786,88 @@ hsa_status_t Runtime::VMemoryGetAccess(const void* va, hsa_access_permission_t*
  return HSA_STATUS_ERROR_INVALID_ALLOCATION;
 }

+hsa_status_t Runtime::VMemoryExportShareableHandle(int* dmabuf_fd,
+                                                   hsa_amd_vmem_alloc_handle_t handle,
+                                                   uint64_t flags) {
+  *dmabuf_fd = -1;
+  auto memoryHandle = memory_handle_map_.find((void*)handle.handle);
+  if (memoryHandle == memory_handle_map_.end()) {
+    debug_warning(false && "Can't find memory handle");
+    return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+  }
+
+  uint64_t offset, ret;
+
+  ret = hsaKmtExportDMABufHandle(memoryHandle->second.thunk_handle, memoryHandle->second.size,
+                                 dmabuf_fd, &offset);
+  if (ret != HSAKMT_STATUS_SUCCESS) return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t Runtime::VMemoryImportShareableHandle(int dmabuf_fd,
+                                                   hsa_amd_vmem_alloc_handle_t* memoryOnlyHandle) {
+  auto lookupRegion = [this](int nodeid, const AMD::MemoryRegion** ret) {
+    auto nodeAgent = agents_by_node_.find(nodeid);
+    if (nodeAgent == agents_by_node_.end()) {
+      *ret = NULL;
+      return;
+    }
+
+    Agent* agent = nodeAgent->second.front();
+    if (agent == nullptr || !agent->IsValid() || agent->device_type() != Agent::kAmdGpuDevice) {
+      *ret = NULL;
+      return;
+    }
+
+    for (const core::MemoryRegion* region : agent->regions()) {
+      const AMD::MemoryRegion* amd_region = reinterpret_cast<const AMD::MemoryRegion*>(region);
+
+      // TODO: Verify that this works on a system with FINE_GRAINED memory.
+      // System's with FINE_GRAINED will have both COARSE and FINE grain... need to get the
+      // rigtht one.
+
+      bool alloc_allowed;
+      hsa_status_t status =
+          amd_region->GetInfo(HSA_REGION_INFO_RUNTIME_ALLOC_ALLOWED, &alloc_allowed);
+      if (status == HSA_STATUS_SUCCESS && alloc_allowed) *ret = amd_region;
+    }
+  };
+
+  HsaGraphicsResourceInfo info;
+  int ret = hsaKmtRegisterGraphicsHandleToNodes(dmabuf_fd, &info, 0, NULL);
+  if (ret) return HSA_STATUS_ERROR_INCOMPATIBLE_ARGUMENTS;
+
+  ThunkHandle thunk_handle = info.MemoryAddress;
+  size_t size = info.SizeInBytes;
+  int gpuid = info.NodeId;
+
+
+  auto memoryHandleIt = memory_handle_map_.find(thunk_handle);
+  if (memoryHandleIt != memory_handle_map_.end()) {
+    /* This handle was already imported, increment ref_count and return */
+    memoryHandleIt->second.ref_count++;
+    *memoryOnlyHandle = MemoryHandle::Convert(thunk_handle);
+    return HSA_STATUS_SUCCESS;
+  }
+
+  const AMD::MemoryRegion* region = NULL;
+  lookupRegion(gpuid, &region);
+  if (!region) return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+  HsaPointerInfo ptrInfo;
+  ret = hsaKmtQueryPointerInfo(info.MemoryAddress, &ptrInfo);
+  if (ret != HSA_STATUS_SUCCESS || ptrInfo.Type == HSA_POINTER_UNKNOWN)
+    return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+  MemoryRegion::AllocateFlags alloc_flag = core::MemoryRegion::AllocateNoFlags;
+  if (ptrInfo.MemFlags.ui32.NoSubstitute) alloc_flag |= core::MemoryRegion::AllocatePinned;
+
+  memory_handle_map_[thunk_handle] = MemoryHandle(region, size, 0, thunk_handle, alloc_flag);
+  *memoryOnlyHandle = MemoryHandle::Convert(thunk_handle);
+
+  return HSA_STATUS_SUCCESS;
+}
+
 }  // namespace core
 }  // namespace rocr
@@ -241,6 +241,8 @@ global:
 	hsa_amd_vmem_unmap;
 	hsa_amd_vmem_set_access;
 	hsa_amd_vmem_get_access;
+	hsa_amd_vmem_export_shareable_handle;
+	hsa_amd_vmem_import_shareable_handle;

 local:
    *;
@@ -202,6 +202,8 @@ struct AmdExtTable {
  decltype(hsa_amd_vmem_unmap)* hsa_amd_vmem_unmap_fn;
  decltype(hsa_amd_vmem_set_access)* hsa_amd_vmem_set_access_fn;
  decltype(hsa_amd_vmem_get_access)* hsa_amd_vmem_get_access_fn;
+  decltype(hsa_amd_vmem_export_shareable_handle)* hsa_amd_vmem_export_shareable_handle_fn;
+  decltype(hsa_amd_vmem_import_shareable_handle)* hsa_amd_vmem_import_shareable_handle_fn;
 };

 // Table to export HSA Core Runtime Apis
@@ -2898,6 +2898,49 @@ hsa_status_t hsa_amd_vmem_set_access(void* va, size_t size,
 */
 hsa_status_t hsa_amd_vmem_get_access(void* va, hsa_access_permission_t* perms,
                                     hsa_agent_t agent_handle);
+
+/*
+ * @brief Get an exportable shareable handle
+ *
+ * Get an exportable shareable handle for a memory_handle. This shareabl handle can then be used to
+ * re-create a virtual memory handle using hsa_amd_vmem_import_shareable_handle. The shareable
+ * handle can be transferred using mechanisms that support posix file descriptors Once all shareable
+ * handles are closed, the memory_handle is released.
+ *
+ * @param[out] dmabuf_fd shareable handle
+ * @param[in] handle previously allocated virtual memory handle
+ * @param[in] flags Currently unsupported
+ *
+ * @retval ::HSA_STATUS_SUCCESS
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ALLOCATION Invalid memory handle
+ *
+ * @retval ::HSA_STATUS_ERROR_OUT_OF_RESOURCES Out of resources
+ *
+ * @retval ::HSA_STATUS_ERROR Unexpected internal error
+ */
+hsa_status_t hsa_amd_vmem_export_shareable_handle(int* dmabuf_fd,
+                                                  hsa_amd_vmem_alloc_handle_t handle,
+                                                  uint64_t flags);
+/*
+ * @brief Import a shareable handle
+ *
+ * Import a shareable handle for a memory handle. Importing a shareable handle that has been closed
+ * and released results in undefined behavior.
+ *
+ * @param[in] dmabuf_fd shareable handle exported with hsa_amd_vmem_export_shareable_handle
+ * @param[out] handle virtual memory handle
+ *
+ * @retval ::HSA_STATUS_SUCCESS
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ALLOCATION Invalid memory handle
+ *
+ * @retval ::HSA_STATUS_ERROR_OUT_OF_RESOURCES Out of resources
+ *
+ * @retval ::HSA_STATUS_ERROR Unexpected internal error
+ */
+hsa_status_t hsa_amd_vmem_import_shareable_handle(int dmabuf_fd,
+                                                  hsa_amd_vmem_alloc_handle_t* handle);
 #ifdef __cplusplus
 }  // end extern "C" block
 #endif