diff --git a/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp b/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
index 86211d252..459f88a50 100644
--- a/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
+++ b/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
@@ -297,7 +297,46 @@ hsa_status_t XdnaDriver::ImportDMABuf(int dmabuf_fd, uint32_t *handle) {
   return HSA_STATUS_SUCCESS;
 }
 
-hsa_status_t XdnaDriver::ReleaseBO(uint32_t handle) {
+__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:
+  default:
+    return PROT_NONE;
+  }
+}
+
+hsa_status_t XdnaDriver::Map(uint32_t handle, void *va, size_t offset,
+                             size_t size, hsa_access_permission_t perms) {
+  // Get fd associated with the handle.
+  drm_prime_handle params = {};
+  params.handle = handle;
+  params.fd = -1;
+  if (ioctl(fd_, DRM_IOCTL_PRIME_HANDLE_TO_FD, &params) < 0)
+    return HSA_STATUS_ERROR;
+
+  // Change permissions; the result should match the already mapped memory.
+  void *mapped_ptr = mmap(va, size, mmap_perm(perms), MAP_FIXED | MAP_SHARED,
+                          params.fd, offset);
+  if (mapped_ptr != va)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t XdnaDriver::Unmap(uint32_t handle, void *va, size_t size) {
+  if (munmap(va, size) != 0)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t XdnaDriver::ReleaseHandle(uint32_t handle) {
   drm_gem_close close_params = {};
   close_params.handle = handle;
   if (ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_params) < 0)
diff --git a/runtime/hsa-runtime/core/inc/amd_xdna_driver.h b/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
index ef8d43354..49b4438d4 100644
--- a/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
+++ b/runtime/hsa-runtime/core/inc/amd_xdna_driver.h
@@ -101,10 +101,27 @@ class XdnaDriver : public core::Driver {
   /// @param[out] handle handle to the imported memory
   hsa_status_t ImportDMABuf(int dmabuf_fd, uint32_t *handle);
 
+  /// @brief Maps the memory associated with the handle.
+  ///
+  /// @param[in] handle handle to the memory object
+  /// @param[in] va virtual address associated with the handle
+  /// @param[in] offset memory offset in bytes
+  /// @param[in] size memory size in bytes
+  /// @param[perms] perms new permissions
+  hsa_status_t Map(uint32_t handle, void *va, size_t offset, size_t size,
+                   hsa_access_permission_t perms);
+
+  /// @brief Unmaps the memory associated with the handle.
+  ///
+  /// @param[in] handle handle to the memory object
+  /// @param[in] va virtual address associated with the handle
+  /// @param[in] size memory size in bytes
+  hsa_status_t Unmap(uint32_t handle, void *va, size_t size);
+
   /// @brief Releases the object associated with the handle.
   ///
-  /// @param[in] handle handle of the object to release.
-  hsa_status_t ReleaseBO(uint32_t handle);
+  /// @param[in] handle handle of the object to release
+  hsa_status_t ReleaseHandle(uint32_t handle);
 
 private:
   hsa_status_t QueryDriverVersion();
diff --git a/runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp
index f6cafd0f2..f69a3e335 100644
--- a/runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp
+++ b/runtime/hsa-runtime/core/runtime/amd_aie_agent.cpp
@@ -301,38 +301,18 @@ hsa_status_t AieAgent::QueueCreate(size_t size, hsa_queue_type32_t queue_type,
   return HSA_STATUS_SUCCESS;
 }
 
-__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:
-  default:
-    return PROT_NONE;
-  }
-}
-
 hsa_status_t AieAgent::Map(void *handle, void *va, size_t offset, size_t size,
                            int fd, hsa_access_permission_t perms) {
-  void *mapped_ptr =
-      mmap(va, size, mmap_perm(perms), MAP_PRIVATE | MAP_FIXED, fd, offset);
-  if (mapped_ptr != va) {
-    perror("MMMAPPPPP: ");
-    return HSA_STATUS_ERROR;
-  }
-
-  return HSA_STATUS_SUCCESS;
+  auto &driver = static_cast<XdnaDriver &>(
+      core::Runtime::runtime_singleton_->AgentDriver(driver_type));
+  return driver.Map(*static_cast<uint32_t *>(handle), va, offset, size, perms);
 }
 
 hsa_status_t AieAgent::Unmap(void *handle, void *va, size_t offset,
                              size_t size) {
-  if (munmap(va, size) != 0)
-    return HSA_STATUS_ERROR;
-
-  return HSA_STATUS_SUCCESS;
+  auto &driver = static_cast<XdnaDriver &>(
+      core::Runtime::runtime_singleton_->AgentDriver(driver_type));
+  return driver.Unmap(*static_cast<uint32_t *>(handle), va, size);
 }
 
 hsa_status_t AieAgent::ImportDMABuf(int dmabuf_fd, void *handle) {
@@ -345,7 +325,7 @@ hsa_status_t AieAgent::ReleaseShareableHandle(void *handle, void *va,
                                               size_t size) {
   auto &driver = static_cast<XdnaDriver &>(
       core::Runtime::runtime_singleton_->AgentDriver(driver_type));
-  return driver.ReleaseBO(*static_cast<uint32_t *>(handle));
+  return driver.ReleaseHandle(*static_cast<uint32_t *>(handle));
 }
 
 void AieAgent::InitRegionList() {