A simple LLVM 12 JIT test with optimizations working

gistfile1.txt

orc: jit.cpp
	clang++ -g -o $@ $^ $(shell llvm-config-12 --cxxflags --ldflags --system-libs --libs core) 

jit.cpp

#include <iostream>
#define LLVM_ENABLE_DUMP

#include "llvm/ADT/Triple.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Analysis/CallGraphSCCPass.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/RegionPass.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/AsmParser/Parser.h"
#include "llvm/CodeGen/CommandFlags.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/Config/llvm-config.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LLVMRemarkStreamer.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/LegacyPassNameParser.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Verifier.h"
#include "llvm/IRReader/IRReader.h"
#include "llvm/InitializePasses.h"
#include "llvm/LinkAllIR.h"
#include "llvm/LinkAllPasses.h"
#include "llvm/Linker/Linker.h"
#include "llvm/MC/SubtargetFeature.h"
#include "llvm/Remarks/HotnessThresholdParser.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Host.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/PluginLoader.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/SystemUtils.h"
#include "llvm/Support/TargetRegistry.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/YAMLTraits.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Transforms/Coroutines.h"
#include "llvm/Transforms/IPO/AlwaysInliner.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "llvm/Transforms/IPO/WholeProgramDevirt.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/Debugify.h"
#include "llvm/ExecutionEngine/JITSymbol.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
#include "llvm/ExecutionEngine/Orc/Core.h"
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassManager.h"

#include "llvm/Transforms/InstCombine/InstCombine.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Scalar/GVN.h"
#include "llvm/Transforms/Scalar/DCE.h"
#include "llvm/Transforms/IPO.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Verifier.h"
#include "llvm/ExecutionEngine/ExecutionEngine.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
#include "llvm/Support/TargetSelect.h"

// Optimizations
#include "llvm/Transforms/Scalar.h"
#include "llvm/Analysis/BasicAliasAnalysis.h"

using namespace llvm;


// https://gist.github.com/ofan/721464

void addOptPasses(
  llvm::legacy::PassManagerBase &passes,
  llvm::legacy::FunctionPassManager &fnPasses,
  llvm::TargetMachine *machine
) {
  llvm::PassManagerBuilder builder;
  builder.OptLevel = 0;
  builder.SizeLevel = 0;
  //builder.Inliner = llvm::createFunctionInliningPass(3, 0, false);
  //builder.LoopVectorize = true;
  //builder.SLPVectorize = true;
  machine->adjustPassManager(builder);

  builder.populateFunctionPassManager(fnPasses);
  builder.populateModulePassManager(passes);
}

void addLinkPasses(llvm::legacy::PassManagerBase &passes) {
  llvm::PassManagerBuilder builder;
  builder.VerifyInput = true;
  //builder.Inliner = llvm::createFunctionInliningPass(3, 0, false);
  //builder.populateLTOPassManager(passes);
}

void optimizeModule(llvm::TargetMachine *machine, llvm::Module *module) {
  module->setTargetTriple(machine->getTargetTriple().str());
  module->setDataLayout(machine->createDataLayout());

  llvm::legacy::PassManager passes;
  passes.add(new llvm::TargetLibraryInfoWrapperPass(machine->getTargetTriple()));
  passes.add(llvm::createTargetTransformInfoWrapperPass(machine->getTargetIRAnalysis()));

  llvm::legacy::FunctionPassManager fnPasses(module);
  fnPasses.add(llvm::createTargetTransformInfoWrapperPass(machine->getTargetIRAnalysis()));
  //fnPasses.add(llvm::createVerifierPass());

  addOptPasses(passes, fnPasses, machine);
  addLinkPasses(passes);

  fnPasses.doInitialization();
  for (llvm::Function &func : *module) {
    fnPasses.run(func);
  }
  fnPasses.doFinalization();

  //passes.run(*module);
}

llvm::Function* createSumFunction(std::string name, Module* module, Function* fn = nullptr) {
    /* Builds the following function:
    
    int sum(int a, int b) {
        int sum1 = 1 + 1;
        int sum2 = sum1 + a;
        int result = sum2 + b;
        return result;
    }
    */

    LLVMContext &context = module->getContext();
    IRBuilder<> builder(context);

    // Define function's signature
    std::vector<Type*> Integers(2, builder.getInt32Ty());
    auto *funcType = FunctionType::get(builder.getInt32Ty(), Integers, false);

    // create the function "sum" and bind it to the module with ExternalLinkage,
    // so we can retrieve it later
    auto *fooFunc = Function::Create(
        funcType, Function::ExternalLinkage, name, module
    );

    // Define the entry block and fill it with an appropriate code
    auto *entry = BasicBlock::Create(context, "entry", fooFunc);
    builder.SetInsertPoint(entry);


    // Retrieve arguments and proceed with further adding...
    auto this_args = fooFunc->arg_begin();
    Value *arg1 = &(*this_args);
    this_args = std::next(this_args);
    Value *arg2 = &(*this_args);

    if(fn) {

        Function::Create(FunctionType::get(builder.getInt32Ty(), Integers, false), 
                         Function::ExternalLinkage, fn->getName(), module);
        auto *added = builder.CreateAdd(arg1, arg2, "result");  
        
        std::vector<Value*> fn_args;
        fn_args.push_back(added);
        fn_args.push_back(added);
        auto *result = builder.CreateCall(fn, fn_args, "calltmp");
        builder.CreateRet(result);
    } else {

        auto *result = builder.CreateAdd(arg1, arg2, "result");  
        
        // ...and return
        builder.CreateRet(result);

    }


    // Verify at the end
    verifyFunction(*fooFunc);
    return fooFunc;
};

int main(int argc, char* argv[]) {

    // Initilaze native target
    llvm::TargetOptions Opts;
    InitializeNativeTarget();
    InitializeNativeTargetAsmPrinter();


    LLVMContext context;

    // Module 1
    auto _M = std::make_unique<Module>("My First JIT", context);
    auto *M = _M.get();

    auto* func = createSumFunction("sum", M);  // create function
    func->addFnAttr(llvm::Attribute::AttrKind::AlwaysInline);


    // Module 2
    auto _M2 = std::make_unique<Module>("My First JIT", context);
    auto* M2 = _M2.get();
    auto* func2 = createSumFunction("module2", M2, func);  // create function

    // Module 3
    auto _M3 = std::make_unique<Module>("My First JIT", context);
    auto* M3 = _M3.get();
    auto* func3 = createSumFunction("module3", M3, func2);  // create function


    // JIT Run
    std::unique_ptr<llvm::RTDyldMemoryManager> MemMgr(new llvm::SectionMemoryManager());

    // Create JIT engine
    llvm::EngineBuilder factory(std::move(_M));

    factory.setEngineKind(llvm::EngineKind::JIT);
    factory.setOptLevel(CodeGenOpt::Aggressive);
    factory.setTargetOptions(Opts);
    factory.setMCJITMemoryManager(std::move(MemMgr));

    auto EE = std::unique_ptr<llvm::ExecutionEngine>(factory.create());
    M->setDataLayout(EE->getDataLayout());
    M2->setDataLayout(EE->getDataLayout());
    M3->setDataLayout(EE->getDataLayout());
    //EE->addGlobalMapping(func, (void* )EE->getPointerToNamedFunction("sum"));
    EE->addModule(std::move(_M2));
    EE->addModule(std::move(_M3));



    //Finalize and run
    /*llvm::Linker L(*M);
    L.linkInModule(std::move(_M2));
    L.linkInModule(std::move(_M3));*/
    
    //M2->print(llvm::errs(), nullptr);

    optimizeModule(EE->getTargetMachine(), M);
    EE->finalizeObject();



    // Get raw pointer
    std::cout << "Finished optimizeation run." << "\n";
    auto* raw_ptr = EE->getPointerToNamedFunction("module3");
    if(raw_ptr == nullptr) {
        std::cout << "pointer to module2 not found" << std::endl;
        return -1;
    }
    auto* func_ptr = (int(*)(int, int))raw_ptr;

    // Execute
    int arg1 = 1;
    int arg2 = 1;
    M->print(llvm::errs(), nullptr);
    M2->print(llvm::errs(), nullptr);
    M3->print(llvm::errs(), nullptr);
    int result = func_ptr(arg1, arg2);
    std::cout << arg1 << " Result: " << result << "\n";

    return 0;
}