Add optimization passes controlled via -Ox command line flag

Author: weliveindetail

CMakeLists.txt

@@ -39,6 +39,7 @@ endif()
 add_executable(JitFromScratch
   main.cpp
   JitFromScratch.cpp
+  SimpleOptimizer.cpp
 )
 
 target_include_directories(JitFromScratch PRIVATE
@@ -65,6 +66,7 @@ set(llvm_libs
   LLVMGlobalISel
   LLVMInstCombine
   LLVMInstrumentation
+  LLVMipo
   LLVMJITLink
   LLVMLinker
   LLVMIRReader
@@ -81,6 +83,7 @@ set(llvm_libs
   LLVMSupport
   LLVMTarget
   LLVMTransformUtils
+  LLVMVectorize
   LLVMX86AsmParser
   LLVMX86CodeGen
   LLVMX86Desc

JitFromScratch.cpp

@@ -1,4 +1,5 @@
 #include "JitFromScratch.h"
+#include "SimpleOptimizer.h"
 
 #include <llvm/ExecutionEngine/Orc/CompileUtils.h>
 #include <llvm/ExecutionEngine/Orc/ExecutionUtils.h>
@@ -7,6 +8,7 @@
 #include <llvm/IR/Mangler.h>
 #include <llvm/Support/Debug.h>
 #include <llvm/Support/raw_ostream.h>
+#include <llvm/Transforms/IPO/PassManagerBuilder.h>
 
 #define DEBUG_TYPE "jitfromscratch"
 
@@ -18,7 +20,8 @@ JitFromScratch::JitFromScratch(ExitOnError ExitOnErr)
       TM(createTargetMachine(ExitOnErr)),
       GDBListener(JITEventListener::createGDBRegistrationListener()),
       ObjLinkingLayer(*ES, createMemoryManagerFtor()),
-      CompileLayer(*ES, ObjLinkingLayer, SimpleCompiler(*TM)) {
+      CompileLayer(*ES, ObjLinkingLayer, SimpleCompiler(*TM)),
+      OptimizeLayer(*ES, CompileLayer) {
   ObjLinkingLayer.setNotifyLoaded(createNotifyLoadedFtor());
   if (auto R = createHostProcessResolver())
     ES->getMainJITDylib().setGenerator(std::move(R));
@@ -90,15 +93,18 @@ Error JitFromScratch::applyDataLayout(Module &M) {
 }
 
 Error JitFromScratch::submitModule(std::unique_ptr<Module> M,
-                                   std::unique_ptr<LLVMContext> C) {
+                                   std::unique_ptr<LLVMContext> C,
+                                   unsigned OptLevel) {
   LLVM_DEBUG(dbgs() << "Submit IR module:\n\n" << *M << "\n\n");
 
   if (auto Err = applyDataLayout(*M))
     return Err;
 
-  return CompileLayer.add(ES->getMainJITDylib(),
-                          ThreadSafeModule(std::move(M), std::move(C)),
-                          ES->allocateVModule());
+  OptimizeLayer.setTransform(SimpleOptimizer(OptLevel));
+
+  return OptimizeLayer.add(ES->getMainJITDylib(),
+                           ThreadSafeModule(std::move(M), std::move(C)),
+                           ES->allocateVModule());
 }
 
 Expected<JITTargetAddress> JitFromScratch::getFunctionAddr(StringRef Name) {

JitFromScratch.h

@@ -36,7 +36,8 @@ class JitFromScratch {
   }
 
   llvm::Error submitModule(std::unique_ptr<llvm::Module> M,
-                           std::unique_ptr<llvm::LLVMContext> C);
+                           std::unique_ptr<llvm::LLVMContext> C,
+                           unsigned OptLevel);
 
   template <class Signature_t>
   llvm::Expected<std::function<Signature_t>> getFunction(llvm::StringRef Name) {
@@ -54,6 +55,7 @@ class JitFromScratch {
 
   llvm::orc::RTDyldObjectLinkingLayer ObjLinkingLayer;
   llvm::orc::IRCompileLayer CompileLayer;
+  llvm::orc::IRTransformLayer OptimizeLayer;
 
   llvm::orc::JITDylib::GeneratorFunction createHostProcessResolver();
 

SimpleOptimizer.cpp

@@ -0,0 +1,32 @@
+#include "SimpleOptimizer.h"
+
+#include <llvm/IR/LegacyPassManager.h>
+#include <llvm/Support/Debug.h>
+#include <llvm/Support/raw_ostream.h>
+
+#define DEBUG_TYPE "jitfromscratch"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+Expected<ThreadSafeModule>
+SimpleOptimizer::operator()(ThreadSafeModule TSM,
+                            const MaterializationResponsibility &) {
+  Module &M = *TSM.getModule();
+
+  legacy::FunctionPassManager FPM(&M);
+  B.populateFunctionPassManager(FPM);
+
+  FPM.doInitialization();
+  for (Function &F : M)
+    FPM.run(F);
+  FPM.doFinalization();
+
+  legacy::PassManager MPM;
+  B.populateModulePassManager(MPM);
+  MPM.run(M);
+
+  LLVM_DEBUG(dbgs() << "Optimized IR module:\n\n" << M << "\n\n");
+
+  return std::move(TSM);
+}

SimpleOptimizer.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <llvm/ExecutionEngine/Orc/Core.h>
+#include <llvm/ExecutionEngine/Orc/ThreadSafeModule.h>
+#include <llvm/Support/Error.h>
+#include <llvm/Transforms/IPO/PassManagerBuilder.h>
+
+class SimpleOptimizer {
+public:
+  SimpleOptimizer(unsigned OptLevel) { B.OptLevel = OptLevel; }
+
+  llvm::Expected<llvm::orc::ThreadSafeModule>
+  operator()(llvm::orc::ThreadSafeModule TSM,
+             const llvm::orc::MaterializationResponsibility &);
+
+private:
+  llvm::PassManagerBuilder B;
+};

main.cpp

@@ -21,6 +21,26 @@
 
 using namespace llvm;
 
+static cl::opt<char>
+    OptLevel("O",
+        cl::desc("Optimization level. [-O0, -O1, -O2, or -O3] "
+                  "(default = '-O2')"),
+        cl::Prefix, cl::ZeroOrMore, cl::init(' '));
+
+Expected<unsigned> getOptLevel() {
+  switch (OptLevel) {
+    case '0': return 0;
+    case '1': return 1;
+    case ' ':
+    case '2': return 2;
+    case '3': return 3;
+    default:
+      return createStringError(inconvertibleErrorCode(),
+                               "Invalid optimization level: -O%c",
+                               (char)OptLevel);
+  }
+}
+
 Expected<std::string> codegenIR(Module &module, unsigned items) {
   LLVMContext &ctx = module.getContext();
   IRBuilder<> B(ctx);
@@ -131,7 +151,9 @@ int main(int argc, char **argv) {
   M->setDataLayout(Jit.getDataLayout());
 
   std::string JitedFnName = ExitOnErr(codegenIR(*M, arrayElements(x)));
-  ExitOnErr(Jit.submitModule(std::move(M), std::move(C)));
+  unsigned OptLevel = ExitOnErr(getOptLevel());
+
+  ExitOnErr(Jit.submitModule(std::move(M), std::move(C), OptLevel));
 
   // Request function; this compiles to machine code and links.
   auto integerDistances =