[ORC] Implement basic reoptimization. (#67050)

Author: sunho

compiler-rt/lib/orc/common.h

@@ -19,9 +19,9 @@
 
 /// This macro should be used to define tags that will be associated with
 /// handlers in the JIT process, and call can be used to define tags f
-#define ORC_RT_JIT_DISPATCH_TAG(X) \
-extern "C" char X; \
-char X = 0;
+#define ORC_RT_JIT_DISPATCH_TAG(X)                                             \
+  ORC_RT_INTERFACE char X;                                                     \
+  char X = 0;
 
 /// Opaque struct for external symbols.
 struct __orc_rt_Opaque {};

compiler-rt/lib/orc/elfnix_platform.cpp

@@ -28,6 +28,7 @@ using namespace __orc_rt;
 using namespace __orc_rt::elfnix;
 
 // Declare function tags for functions in the JIT process.
+ORC_RT_JIT_DISPATCH_TAG(__orc_rt_reoptimize_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_get_initializers_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_get_deinitializers_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_symbol_lookup_tag)

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h

@@ -21,6 +21,7 @@
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
@@ -48,6 +49,7 @@ class KaleidoscopeJIT {
   RTDyldObjectLinkingLayer ObjectLayer;
   IRCompileLayer CompileLayer;
   IRTransformLayer OptimizeLayer;
+  IRPartitionLayer IPLayer;
   CompileOnDemandLayer CODLayer;
 
   JITDylib &MainJD;
@@ -68,8 +70,8 @@ class KaleidoscopeJIT {
         CompileLayer(*this->ES, ObjectLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
-        CODLayer(*this->ES, OptimizeLayer,
-                 this->EPCIU->getLazyCallThroughManager(),
+        IPLayer(*this->ES, OptimizeLayer),
+        CODLayer(*this->ES, IPLayer, this->EPCIU->getLazyCallThroughManager(),
                  [this] { return this->EPCIU->createIndirectStubsManager(); }),
         MainJD(this->ES->createBareJITDylib("<main>")) {
     MainJD.addGenerator(

llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp

@@ -3,6 +3,7 @@
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
@@ -109,13 +110,14 @@ class SpeculativeJIT {
       IndirectStubsManagerBuilderFunction ISMBuilder,
       std::unique_ptr<DynamicLibrarySearchGenerator> ProcessSymbolsGenerator)
       : ES(std::move(ES)), DL(std::move(DL)),
-        MainJD(this->ES->createBareJITDylib("<main>")), LCTMgr(std::move(LCTMgr)),
+        MainJD(this->ES->createBareJITDylib("<main>")),
+        LCTMgr(std::move(LCTMgr)),
         CompileLayer(*this->ES, ObjLayer,
                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
         S(Imps, *this->ES),
         SpeculateLayer(*this->ES, CompileLayer, S, Mangle, BlockFreqQuery()),
-        CODLayer(*this->ES, SpeculateLayer, *this->LCTMgr,
-                 std::move(ISMBuilder)) {
+        IPLayer(*this->ES, SpeculateLayer),
+        CODLayer(*this->ES, IPLayer, *this->LCTMgr, std::move(ISMBuilder)) {
     MainJD.addGenerator(std::move(ProcessSymbolsGenerator));
     this->CODLayer.setImplMap(&Imps);
     ExitOnErr(S.addSpeculationRuntime(MainJD, Mangle));
@@ -141,6 +143,7 @@ class SpeculativeJIT {
   Speculator S;
   RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr};
   IRSpeculationLayer SpeculateLayer;
+  IRPartitionLayer IPLayer;
   CompileOnDemandLayer CODLayer;
 };
 

llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

@@ -53,37 +53,15 @@ namespace llvm {
 namespace orc {
 
 class CompileOnDemandLayer : public IRLayer {
-  friend class PartitioningIRMaterializationUnit;
-
 public:
   /// Builder for IndirectStubsManagers.
   using IndirectStubsManagerBuilder =
       std::function<std::unique_ptr<IndirectStubsManager>()>;
 
-  using GlobalValueSet = std::set<const GlobalValue *>;
-
-  /// Partitioning function.
-  using PartitionFunction =
-      std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>;
-
-  /// Off-the-shelf partitioning which compiles all requested symbols (usually
-  /// a single function at a time).
-  static std::optional<GlobalValueSet>
-  compileRequested(GlobalValueSet Requested);
-
-  /// Off-the-shelf partitioning which compiles whole modules whenever any
-  /// symbol in them is requested.
-  static std::optional<GlobalValueSet>
-  compileWholeModule(GlobalValueSet Requested);
-
   /// Construct a CompileOnDemandLayer.
   CompileOnDemandLayer(ExecutionSession &ES, IRLayer &BaseLayer,
-                        LazyCallThroughManager &LCTMgr,
-                        IndirectStubsManagerBuilder BuildIndirectStubsManager);
-
-  /// Sets the partition function.
-  void setPartitionFunction(PartitionFunction Partition);
-
+                       LazyCallThroughManager &LCTMgr,
+                       IndirectStubsManagerBuilder BuildIndirectStubsManager);
   /// Sets the ImplSymbolMap
   void setImplMap(ImplSymbolMap *Imp);
 
@@ -110,22 +88,12 @@ class CompileOnDemandLayer : public IRLayer {
 
   PerDylibResources &getPerDylibResources(JITDylib &TargetD);
 
-  void cleanUpModule(Module &M);
-
-  void expandPartition(GlobalValueSet &Partition);
-
-  void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
-                     ThreadSafeModule TSM,
-                     IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
-
   mutable std::mutex CODLayerMutex;
 
   IRLayer &BaseLayer;
   LazyCallThroughManager &LCTMgr;
   IndirectStubsManagerBuilder BuildIndirectStubsManager;
   PerDylibResourcesMap DylibResources;
-  PartitionFunction Partition = compileRequested;
-  SymbolLinkagePromoter PromoteSymbols;
   ImplSymbolMap *AliaseeImpls = nullptr;
 };
 

llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h

@@ -0,0 +1,85 @@
+//===- IRPartitionLayer.h - Partition IR module on lookup -------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// JIT layer for breaking up modules into smaller submodules that only contains
+// looked up symbols.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
+#define LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
+
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalAlias.h"
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Mangler.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+
+namespace llvm {
+namespace orc {
+
+/// A layer that breaks up IR modules into smaller submodules that only contains
+/// looked up symbols.
+class IRPartitionLayer : public IRLayer {
+  friend class PartitioningIRMaterializationUnit;
+
+public:
+  using GlobalValueSet = std::set<const GlobalValue *>;
+
+  /// Partitioning function.
+  using PartitionFunction =
+      std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>;
+
+  /// Construct a IRPartitionLayer.
+  IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer);
+
+  /// Off-the-shelf partitioning which compiles all requested symbols (usually
+  /// a single function at a time).
+  static std::optional<GlobalValueSet>
+  compileRequested(GlobalValueSet Requested);
+
+  /// Off-the-shelf partitioning which compiles whole modules whenever any
+  /// symbol in them is requested.
+  static std::optional<GlobalValueSet>
+  compileWholeModule(GlobalValueSet Requested);
+
+  /// Sets the partition function.
+  void setPartitionFunction(PartitionFunction Partition);
+
+  /// Emits the given module. This should not be called by clients: it will be
+  /// called by the JIT when a definition added via the add method is requested.
+  void emit(std::unique_ptr<MaterializationResponsibility> R,
+            ThreadSafeModule TSM) override;
+
+private:
+  void cleanUpModule(Module &M);
+
+  void expandPartition(GlobalValueSet &Partition);
+
+  void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
+                     ThreadSafeModule TSM,
+                     IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
+
+  IRLayer &BaseLayer;
+  PartitionFunction Partition = compileRequested;
+  SymbolLinkagePromoter PromoteSymbols;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif

llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h

@@ -0,0 +1,107 @@
+//===- JITLinkRedirectableSymbolManager.h - JITLink redirection -*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Redirectable Symbol Manager implementation using JITLink
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_JITLINKREDIRECABLEMANAGER_H
+#define LLVM_EXECUTIONENGINE_ORC_JITLINKREDIRECABLEMANAGER_H
+
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
+#include "llvm/Support/StringSaver.h"
+
+namespace llvm {
+namespace orc {
+
+class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
+                                         public ResourceManager {
+public:
+  /// Create redirection manager that uses JITLink based implementaion.
+  static Expected<std::unique_ptr<RedirectableSymbolManager>>
+  Create(ObjectLinkingLayer &ObjLinkingLayer, JITDylib &JD) {
+    Error Err = Error::success();
+    auto RM = std::unique_ptr<RedirectableSymbolManager>(
+        new JITLinkRedirectableSymbolManager(ObjLinkingLayer, JD, Err));
+    if (Err)
+      return Err;
+    return std::move(RM);
+  }
+
+  void emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> R,
+                               const SymbolAddrMap &InitialDests) override;
+
+  Error redirect(JITDylib &TargetJD, const SymbolAddrMap &NewDests) override;
+
+  Error handleRemoveResources(JITDylib &TargetJD, ResourceKey K) override;
+
+  void handleTransferResources(JITDylib &TargetJD, ResourceKey DstK,
+                               ResourceKey SrcK) override;
+
+private:
+  using StubHandle = unsigned;
+  constexpr static unsigned StubBlockSize = 256;
+  constexpr static StringRef JumpStubPrefix = "$__IND_JUMP_STUBS";
+  constexpr static StringRef StubPtrPrefix = "$IND_JUMP_PTR_";
+  constexpr static StringRef JumpStubTableName = "$IND_JUMP_";
+  constexpr static StringRef StubPtrTableName = "$__IND_JUMP_PTRS";
+
+  JITLinkRedirectableSymbolManager(ObjectLinkingLayer &ObjLinkingLayer,
+                                   JITDylib &JD, Error &Err)
+      : ObjLinkingLayer(ObjLinkingLayer), JD(JD),
+        AnonymousPtrCreator(jitlink::getAnonymousPointerCreator(
+            ObjLinkingLayer.getExecutionSession().getTargetTriple())),
+        PtrJumpStubCreator(jitlink::getPointerJumpStubCreator(
+            ObjLinkingLayer.getExecutionSession().getTargetTriple())) {
+    if (!AnonymousPtrCreator || !PtrJumpStubCreator)
+      Err = make_error<StringError>("Architecture not supported",
+                                    inconvertibleErrorCode());
+    if (Err)
+      return;
+    ObjLinkingLayer.getExecutionSession().registerResourceManager(*this);
+  }
+
+  ~JITLinkRedirectableSymbolManager() {
+    ObjLinkingLayer.getExecutionSession().deregisterResourceManager(*this);
+  }
+
+  StringRef JumpStubSymbolName(unsigned I) {
+    return *ObjLinkingLayer.getExecutionSession().intern(
+        (JumpStubPrefix + Twine(I)).str());
+  }
+
+  StringRef StubPtrSymbolName(unsigned I) {
+    return *ObjLinkingLayer.getExecutionSession().intern(
+        (StubPtrPrefix + Twine(I)).str());
+  }
+
+  unsigned GetNumAvailableStubs() const { return AvailableStubs.size(); }
+
+  Error redirectInner(JITDylib &TargetJD, const SymbolAddrMap &NewDests);
+  Error grow(unsigned Need);
+
+  ObjectLinkingLayer &ObjLinkingLayer;
+  JITDylib &JD;
+  jitlink::AnonymousPointerCreator AnonymousPtrCreator;
+  jitlink::PointerJumpStubCreator PtrJumpStubCreator;
+
+  std::vector<StubHandle> AvailableStubs;
+  using SymbolToStubMap = DenseMap<SymbolStringPtr, StubHandle>;
+  DenseMap<JITDylib *, SymbolToStubMap> SymbolToStubs;
+  std::vector<ExecutorSymbolDef> JumpStubs;
+  std::vector<ExecutorSymbolDef> StubPointers;
+  DenseMap<ResourceKey, std::vector<SymbolStringPtr>> TrackedResources;
+
+  std::mutex Mutex;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif

llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h

@@ -17,6 +17,7 @@
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
@@ -270,9 +271,8 @@ class LLLazyJIT : public LLJIT {
 public:
 
   /// Sets the partition function.
-  void
-  setPartitionFunction(CompileOnDemandLayer::PartitionFunction Partition) {
-    CODLayer->setPartitionFunction(std::move(Partition));
+  void setPartitionFunction(IRPartitionLayer::PartitionFunction Partition) {
+    IPLayer->setPartitionFunction(std::move(Partition));
   }
 
   /// Returns a reference to the on-demand layer.
@@ -292,6 +292,7 @@ class LLLazyJIT : public LLJIT {
   LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
 
   std::unique_ptr<LazyCallThroughManager> LCTMgr;
+  std::unique_ptr<IRPartitionLayer> IPLayer;
   std::unique_ptr<CompileOnDemandLayer> CODLayer;
 };