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

1 files changed, 15 insertions, 8 deletions

diff --git a/libc/utils/gpu/loader/Loader.h b/libc/utils/gpu/loader/Loader.h
index dad96c10d37c..2c3059307584 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;
 };