[llvm-exegesis] Switch from MCJIT to LLJIT

[boomanaiden154]

This patch switches from using MCJIT to LLJIT as MCJIT is going to be deprecated soon.

Fixes #41764.

[llvmbot]

@llvm/pr-subscribers-tools-llvm-exegesis

Author: Aiden Grossman (boomanaiden154)

Changes

This patch switches from using MCJIT to LLJIT as MCJIT is going to be deprecated soon.

---

Full diff: https://github.com/llvm/llvm-project/pull/72838.diff

3 Files Affected:

• (modified) llvm/tools/llvm-exegesis/lib/Assembler.cpp (+28-49)
• (modified) llvm/tools/llvm-exegesis/lib/Assembler.h (+3-4)
• (modified) llvm/tools/llvm-exegesis/lib/CMakeLists.txt (+1-1)

diff --git a/llvm/tools/llvm-exegesis/lib/Assembler.cpp b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
index 9ff33258e965f7e..54ba29ed04c6f95 100644
--- a/llvm/tools/llvm-exegesis/lib/Assembler.cpp
+++ b/llvm/tools/llvm-exegesis/lib/Assembler.cpp
@@ -21,11 +21,12 @@
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
-#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/MC/MCInstrInfo.h"
+#include "llvm/Object/SymbolSize.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/raw_ostream.h"
@@ -105,7 +106,7 @@ MachineFunction &createVoidVoidPtrMachineFunction(StringRef FunctionName,
   FunctionType *FunctionType =
       FunctionType::get(ReturnType, {MemParamType}, false);
   Function *const F = Function::Create(
-      FunctionType, GlobalValue::InternalLinkage, FunctionName, Module);
+      FunctionType, GlobalValue::ExternalLinkage, FunctionName, Module);
   BasicBlock *BB = BasicBlock::Create(Module->getContext(), "", F);
   new UnreachableInst(Module->getContext(), BB);
   return MMI->getOrCreateMachineFunction(*F);
@@ -324,66 +325,44 @@ object::OwningBinary<object::ObjectFile> getObjectFromFile(StringRef Filename) {
   return cantFail(object::ObjectFile::createObjectFile(Filename));
 }
 
-namespace {
-
-// Implementation of this class relies on the fact that a single object with a
-// single function will be loaded into memory.
-class TrackingSectionMemoryManager : public SectionMemoryManager {
-public:
-  explicit TrackingSectionMemoryManager(uintptr_t *CodeSize)
-      : CodeSize(CodeSize) {}
-
-  uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment,
-                               unsigned SectionID,
-                               StringRef SectionName) override {
-    *CodeSize = Size;
-    return SectionMemoryManager::allocateCodeSection(Size, Alignment, SectionID,
-                                                     SectionName);
-  }
-
-private:
-  uintptr_t *const CodeSize = nullptr;
-};
-
-} // namespace
-
 Expected<ExecutableFunction> ExecutableFunction::create(
     std::unique_ptr<LLVMTargetMachine> TM,
     object::OwningBinary<object::ObjectFile> &&ObjectFileHolder) {
   assert(ObjectFileHolder.getBinary() && "cannot create object file");
   std::unique_ptr<LLVMContext> Ctx = std::make_unique<LLVMContext>();
-  // Initializing the execution engine.
-  // We need to use the JIT EngineKind to be able to add an object file.
-  LLVMLinkInMCJIT();
-  uintptr_t CodeSize = 0;
-  std::string Error;
-  std::unique_ptr<ExecutionEngine> EE(
-      EngineBuilder(createModule(Ctx, TM->createDataLayout()))
-          .setErrorStr(&Error)
-          .setMCPU(TM->getTargetCPU())
-          .setEngineKind(EngineKind::JIT)
-          .setMCJITMemoryManager(
-              std::make_unique<TrackingSectionMemoryManager>(&CodeSize))
-          .create(TM.release()));
-  if (!EE)
-    return make_error<StringError>(Twine(Error), inconvertibleErrorCode());
-  // Adding the generated object file containing the assembled function.
-  // The ExecutionEngine makes sure the object file is copied into an
-  // executable page.
-  EE->addObjectFile(std::move(ObjectFileHolder));
-  // Fetching function bytes.
-  const uint64_t FunctionAddress = EE->getFunctionAddress(FunctionID);
+
+  auto SymbolSizes = object::computeSymbolSizes(*ObjectFileHolder.getBinary());
+  assert(SymbolSizes.size() == 3);
+  uintptr_t CodeSize = std::get<1>(SymbolSizes[2]);
+
+  auto EJITOrErr = orc::LLJITBuilder().create();
+  if (!EJITOrErr)
+    return EJITOrErr.takeError();
+
+  auto EJIT = std::move(*EJITOrErr);
+
+  if (auto ObjErr =
+          EJIT->addObjectFile(std::get<1>(ObjectFileHolder.takeBinary())))
+    return ObjErr;
+
+  auto FunctionAddressOrErr = EJIT->lookup(FunctionID);
+  if (!FunctionAddressOrErr)
+    return FunctionAddressOrErr.takeError();
+
+  const uint64_t FunctionAddress = FunctionAddressOrErr->getValue();
+
   assert(isAligned(kFunctionAlignment, FunctionAddress) &&
          "function is not properly aligned");
+
   StringRef FBytes =
       StringRef(reinterpret_cast<const char *>(FunctionAddress), CodeSize);
-  return ExecutableFunction(std::move(Ctx), std::move(EE), FBytes);
+  return ExecutableFunction(std::move(Ctx), std::move(EJIT), FBytes);
 }
 
 ExecutableFunction::ExecutableFunction(std::unique_ptr<LLVMContext> Ctx,
-                                       std::unique_ptr<ExecutionEngine> EE,
+                                       std::unique_ptr<orc::LLJIT> EJIT,
                                        StringRef FB)
-    : FunctionBytes(FB), Context(std::move(Ctx)), ExecEngine(std::move(EE)) {}
+    : FunctionBytes(FB), Context(std::move(Ctx)), ExecJIT(std::move(EJIT)) {}
 
 Error getBenchmarkFunctionBytes(const StringRef InputData,
                                 std::vector<uint8_t> &Bytes) {
diff --git a/llvm/tools/llvm-exegesis/lib/Assembler.h b/llvm/tools/llvm-exegesis/lib/Assembler.h
index 5f1bf8cdfb7ad6c..abc5aa7be8cfeef 100644
--- a/llvm/tools/llvm-exegesis/lib/Assembler.h
+++ b/llvm/tools/llvm-exegesis/lib/Assembler.h
@@ -23,7 +23,7 @@
 #include "llvm/ADT/BitVector.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
-#include "llvm/ExecutionEngine/ExecutionEngine.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/MC/MCInst.h"
@@ -124,11 +124,10 @@ class ExecutableFunction {
 
 private:
   ExecutableFunction(std::unique_ptr<LLVMContext> Ctx,
-                     std::unique_ptr<ExecutionEngine> EE,
-                     StringRef FunctionBytes);
+                     std::unique_ptr<orc::LLJIT> EJIT, StringRef FunctionBytes);
 
   std::unique_ptr<LLVMContext> Context;
-  std::unique_ptr<ExecutionEngine> ExecEngine;
+  std::unique_ptr<orc::LLJIT> ExecJIT;
 };
 
 // Copies benchmark function's bytes from benchmark object.
diff --git a/llvm/tools/llvm-exegesis/lib/CMakeLists.txt b/llvm/tools/llvm-exegesis/lib/CMakeLists.txt
index 92135e6aaf673be..7312811989dd800 100644
--- a/llvm/tools/llvm-exegesis/lib/CMakeLists.txt
+++ b/llvm/tools/llvm-exegesis/lib/CMakeLists.txt
@@ -29,10 +29,10 @@ set(LLVM_LINK_COMPONENTS
   MC
   MCA
   MCDisassembler
-  MCJIT
   MCParser
   Object
   ObjectYAML
+  OrcJIT
   RuntimeDyld
   Support
   TargetParser

[boomanaiden154]

@legrosbuffle can you review this when you get a chance? It's a stacked patch against #72837 and should be relatively mechanical maintenance.

[zmodem]

This broke lit tests on Mac, see https://green.lab.llvm.org/green/job/clang-stage1-cmake-RA-incremental/38602/

[boomanaiden154]

This broke lit tests on Mac, see https://green.lab.llvm.org/green/job/clang-stage1-cmake-RA-incremental/38602/

Thanks for the revert. That test shouldn't be running on Darwin or Windows, so I've specified it requires Linux, which should fix the issue. I'll keep an eye on that bot and make sure everything goes through.

[lhames]

This patch switches from using MCJIT to LLJIT as MCJIT is going to be deprecated soon.

Late to the party here, but thanks very much for this -- it's great to see LLVM tools moving over to ORC! :)