[libc-commits] [libc] dfc162a - [libc] Free the GPU memory allocated in the device loaders

[Joseph Huber via libc-commits]

Author: Joseph Huber
Date: 2023-04-03T11:55:32-05:00
New Revision: dfc162ad3fcb9636685ee99360cb1c4f5a56cd6b

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

LOG: [libc] Free the GPU memory allocated in the device loaders

Summary:
This part was ignored and we just hoped that shutting down the runtime
freed these correctly. But it's best to be specific and free the memory
we've allocated.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/libc/utils/gpu/loader/amdgpu/Loader.cpp b/libc/utils/gpu/loader/amdgpu/Loader.cpp
index d090b98c5a2e..87dd3ce48d82 100644
--- a/libc/utils/gpu/loader/amdgpu/Loader.cpp
+++ b/libc/utils/gpu/loader/amdgpu/Loader.cpp
@@ -417,6 +417,22 @@ int load(int argc, char **argv, char **envp, void *image, size_t size) {
   // Save the return value and perform basic clean-up.
   int ret = *static_cast<int *>(host_ret);
 
+  // Free the memory allocated for the device.
+  if (hsa_status_t err = hsa_amd_memory_pool_free(args))
+    handle_error(err);
+  if (hsa_status_t err = hsa_amd_memory_pool_free(dev_argv))
+    handle_error(err);
+  if (hsa_status_t err = hsa_amd_memory_pool_free(dev_ret))
+    handle_error(err);
+  if (hsa_status_t err = hsa_amd_memory_pool_free(server_inbox))
+    handle_error(err);
+  if (hsa_status_t err = hsa_amd_memory_pool_free(server_outbox))
+    handle_error(err);
+  if (hsa_status_t err = hsa_amd_memory_pool_free(buffer))
+    handle_error(err);
+  if (hsa_status_t err = hsa_amd_memory_pool_free(host_ret))
+    handle_error(err);
+
   if (hsa_status_t err = hsa_signal_destroy(memory_signal))
     handle_error(err);
 

diff  --git a/libc/utils/gpu/loader/nvptx/Loader.cpp b/libc/utils/gpu/loader/nvptx/Loader.cpp
index 88ef1701b158..ed8b8d018c6a 100644
--- a/libc/utils/gpu/loader/nvptx/Loader.cpp
+++ b/libc/utils/gpu/loader/nvptx/Loader.cpp
@@ -176,6 +176,18 @@ int load(int argc, char **argv, char **envp, void *image, size_t size) {
   if (CUresult err = cuStreamSynchronize(stream))
     handle_error(err);
 
+  // Free the memory allocated for the device.
+  if (CUresult err = cuMemFree(dev_ret))
+    handle_error(err);
+  if (CUresult err = cuMemFreeHost(dev_argv))
+    handle_error(err);
+  if (CUresult err = cuMemFreeHost(server_inbox))
+    handle_error(err);
+  if (CUresult err = cuMemFreeHost(server_outbox))
+    handle_error(err);
+  if (CUresult err = cuMemFreeHost(buffer))
+    handle_error(err);
+
   // Destroy the context and the loaded binary.
   if (CUresult err = cuModuleUnload(binary))
     handle_error(err);