Merge commit 'd1b294029004' from llvm.org/main into next

Author: git apple-llvm automerger

libc/utils/gpu/loader/Main.cpp

@@ -16,13 +16,18 @@
 #include "llvm/BinaryFormat/Magic.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/WithColor.h"
 
+#include <cerrno>
 #include <cstdio>
 #include <cstdlib>
+#include <cstring>
 #include <string>
+#include <sys/file.h>
 
 using namespace llvm;
 
@@ -62,6 +67,12 @@ static cl::opt<bool>
                          cl::desc("Output resource usage of launched kernels"),
                          cl::init(false), cl::cat(loader_category));
 
+static cl::opt<bool>
+    no_parallelism("no-parallelism",
+                   cl::desc("Allows only a single process to use the GPU at a "
+                            "time. Useful to suppress out-of-resource errors"),
+                   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));
@@ -75,6 +86,12 @@ static cl::list<std::string> args(cl::ConsumeAfter,
   exit(EXIT_FAILURE);
 }
 
+std::string get_main_executable(const char *name) {
+  void *ptr = (void *)(intptr_t)&get_main_executable;
+  auto cow_path = sys::fs::getMainExecutable(name, ptr);
+  return sys::path::parent_path(cow_path).str();
+}
+
 int main(int argc, const char **argv, const char **envp) {
   sys::PrintStackTraceOnErrorSignal(argv[0]);
   cl::HideUnrelatedOptions(loader_category);
@@ -98,12 +115,28 @@ int main(int argc, const char **argv, const char **envp) {
   llvm::transform(args, std::back_inserter(new_argv),
                   [](const std::string &arg) { return arg.c_str(); });
 
+  // Claim a file lock on the executable so only a single process can enter this
+  // region if requested. This prevents the loader from spurious failures.
+  int fd = -1;
+  if (no_parallelism) {
+    fd = open(get_main_executable(argv[0]).c_str(), O_RDONLY);
+    if (flock(fd, LOCK_EX) == -1)
+      report_error(createStringError("Failed to lock '%s': %s", argv[0],
+                                     strerror(errno)));
+  }
+
   // Drop the loader from the program arguments.
   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);
 
+  if (no_parallelism) {
+    if (flock(fd, LOCK_UN) == -1)
+      report_error(createStringError("Failed to unlock '%s': %s", argv[0],
+                                     strerror(errno)));
+  }
+
   return ret;
 }