Break free from LLJIT

Author: weliveindetail

JitFromScratch.cpp

@@ -1,36 +1,39 @@
 #include "JitFromScratch.h"
 
+#include <llvm/ExecutionEngine/Orc/CompileUtils.h>
+#include <llvm/ExecutionEngine/Orc/ExecutionUtils.h>
 #include <llvm/ExecutionEngine/Orc/ThreadSafeModule.h>
+#include <llvm/ExecutionEngine/SectionMemoryManager.h>
+#include <llvm/IR/Mangler.h>
 #include <llvm/Support/Debug.h>
+#include <llvm/Support/raw_ostream.h>
 
 #define DEBUG_TYPE "jitfromscratch"
 
 using namespace llvm;
 using namespace llvm::orc;
 
-JitFromScratch::JitFromScratch(ExitOnError ExitOnErr) {
-  LLJITBuilder Builder;
-  ExitOnErr(Builder.prepareForConstruction());
-  TT = Builder.JTMB->getTargetTriple();
-  LLJIT = ExitOnErr(Builder.create());
-
-  if (auto R = createHostProcessResolver(LLJIT->getDataLayout()))
-    LLJIT->getMainJITDylib().setGenerator(std::move(R));
+JitFromScratch::JitFromScratch(ExitOnError ExitOnErr)
+    : ES(std::make_unique<ExecutionSession>()),
+      TM(createTargetMachine(ExitOnErr)),
+      ObjLinkingLayer(*ES, createMemoryManagerFtor()),
+      CompileLayer(*ES, ObjLinkingLayer, SimpleCompiler(*TM)) {
+  if (auto R = createHostProcessResolver())
+    ES->getMainJITDylib().setGenerator(std::move(R));
 }
 
-JITDylib::GeneratorFunction
-JitFromScratch::createHostProcessResolver(DataLayout DL) {
-  char Prefix = DL.getGlobalPrefix();
+JITDylib::GeneratorFunction JitFromScratch::createHostProcessResolver() {
+  char Prefix = TM->createDataLayout().getGlobalPrefix();
   Expected<JITDylib::GeneratorFunction> R =
       DynamicLibrarySearchGenerator::GetForCurrentProcess(Prefix);
 
   if (!R) {
-    LLJIT->getExecutionSession().reportError(R.takeError());
+    ES->reportError(R.takeError());
     return nullptr;
   }
 
   if (!*R) {
-    LLJIT->getExecutionSession().reportError(createStringError(
+    ES->reportError(createStringError(
         inconvertibleErrorCode(),
         "Generator function for host process symbols must not be null"));
     return nullptr;
@@ -39,14 +42,61 @@ JitFromScratch::createHostProcessResolver(DataLayout DL) {
   return *R;
 }
 
+std::unique_ptr<TargetMachine>
+JitFromScratch::createTargetMachine(ExitOnError ExitOnErr) {
+  auto JTMB = ExitOnErr(JITTargetMachineBuilder::detectHost());
+  return ExitOnErr(JTMB.createTargetMachine());
+}
+
+using GetMemoryManagerFunction =
+    RTDyldObjectLinkingLayer::GetMemoryManagerFunction;
+
+GetMemoryManagerFunction JitFromScratch::createMemoryManagerFtor() {
+  return []() -> GetMemoryManagerFunction::result_type {
+    return std::make_unique<SectionMemoryManager>();
+  };
+}
+
+std::string JitFromScratch::mangle(StringRef UnmangledName) {
+  std::string MangledName;
+  {
+    DataLayout DL = TM->createDataLayout();
+    raw_string_ostream MangledNameStream(MangledName);
+    Mangler::getNameWithPrefix(MangledNameStream, UnmangledName, DL);
+  }
+  return MangledName;
+}
+
+Error JitFromScratch::applyDataLayout(Module &M) {
+  DataLayout DL = TM->createDataLayout();
+  if (M.getDataLayout().isDefault())
+    M.setDataLayout(DL);
+
+  if (M.getDataLayout() != DL)
+    return make_error<StringError>(
+        "Added modules have incompatible data layouts",
+        inconvertibleErrorCode());
+
+  return Error::success();
+}
+
 Error JitFromScratch::submitModule(std::unique_ptr<Module> M,
                                    std::unique_ptr<LLVMContext> C) {
   LLVM_DEBUG(dbgs() << "Submit IR module:\n\n" << *M << "\n\n");
-  return LLJIT->addIRModule(ThreadSafeModule(std::move(M), std::move(C)));
+
+  if (auto Err = applyDataLayout(*M))
+    return Err;
+
+  return CompileLayer.add(ES->getMainJITDylib(),
+                          ThreadSafeModule(std::move(M), std::move(C)),
+                          ES->allocateVModule());
 }
 
 Expected<JITTargetAddress> JitFromScratch::getFunctionAddr(StringRef Name) {
-  Expected<JITEvaluatedSymbol> S = LLJIT->lookup(Name);
+  SymbolStringPtr NamePtr = ES->intern(mangle(Name));
+  JITDylibSearchList JDs{{&ES->getMainJITDylib(), true}};
+
+  Expected<JITEvaluatedSymbol> S = ES->lookup(JDs, NamePtr);
   if (!S)
     return S.takeError();
 

JitFromScratch.h

@@ -1,16 +1,20 @@
 #pragma once
 
 #include <llvm/ADT/StringRef.h>
-#include <llvm/ADT/Triple.h>
 #include <llvm/ExecutionEngine/JITSymbol.h>
-#include <llvm/ExecutionEngine/Orc/LLJIT.h>
+#include <llvm/ExecutionEngine/Orc/Core.h>
+#include <llvm/ExecutionEngine/Orc/IRCompileLayer.h>
+#include <llvm/ExecutionEngine/Orc/IRTransformLayer.h>
+#include <llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h>
 #include <llvm/IR/DataLayout.h>
 #include <llvm/IR/LLVMContext.h>
 #include <llvm/IR/Module.h>
 #include <llvm/Support/Error.h>
+#include <llvm/Target/TargetMachine.h>
 
 #include <functional>
 #include <memory>
+#include <string>
 
 class JitFromScratch {
 public:
@@ -23,11 +27,11 @@ class JitFromScratch {
   JitFromScratch &operator=(JitFromScratch &&) = delete;
 
   llvm::DataLayout getDataLayout() const {
-    return LLJIT->getDataLayout();
+    return TM->createDataLayout();
   }
 
   const llvm::Triple &getTargetTriple() const {
-    return TT;
+    return TM->getTargetTriple();
   }
 
   llvm::Error submitModule(std::unique_ptr<llvm::Module> M,
@@ -43,11 +47,22 @@ class JitFromScratch {
   }
 
 private:
-  std::unique_ptr<llvm::orc::LLJIT> LLJIT;
-  llvm::Triple TT;
+  std::unique_ptr<llvm::orc::ExecutionSession> ES;
+  std::unique_ptr<llvm::TargetMachine> TM;
 
-  llvm::orc::JITDylib::GeneratorFunction
-  createHostProcessResolver(llvm::DataLayout DL);
+  llvm::orc::RTDyldObjectLinkingLayer ObjLinkingLayer;
+  llvm::orc::IRCompileLayer CompileLayer;
+
+  llvm::orc::JITDylib::GeneratorFunction createHostProcessResolver();
+
+  std::unique_ptr<llvm::TargetMachine>
+  createTargetMachine(llvm::ExitOnError ExitOnErr);
+
+  llvm::orc::RTDyldObjectLinkingLayer::GetMemoryManagerFunction
+  createMemoryManagerFtor();
+
+  std::string mangle(llvm::StringRef UnmangledName);
+  llvm::Error applyDataLayout(llvm::Module &M);
 
   llvm::Expected<llvm::JITTargetAddress> getFunctionAddr(llvm::StringRef Name);
 };