fix llvm-kompile-codegen so it optimizes code properly (#1097)

When we migrated the llvm backend optimization pipeline from being
applied by `opt` to being applied by `llvm-kompile-codegen`, we made an
error. The code generator was not, in fact, applying the `opt`
transformation pipeline correctly. This means that `kompile -O2` will
not, in fact, apply `llc -O2` to the bitcode generated by the llvm
backend, leading to a significant performance regression.

While this PR does not entirely fix the issue (an upstream change to the
K frontend is also required), it fixes the issue within this repository.
We see a roughly 1.5x speedup in runtime when `-O2` is passed to
kompile. We do lose something in compilation time, however. This is
normal; compiling with optimizations is more expensive, we just weren't
doing it correctly before.

The main changes in this pull request are threefold:

1. Convert optimizer code to new pass manager.
2. Make sure to run the middle-end optimizer.
3. Convert `TailCallElimination` and `Mem2Reg` from required to optional
passes on -O0. In order to make the tail call optimization still work
correctly without `TailCallElimination` manually marking tail calls as
`tail`, we instead explicitly mark them as `musttail` in the IR
generated by the code generator.

Author: Dwight Guth

.github/workflows/ci-tests.sh

@@ -17,7 +17,7 @@ make -j$(nproc) install
 popd
 
 pushd matching
-mvn package
+mvn -B package
 popd
 
 export PATH="$(realpath ./build/install/bin):$(realpath ./build/bin):/usr/lib/llvm-${llvm_version}/bin:$PATH"

include/kllvm/codegen/Options.h

@@ -3,8 +3,6 @@
 
 #include "llvm/Support/CommandLine.h"
 
-enum class opt_level { O0, O1, O2, O3 };
-
 extern llvm::cl::OptionCategory codegen_lib_cat;
 
 extern llvm::cl::opt<bool> debug;
@@ -15,7 +13,7 @@ extern llvm::cl::opt<bool> force_binary;
 extern llvm::cl::opt<bool> proof_hint_instrumentation;
 extern llvm::cl::opt<bool> proof_hint_instrumentation_slow;
 extern llvm::cl::opt<bool> keep_frame_pointer;
-extern llvm::cl::opt<opt_level> optimization_level;
+extern llvm::cl::opt<char> optimization_level;
 
 namespace kllvm {
 

include/kllvm/codegen/SetVisibilityHidden.h

@@ -31,6 +31,9 @@ struct set_visibility_hidden : llvm::PassInfoMixin<set_visibility_hidden> {
     }
     return llvm::PreservedAnalyses::none();
   }
+
+  // NOLINTNEXTLINE(*-identifier-naming)
+  static bool isRequired() { return true; }
 };
 
 } // namespace kllvm

lib/codegen/ApplyPasses.cpp

@@ -37,23 +37,65 @@ namespace kllvm {
 
 auto get_opt_level() {
   switch (optimization_level) {
-  case opt_level::O0: return CODEGEN_OPT_LEVEL::None;
-  case opt_level::O1: return CODEGEN_OPT_LEVEL::Less;
-  case opt_level::O2: return CODEGEN_OPT_LEVEL::Default;
-  case opt_level::O3: return CODEGEN_OPT_LEVEL::Aggressive;
+  case '0': return CODEGEN_OPT_LEVEL::None;
+  case '1': return CODEGEN_OPT_LEVEL::Less;
+  case '2': return CODEGEN_OPT_LEVEL::Default;
+  case '3': return CODEGEN_OPT_LEVEL::Aggressive;
+  default:
+    throw std::runtime_error(
+        fmt::format("Invalid optimization level: {}", optimization_level));
   }
 }
 
-void apply_kllvm_opt_passes(llvm::Module &mod, bool hidden_visibility) {
-  auto pm = legacy::PassManager();
-
-  pm.add(createPromoteMemoryToRegisterPass());
-  pm.add(createTailCallEliminationPass());
-  if (hidden_visibility) {
-    pm.add(new legacy_set_visibility_hidden());
+auto get_pass_opt_level() {
+  switch (optimization_level) {
+  case '0': return OptimizationLevel::O0;
+  case '1': return OptimizationLevel::O1;
+  case '2': return OptimizationLevel::O2;
+  case '3': return OptimizationLevel::O3;
+  default:
+    throw std::runtime_error(
+        fmt::format("Invalid optimization level: {}", optimization_level));
   }
+}
 
-  pm.run(mod);
+void apply_kllvm_opt_passes(llvm::Module &mod, bool hidden_visibility) {
+  // Create the analysis managers.
+  // These must be declared in this order so that they are destroyed in the
+  // correct order due to inter-analysis-manager references.
+  LoopAnalysisManager lam;
+  FunctionAnalysisManager fam;
+  CGSCCAnalysisManager cgam;
+  ModuleAnalysisManager mam;
+
+  // Create the new pass manager builder.
+  // Take a look at the PassBuilder constructor parameters for more
+  // customization, e.g. specifying a TargetMachine or various debugging
+  // options.
+  PassBuilder pb;
+
+  // Register all the basic analyses with the managers.
+  pb.registerModuleAnalyses(mam);
+  pb.registerCGSCCAnalyses(cgam);
+  pb.registerFunctionAnalyses(fam);
+  pb.registerLoopAnalyses(lam);
+  pb.crossRegisterProxies(lam, fam, cgam, mam);
+
+  // register custom passes
+  pb.registerPipelineStartEPCallback(
+      [hidden_visibility](
+          llvm::ModulePassManager &pm, OptimizationLevel level) {
+        if (hidden_visibility) {
+          pm.addPass(set_visibility_hidden());
+        }
+      });
+
+  // Create the pass manager.
+  ModulePassManager mpm
+      = pb.buildPerModuleDefaultPipeline(get_pass_opt_level());
+
+  // Optimize the IR!
+  mpm.run(mod, mam);
 }
 
 void generate_object_file(llvm::Module &mod, llvm::raw_ostream &os) {

lib/codegen/CreateTerm.cpp

@@ -795,6 +795,7 @@ llvm::Value *create_term::create_function_call(
   set_debug_loc(call);
   if (tailcc) {
     call->setCallingConv(llvm::CallingConv::Tail);
+    call->setTailCall();
   }
   if (sret) {
     llvm::Attribute sret_attr
@@ -1138,7 +1139,20 @@ bool make_function(
     auto *call = llvm::CallInst::Create(step, {retval}, "", current_block);
     set_debug_loc(call);
     call->setCallingConv(llvm::CallingConv::Tail);
+    call->setTailCallKind(llvm::CallInst::TCK_MustTail);
     retval = call;
+  } else {
+    if (auto *call = llvm::dyn_cast<llvm::CallInst>(retval)) {
+      // check that musttail requirements are met:
+      // 1. Call is in tail position (guaranteed)
+      // 2. Return returns return value of call (guaranteed)
+      // 3. Calling convention is tailcc
+      // 4. Function is not sret (here approximated by checking if return type is void)
+      if (call->getCallingConv() == llvm::CallingConv::Tail
+          && call->getType() != llvm::Type::getVoidTy(module->getContext())) {
+        call->setTailCallKind(llvm::CallInst::TCK_MustTail);
+      }
+    }
   }
   auto *ret
       = llvm::ReturnInst::Create(module->getContext(), retval, current_block);
@@ -1262,6 +1276,7 @@ std::string make_apply_rule_function(
       = llvm::CallInst::Create(step, args, "", creator.get_current_block());
   set_debug_loc(retval);
   retval->setCallingConv(llvm::CallingConv::Tail);
+  retval->setTailCallKind(llvm::CallInst::TCK_MustTail);
   llvm::ReturnInst::Create(
       module->getContext(), retval, creator.get_current_block());
   return name;
@@ -1277,7 +1292,7 @@ std::string make_side_condition_function(
   }
   std::string name = "side_condition_" + std::to_string(axiom->get_ordinal());
   if (make_function(
-          name, pattern, definition, module, false, false, false, axiom,
+          name, pattern, definition, module, true, false, false, axiom,
           ".sc")) {
     return name;
   }

lib/codegen/Decision.cpp

@@ -441,7 +441,8 @@ void function_node::codegen(decision *d) {
   create_term creator(
       final_subst, d->definition_, d->current_block_, d->module_, false);
   auto *call = creator.create_function_call(
-      function_, cat_, args, function_.substr(0, 5) == "hook_", false);
+      function_, cat_, args, function_.substr(0, 5) == "hook_",
+      is_side_condition);
   call->setName(name_.substr(0, max_name_length));
   d->store(std::make_pair(name_, type_), call);
 
@@ -625,6 +626,7 @@ void leaf_node::codegen(decision *d) {
   call->setCallingConv(llvm::CallingConv::Tail);
 
   if (child_ == nullptr) {
+    call->setTailCallKind(llvm::CallInst::TCK_MustTail);
     llvm::ReturnInst::Create(d->ctx_, call, d->current_block_);
   } else {
     new llvm::StoreInst(

lib/codegen/Options.cpp

@@ -20,14 +20,11 @@ cl::opt<bool> keep_frame_pointer(
     cl::desc("Keep frame pointer in compiled code for debugging purposes"),
     cl::cat(codegen_lib_cat));
 
-cl::opt<opt_level> optimization_level(
-    cl::desc("Choose optimization level"),
-    cl::values(
-        clEnumVal(opt_level::O0, "No optimizations"),
-        clEnumVal(opt_level::O1, "Enable trivial optimizations"),
-        clEnumVal(opt_level::O2, "Enable default optimizations"),
-        clEnumVal(opt_level::O3, "Enable expensive optimizations")),
-    cl::cat(codegen_lib_cat));
+cl::opt<char> optimization_level(
+    "O",
+    cl::desc("Optimization level. [-O0, -O1, -O2, or -O3] "
+             "(default = '-O0')"),
+    cl::Prefix, cl::init('0'), cl::cat(codegen_lib_cat));
 
 cl::opt<bool> debug(
     "debug", cl::desc("Enable debug information"), cl::ZeroOrMore,