[libc-commits] [libc] cd4e888 - [libc] Improve copying system vectors to the GPU

[Joseph Huber via libc-commits]

Author: Joseph Huber
Date: 2023-04-03T10:10:43-05:00
New Revision: cd4e8884edefd37d3c3fd42c7740ef2a82a6592b

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

LOG: [libc] Improve copying system vectors to the GPU

Summary:
This implementation was buggy and inefficient. Fine-grained memory can
only be allocated on a page-level granularity. Which means that each
allocated string used about 4096 bytes. This is wasteful in general, and
also allowed for buggy behaviour. The previous copying of the
environment vector only worked because the large buffer size meant that
we would typically have a null byte after the allocated memory. However
this would break if the vector was larger than a page. This patch
allocates everything into a single buffer. It makes it easier to free,
use, and it more correct.

Added: 
    

Modified: 
    libc/utils/gpu/loader/Loader.h

Removed: 
    


################################################################################
diff  --git a/libc/utils/gpu/loader/Loader.h b/libc/utils/gpu/loader/Loader.h
index dad96c10d37c8..2c30593075845 100644
--- a/libc/utils/gpu/loader/Loader.h
+++ b/libc/utils/gpu/loader/Loader.h
@@ -9,6 +9,7 @@
 #ifndef LLVM_LIBC_UTILS_GPU_LOADER_LOADER_H
 #define LLVM_LIBC_UTILS_GPU_LOADER_LOADER_H
 
+#include <cstdint>
 #include <cstring>
 #include <stddef.h>
 
@@ -20,21 +21,27 @@ int load(int argc, char **argv, char **evnp, void *image, size_t size);
 /// Copy the system's argument vector to GPU memory allocated using \p alloc.
 template <typename Allocator>
 void *copy_argument_vector(int argc, char **argv, Allocator alloc) {
-  void *dev_argv = alloc(argc * sizeof(char *));
-  if (dev_argv == nullptr)
+  size_t argv_size = sizeof(char *) * (argc + 1);
+  size_t str_size = 0;
+  for (int i = 0; i < argc; ++i)
+    str_size += strlen(argv[i]) + 1;
+
+  // We allocate enough space for a null terminated array and all the strings.
+  void *dev_argv = alloc(argv_size + str_size);
+  if (!dev_argv)
     return nullptr;
 
+  // Store the strings linerally in the same memory buffer.
+  void *dev_str = reinterpret_cast<uint8_t *>(dev_argv) + argv_size;
   for (int i = 0; i < argc; ++i) {
     size_t size = strlen(argv[i]) + 1;
-    void *dev_str = alloc(size);
-    if (dev_str == nullptr)
-      return nullptr;
-
-    // Load the host memory buffer with the pointer values of the newly
-    // allocated strings.
     std::memcpy(dev_str, argv[i], size);
     static_cast<void **>(dev_argv)[i] = dev_str;
+    dev_str = reinterpret_cast<uint8_t *>(dev_str) + size;
   }
+
+  // Ensure the vector is null terminated.
+  reinterpret_cast<void **>(dev_argv)[argv_size] = nullptr;
   return dev_argv;
 };