Search code examples
c++compilationllvmjit

LLVM: simple example of a just-in-time compilation


I'm learning LLVM and trying to compile a simple function:

int sum(int a, int b) {
    return a+b;
};

on the fly.

So here's the code I have so far:

#include <string>
#include <vector>
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Verifier.h"

using namespace llvm;

static LLVMContext &Context = getGlobalContext();
static std::unique_ptr<Module> MyModule = make_unique<Module>("my compiler", Context);

Function *createFunc(IRBuilder<> &Builder, std::string Name) {
    std::vector<Type*> Integers(2, Builder.getInt32Ty());
    auto *funcType = FunctionType::get(Builder.getInt32Ty(), Integers, false);
    auto *fooFunc = Function::Create(funcType, Function::ExternalLinkage, Name, MyModule.get());
    return fooFunc;
};

int main(int argc, char* argv[]) {
    static IRBuilder<> Builder(Context);
    auto *fooFunc = createFunc(Builder, "sum");
    auto *entry = BasicBlock::Create(Context, "entry", fooFunc);
    Builder.SetInsertPoint(entry);

    // Fill the function body
    auto args = fooFunc->arg_begin();
    Value *arg1 = &(*args);
    args = std::next(args);
    Value *arg2 = &(*args);
    auto *sum = Builder.CreateAdd(arg1, arg2, "tmp");
    Builder.CreateRet(sum);
    verifyFunction(*fooFunc);

    // TODO: compile and run it
    MyModule->dump();
    return 0;
}

This compiles and when I run it I get the expected output:

; ModuleID = 'my compiler'

define i32 @sum(i32, i32) {
entry:
  %tmp = add i32 %0, %1
  ret i32 %tmp
}

just like in the tutorial.

But now I want to compile this function and run it from C++. I'm looking for the easiest way to do something like that:

auto compiledStuff = ...;
auto compiledFn = (int (*)(int, int))compiledStuff;
auto result = compiledFn(3, 8);

I've been digging through the official Kaleidoscope tutorial but the JIT tutorial is really complicated and seems to focus on optimizations and laziness while I still can't figure out how to easily compile a module and call a function from it.

Any help?


Solution

  • So, I've digged through the KaleidoscopeJIT and retrieved the most important pieces. First of all note that I'm using llvm-4.0. I've had lots of issues by not realizing how really incompatible 4.0 and lower versions are.

    The code works with C++11. I'm using clang++-4.0 with following compilation flags:

    llvm-config-4.0 --cxxflags --ldflags --system-libs --libs core engine
    

    And now the entire code:

    #include <string>
    #include <vector>
    #include <iostream>
    #include "llvm/IR/LLVMContext.h"
    #include "llvm/IR/Module.h"
    #include "llvm/IR/IRBuilder.h"
    #include "llvm/IR/Verifier.h"
    #include "llvm/ADT/iterator_range.h"
    #include "llvm/ADT/STLExtras.h"
    #include "llvm/ExecutionEngine/ExecutionEngine.h"
    #include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
    #include "llvm/ExecutionEngine/RuntimeDyld.h"
    #include "llvm/ExecutionEngine/SectionMemoryManager.h"
    #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
    #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
    #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
    #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
    #include "llvm/ExecutionEngine/JITSymbol.h"
    #include "llvm/IR/DataLayout.h"
    #include "llvm/IR/Mangler.h"
    #include "llvm/Support/DynamicLibrary.h"
    #include "llvm/Support/raw_ostream.h"
    #include "llvm/Support/TargetSelect.h"
    
    using namespace llvm;
    typedef orc::ObjectLinkingLayer<> ObjLayerT;
    typedef orc::IRCompileLayer<ObjLayerT> CompileLayerT;
    
    static LLVMContext Context;
    static auto MyModule = make_unique<Module>("my compiler", Context);
    
    Function *createFunc(IRBuilder<> &Builder, std::string Name) {
        std::vector<Type*> Integers(2, Builder.getInt32Ty());
        auto *funcType = FunctionType::get(Builder.getInt32Ty(), Integers, false);
        auto *fooFunc = Function::Create(funcType, Function::ExternalLinkage, Name, MyModule.get());
        return fooFunc;
    };
    
    void updateBody(Function *fooFunc, IRBuilder<> &Builder) {
        auto *entry = BasicBlock::Create(Context, "entry", fooFunc);
        Builder.SetInsertPoint(entry);
        auto args = fooFunc->arg_begin();
        Value *arg1 = &(*args);
        args = std::next(args);
        Value *arg2 = &(*args);
        auto *sum = Builder.CreateAdd(arg1, arg2, "tmp");
        Builder.CreateRet(sum);
    };
    
    int main(int argc, char* argv[]) {
        // Prepare the module
        static IRBuilder<> Builder(Context);
        auto *fooFunc = createFunc(Builder, "sum");
        updateBody(fooFunc, Builder);
        verifyFunction(*fooFunc);
    
        // Initilaze native target
        InitializeNativeTarget();
        InitializeNativeTargetAsmPrinter();
        InitializeNativeTargetAsmParser();
    
        // Prepare jit layer
        ObjLayerT ObjectLayer;
        std::unique_ptr<TargetMachine> TM(EngineBuilder().selectTarget());
        DataLayout DL(TM->createDataLayout());
        CompileLayerT CompileLayer(ObjectLayer, orc::SimpleCompiler(*TM));
        auto Resolver = orc::createLambdaResolver(
            [&](const std::string &Name) {
                if (auto Sym = CompileLayer.findSymbol(Name, false))
                    return Sym;
                return JITSymbol(nullptr);
            },
            [](const std::string &S) { return nullptr; }
        );
    
        // Add MyModule to the jit layer
        std::vector<std::unique_ptr<Module>> Modules;
        Modules.push_back(std::move(MyModule));
        CompileLayer.addModuleSet(
            std::move(Modules),
            make_unique<SectionMemoryManager>(),
            std::move(Resolver)
        );
    
        // Retrieve the foo symbol
        std::string MangledName;
        raw_string_ostream MangledNameStream(MangledName);
        Mangler::getNameWithPrefix(MangledNameStream, "sum", DL);
        auto Sym = CompileLayer.findSymbol(MangledNameStream.str(), true);
    
        // Cast to function
        auto func = (int(*)(int, int))Sym.getAddress();
    
        // Try it
        std::cout << func(5, 7) << std::endl;
    
        return 0;
    }
    

    I'm not sure if all includes are needed but anyway it works like a charm. Although I'm looking forward for any comment on how to improve it. :)