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

[Joseph Huber via libc-commits]

https://github.com/jhuber6 created https://github.com/llvm/llvm-project/pull/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.


>From e80daaa581edda49aa406e7798b819feb0566979 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Thu, 30 Nov 2023 13:28:44 -0600
Subject: [PATCH] [libc] Explicitly pin memory for the HSA memory transfer

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.
---
 libc/utils/gpu/loader/amdgpu/Loader.cpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

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);