libc/utils/gpu/loader/Main.cpp

   1 //===-- Main entry into the loader interface ------------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file opens a device image passed on the command line and passes it to
  10 // one of the loader implementations for launch.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "Loader.h"
  15
  16 #include "llvm/BinaryFormat/Magic.h"
  17 #include "llvm/Support/CommandLine.h"
  18 #include "llvm/Support/Error.h"
  19 #include "llvm/Support/FileSystem.h"
  20 #include "llvm/Support/MemoryBuffer.h"
  21 #include "llvm/Support/Path.h"
  22 #include "llvm/Support/Signals.h"
  23 #include "llvm/Support/WithColor.h"
  24
  25 #include <cerrno>
  26 #include <cstdio>
  27 #include <cstdlib>
  28 #include <cstring>
  29 #include <string>
  30 #include <sys/file.h>
  31
  32 using namespace llvm;
  33
  34 static cl::OptionCategory loader_category("loader options");
  35
  36 static cl::opt<bool> help("h", cl::desc("Alias for -help"), cl::Hidden,
  37                           cl::cat(loader_category));
  38
  39 static cl::opt<unsigned>
  40     threads_x("threads-x", cl::desc("Number of threads in the 'x' dimension"),
  41               cl::init(1), cl::cat(loader_category));
  42 static cl::opt<unsigned>
  43     threads_y("threads-y", cl::desc("Number of threads in the 'y' dimension"),
  44               cl::init(1), cl::cat(loader_category));
  45 static cl::opt<unsigned>
  46     threads_z("threads-z", cl::desc("Number of threads in the 'z' dimension"),
  47               cl::init(1), cl::cat(loader_category));
  48 static cl::alias threads("threads", cl::aliasopt(threads_x),
  49                          cl::desc("Alias for --threads-x"),
  50                          cl::cat(loader_category));
  51
  52 static cl::opt<unsigned>
  53     blocks_x("blocks-x", cl::desc("Number of blocks in the 'x' dimension"),
  54              cl::init(1), cl::cat(loader_category));
  55 static cl::opt<unsigned>
  56     blocks_y("blocks-y", cl::desc("Number of blocks in the 'y' dimension"),
  57              cl::init(1), cl::cat(loader_category));
  58 static cl::opt<unsigned>
  59     blocks_z("blocks-z", cl::desc("Number of blocks in the 'z' dimension"),
  60              cl::init(1), cl::cat(loader_category));
  61 static cl::alias blocks("blocks", cl::aliasopt(blocks_x),
  62                         cl::desc("Alias for --blocks-x"),
  63                         cl::cat(loader_category));
  64
  65 static cl::opt<bool>
  66     print_resource_usage("print-resource-usage",
  67                          cl::desc("Output resource usage of launched kernels"),
  68                          cl::init(false), cl::cat(loader_category));
  69
  70 static cl::opt<bool>
  71     no_parallelism("no-parallelism",
  72                    cl::desc("Allows only a single process to use the GPU at a "
  73                             "time. Useful to suppress out-of-resource errors"),
  74                    cl::init(false), cl::cat(loader_category));
  75
  76 static cl::opt<std::string> file(cl::Positional, cl::Required,
  77                                  cl::desc("<gpu executable>"),
  78                                  cl::cat(loader_category));
  79 static cl::list<std::string> args(cl::ConsumeAfter,
  80                                   cl::desc("<program arguments>..."),
  81                                   cl::cat(loader_category));
  82
  83 [[noreturn]] void report_error(Error E) {
  84   outs().flush();
  85   logAllUnhandledErrors(std::move(E), WithColor::error(errs(), "loader"));
  86   exit(EXIT_FAILURE);
  87 }
  88
  89 std::string get_main_executable(const char *name) {
  90   void *ptr = (void *)(intptr_t)&get_main_executable;
  91   auto cow_path = sys::fs::getMainExecutable(name, ptr);
  92   return sys::path::parent_path(cow_path).str();
  93 }
  94
  95 int main(int argc, const char **argv, const char **envp) {
  96   sys::PrintStackTraceOnErrorSignal(argv[0]);
  97   cl::HideUnrelatedOptions(loader_category);
  98   cl::ParseCommandLineOptions(
  99       argc, argv,
 100       "A utility used to launch unit tests built for a GPU target. This is\n"
 101       "intended to provide an intrface simular to cross-compiling emulators\n");
 102
 103   if (help) {
 104     cl::PrintHelpMessage();
 105     return EXIT_SUCCESS;
 106   }
 107
 108   ErrorOr<std::unique_ptr<MemoryBuffer>> image_or_err =
 109       MemoryBuffer::getFileOrSTDIN(file);
 110   if (std::error_code ec = image_or_err.getError())
 111     report_error(errorCodeToError(ec));
 112   MemoryBufferRef image = **image_or_err;
 113
 114   SmallVector<const char *> new_argv = {file.c_str()};
 115   llvm::transform(args, std::back_inserter(new_argv),
 116                   [](const std::string &arg) { return arg.c_str(); });
 117
 118   // Claim a file lock on the executable so only a single process can enter this
 119   // region if requested. This prevents the loader from spurious failures.
 120   int fd = -1;
 121   if (no_parallelism) {
 122     fd = open(get_main_executable(argv[0]).c_str(), O_RDONLY);
 123     if (flock(fd, LOCK_EX) == -1)
 124       report_error(createStringError("Failed to lock '%s': %s", argv[0],
 125                                      strerror(errno)));
 126   }
 127
 128   // Drop the loader from the program arguments.
 129   LaunchParameters params{threads_x, threads_y, threads_z,
 130                           blocks_x,  blocks_y,  blocks_z};
 131   int ret = load(new_argv.size(), new_argv.data(), envp,
 132                  const_cast<char *>(image.getBufferStart()),
 133                  image.getBufferSize(), params, print_resource_usage);
 134
 135   if (no_parallelism) {
 136     if (flock(fd, LOCK_UN) == -1)
 137       report_error(createStringError("Failed to unlock '%s': %s", argv[0],
 138                                      strerror(errno)));
 139   }
 140
 141   return ret;
 142 }