[libc-commits] [libc] 0584e6c - [libc] Explicitly pin memory for the HSA memory transfer (#73973)

[via libc-commits]

Author: Joseph Huber
Date: 2023-11-30T13:46:52-06:00
New Revision: 0584e6c1669dbfc6252f5b069e2a895389660120

URL: https://github.com/llvm/llvm-project/commit/0584e6c1669dbfc6252f5b069e2a895389660120
DIFF: https://github.com/llvm/llvm-project/commit/0584e6c1669dbfc6252f5b069e2a895389660120.diff

LOG: [libc] Explicitly pin memory for the HSA memory transfer (#73973)

Summary:
This portion of code handles mapping the RPC client memory over to the
device. HSA copies need to be between two slices of memory that HSA has
allocated. Previously we used coarse-grained memory to act as the host
source. However, the support for this varies depending on the kernel and
version and should not be relied upon. This patch changes that handling
to use the `hsa_amd_memory_lock` API to explicitly pin memory to a
location sufficient for a DMA transfer to the GPU.

Added: 
    

Modified: 
    libc/utils/gpu/loader/amdgpu/Loader.cpp

Removed: 
    


################################################################################
diff  --git a/libc/utils/gpu/loader/amdgpu/Loader.cpp b/libc/utils/gpu/loader/amdgpu/Loader.cpp
index 4272cf40bfd1f67..5c28607b2f290a9 100644
--- a/libc/utils/gpu/loader/amdgpu/Loader.cpp
+++ b/libc/utils/gpu/loader/amdgpu/Loader.cpp
@@ -482,12 +482,11 @@ int load(int argc, char **argv, char **envp, void *image, size_t size,
     handle_error(err);
 
   void *rpc_client_buffer;
-  if (hsa_status_t err = hsa_amd_memory_pool_allocate(
-          coarsegrained_pool, rpc_get_client_size(),
-          /*flags=*/0, &rpc_client_buffer))
+  if (hsa_status_t err = hsa_amd_memory_lock(
+          const_cast<void *>(rpc_get_client_buffer(device_id)),
+          rpc_get_client_size(),
+          /*agents=*/nullptr, 0, &rpc_client_buffer))
     handle_error(err);
-  std::memcpy(rpc_client_buffer, rpc_get_client_buffer(device_id),
-              rpc_get_client_size());
 
   // Copy the RPC client buffer to the address pointed to by the symbol.
   if (hsa_status_t err =
@@ -495,7 +494,8 @@ int load(int argc, char **argv, char **envp, void *image, size_t size,
                      rpc_client_buffer, host_agent, rpc_get_client_size()))
     handle_error(err);
 
-  if (hsa_status_t err = hsa_amd_memory_pool_free(rpc_client_buffer))
+  if (hsa_status_t err = hsa_amd_memory_unlock(
+          const_cast<void *>(rpc_get_client_buffer(device_id))))
     handle_error(err);
   if (hsa_status_t err = hsa_amd_memory_pool_free(rpc_client_host))
     handle_error(err);