From 2911ead7a921fd5852691f48defbbe9c0be5348a Mon Sep 17 00:00:00 2001
From: Sameer Rahmani <lxsameer@gnu.org>
Date: Thu, 4 Nov 2021 00:57:00 +0000
Subject: [PATCH] Add the compile on demand layer to the jit

---
 include/serene/context.h     |  3 +++
 include/serene/jit/engine.h  | 23 ++++++++++++++++----
 src/libserene/jit/engine.cpp | 41 ++++++++++++++++++++++++++++++++----
 3 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/include/serene/context.h b/include/serene/context.h
index 3d3205a..d6c3bb0 100644
--- a/include/serene/context.h
+++ b/include/serene/context.h
@@ -99,6 +99,9 @@ public:
 
   std::string targetTriple;
 
+  // TODO: Replace target Triple with this one
+  llvm::Triple triple;
+
   /// Insert the given `ns` into the context. The Context object is
   /// the owner of all the namespaces. The `ns` will overwrite any
   /// namespace with the same name.
diff --git a/include/serene/jit/engine.h b/include/serene/jit/engine.h
index ab643c7..2087703 100644
--- a/include/serene/jit/engine.h
+++ b/include/serene/jit/engine.h
@@ -20,6 +20,9 @@
  * Commentary:
  */
 
+// TODO: Look at the LLJITBuilderState::prepareForConstruction to setup the
+//       object linking layer
+
 #ifndef SERENE_JIT_ENGINE_H
 #define SERENE_JIT_ENGINE_H
 
@@ -35,6 +38,7 @@
 #include <llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h>
 #include <llvm/ExecutionEngine/SectionMemoryManager.h>
 #include <llvm/IR/DataLayout.h>
+#include <llvm/Support/raw_ostream.h>
 
 #include <memory>
 
@@ -73,6 +77,12 @@ class SereneJIT {
   /// for us
   orc::IRCompileLayer compileLayer;
 
+  // TODO: Enable these two layers when we add the `Platform support`
+  // std::unique_ptr<orc::IRTransformLayer> transformLayer;
+  // std::unique_ptr<orc::IRTransformLayer> initHelperTransformLayer;
+
+  orc::CompileOnDemandLayer compileOnDemandLayer;
+
   /// Transform layaer is responsible for running a pass pipeline on the AST
   /// and generate LLVM IR
   // orc::IRTransformLayer transformLayer;
@@ -93,13 +103,18 @@ public:
   SereneJIT(serene::SereneContext &ctx,
             std::unique_ptr<orc::ExecutionSession> es,
             std::unique_ptr<orc::EPCIndirectionUtils> epciu,
-            orc::JITTargetMachineBuilder jtmb, llvm::DataLayout &&dl);
+            orc::JITTargetMachineBuilder jtmb, llvm::DataLayout &&dl,
+            unsigned numCompileThreads = 0);
 
   ~SereneJIT() {
-    if (auto Err = es->endSession()) {
-      es->reportError(std::move(Err));
+    // if (compileThreads) {
+    //   compileThreads->wait();
+    // }
+
+    if (auto err = es->endSession()) {
+      es->reportError(std::move(err));
     }
-  }
+  };
 
   const llvm::DataLayout &getDataLayout() const { return dl; }
 
diff --git a/src/libserene/jit/engine.cpp b/src/libserene/jit/engine.cpp
index 7f565b2..872a1dc 100644
--- a/src/libserene/jit/engine.cpp
+++ b/src/libserene/jit/engine.cpp
@@ -22,12 +22,24 @@
 #include "serene/jit/layers.h"
 #include "serene/utils.h"
 
+#include <llvm/ExecutionEngine/JITSymbol.h>
+
+#include <memory>
+
 namespace serene::jit {
 
+static void handleLazyCallThroughError() {
+  // TODO: Report to the diag engine
+  llvm::errs() << "LazyCallThrough error: Could not find function body";
+  // TODO: terminate ?
+}
+
 SereneJIT::SereneJIT(serene::SereneContext &ctx,
                      std::unique_ptr<orc::ExecutionSession> es,
                      std::unique_ptr<orc::EPCIndirectionUtils> epciu,
-                     orc::JITTargetMachineBuilder jtmb, llvm::DataLayout &&dl)
+                     orc::JITTargetMachineBuilder jtmb, llvm::DataLayout &&dl,
+                     unsigned numCompileThreads)
+
     : es(std::move(es)), epciu(std::move(epciu)), dl(dl),
       mangler(*this->es, this->dl),
       objectLayer(
@@ -36,12 +48,21 @@ SereneJIT::SereneJIT(serene::SereneContext &ctx,
       compileLayer(
           *this->es, objectLayer,
           std::make_unique<orc::ConcurrentIRCompiler>(std::move(jtmb))),
-      nsLayer(ctx, compileLayer, mangler, dl),
+      // TODO: Change compileOnDemandLayer to use an optimization layer
+      //       as the parent
+      compileOnDemandLayer(
+          *this->es, compileLayer, this->epciu->getLazyCallThroughManager(),
+          [this] { return this->epciu->createIndirectStubsManager(); }),
+      nsLayer(ctx, compileOnDemandLayer, mangler, dl),
       mainJD(this->es->createBareJITDylib(ctx.getCurrentNS().name)), ctx(ctx) {
   UNUSED(this->ctx);
   mainJD.addGenerator(
       cantFail(orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
           dl.getGlobalPrefix())));
+
+  if (numCompileThreads > 0) {
+    compileOnDemandLayer.setCloneToNewContextOnEmit(true);
+  }
 }
 
 llvm::Error SereneJIT::addNS(llvm::StringRef nsname,
@@ -60,6 +81,18 @@ makeSereneJIT(serene::SereneContext &ctx) {
     return epc.takeError();
   }
   auto es = std::make_unique<orc::ExecutionSession>(std::move(*epc));
+  auto epciu =
+      orc::EPCIndirectionUtils::Create(es->getExecutorProcessControl());
+  if (!epciu) {
+    return epciu.takeError();
+  }
+
+  (*epciu)->createLazyCallThroughManager(
+      *es, llvm::pointerToJITTargetAddress(&handleLazyCallThroughError));
+
+  if (auto err = setUpInProcessLCTMReentryViaEPCIU(**epciu)) {
+    return std::move(err);
+  }
 
   orc::JITTargetMachineBuilder jtmb(
       es->getExecutorProcessControl().getTargetTriple());
@@ -69,7 +102,7 @@ makeSereneJIT(serene::SereneContext &ctx) {
     return dl.takeError();
   }
 
-  return std::make_unique<SereneJIT>(ctx, std::move(es), std::move(jtmb),
-                                     std::move(*dl));
+  return std::make_unique<SereneJIT>(ctx, std::move(es), std::move(*epciu),
+                                     std::move(jtmb), std::move(*dl));
 };
 } // namespace serene::jit