[llvm] Add ability to run -O3 and -Oz optimizations.

compiler_gym/envs/llvm/service/Benchmark.cc

@@ -297,4 +297,8 @@ Status Benchmark::compile() {
  return Status::OK;
 }
 
+bool Benchmark::applyBaselineOptimizations(unsigned optLevel, unsigned sizeLevel) {
+ return applyBaselineOptimizationsToModule(&module(), optLevel, sizeLevel);
+}
+
 } // namespace compiler_gym::llvm_service

compiler_gym/envs/llvm/service/Benchmark.h

@@ -130,6 +130,15 @@ class Benchmark {
 
  grpc::Status compile();
 
+ /**
+ * Apply the given baseline optimizations.
+ *
+ * @param optLevel The runtime optimization level.
+ * @param sizeLevel The size optimization level
+ * @return Whether the baseline optimizations modified the module.
+ */
+ bool applyBaselineOptimizations(unsigned optLevel, unsigned sizeLevel);
+
  inline bool isBuildable() const { return isBuildable_; }
  inline bool isRunnable() const { return isRunnable_; }
 

compiler_gym/envs/llvm/service/Cost.cc

@@ -28,31 +28,6 @@ namespace compiler_gym::llvm_service {
 
 namespace {
 
-// Apply the optimizations from a given LLVM optimization level.
-bool applyBaselineOptimizations(llvm::Module* module, unsigned optLevel, unsigned sizeLevel) {
- llvm::legacy::PassManager passManager;
- llvm::legacy::FunctionPassManager functionPassManager(module);
-
- llvm::PassManagerBuilder builder;
- builder.OptLevel = optLevel;
- builder.SizeLevel = sizeLevel;
- if (optLevel > 1) {
- builder.Inliner = llvm::createFunctionInliningPass(optLevel, sizeLevel, false);
- }
-
- builder.populateFunctionPassManager(functionPassManager);
- builder.populateModulePassManager(passManager);
-
- bool changed = passManager.run(*module);
- changed |= (functionPassManager.doInitialization() ? 1 : 0);
- for (auto& function : *module) {
- changed |= (functionPassManager.run(function) ? 1 : 0);
- }
- changed |= (functionPassManager.doFinalization() ? 1 : 0);
-
- return changed;
-}
-
 // Serialize the module to a string.
 std::string moduleToString(llvm::Module& module) {
  std::string str;
@@ -144,6 +119,39 @@ inline size_t getBaselineCostIndex(LlvmBaselinePolicy policy, LlvmCostFunction c
 
 } // anonymous namespace
 
+/**
+ * Apply the given baseline optimizations.
+ *
+ * @param module The module to optimize.
+ * @param optLevel The runtime optimization level.
+ * @param sizeLevel The size optimization level
+ * @return Whether the baseline optimizations modified the module.
+ */
+bool applyBaselineOptimizationsToModule(llvm::Module* module, unsigned optLevel,
+ unsigned sizeLevel) {
+ llvm::legacy::PassManager passManager;
+ llvm::legacy::FunctionPassManager functionPassManager(module);
+
+ llvm::PassManagerBuilder builder;
+ builder.OptLevel = optLevel;
+ builder.SizeLevel = sizeLevel;
+ if (optLevel > 1) {
+ builder.Inliner = llvm::createFunctionInliningPass(optLevel, sizeLevel, false);
+ }
+
+ builder.populateFunctionPassManager(functionPassManager);
+ builder.populateModulePassManager(passManager);
+
+ bool changed = passManager.run(*module);
+ changed |= (functionPassManager.doInitialization() ? 1 : 0);
+ for (auto& function : *module) {
+ changed |= (functionPassManager.run(function) ? 1 : 0);
+ }
+ changed |= (functionPassManager.doFinalization() ? 1 : 0);
+
+ return changed;
+}
+
 Status setCost(const LlvmCostFunction& costFunction, llvm::Module& module,
  const fs::path& workingDirectory, double* cost) {
  switch (costFunction) {
@@ -195,10 +203,10 @@ Status setBaselineCosts(llvm::Module& unoptimizedModule, BaselineCosts* baseline
  // Create a copy of the unoptimized module and apply the default set of LLVM
  // optimizations.
  std::unique_ptr<llvm::Module> moduleOz = llvm::CloneModule(unoptimizedModule);
- applyBaselineOptimizations(moduleOz.get(), /*optLevel=*/2, /*sizeLevel=*/2);
+ applyBaselineOptimizationsToModule(moduleOz.get(), /*optLevel=*/2, /*sizeLevel=*/2);
 
  std::unique_ptr<llvm::Module> moduleO3 = llvm::CloneModule(unoptimizedModule);
- applyBaselineOptimizations(moduleO3.get(), /*optLevel=*/3, /*sizeLevel=*/0);
+ applyBaselineOptimizationsToModule(moduleO3.get(), /*optLevel=*/3, /*sizeLevel=*/0);
 
  for (const auto policy : magic_enum::enum_values<LlvmBaselinePolicy>()) {
  // Set the baseline module.

compiler_gym/envs/llvm/service/Cost.h

@@ -51,6 +51,17 @@ constexpr size_t numBaselineCosts = magic_enum::enum_count<LlvmBaselinePolicy>()
 using BaselineCosts = std::array<double, numBaselineCosts>;
 using PreviousCosts = std::array<std::optional<double>, numCosts>;
 
+/**
+ * Apply the given baseline optimizations.
+ *
+ * @param module The module to optimize.
+ * @param optLevel The runtime optimization level.
+ * @param sizeLevel The size optimization level
+ * @return Whether the baseline optimizations modified the module.
+ */
+bool applyBaselineOptimizationsToModule(llvm::Module* module, unsigned optLevel,
+ unsigned sizeLevel);
+
 /**
  * Compute the cost using a given cost function. A lower cost is better.
  *

compiler_gym/envs/llvm/service/LlvmSession.cc

@@ -172,6 +172,17 @@ Status LlvmSession::handleSessionParameter(const std::string& key, const std::st
  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;
 }