[ORC] Add custom IR compiler configuration to LLJITBuilder to enable obj caches.
authorLang Hames <lhames@gmail.com>
Wed, 10 Jul 2019 17:24:24 +0000 (17:24 +0000)
committerLang Hames <lhames@gmail.com>
Wed, 10 Jul 2019 17:24:24 +0000 (17:24 +0000)
LLJITBuilder now has a setCompileFunctionCreator method which can be used to
construct a CompileFunction for the LLJIT instance being created. The motivating
use-case for this is supporting ObjectCaches, which can now be set up at
compile-function construction time. To demonstrate this an example project,
LLJITWithObjectCache, is included.

llvm-svn: 365671

llvm/examples/CMakeLists.txt
llvm/examples/HowToUseLLJIT/HowToUseLLJIT.cpp
llvm/examples/LLJITExamples/CMakeLists.txt [new file with mode: 0644]
llvm/examples/LLJITExamples/ExampleModules.h [new file with mode: 0644]
llvm/examples/LLJITExamples/LLJITWithObjectCache/CMakeLists.txt [new file with mode: 0644]
llvm/examples/LLJITExamples/LLJITWithObjectCache/LLJITWithObjectCache.cpp [new file with mode: 0644]
llvm/include/llvm/ExecutionEngine/Orc/CompileUtils.h
llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

index f8d4ee9..939bac1 100644 (file)
@@ -2,6 +2,7 @@ add_subdirectory(BrainF)
 add_subdirectory(Fibonacci)
 add_subdirectory(HowToUseJIT)
 add_subdirectory(HowToUseLLJIT)
+add_subdirectory(LLJITExamples)
 add_subdirectory(Kaleidoscope)
 add_subdirectory(ModuleMaker)
 
index 372643d..0b77ff0 100644 (file)
@@ -1,3 +1,11 @@
+//===-- examples/HowToUseJIT/HowToUseJIT.cpp - An example use of the JIT --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
diff --git a/llvm/examples/LLJITExamples/CMakeLists.txt b/llvm/examples/LLJITExamples/CMakeLists.txt
new file mode 100644 (file)
index 0000000..2ca888e
--- /dev/null
@@ -0,0 +1 @@
+add_subdirectory(LLJITWithObjectCache)
diff --git a/llvm/examples/LLJITExamples/ExampleModules.h b/llvm/examples/LLJITExamples/ExampleModules.h
new file mode 100644 (file)
index 0000000..aa6b2b9
--- /dev/null
@@ -0,0 +1,54 @@
+//===----- ExampleModules.h - IR modules for LLJIT examples -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Example modules for LLJIT examples
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXAMPLES_HOWTOUSELLJIT_EXAMPLEMODULES_H
+#define LLVM_EXAMPLES_HOWTOUSELLJIT_EXAMPLEMODULES_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/SourceMgr.h"
+
+const llvm::StringRef Add1Example =
+    R"(
+  define i32 @add1(i32 %x) {
+  entry:
+    %r = add nsw i32 %x, 1
+    ret i32 %r
+  }
+)";
+
+inline llvm::Expected<llvm::orc::ThreadSafeModule>
+parseExampleModule(llvm::StringRef Source, llvm::StringRef Name) {
+  using namespace llvm;
+  using namespace llvm::orc;
+
+  auto Ctx = llvm::make_unique<LLVMContext>();
+  SMDiagnostic Err;
+  auto M = parseIR(MemoryBufferRef(Source, Name), Err, *Ctx);
+
+  if (!M) {
+    std::string ErrMsg;
+    {
+      raw_string_ostream ErrStream(ErrMsg);
+      Err.print("", ErrStream);
+    }
+    return make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode());
+  }
+
+  return ThreadSafeModule(std::move(M), std::move(Ctx));
+}
+
+#endif // LLVM_EXAMPLES_HOWTOUSELLJIT_EXAMPLEMODULES_H
diff --git a/llvm/examples/LLJITExamples/LLJITWithObjectCache/CMakeLists.txt b/llvm/examples/LLJITExamples/LLJITWithObjectCache/CMakeLists.txt
new file mode 100644 (file)
index 0000000..4c9680a
--- /dev/null
@@ -0,0 +1,11 @@
+set(LLVM_LINK_COMPONENTS
+  Core
+  IRReader
+  OrcJIT
+  Support
+  nativecodegen
+  )
+
+add_llvm_example(LLJITWithObjectCache
+  LLJITWithObjectCache.cpp
+  )
diff --git a/llvm/examples/LLJITExamples/LLJITWithObjectCache/LLJITWithObjectCache.cpp b/llvm/examples/LLJITExamples/LLJITWithObjectCache/LLJITWithObjectCache.cpp
new file mode 100644 (file)
index 0000000..377babb
--- /dev/null
@@ -0,0 +1,95 @@
+//===--- LLJITWithObjectCache.cpp - An LLJIT example with an ObjectCache --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ExecutionEngine/ObjectCache.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "../ExampleModules.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+ExitOnError ExitOnErr;
+
+class MyObjectCache : public ObjectCache {
+public:
+  void notifyObjectCompiled(const Module *M,
+                            MemoryBufferRef ObjBuffer) override {
+    CachedObjects[M->getModuleIdentifier()] = MemoryBuffer::getMemBufferCopy(
+        ObjBuffer.getBuffer(), ObjBuffer.getBufferIdentifier());
+  }
+
+  std::unique_ptr<MemoryBuffer> getObject(const Module *M) override {
+    auto I = CachedObjects.find(M->getModuleIdentifier());
+    if (I == CachedObjects.end()) {
+      dbgs() << "No object for " << M->getModuleIdentifier()
+             << " in cache. Compiling.\n";
+      return nullptr;
+    }
+
+    dbgs() << "Object for " << M->getModuleIdentifier()
+           << " loaded from cache.\n";
+    return MemoryBuffer::getMemBuffer(I->second->getMemBufferRef());
+  }
+
+private:
+  StringMap<std::unique_ptr<MemoryBuffer>> CachedObjects;
+};
+
+void runJITWithCache(ObjectCache &ObjCache) {
+
+  // Create an LLJIT instance with a custom CompileFunction.
+  auto J = ExitOnErr(
+      LLJITBuilder()
+          .setCompileFunctionCreator(
+              [&](JITTargetMachineBuilder JTMB)
+                  -> Expected<IRCompileLayer::CompileFunction> {
+                auto TM = JTMB.createTargetMachine();
+                if (!TM)
+                  return TM.takeError();
+                return IRCompileLayer::CompileFunction(
+                    TMOwningSimpleCompiler(std::move(*TM), &ObjCache));
+              })
+          .create());
+
+  auto M = ExitOnErr(parseExampleModule(Add1Example, "add1"));
+
+  ExitOnErr(J->addIRModule(std::move(M)));
+
+  // Look up the JIT'd function, cast it to a function pointer, then call it.
+  auto Add1Sym = ExitOnErr(J->lookup("add1"));
+  int (*Add1)(int) = (int (*)(int))Add1Sym.getAddress();
+
+  int Result = Add1(42);
+  outs() << "add1(42) = " << Result << "\n";
+}
+
+int main(int argc, char *argv[]) {
+  // Initialize LLVM.
+  InitLLVM X(argc, argv);
+
+  InitializeNativeTarget();
+  InitializeNativeTargetAsmPrinter();
+
+  cl::ParseCommandLineOptions(argc, argv, "HowToUseLLJIT");
+  ExitOnErr.setBanner(std::string(argv[0]) + ": ");
+
+  MyObjectCache MyCache;
+
+  runJITWithCache(MyCache);
+  runJITWithCache(MyCache);
+
+  return 0;
+}
index 1585925..eb6d84e 100644 (file)
@@ -18,7 +18,6 @@
 
 namespace llvm {
 
-class JITTargetMachineBuilder;
 class MCContext;
 class MemoryBuffer;
 class Module;
@@ -27,6 +26,8 @@ class TargetMachine;
 
 namespace orc {
 
+class JITTargetMachineBuilder;
+
 /// Simple compile functor: Takes a single IR module and returns an ObjectFile.
 /// This compiler supports a single compilation thread and LLVMContext only.
 /// For multithreaded compilation, use ConcurrentIRCompiler below.
@@ -52,6 +53,22 @@ private:
   ObjectCache *ObjCache = nullptr;
 };
 
+/// A SimpleCompiler that owns its TargetMachine.
+///
+/// This convenient for clients who don't want to own their TargetMachines,
+/// e.g. LLJIT.
+class TMOwningSimpleCompiler : public SimpleCompiler {
+public:
+  TMOwningSimpleCompiler(std::unique_ptr<TargetMachine> TM,
+                         ObjectCache *ObjCache = nullptr)
+      : SimpleCompiler(*TM, ObjCache), TM(std::move(TM)) {}
+
+private:
+  // FIXME: shared because std::functions (and consequently
+  // IRCompileLayer::CompileFunction) are not moveable.
+  std::shared_ptr<llvm::TargetMachine> TM;
+};
+
 /// A thread-safe version of SimpleCompiler.
 ///
 /// This class creates a new TargetMachine and SimpleCompiler instance for each
index b54c7d8..0aac191 100644 (file)
@@ -121,6 +121,9 @@ protected:
   static std::unique_ptr<ObjectLayer>
   createObjectLinkingLayer(LLJITBuilderState &S, ExecutionSession &ES);
 
+  static Expected<IRCompileLayer::CompileFunction>
+  createCompileFunction(LLJITBuilderState &S, JITTargetMachineBuilder JTMB);
+
   /// Create an LLJIT instance with a single compile thread.
   LLJIT(LLJITBuilderState &S, Error &Err);
 
@@ -181,12 +184,17 @@ private:
 
 class LLJITBuilderState {
 public:
-  using CreateObjectLinkingLayerFunction =
+  using ObjectLinkingLayerCreator =
       std::function<std::unique_ptr<ObjectLayer>(ExecutionSession &)>;
 
+  using CompileFunctionCreator =
+      std::function<Expected<IRCompileLayer::CompileFunction>(
+          JITTargetMachineBuilder JTMB)>;
+
   std::unique_ptr<ExecutionSession> ES;
   Optional<JITTargetMachineBuilder> JTMB;
-  CreateObjectLinkingLayerFunction CreateObjectLinkingLayer;
+  ObjectLinkingLayerCreator CreateObjectLinkingLayer;
+  CompileFunctionCreator CreateCompileFunction;
   unsigned NumCompileThreads = 0;
 
   /// Called prior to JIT class construcion to fix up defaults.
@@ -215,13 +223,24 @@ public:
   ///
   /// If this method is not called, a default creation function will be used
   /// that will construct an RTDyldObjectLinkingLayer.
-  SetterImpl &setCreateObjectLinkingLayer(
-      LLJITBuilderState::CreateObjectLinkingLayerFunction
-          CreateObjectLinkingLayer) {
+  SetterImpl &setObjectLinkingLayerCreator(
+      LLJITBuilderState::ObjectLinkingLayerCreator CreateObjectLinkingLayer) {
     impl().CreateObjectLinkingLayer = std::move(CreateObjectLinkingLayer);
     return impl();
   }
 
+  /// Set a CompileFunctionCreator.
+  ///
+  /// If this method is not called, a default creation function wil be used
+  /// that will construct a basic IR compile function that is compatible with
+  /// the selected number of threads (SimpleCompiler for '0' compile threads,
+  /// ConcurrentIRCompiler otherwise).
+  SetterImpl &setCompileFunctionCreator(
+      LLJITBuilderState::CompileFunctionCreator CreateCompileFunction) {
+    impl().CreateCompileFunction = std::move(CreateCompileFunction);
+    return impl();
+  }
+
   /// Set the number of compile threads to use.
   ///
   /// If set to zero, compilation will be performed on the execution thread when
index ae8a66f..b120691 100644 (file)
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/Mangler.h"
 
-namespace {
-
-  // A SimpleCompiler that owns its TargetMachine.
-  class TMOwningSimpleCompiler : public llvm::orc::SimpleCompiler {
-  public:
-    TMOwningSimpleCompiler(std::unique_ptr<llvm::TargetMachine> TM)
-      : llvm::orc::SimpleCompiler(*TM), TM(std::move(TM)) {}
-  private:
-    // FIXME: shared because std::functions (and thus
-    // IRCompileLayer::CompileFunction) are not moveable.
-    std::shared_ptr<llvm::TargetMachine> TM;
-  };
-
-} // end anonymous namespace
-
 namespace llvm {
 namespace orc {
 
@@ -86,6 +71,26 @@ LLJIT::createObjectLinkingLayer(LLJITBuilderState &S, ExecutionSession &ES) {
   return llvm::make_unique<RTDyldObjectLinkingLayer>(ES, std::move(GetMemMgr));
 }
 
+Expected<IRCompileLayer::CompileFunction>
+LLJIT::createCompileFunction(LLJITBuilderState &S,
+                             JITTargetMachineBuilder JTMB) {
+
+  /// If there is a custom compile function creator set then use it.
+  if (S.CreateCompileFunction)
+    return S.CreateCompileFunction(std::move(JTMB));
+
+  // Otherwise default to creating a SimpleCompiler, or ConcurrentIRCompiler,
+  // depending on the number of threads requested.
+  if (S.NumCompileThreads > 0)
+    return ConcurrentIRCompiler(std::move(JTMB));
+
+  auto TM = JTMB.createTargetMachine();
+  if (!TM)
+    return TM.takeError();
+
+  return TMOwningSimpleCompiler(std::move(*TM));
+}
+
 LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
     : ES(S.ES ? std::move(S.ES) : llvm::make_unique<ExecutionSession>()),
       Main(this->ES->getMainJITDylib()), DL(""), CtorRunner(Main),
@@ -95,25 +100,25 @@ LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
 
   ObjLinkingLayer = createObjectLinkingLayer(S, *ES);
 
-  if (S.NumCompileThreads > 0) {
-
-    // Configure multi-threaded.
+  if (auto DLOrErr = S.JTMB->getDefaultDataLayoutForTarget())
+    DL = std::move(*DLOrErr);
+  else {
+    Err = DLOrErr.takeError();
+    return;
+  }
 
-    if (auto DLOrErr = S.JTMB->getDefaultDataLayoutForTarget())
-      DL = std::move(*DLOrErr);
-    else {
-      Err = DLOrErr.takeError();
+  {
+    auto CompileFunction = createCompileFunction(S, std::move(*S.JTMB));
+    if (!CompileFunction) {
+      Err = CompileFunction.takeError();
       return;
     }
+    CompileLayer = llvm::make_unique<IRCompileLayer>(
+        *ES, *ObjLinkingLayer, std::move(*CompileFunction));
+  }
 
-    {
-      auto TmpCompileLayer = llvm::make_unique<IRCompileLayer>(
-          *ES, *ObjLinkingLayer, ConcurrentIRCompiler(std::move(*S.JTMB)));
-
-      TmpCompileLayer->setCloneToNewContextOnEmit(true);
-      CompileLayer = std::move(TmpCompileLayer);
-    }
-
+  if (S.NumCompileThreads > 0) {
+    CompileLayer->setCloneToNewContextOnEmit(true);
     CompileThreads = llvm::make_unique<ThreadPool>(S.NumCompileThreads);
     ES->setDispatchMaterialization(
         [this](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
@@ -122,20 +127,6 @@ LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
           auto Work = [SharedMU, &JD]() { SharedMU->doMaterialize(JD); };
           CompileThreads->async(std::move(Work));
         });
-  } else {
-
-    // Configure single-threaded.
-
-    auto TM = S.JTMB->createTargetMachine();
-    if (!TM) {
-      Err = TM.takeError();
-      return;
-    }
-
-    DL = (*TM)->createDataLayout();
-
-    CompileLayer = llvm::make_unique<IRCompileLayer>(
-        *ES, *ObjLinkingLayer, TMOwningSimpleCompiler(std::move(*TM)));
   }
 }