-
Notifications
You must be signed in to change notification settings - Fork 123
/
LlvmSession.cc
320 lines (270 loc) · 11.6 KB
/
LlvmSession.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
// Copyright (c) Facebook, Inc. and its affiliates.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.
#include "compiler_gym/envs/llvm/service/LlvmSession.h"
#include <cpuinfo.h>
#include <fmt/format.h>
#include <glog/logging.h>
#include <boost/process.hpp>
#include <chrono>
#include <future>
#include <iomanip>
#include <optional>
#include <string>
#include "boost/filesystem.hpp"
#include "compiler_gym/envs/llvm/service/ActionSpace.h"
#include "compiler_gym/envs/llvm/service/Benchmark.h"
#include "compiler_gym/envs/llvm/service/BenchmarkFactory.h"
#include "compiler_gym/envs/llvm/service/Cost.h"
#include "compiler_gym/envs/llvm/service/Observation.h"
#include "compiler_gym/envs/llvm/service/ObservationSpaces.h"
#include "compiler_gym/envs/llvm/service/passes/ActionHeaders.h"
#include "compiler_gym/envs/llvm/service/passes/ActionSwitch.h"
#include "compiler_gym/third_party/autophase/InstCount.h"
#include "compiler_gym/third_party/llvm/InstCount.h"
#include "compiler_gym/util/EnumUtil.h"
#include "compiler_gym/util/GrpcStatusMacros.h"
#include "compiler_gym/util/RunfilesPath.h"
#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/CodeGen/Passes.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Metadata.h"
#include "llvm/IR/Module.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/raw_ostream.h"
#include "nlohmann/json.hpp"
#include "programl/graph/format/node_link_graph.h"
#include "programl/ir/llvm/llvm.h"
namespace fs = boost::filesystem;
namespace bp = boost::process;
namespace compiler_gym::llvm_service {
using grpc::Status;
using grpc::StatusCode;
using nlohmann::json;
using BenchmarkProto = compiler_gym::Benchmark;
using ActionSpaceProto = compiler_gym::ActionSpace;
namespace {
// Return the target library information for a module.
llvm::TargetLibraryInfoImpl getTargetLibraryInfo(llvm::Module& module) {
llvm::Triple triple(module.getTargetTriple());
return llvm::TargetLibraryInfoImpl(triple);
}
} // anonymous namespace
std::string LlvmSession::getCompilerVersion() const {
std::stringstream ss;
ss << LLVM_VERSION_STRING << " " << llvm::Triple::normalize(LLVM_DEFAULT_TARGET_TRIPLE);
return ss.str();
}
std::vector<ActionSpace> LlvmSession::getActionSpaces() const { return getLlvmActionSpaceList(); }
std::vector<ObservationSpace> LlvmSession::getObservationSpaces() const {
return getLlvmObservationSpaceList();
}
LlvmSession::LlvmSession(const boost::filesystem::path& workingDirectory)
: CompilationSession(workingDirectory),
observationSpaceNames_(util::createPascalCaseToEnumLookupTable<LlvmObservationSpace>()) {
cpuinfo_initialize();
}
Status LlvmSession::init(const ActionSpace& actionSpace, const BenchmarkProto& benchmark) {
BenchmarkFactory& benchmarkFactory = BenchmarkFactory::getSingleton(workingDirectory());
// Get the benchmark or return an error.
std::unique_ptr<Benchmark> llvmBenchmark;
RETURN_IF_ERROR(benchmarkFactory.getBenchmark(benchmark, &llvmBenchmark));
// Verify the benchmark now to catch errors early.
RETURN_IF_ERROR(llvmBenchmark->verify_module());
LlvmActionSpace actionSpaceEnum;
RETURN_IF_ERROR(util::pascalCaseToEnum(actionSpace.name(), &actionSpaceEnum));
return init(actionSpaceEnum, std::move(llvmBenchmark));
}
Status LlvmSession::init(CompilationSession* other) {
// TODO: Static cast?
auto llvmOther = static_cast<LlvmSession*>(other);
return init(llvmOther->actionSpace(), llvmOther->benchmark().clone(workingDirectory()));
}
Status LlvmSession::init(const LlvmActionSpace& actionSpace, std::unique_ptr<Benchmark> benchmark) {
benchmark_ = std::move(benchmark);
actionSpace_ = actionSpace;
tlii_ = getTargetLibraryInfo(benchmark_->module());
return Status::OK;
}
Status LlvmSession::applyAction(const Action& action, bool& endOfEpisode,
std::optional<ActionSpace>& newActionSpace,
bool& actionHadNoEffect) {
DCHECK(benchmark_) << "Calling applyAction() before init()";
// Apply the requested action.
switch (actionSpace()) {
case LlvmActionSpace::PASSES_ALL:
LlvmAction actionEnum;
if (action.choice_size() != 1) {
return Status(
StatusCode::INVALID_ARGUMENT,
fmt::format("Invalid choice count. Expected 1, received {}", action.choice_size()));
}
RETURN_IF_ERROR(util::intToEnum(action.choice(0).named_discrete_value_index(), &actionEnum));
RETURN_IF_ERROR(applyPassAction(actionEnum, actionHadNoEffect));
}
return Status::OK;
}
Status LlvmSession::endOfStep(bool actionHadNoEffect, bool& endOfEpisode,
std::optional<ActionSpace>& newActionSpace) {
if (actionHadNoEffect) {
return Status::OK;
} else {
return benchmark().verify_module();
}
}
Status LlvmSession::computeObservation(const ObservationSpace& observationSpace,
Observation& observation) {
DCHECK(benchmark_) << "Calling computeObservation() before init()";
const auto& it = observationSpaceNames_.find(observationSpace.name());
if (it == observationSpaceNames_.end()) {
return Status(
StatusCode::INVALID_ARGUMENT,
fmt::format("Could not interpret observation space name: {}", observationSpace.name()));
}
const LlvmObservationSpace observationSpaceEnum = it->second;
return setObservation(observationSpaceEnum, workingDirectory(), benchmark(), observation);
}
Status LlvmSession::handleSessionParameter(const std::string& key, const std::string& value,
std::optional<std::string>& reply) {
if (key == "llvm.set_runtimes_per_observation_count") {
const int ivalue = std::stoi(value);
if (ivalue < 1) {
return Status(
StatusCode::INVALID_ARGUMENT,
fmt::format("runtimes_per_observation_count must be >= 1. Received: {}", ivalue));
}
benchmark().setRuntimesPerObservationCount(ivalue);
reply = value;
} else if (key == "llvm.get_runtimes_per_observation_count") {
reply = fmt::format("{}", benchmark().getRuntimesPerObservationCount());
} else if (key == "llvm.set_warmup_runs_count_per_runtime_observation") {
const int ivalue = std::stoi(value);
if (ivalue < 0) {
return Status(
StatusCode::INVALID_ARGUMENT,
fmt::format("warmup_runs_count_per_runtime_observation must be >= 0. Received: {}",
ivalue));
}
benchmark().setWarmupRunsPerRuntimeObservationCount(ivalue);
reply = value;
} else if (key == "llvm.get_warmup_runs_count_per_runtime_observation") {
reply = fmt::format("{}", benchmark().getWarmupRunsPerRuntimeObservationCount());
} else if (key == "llvm.set_buildtimes_per_observation_count") {
const int ivalue = std::stoi(value);
if (ivalue < 1) {
return Status(
StatusCode::INVALID_ARGUMENT,
fmt::format("buildtimes_per_observation_count must be >= 1. Received: {}", ivalue));
}
benchmark().setBuildtimesPerObservationCount(ivalue);
reply = value;
} else if (key == "llvm.get_buildtimes_per_observation_count") {
reply = fmt::format("{}", benchmark().getBuildtimesPerObservationCount());
} else if (key == "llvm.apply_baseline_optimizations") {
if (value == "-Oz") {
bool changed = benchmark().applyBaselineOptimizations(/*optLevel=*/2, /*sizeLevel=*/2);
reply = changed ? "1" : "0";
} else if (value == "-O3") {
bool changed = benchmark().applyBaselineOptimizations(/*optLevel=*/3, /*sizeLevel=*/0);
reply = changed ? "1" : "0";
} else {
return Status(StatusCode::INVALID_ARGUMENT,
fmt::format("Invalid value for llvm.apply_baseline_optimizations: {}", value));
}
}
return Status::OK;
}
Status LlvmSession::applyPassAction(LlvmAction action, bool& actionHadNoEffect) {
#ifdef EXPERIMENTAL_UNSTABLE_GVN_SINK_PASS
// NOTE(https://github.com/facebookresearch/CompilerGym/issues/46): The
// -gvn-sink pass has been found to have nondeterministic behavior so has
// been disabled in compiler_gym/envs/llvm/service/pass/config.py. Invoking
// the command line was found to produce more stable results.
if (action == LlvmAction::GVNSINK_PASS) {
RETURN_IF_ERROR(runOptWithArgs({"-gvn-sink"}));
actionHadNoEffect = true;
return Status::OK;
}
#endif
// Use the generated HANDLE_PASS() switch statement to dispatch to runPass().
#define HANDLE_PASS(pass) actionHadNoEffect = !runPass(pass);
HANDLE_ACTION(action, HANDLE_PASS)
#undef HANDLE_PASS
if (!actionHadNoEffect) {
benchmark().markModuleModified();
}
return Status::OK;
}
bool LlvmSession::runPass(llvm::Pass* pass) {
llvm::legacy::PassManager passManager;
setupPassManager(&passManager, pass);
return passManager.run(benchmark().module());
}
bool LlvmSession::runPass(llvm::FunctionPass* pass) {
llvm::legacy::FunctionPassManager passManager(&benchmark().module());
setupPassManager(&passManager, pass);
bool changed = passManager.doInitialization();
for (auto& function : benchmark().module()) {
changed |= (passManager.run(function) ? 1 : 0);
}
changed |= (passManager.doFinalization() ? 1 : 0);
return changed;
}
Status LlvmSession::runOptWithArgs(const std::vector<std::string>& optArgs) {
// Create temporary files for `opt` to read from and write to.
const auto before_path = fs::unique_path(workingDirectory() / "module-%%%%%%%%.bc");
const auto after_path = fs::unique_path(workingDirectory() / "module-%%%%%%%%.bc");
RETURN_IF_ERROR(writeBitcodeFile(benchmark().module(), before_path));
// Build a command line invocation: `opt input.bc -o output.bc <optArgs...>`.
const auto optPath = util::getSiteDataPath("llvm-v0/bin/opt");
if (!fs::exists(optPath)) {
return Status(StatusCode::INTERNAL, fmt::format("File not found: {}", optPath.string()));
}
std::string optCmd =
fmt::format("{} {} -o {}", optPath.string(), before_path.string(), after_path.string());
for (const auto& arg : optArgs) {
optCmd += " " + arg;
}
// Run the opt command line.
try {
boost::asio::io_context optStderrStream;
std::future<std::string> optStderrFuture;
bp::child opt(optCmd, bp::std_in.close(), bp::std_out > bp::null, bp::std_err > optStderrFuture,
optStderrStream);
if (!util::wait_for(opt, std::chrono::seconds(60))) {
return Status(StatusCode::DEADLINE_EXCEEDED,
fmt::format("Failed to run opt within 60 seconds: {}", optCmd));
}
optStderrStream.run();
if (opt.exit_code()) {
const std::string stderr = optStderrFuture.get();
return Status(StatusCode::INTERNAL,
fmt::format("Opt command '{}' failed with return code {}: {}", optCmd,
opt.exit_code(), stderr));
}
fs::remove(before_path);
} catch (bp::process_error& e) {
fs::remove(before_path);
return Status(StatusCode::INTERNAL,
fmt::format("Failed to run opt command '{}': {}", optCmd, e.what()));
}
if (!fs::exists(after_path)) {
return Status(StatusCode::INTERNAL, "Failed to generate output file");
}
// Read the bitcode file generated by `opt`.
Bitcode bitcode;
auto status = readBitcodeFile(after_path, &bitcode);
fs::remove(after_path);
if (!status.ok()) {
return status;
}
// Replace the benchmark's module with the one generated by `opt`.
auto module = makeModule(benchmark().context(), bitcode, benchmark().name(), &status);
RETURN_IF_ERROR(status);
benchmark().replaceModule(std::move(module));
return Status::OK;
}
} // namespace compiler_gym::llvm_service