Merge pull request #9085 from edgargabriel/pr/hipHostMalloc-fix-1.14

UCT/ROCM: don't try to lock host pinned memory -v1.14
openucx · May 18, 2023 · 1adefd8 · 1adefd8
2 parents 315c0e8 + 38a7987
commit 1adefd8
Showing 1 changed file with 14 additions and 8 deletions.
diff --git a/src/uct/rocm/copy/rocm_copy_ep.c b/src/uct/rocm/copy/rocm_copy_ep.c
@@ -119,11 +119,14 @@ ucs_status_t uct_rocm_copy_ep_zcopy(uct_ep_h tl_ep,
         return UCS_ERR_IO_ERROR;
     }
 
-    if ((remote_addr_mem_type == HSA_EXT_POINTER_TYPE_HSA) &&
-        (dev_type == HSA_DEVICE_TYPE_GPU)) {
-        /* UCS_MEMORY_TYPE_ROCM */
+    if (remote_addr_mem_type == HSA_EXT_POINTER_TYPE_HSA) {
         remote_addr_mod = (void*)remote_addr;
-        agent           = tmp_agent;
+        if (dev_type == HSA_DEVICE_TYPE_GPU) {
+            /* UCS_MEMORY_TYPE_ROCM */
+            agent = tmp_agent;
+        } else {
+            remote_addr_is_host = 1;
+        }
     } else if ((remote_addr_mem_type == HSA_EXT_POINTER_TYPE_LOCKED) &&
                (size == dev_size)) {
         /* locked host memory, e.g. hipHostRegister, OR previously registered */
@@ -152,11 +155,14 @@ ucs_status_t uct_rocm_copy_ep_zcopy(uct_ep_h tl_ep,
         return UCS_ERR_IO_ERROR;
     }
 
-    if ((iov_buffer_mem_type == HSA_EXT_POINTER_TYPE_HSA) &&
-        (dev_type == HSA_DEVICE_TYPE_GPU)) {
-        /* UCS_MEMORY_TYPE_ROCM */
+    if (iov_buffer_mem_type == HSA_EXT_POINTER_TYPE_HSA) {
         iov_buffer_mod = iov->buffer;
-        agent          = tmp_agent;
+        if (dev_type == HSA_DEVICE_TYPE_GPU) {
+            /* UCS_MEMORY_TYPE_ROCM */
+            agent = tmp_agent;
+        } else {
+            iov_buffer_is_host = 1;
+        }
     } else if ((iov_buffer_mem_type == HSA_EXT_POINTER_TYPE_LOCKED) &&
                (size == dev_size)) {
         /* locked host memory (e.g. hipHostRegister) OR previously registered */