[libc] Use LLVM CommandLine for loader tool (#101501)

Summary:
This patch removes the ad-hoc parsing that I used previously and
replaces it with the LLVM CommnadLine interface. This doesn't change any
functionality, but makes it easier to maintain.

Author: jhuber6

libc/utils/gpu/loader/CMakeLists.txt

@@ -7,6 +7,9 @@ target_include_directories(gpu_loader PUBLIC
   ${LLVM_MAIN_INCLUDE_DIR}
   ${LLVM_BINARY_DIR}/include
 )
+if(NOT LLVM_ENABLE_RTTI)
+  target_compile_options(gpu_loader PUBLIC -fno-rtti)
+endif()
 
 find_package(hsa-runtime64 QUIET 1.2.0 HINTS ${CMAKE_INSTALL_PREFIX} PATHS /opt/rocm)
 if(hsa-runtime64_FOUND)

libc/utils/gpu/loader/Loader.h

@@ -53,8 +53,9 @@ struct end_args_t {
 /// 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, char **evnp, void *image, size_t size,
-         const LaunchParameters &params, bool print_resource_usage);
+int load(int argc, const char **argv, const char **evnp, void *image,
+         size_t size, const LaunchParameters &params,
+         bool print_resource_usage);
 
 /// Return \p V aligned "upwards" according to \p Align.
 template <typename V, typename A> inline V align_up(V val, A align) {
@@ -63,7 +64,7 @@ template <typename V, typename A> inline V align_up(V val, A align) {
 
 /// 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 *copy_argument_vector(int argc, const char **argv, Allocator alloc) {
   size_t argv_size = sizeof(char *) * (argc + 1);
   size_t str_size = 0;
   for (int i = 0; i < argc; ++i)
@@ -90,9 +91,9 @@ void *copy_argument_vector(int argc, char **argv, Allocator alloc) {
 
 /// Copy the system's environment to GPU memory allocated using \p alloc.
 template <typename Allocator>
-void *copy_environment(char **envp, Allocator alloc) {
+void *copy_environment(const char **envp, Allocator alloc) {
   int envc = 0;
-  for (char **env = envp; *env != 0; ++env)
+  for (const char **env = envp; *env != 0; ++env)
     ++envc;
 
   return copy_argument_vector(envc, envp, alloc);

libc/utils/gpu/loader/Main.cpp

@@ -13,88 +13,97 @@
 
 #include "Loader.h"
 
+#include "llvm/BinaryFormat/Magic.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Signals.h"
+#include "llvm/Support/WithColor.h"
+
 #include <cstdio>
 #include <cstdlib>
 #include <string>
-#include <vector>
 
-int main(int argc, char **argv, char **envp) {
-  if (argc < 2) {
-    printf("USAGE: ./loader [--threads <n>, --blocks <n>, "
-           "--print-resource-usage] <device_image> "
-           "<args>, ...\n");
-    return EXIT_SUCCESS;
-  }
+using namespace llvm;
 
-  int offset = 0;
-  FILE *file = nullptr;
-  char *ptr;
-  LaunchParameters params = {1, 1, 1, 1, 1, 1};
-  bool print_resource_usage = false;
-  while (!file && ++offset < argc) {
-    if (argv[offset] == std::string("--threads") ||
-        argv[offset] == std::string("--threads-x")) {
-      params.num_threads_x =
-          offset + 1 < argc ? strtoul(argv[offset + 1], &ptr, 10) : 1;
-      offset++;
-      continue;
-    } else if (argv[offset] == std::string("--threads-y")) {
-      params.num_threads_y =
-          offset + 1 < argc ? strtoul(argv[offset + 1], &ptr, 10) : 1;
-      offset++;
-      continue;
-    } else if (argv[offset] == std::string("--threads-z")) {
-      params.num_threads_z =
-          offset + 1 < argc ? strtoul(argv[offset + 1], &ptr, 10) : 1;
-      offset++;
-      continue;
-    } else if (argv[offset] == std::string("--blocks") ||
-               argv[offset] == std::string("--blocks-x")) {
-      params.num_blocks_x =
-          offset + 1 < argc ? strtoul(argv[offset + 1], &ptr, 10) : 1;
-      offset++;
-      continue;
-    } else if (argv[offset] == std::string("--blocks-y")) {
-      params.num_blocks_y =
-          offset + 1 < argc ? strtoul(argv[offset + 1], &ptr, 10) : 1;
-      offset++;
-      continue;
-    } else if (argv[offset] == std::string("--blocks-z")) {
-      params.num_blocks_z =
-          offset + 1 < argc ? strtoul(argv[offset + 1], &ptr, 10) : 1;
-      offset++;
-      continue;
-    } else if (argv[offset] == std::string("--print-resource-usage")) {
-      print_resource_usage = true;
-      continue;
-    } else {
-      file = fopen(argv[offset], "r");
-      if (!file) {
-        fprintf(stderr, "Failed to open image file '%s'\n", argv[offset]);
-        return EXIT_FAILURE;
-      }
-      break;
-    }
-  }
+static cl::OptionCategory loader_category("loader options");
+
+static cl::opt<bool> help("h", cl::desc("Alias for -help"), cl::Hidden,
+                          cl::cat(loader_category));
+
+static cl::opt<unsigned>
+    threads_x("threads-x", cl::desc("Number of threads in the 'x' dimension"),
+              cl::init(1), cl::cat(loader_category));
+static cl::opt<unsigned>
+    threads_y("threads-y", cl::desc("Number of threads in the 'y' dimension"),
+              cl::init(1), cl::cat(loader_category));
+static cl::opt<unsigned>
+    threads_z("threads-z", cl::desc("Number of threads in the 'z' dimension"),
+              cl::init(1), cl::cat(loader_category));
+static cl::alias threads("threads", cl::aliasopt(threads_x),
+                         cl::desc("Alias for --threads-x"),
+                         cl::cat(loader_category));
+
+static cl::opt<unsigned>
+    blocks_x("blocks-x", cl::desc("Number of blocks in the 'x' dimension"),
+             cl::init(1), cl::cat(loader_category));
+static cl::opt<unsigned>
+    blocks_y("blocks-y", cl::desc("Number of blocks in the 'y' dimension"),
+             cl::init(1), cl::cat(loader_category));
+static cl::opt<unsigned>
+    blocks_z("blocks-z", cl::desc("Number of blocks in the 'z' dimension"),
+             cl::init(1), cl::cat(loader_category));
+static cl::alias blocks("blocks", cl::aliasopt(blocks_x),
+                        cl::desc("Alias for --blocks-x"),
+                        cl::cat(loader_category));
 
-  if (!file) {
-    fprintf(stderr, "No image file provided\n");
-    return EXIT_FAILURE;
+static cl::opt<bool>
+    print_resource_usage("print-resource-usage",
+                         cl::desc("Output resource usage of launched kernels"),
+                         cl::init(false), cl::cat(loader_category));
+
+static cl::opt<std::string> file(cl::Positional, cl::Required,
+                                 cl::desc("<gpu executable>"),
+                                 cl::cat(loader_category));
+static cl::list<std::string> args(cl::ConsumeAfter,
+                                  cl::desc("<program arguments>..."),
+                                  cl::cat(loader_category));
+
+[[noreturn]] void report_error(Error E) {
+  outs().flush();
+  logAllUnhandledErrors(std::move(E), WithColor::error(errs(), "loader"));
+  exit(EXIT_FAILURE);
+}
+
+int main(int argc, const char **argv, const char **envp) {
+  sys::PrintStackTraceOnErrorSignal(argv[0]);
+  cl::HideUnrelatedOptions(loader_category);
+  cl::ParseCommandLineOptions(
+      argc, argv,
+      "A utility used to launch unit tests built for a GPU target. This is\n"
+      "intended to provide an intrface simular to cross-compiling emulators\n");
+
+  if (help) {
+    cl::PrintHelpMessage();
+    return EXIT_SUCCESS;
   }
 
-  // TODO: We should perform some validation on the file.
-  fseek(file, 0, SEEK_END);
-  const auto size = ftell(file);
-  fseek(file, 0, SEEK_SET);
+  ErrorOr<std::unique_ptr<MemoryBuffer>> image_or_err =
+      MemoryBuffer::getFileOrSTDIN(file);
+  if (std::error_code ec = image_or_err.getError())
+    report_error(errorCodeToError(ec));
+  MemoryBufferRef image = **image_or_err;
 
-  void *image = malloc(size * sizeof(char));
-  fread(image, sizeof(char), size, file);
-  fclose(file);
+  SmallVector<const char *> new_argv = {file.c_str()};
+  llvm::transform(args, std::back_inserter(new_argv),
+                  [](const std::string &arg) { return arg.c_str(); });
 
   // Drop the loader from the program arguments.
-  int ret = load(argc - offset, &argv[offset], envp, image, size, params,
-                 print_resource_usage);
+  LaunchParameters params{threads_x, threads_y, threads_z,
+                          blocks_x,  blocks_y,  blocks_z};
+  int ret = load(new_argv.size(), new_argv.data(), envp,
+                 const_cast<char *>(image.getBufferStart()),
+                 image.getBufferSize(), params, print_resource_usage);
 
-  free(image);
   return ret;
 }

libc/utils/gpu/loader/amdgpu/amdhsa-loader.cpp

@@ -334,8 +334,9 @@ static hsa_status_t hsa_memcpy(void *dst, hsa_agent_t dst_agent,
   return HSA_STATUS_SUCCESS;
 }
 
-int load(int argc, char **argv, char **envp, void *image, size_t size,
-         const LaunchParameters &params, bool print_resource_usage) {
+int load(int argc, const char **argv, const char **envp, void *image,
+         size_t size, const LaunchParameters &params,
+         bool print_resource_usage) {
   // Initialize the HSA runtime used to communicate with the device.
   if (hsa_status_t err = hsa_init())
     handle_error(err);

libc/utils/gpu/loader/nvptx/nvptx-loader.cpp

@@ -245,8 +245,9 @@ CUresult launch_kernel(CUmodule binary, CUstream stream,
   return CUDA_SUCCESS;
 }
 
-int load(int argc, char **argv, char **envp, void *image, size_t size,
-         const LaunchParameters &params, bool print_resource_usage) {
+int load(int argc, const char **argv, const char **envp, void *image,
+         size_t size, const LaunchParameters &params,
+         bool print_resource_usage) {
   if (CUresult err = cuInit(0))
     handle_error(err);
   // Obtain the first device found on the system.