rocr/aie: Add dma-buf import support for AIEAgents via the Driver interface

Change-Id: I70f8d8772dda7c06944d75042cb3034ddd89aff4

runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp

@@ -3,7 +3,7 @@
 // The University of Illinois/NCSA
 // Open Source License (NCSA)
 //
-// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 //
 // Developed by:
 //
@@ -45,11 +45,13 @@
 #include <memory>
 #include <string>
 
+#include <amdgpu_drm.h>
 #include <link.h>
 #include <sys/ioctl.h>
 
 #include "hsakmt/hsakmt.h"
 
+#include "core/inc/amd_gpu_agent.h"
 #include "core/inc/amd_memory_region.h"
 #include "core/inc/runtime.h"
 
@@ -58,6 +60,29 @@ extern r_debug _amdgpu_r_debug;
 namespace rocr {
 namespace AMD {
 
+static_assert(
+    (sizeof(core::ShareableHandle::handle) >= sizeof(amdgpu_bo_handle)) &&
+        (alignof(core::ShareableHandle::handle) >= alignof(amdgpu_bo_handle)),
+    "ShareableHandle cannot store a amdgpu_bo_handle");
+
+namespace {
+
+__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:
+  default:
+    return 0;
+  }
+}
+
+} // namespace
+
 KfdDriver::KfdDriver(std::string devnode_name)
     : core::Driver(core::DriverType::KFD, devnode_name) {}
 
@@ -322,6 +347,73 @@ hsa_status_t KfdDriver::DestroyQueue(core::Queue &queue) const {
   return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t KfdDriver::ExportDMABuf(void *mem, size_t size, int *dmabuf_fd,
+                                     size_t *offset) {
+  int dmabuf_fd_res = -1;
+  size_t offset_res = 0;
+  if (hsaKmtExportDMABufHandle(mem, size, &dmabuf_fd_res, &offset_res) !=
+      HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+
+  *dmabuf_fd = dmabuf_fd_res;
+  *offset = offset_res;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdDriver::ImportDMABuf(int dmabuf_fd, core::Agent &agent,
+                                     core::ShareableHandle &handle) {
+  auto &gpu_agent = static_cast<GpuAgent &>(agent);
+  amdgpu_bo_import_result res;
+  auto ret = amdgpu_bo_import(
+      gpu_agent.libDrmDev(), amdgpu_bo_handle_type_dma_buf_fd, dmabuf_fd, &res);
+  if (ret)
+    return HSA_STATUS_ERROR;
+
+  handle.handle = reinterpret_cast<uint64_t>(res.buf_handle);
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdDriver::Map(core::ShareableHandle handle, void *mem,
+                            size_t offset, size_t size,
+                            hsa_access_permission_t perms) {
+  const auto ldrm_bo = reinterpret_cast<amdgpu_bo_handle>(handle.handle);
+  if (!ldrm_bo)
+    return HSA_STATUS_ERROR;
+
+  if (amdgpu_bo_va_op(ldrm_bo, offset, size, reinterpret_cast<uint64_t>(mem),
+                      drm_perm(perms), AMDGPU_VA_OP_MAP) != 0)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdDriver::Unmap(core::ShareableHandle handle, void *mem,
+                              size_t offset, size_t size) {
+  const auto ldrm_bo = reinterpret_cast<amdgpu_bo_handle>(handle.handle);
+  if (!ldrm_bo)
+    return HSA_STATUS_ERROR;
+
+  if (amdgpu_bo_va_op(ldrm_bo, offset, size, reinterpret_cast<uint64_t>(mem), 0,
+                      AMDGPU_VA_OP_UNMAP) != 0)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdDriver::ReleaseShareableHandle(core::ShareableHandle &handle) {
+  const auto ldrm_bo = reinterpret_cast<amdgpu_bo_handle>(handle.handle);
+  if (!ldrm_bo)
+    return HSA_STATUS_ERROR;
+
+  const auto ret = amdgpu_bo_free(ldrm_bo);
+  if (ret)
+    return HSA_STATUS_ERROR;
+
+  handle = {};
+  return HSA_STATUS_SUCCESS;
+}
+
 void *KfdDriver::AllocateKfdMemory(const HsaMemFlags &flags, uint32_t node_id,
                                    size_t size) {
   void *mem = nullptr;

runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp

@@ -3,7 +3,7 @@
 // The University of Illinois/NCSA
 // Open Source License (NCSA)
 //
-// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 //
 // Developed by:
 //
@@ -52,12 +52,17 @@
 #include "core/inc/amd_aie_aql_queue.h"
 #include "core/inc/amd_memory_region.h"
 #include "core/inc/runtime.h"
+#include "core/util/memory.h"
 #include "core/util/utils.h"
 #include "uapi/amdxdna_accel.h"
 
 namespace rocr {
 namespace AMD {
 
+static_assert((sizeof(core::ShareableHandle::handle) >= sizeof(uint32_t)) &&
+                  (alignof(core::ShareableHandle::handle) >= alignof(uint32_t)),
+              "ShareableHandle cannot store a XDNA handle");
+
 XdnaDriver::XdnaDriver(std::string devnode_name)
     : core::Driver(core::DriverType::XDNA, devnode_name) {}
 
@@ -292,6 +297,62 @@ hsa_status_t XdnaDriver::DestroyQueue(core::Queue &queue) const {
   return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t XdnaDriver::ExportDMABuf(void *mem, size_t size, int *dmabuf_fd,
+                                      size_t *offset) {
+  // Not implemented yet.
+  return HSA_STATUS_ERROR;
+}
+
+hsa_status_t XdnaDriver::ImportDMABuf(int dmabuf_fd, core::Agent &agent,
+                                      core::ShareableHandle &handle) {
+  drm_prime_handle import_params = {};
+  import_params.handle = AMDXDNA_INVALID_BO_HANDLE;
+  import_params.fd = dmabuf_fd;
+  if (ioctl(fd_, DRM_IOCTL_PRIME_FD_TO_HANDLE, &import_params) < 0)
+    return HSA_STATUS_ERROR;
+
+  handle.handle = import_params.handle;
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t XdnaDriver::Map(core::ShareableHandle handle, void *mem,
+                             size_t offset, size_t size,
+                             hsa_access_permission_t perms) {
+  // Get fd associated with the handle.
+  drm_prime_handle params = {};
+  params.handle = handle.handle;
+  params.fd = -1;
+  if (ioctl(fd_, DRM_IOCTL_PRIME_HANDLE_TO_FD, &params) < 0)
+    return HSA_STATUS_ERROR;
+
+  // Change permissions.
+  void *mapped_ptr = mmap(mem, size, PermissionsToMmapFlags(perms),
+                          MAP_FIXED | MAP_SHARED, params.fd, offset);
+  if (mapped_ptr == MAP_FAILED)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t XdnaDriver::Unmap(core::ShareableHandle handle, void *mem,
+                               size_t offset, size_t size) {
+  if (munmap(mem, size) != 0)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t XdnaDriver::ReleaseShareableHandle(core::ShareableHandle &handle) {
+  drm_gem_close close_params = {};
+  close_params.handle = handle.handle;
+  if (ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_params) < 0)
+    return HSA_STATUS_ERROR;
+
+  handle = {};
+
+  return HSA_STATUS_SUCCESS;
+}
+
 hsa_status_t XdnaDriver::QueryDriverVersion() {
   amdxdna_drm_query_aie_version aie_version{0, 0};
   amdxdna_drm_get_info args{DRM_AMDXDNA_QUERY_AIE_VERSION, sizeof(aie_version),