[libc] Add environment variables to GPU libc test for AMDGPU

[jhuber6]

Summary

This patch performs the same operation to copy over the argv array to
the envp array. This allows the GPU tests to use environment
variables.

[sivachandra]

This is essentially a GPU only change. Stamping anyway.

libc/startup/gpu/amdgpu/start.cpp

 //===-- Implementation of crt for amdgpu ----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "src/__support/RPC/rpc_client.h"
 
-extern "C" int main(int argc, char **argv);
+extern "C" int main(int argc, char **argv, char **envp);
 
 extern "C" [[gnu::visibility("protected"), clang::amdgpu_kernel]] void
-_start(int argc, char **argv, int *ret, void *in, void *out, void *buffer) {
+_start(int argc, char **argv, char **envp, int *ret, void *in, void *out,
+       void *buffer) {
   __llvm_libc::rpc::client.reset(in, out, buffer);
 
-  __atomic_fetch_or(ret, main(argc, argv), __ATOMIC_RELAXED);
+  __atomic_fetch_or(ret, main(argc, argv, envp), __ATOMIC_RELAXED);
 }

libc/test/integration/startup/gpu/CMakeLists.txt

 add_custom_target(libc-startup-tests)
 add_dependencies(libc-integration-tests libc-startup-tests)
 
 add_integration_test(
   startup_args_test
   SUITE libc-startup-tests
   SRCS
     args_test.cpp
   ARGS
     1 2 3
+  ENV
+    FRANCE=Paris
+    GERMANY=Berlin
 )

libc/test/integration/startup/gpu/args_test.cpp

static bool my_streq(const char *lhs, const char *rhs) {
   const char *l, *r;
   for (l = lhs, r = rhs; *l != '\0' && *r != '\0'; ++l, ++r)
     if (*l != *r)
       return false;
 
   return *l == '\0' && *r == '\0';
 }
 
-TEST_MAIN(int argc, char **argv) {
+TEST_MAIN(int argc, char **argv, char **envp) {
   ASSERT_TRUE(argc == 4);
   ASSERT_TRUE(my_streq(argv[1], "1"));
   ASSERT_TRUE(my_streq(argv[2], "2"));
   ASSERT_TRUE(my_streq(argv[3], "3"));
 
+  bool found_france = false;
+  bool found_germany = false;
+  for (; *envp != nullptr; ++envp) {
+    if (my_streq(*envp, "FRANCE=Paris"))
+      found_france = true;
+    if (my_streq(*envp, "GERMANY=Berlin"))
+      found_germany = true;
+  }
+
+  ASSERT_TRUE(found_france && found_germany);
+
   return 0;
 }

libc/utils/gpu/loader/Loader.h

 //===-- Generic device loader interface -----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include <stddef.h>
 
 /// Generic interface to load the \p image and launch execution of the _start
 /// kernel on the target device. Copies \p argc and \p argv to the device.
 /// Returns the final value of the `main` function on the device.
-int load(int argc, char **argv, void *image, size_t size);
+int load(int argc, char **argv, char **evnp, void *image, size_t size);

libc/utils/gpu/loader/Main.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "Loader.h"
 
 #include <cstdio>
 #include <cstdlib>
 
-int main(int argc, char **argv) {
+int main(int argc, char **argv, char **envp) {
   if (argc < 2) {
     printf("USAGE: ./loader <device_image> <args>, ...\n");
     return EXIT_SUCCESS;
   }
 
   // TODO: We should perform some validation on the file.
   FILE *file = fopen(argv[1], "r");
 
   if (!file) {
     fprintf(stderr, "Failed to open image file %s\n", argv[1]);
     return EXIT_FAILURE;
   }
 
   fseek(file, 0, SEEK_END);
   const auto size = ftell(file);
   fseek(file, 0, SEEK_SET);
 
   void *image = malloc(size * sizeof(char));
   fread(image, sizeof(char), size, file);
   fclose(file);
 
   // Drop the loader from the program arguments.
-  int ret = load(argc - 1, &argv[1], image, size);
+  int ret = load(argc - 1, &argv[1], envp, image, size);
 
   free(image);
   return ret;
 }

libc/utils/gpu/loader/amdgpu/Loader.cpp

 
 /// The name of the kernel we will launch. All AMDHSA kernels end with '.kd'.
 constexpr const char *KERNEL_START = "_start.kd";
 
 /// The arguments to the '_start' kernel.
 struct kernel_args_t {
   int argc;
   void *argv;
+  void *envp;
   void *ret;
   void *inbox;
   void *outbox;
   void *buffer;
 };
 
 static __llvm_libc::rpc::Server server;
 
auto cb = [&](hsa_amd_memory_pool_t memory_pool) {
     if (flags & flag)
       *output_pool = memory_pool;
 
     return HSA_STATUS_SUCCESS;
   };
   return iterate_agent_memory_pools(agent, cb);
 }
 
-int load(int argc, char **argv, void *image, size_t size) {
+int load(int argc, char **argv, char **envp, void *image, size_t size) {
   // Initialize the HSA runtime used to communicate with the device.
   if (hsa_status_t err = hsa_init())
     handle_error(err);
 
   // Register a callback when the device encounters a memory fault.
   if (hsa_status_t err = hsa_amd_register_system_event_handler(
           [](const hsa_amd_event_t *event, void *) -> hsa_status_t {
             if (event->event_type == HSA_AMD_GPU_MEMORY_FAULT_EVENT)
if (hsa_status_t err = hsa_amd_memory_pool_allocate(finegrained_pool, size,
       handle_error(err);
     hsa_amd_agents_allow_access(1, &dev_agent, nullptr, dev_str);
     // 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;
   }
 
+  // Allocate fine-grained memory on the host to hold the pointer array for the
+  // copied environment array and allow the GPU agent to access it.
+  int envc = 0;
+  for (char **env = envp; *env != 0; ++env)
+    ++envc;
+  void *dev_envp;
+  if (hsa_status_t err =
+          hsa_amd_memory_pool_allocate(finegrained_pool, envc * sizeof(char *),
+                                       /*flags=*/0, &dev_envp))
+    handle_error(err);
+  hsa_amd_agents_allow_access(1, &dev_agent, nullptr, dev_envp);
+  for (int i = 0; i < envc; ++i) {
+    size_t size = strlen(envp[i]) + 1;
+    void *dev_str;
+    if (hsa_status_t err = hsa_amd_memory_pool_allocate(finegrained_pool, size,
+                                                        /*flags=*/0, &dev_str))
+      handle_error(err);
+    hsa_amd_agents_allow_access(1, &dev_agent, nullptr, dev_str);
+    // Load the host memory buffer with the pointer values of the newly
+    // allocated strings.
+    std::memcpy(dev_str, envp[i], size);
+    static_cast<void **>(dev_envp)[i] = dev_str;
+  }
+
   // Allocate space for the return pointer and initialize it to zero.
   void *dev_ret;
   if (hsa_status_t err =
           hsa_amd_memory_pool_allocate(coarsegrained_pool, sizeof(int),
                                        /*flags=*/0, &dev_ret))
     handle_error(err);
   hsa_amd_memory_fill(dev_ret, 0, sizeof(int));
 
int load(int argc, char **argv, char **envp, void *image, size_t size) {
   hsa_amd_agents_allow_access(1, &dev_agent, nullptr, buffer);
 
   // Initialie all the arguments (explicit and implicit) to zero, then set the
   // explicit arguments to the values created above.
   std::memset(args, 0, args_size);
   kernel_args_t *kernel_args = reinterpret_cast<kernel_args_t *>(args);
   kernel_args->argc = argc;
   kernel_args->argv = dev_argv;
+  kernel_args->envp = dev_envp;
   kernel_args->ret = dev_ret;
   kernel_args->inbox = server_outbox;
   kernel_args->outbox = server_inbox;
   kernel_args->buffer = buffer;
 
   // Obtain a packet from the queue.
   uint64_t packet_id = hsa_queue_add_write_index_relaxed(queue, 1);
   while (packet_id - hsa_queue_load_read_index_scacquire(queue) >= queue_size)