|
20 | 20 | #include "llvm/Support/Signals.h"
|
21 | 21 | #include "llvm/Support/WithColor.h"
|
22 | 22 |
|
| 23 | +#include <cerrno> |
23 | 24 | #include <cstdio>
|
24 | 25 | #include <cstdlib>
|
| 26 | +#include <cstring> |
25 | 27 | #include <string>
|
| 28 | +#include <sys/file.h> |
26 | 29 |
|
27 | 30 | using namespace llvm;
|
28 | 31 |
|
@@ -62,6 +65,12 @@ static cl::opt<bool>
|
62 | 65 | cl::desc("Output resource usage of launched kernels"),
|
63 | 66 | cl::init(false), cl::cat(loader_category));
|
64 | 67 |
|
| 68 | +static cl::opt<bool> |
| 69 | + no_parallelism("no-parallelism", |
| 70 | + cl::desc("Allows only a single process to use the GPU at a " |
| 71 | + "time. Useful to suppress out-of-resource errors"), |
| 72 | + cl::init(false), cl::cat(loader_category)); |
| 73 | + |
65 | 74 | static cl::opt<std::string> file(cl::Positional, cl::Required,
|
66 | 75 | cl::desc("<gpu executable>"),
|
67 | 76 | cl::cat(loader_category));
|
@@ -98,12 +107,28 @@ int main(int argc, const char **argv, const char **envp) {
|
98 | 107 | llvm::transform(args, std::back_inserter(new_argv),
|
99 | 108 | [](const std::string &arg) { return arg.c_str(); });
|
100 | 109 |
|
| 110 | + // Claim a file lock on the executable so only a single process can enter this |
| 111 | + // region if requested. This prevents the loader from spurious failures. |
| 112 | + int fd = -1; |
| 113 | + if (no_parallelism) { |
| 114 | + fd = open(argv[0], O_RDONLY); |
| 115 | + if (flock(fd, LOCK_EX) == -1) |
| 116 | + report_error(createStringError("Failed to lock '%s': %s", argv[0], |
| 117 | + strerror(errno))); |
| 118 | + } |
| 119 | + |
101 | 120 | // Drop the loader from the program arguments.
|
102 | 121 | LaunchParameters params{threads_x, threads_y, threads_z,
|
103 | 122 | blocks_x, blocks_y, blocks_z};
|
104 | 123 | int ret = load(new_argv.size(), new_argv.data(), envp,
|
105 | 124 | const_cast<char *>(image.getBufferStart()),
|
106 | 125 | image.getBufferSize(), params, print_resource_usage);
|
107 | 126 |
|
| 127 | + if (no_parallelism) { |
| 128 | + if (flock(fd, LOCK_UN) == -1) |
| 129 | + report_error(createStringError("Failed to unlock '%s': %s", argv[0], |
| 130 | + strerror(errno))); |
| 131 | + } |
| 132 | + |
108 | 133 | return ret;
|
109 | 134 | }
|
0 commit comments