jit.hpp

#pragma once

#include "llvm/ExecutionEngine/ExecutionEngine.h"
#include "llvm/ExecutionEngine/RTDyldMemoryManager.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/IR/Mangler.h"
#include "llvm/Support/DynamicLibrary.h"

class JIT
{
public:
    // This sample doesn't implement on-request or lazy compilation. It therefore
    // uses Orc's eager compilation layer directly - IRCompileLayer. It also uses
    // the basis object layer - ObjectLinkingLayer - directly.
    // Orc's SimpleCompiler is used to actually compile the module; it runs LLVM's
    // codegen and MC on the module, producing an object file in memory. No
    // IR-level optimizations are run by the JIT.
    typedef llvm::orc::ObjectLinkingLayer<> ObjLayerT;
    typedef llvm::orc::IRCompileLayer<ObjLayerT> CompileLayerT;
    typedef CompileLayerT::ModuleSetHandleT ModuleHandleT;

    JIT() : TM(llvm::EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
        CompileLayer(ObjectLayer, llvm::orc::SimpleCompiler(*TM))
    {
    }

    // A simple SymbolResolver that doesn't support linking by always returning
    // nullptr.
    struct NoLinkingResolver : public llvm::RuntimeDyld::SymbolResolver
    {
        llvm::RuntimeDyld::SymbolInfo findSymbol(const std::string &Name) override
        {
            return nullptr;
        }

        llvm::RuntimeDyld::SymbolInfo findSymbolInLogicalDylib(const std::string &Name) override
        {
            return nullptr;
        }
    };

    llvm::TargetMachine& getTargetMachine()
    {
        return *TM;
    }

    // Add a module to the JIT.
    ModuleHandleT addModule(std::unique_ptr<llvm::Module> M)
    {
        std::vector<std::unique_ptr<llvm::Module>> MS;
        MS.push_back(std::move(M));
        auto H = CompileLayer.addModuleSet(std::move(MS),
                                           llvm::make_unique<llvm::SectionMemoryManager>(),
                                           llvm::make_unique<NoLinkingResolver>());
        ModuleHandles.push_back(H);
        return H;
    }

    // Remove a module from the JIT.
    void removeModule(ModuleHandleT H)
    {
        ModuleHandles.erase(
            std::find(ModuleHandles.begin(), ModuleHandles.end(), H));
        CompileLayer.removeModuleSet(H);
    }

    // Get the runtime address of the compiled symbol whose name is given.
    llvm::orc::JITSymbol findSymbol(const std::string Name)
    {
        std::string MangledName;
        {
            llvm::raw_string_ostream MangledNameStream(MangledName);
            llvm::Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
        }

        for (auto H : llvm::make_range(ModuleHandles.rbegin(), ModuleHandles.rend()))
        {
            if (auto Sym = CompileLayer.findSymbolIn(H, MangledName, true))
            {
                return Sym;
            }
        }

        return nullptr;
    }

private:
    std::unique_ptr<llvm::TargetMachine> TM;
    const llvm::DataLayout DL;
    ObjLayerT ObjectLayer;
    CompileLayerT CompileLayer;
    std::vector<ModuleHandleT> ModuleHandles;
};

main.cpp

#include <iostream>

#include <llvm/ADT/Triple.h>
#include <llvm/IR/DataLayout.h>
#include <llvm/IR/DerivedTypes.h>
#include <llvm/IR/DIBuilder.h>
#include <llvm/IR/IRBuilder.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Verifier.h>
#include <llvm/Support/TargetSelect.h>
#include <llvm/Transforms/Scalar.h>
#include <llvm/Support/Host.h>

#include "jit.hpp"

std::string GetMCJITProcessTriple()
{
    auto triple = llvm::Triple(llvm::sys::getProcessTriple());
#ifdef _MSC_VER
    //triple.setOS(llvm::Triple::OSType::MinGW32);
    triple.setEnvironment(llvm::Triple::EnvironmentType::Itanium);
#endif
    auto stringtrip = triple.getTriple();
#ifdef _MSC_VER
    // MCJIT can't handle non-ELF on Windows for some reason.
    stringtrip += "-elf";
#endif
    return stringtrip;
}

int main(int, char**)
{
    llvm::InitializeNativeTarget();
    llvm::InitializeNativeTargetAsmPrinter();
    llvm::InitializeNativeTargetAsmParser();
    llvm::LLVMContext& context = llvm::getGlobalContext();

    JIT jit;

    llvm::Module* module = new llvm::Module("top", context);
    module->setTargetTriple(GetMCJITProcessTriple());
    module->setDataLayout(jit.getTargetMachine().createDataLayout());
    module->addModuleFlag(llvm::Module::Warning, "Debug Info Version", llvm::DEBUG_METADATA_VERSION);

    llvm::IRBuilder<> builder(context);

    // Create main function
    {
        llvm::FunctionType* funcType = llvm::FunctionType::get(builder.getVoidTy(), false);
        llvm::Function* mainFunc = llvm::Function::Create(funcType, llvm::Function::ExternalLinkage, "main", module);

        llvm::BasicBlock* entry = llvm::BasicBlock::Create(context, "entrypoint", mainFunc);
        builder.SetInsertPoint(entry);

        llvm::Value* helloWorld = builder.CreateGlobalStringPtr("hello world!\n");

        std::vector<llvm::Type*> putsArgs;
        putsArgs.push_back(builder.getInt8Ty()->getPointerTo());
        llvm::ArrayRef<llvm::Type*> argsRef(putsArgs);

        llvm::FunctionType* putsType = llvm::FunctionType::get(builder.getInt32Ty(), argsRef, false);
        llvm::Constant* putsFunc = module->getOrInsertFunction("puts", putsType);

        builder.CreateCall(putsFunc, helloWorld);
        builder.CreateRetVoid();

        llvm::verifyFunction(*mainFunc);
    }

    module->dump();

    jit.addModule(std::unique_ptr<llvm::Module>(module));

    typedef void(*PFN)();
    PFN pfn = reinterpret_cast<PFN>(jit.findSymbol("main").getAddress());
    pfn();

    std::cin.get();
    return 0;
}