[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.

2 files changed, 28 insertions, 0 deletions

diff --git a/libc/utils/gpu/loader/amdgpu/Loader.cpp b/libc/utils/gpu/loader/amdgpu/Loader.cpp
index d090b98..87dd3ce4 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 88ef170..ed8b8d0 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);