[clang] Karikari (PR #67253)

Sunho Kim via cfe-commits cfe-commits at lists.llvm.org
Sat Sep 23 20:51:59 PDT 2023


https://github.com/sunho created https://github.com/llvm/llvm-project/pull/67253

None

>From 9d759677e725c244a7a3eb06864d6474253b6d86 Mon Sep 17 00:00:00 2001
From: Sunho Kim <ksunhokim123 at gmail.com>
Date: Mon, 21 Aug 2023 15:48:26 +0900
Subject: [PATCH 1/5] [ORC] Introduce RedirectionManager interface and
 implementation using JITLink.

---
 .../Orc/JITLinkRedirectableSymbolManager.h    | 106 +++++++++++
 .../ExecutionEngine/Orc/RedirectionManager.h  | 102 ++++++++++
 llvm/lib/ExecutionEngine/Orc/CMakeLists.txt   |   2 +
 .../Orc/JITLinkRedirectableSymbolManager.cpp  | 179 ++++++++++++++++++
 .../Orc/RedirectionManager.cpp                |  28 +++
 .../ExecutionEngine/Orc/CMakeLists.txt        |   1 +
 .../Orc/JITLinkRedirectionManagerTest.cpp     | 100 ++++++++++
 7 files changed, 518 insertions(+)
 create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
 create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
 create mode 100644 llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
 create mode 100644 llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp
 create mode 100644 llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
new file mode 100644
index 000000000000000..5de0da1f52d0db6
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
@@ -0,0 +1,106 @@
+//===- 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(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
+         JITDylib &JD) {
+    Error Err = Error::success();
+    auto RM = std::unique_ptr<RedirectableSymbolManager>(
+        new JITLinkRedirectableSymbolManager(ES, 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(ExecutionSession &ES,
+                                   ObjectLinkingLayer &ObjLinkingLayer,
+                                   JITDylib &JD, Error &Err)
+      : ES(ES), ObjLinkingLayer(ObjLinkingLayer), JD(JD),
+        AnonymousPtrCreator(
+            jitlink::getAnonymousPointerCreator(ES.getTargetTriple())),
+        PtrJumpStubCreator(
+            jitlink::getPointerJumpStubCreator(ES.getTargetTriple())) {
+    if (!AnonymousPtrCreator || !PtrJumpStubCreator)
+      Err = make_error<StringError>("Architecture not supported",
+                                    inconvertibleErrorCode());
+    if (Err)
+      return;
+    ES.registerResourceManager(*this);
+  }
+
+  ~JITLinkRedirectableSymbolManager() { ES.deregisterResourceManager(*this); }
+
+  StringRef JumpStubSymbolName(unsigned I) {
+    return *ES.intern((JumpStubPrefix + Twine(I)).str());
+  }
+
+  StringRef StubPtrSymbolName(unsigned I) {
+    return *ES.intern((StubPtrPrefix + Twine(I)).str());
+  }
+
+  unsigned GetNumAvailableStubs() const { return AvailableStubs.size(); }
+
+  Error redirectInner(JITDylib &TargetJD, const SymbolAddrMap &NewDests);
+  Error grow(unsigned Need);
+
+  ExecutionSession &ES;
+  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
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h b/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
new file mode 100644
index 000000000000000..6116102f4fa4d4a
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/RedirectionManager.h
@@ -0,0 +1,102 @@
+//===- RedirectionManager.h - Redirection manager interface -----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Redirection manager interface that redirects a call to symbol to another.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_REDIRECTIONMANAGER_H
+#define LLVM_EXECUTIONENGINE_ORC_REDIRECTIONMANAGER_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+
+namespace llvm {
+namespace orc {
+
+/// Base class for performing redirection of call to symbol to another symbol in
+/// runtime.
+class RedirectionManager {
+public:
+  /// Symbol name to symbol definition map.
+  using SymbolAddrMap = DenseMap<SymbolStringPtr, ExecutorSymbolDef>;
+
+  virtual ~RedirectionManager() = default;
+  /// Change the redirection destination of given symbols to new destination
+  /// symbols.
+  virtual Error redirect(JITDylib &JD, const SymbolAddrMap &NewDests) = 0;
+
+  /// Change the redirection destination of given symbol to new destination
+  /// symbol.
+  virtual Error redirect(JITDylib &JD, SymbolStringPtr Symbol,
+                         ExecutorSymbolDef NewDest) {
+    return redirect(JD, {{Symbol, NewDest}});
+  }
+
+private:
+  virtual void anchor();
+};
+
+/// Base class for managing redirectable symbols in which a call
+/// gets redirected to another symbol in runtime.
+class RedirectableSymbolManager : public RedirectionManager {
+public:
+  /// Create redirectable symbols with given symbol names and initial
+  /// desitnation symbol addresses.
+  Error createRedirectableSymbols(ResourceTrackerSP RT,
+                                  const SymbolMap &InitialDests);
+
+  /// Create a single redirectable symbol with given symbol name and initial
+  /// desitnation symbol address.
+  Error createRedirectableSymbol(ResourceTrackerSP RT, SymbolStringPtr Symbol,
+                                 ExecutorSymbolDef InitialDest) {
+    return createRedirectableSymbols(RT, {{Symbol, InitialDest}});
+  }
+
+  /// Emit redirectable symbol
+  virtual void
+  emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> MR,
+                          const SymbolMap &InitialDests) = 0;
+};
+
+class RedirectableMaterializationUnit : public MaterializationUnit {
+public:
+  RedirectableMaterializationUnit(RedirectableSymbolManager &RM,
+                                  const SymbolMap &InitialDests)
+      : MaterializationUnit(convertToFlags(InitialDests)), RM(RM),
+        InitialDests(InitialDests) {}
+
+  StringRef getName() const override {
+    return "RedirectableSymbolMaterializationUnit";
+  }
+
+  void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
+    RM.emitRedirectableSymbols(std::move(R), std::move(InitialDests));
+  }
+
+  void discard(const JITDylib &JD, const SymbolStringPtr &Name) override {
+    InitialDests.erase(Name);
+  }
+
+private:
+  static MaterializationUnit::Interface
+  convertToFlags(const SymbolMap &InitialDests) {
+    SymbolFlagsMap Flags;
+    for (auto [K, V] : InitialDests) {
+      Flags[K] = V.getFlags();
+    }
+    return MaterializationUnit::Interface(Flags, {});
+  }
+
+  RedirectableSymbolManager &RM;
+  SymbolMap InitialDests;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif
diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index 3256ed8b7362c66..17e611d5795cec3 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
@@ -49,6 +49,8 @@ add_llvm_component_library(LLVMOrcJIT
   ExecutorProcessControl.cpp
   TaskDispatch.cpp
   ThreadSafeModule.cpp
+  RedirectionManager.cpp
+  JITLinkRedirectableSymbolManager.cpp
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc
 
diff --git a/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
new file mode 100644
index 000000000000000..6d76d50271b9f2d
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
@@ -0,0 +1,179 @@
+//===-- JITLinkRedirectableSymbolManager.cpp - JITLink redirection in Orc -===//
+//
+// 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/JITLinkRedirectableSymbolManager.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+
+#define DEBUG_TYPE "orc"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
+    std::unique_ptr<MaterializationResponsibility> R,
+    const SymbolAddrMap &InitialDests) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  if (GetNumAvailableStubs() < InitialDests.size())
+    if (auto Err = grow(InitialDests.size() - GetNumAvailableStubs())) {
+      ES.reportError(std::move(Err));
+      R->failMaterialization();
+      return;
+    }
+
+  JITDylib &TargetJD = R->getTargetJITDylib();
+  SymbolMap NewSymbolDefs;
+  std::vector<SymbolStringPtr> Symbols;
+  for (auto &[K, V] : InitialDests) {
+    StubHandle StubID = AvailableStubs.back();
+    if (SymbolToStubs[&TargetJD].count(K)) {
+      ES.reportError(make_error<StringError>(
+          "Tried to create duplicate redirectable symbols",
+          inconvertibleErrorCode()));
+      R->failMaterialization();
+      return;
+    }
+    dbgs() << *K << "\n";
+    SymbolToStubs[&TargetJD][K] = StubID;
+    NewSymbolDefs[K] = JumpStubs[StubID];
+    NewSymbolDefs[K].setFlags(V.getFlags());
+    Symbols.push_back(K);
+    AvailableStubs.pop_back();
+  }
+
+  if (auto Err = R->replace(absoluteSymbols(NewSymbolDefs))) {
+    ES.reportError(std::move(Err));
+    R->failMaterialization();
+    return;
+  }
+
+  if (auto Err = redirectInner(TargetJD, InitialDests)) {
+    ES.reportError(std::move(Err));
+    R->failMaterialization();
+    return;
+  }
+
+  auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
+    TrackedResources[Key].insert(TrackedResources[Key].end(), Symbols.begin(),
+                                 Symbols.end());
+  });
+  if (Err) {
+    ES.reportError(std::move(Err));
+    R->failMaterialization();
+    return;
+  }
+}
+
+Error JITLinkRedirectableSymbolManager::redirect(
+    JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  return redirectInner(TargetJD, NewDests);
+}
+
+Error JITLinkRedirectableSymbolManager::redirectInner(
+    JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
+  std::vector<tpctypes::PointerWrite> PtrWrites;
+  for (auto &[K, V] : NewDests) {
+    if (!SymbolToStubs[&TargetJD].count(K))
+      return make_error<StringError>(
+          "Tried to redirect non-existent redirectalbe symbol",
+          inconvertibleErrorCode());
+    StubHandle StubID = SymbolToStubs[&TargetJD].at(K);
+    PtrWrites.push_back({StubPointers[StubID].getAddress(), V.getAddress()});
+  }
+  if (auto Err = ES.getExecutorProcessControl().getMemoryAccess().writePointers(
+          PtrWrites))
+    return Err;
+  return Error::success();
+}
+
+Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
+  unsigned OldSize = JumpStubs.size();
+  unsigned NumNewStubs = alignTo(Need, StubBlockSize);
+  unsigned NewSize = OldSize + NumNewStubs;
+
+  JumpStubs.resize(NewSize);
+  StubPointers.resize(NewSize);
+  AvailableStubs.reserve(NewSize);
+
+  SymbolLookupSet LookupSymbols;
+  DenseMap<SymbolStringPtr, ExecutorSymbolDef *> NewDefsMap;
+
+  Triple TT = ES.getTargetTriple();
+  auto G = std::make_unique<jitlink::LinkGraph>(
+      "<INDIRECT STUBS>", TT, TT.isArch64Bit() ? 8 : 4,
+      TT.isLittleEndian() ? support::little : support::big,
+      jitlink::getGenericEdgeKindName);
+  auto &PointerSection =
+      G->createSection(StubPtrTableName, MemProt::Write | MemProt::Read);
+  auto &StubsSection =
+      G->createSection(JumpStubTableName, MemProt::Exec | MemProt::Read);
+
+  for (size_t I = OldSize; I < NewSize; I++) {
+    auto Pointer = AnonymousPtrCreator(*G, PointerSection, nullptr, 0);
+    if (auto Err = Pointer.takeError())
+      return Err;
+
+    StringRef PtrSymName = StubPtrSymbolName(I);
+    Pointer->setName(PtrSymName);
+    Pointer->setScope(jitlink::Scope::Default);
+    LookupSymbols.add(ES.intern(PtrSymName));
+    NewDefsMap[ES.intern(PtrSymName)] = &StubPointers[I];
+
+    auto Stub = PtrJumpStubCreator(*G, StubsSection, *Pointer);
+    if (auto Err = Stub.takeError())
+      return Err;
+
+    StringRef JumpStubSymName = JumpStubSymbolName(I);
+    Stub->setName(JumpStubSymName);
+    Stub->setScope(jitlink::Scope::Default);
+    LookupSymbols.add(ES.intern(JumpStubSymName));
+    NewDefsMap[ES.intern(JumpStubSymName)] = &JumpStubs[I];
+  }
+
+  if (auto Err = ObjLinkingLayer.add(JD, std::move(G)))
+    return Err;
+
+  auto LookupResult = ES.lookup(makeJITDylibSearchOrder(&JD), LookupSymbols);
+  if (auto Err = LookupResult.takeError())
+    return Err;
+
+  for (auto &[K, V] : *LookupResult)
+    *NewDefsMap.at(K) = V;
+
+  for (size_t I = OldSize; I < NewSize; I++)
+    AvailableStubs.push_back(I);
+
+  return Error::success();
+}
+
+Error JITLinkRedirectableSymbolManager::handleRemoveResources(
+    JITDylib &TargetJD, ResourceKey K) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  for (auto &Symbol : TrackedResources[K]) {
+    if (!SymbolToStubs[&TargetJD].count(Symbol))
+      return make_error<StringError>(
+          "Tried to remove non-existent redirectable symbol",
+          inconvertibleErrorCode());
+    AvailableStubs.push_back(SymbolToStubs[&TargetJD].at(Symbol));
+    SymbolToStubs[&TargetJD].erase(Symbol);
+    if (SymbolToStubs[&TargetJD].empty())
+      SymbolToStubs.erase(&TargetJD);
+  }
+  TrackedResources.erase(K);
+
+  return Error::success();
+}
+
+void JITLinkRedirectableSymbolManager::handleTransferResources(
+    JITDylib &TargetJD, ResourceKey DstK, ResourceKey SrcK) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  TrackedResources[DstK].insert(TrackedResources[DstK].end(),
+                                TrackedResources[SrcK].begin(),
+                                TrackedResources[SrcK].end());
+  TrackedResources.erase(SrcK);
+}
diff --git a/llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp b/llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp
new file mode 100644
index 000000000000000..61d765066c8e4cd
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/RedirectionManager.cpp
@@ -0,0 +1,28 @@
+//===---- RedirectionManager.cpp - Redirection manager interface in Orc ---===//
+//
+// 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/RedirectionManager.h"
+
+#define DEBUG_TYPE "orc"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+void RedirectionManager::anchor() {}
+
+Error RedirectableSymbolManager::createRedirectableSymbols(
+    ResourceTrackerSP RT, const SymbolMap &InitialDests) {
+  auto &JD = RT->getJITDylib();
+
+  if (auto Err = JD.define(std::make_unique<RedirectableMaterializationUnit>(
+                               *this, InitialDests),
+                           RT))
+    return Err;
+
+  return Error::success();
+}
diff --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
index 37768e91fd44729..e932ff94361a7b7 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
@@ -40,6 +40,7 @@ add_llvm_unittest(OrcJITTests
   TaskDispatchTest.cpp
   ThreadSafeModuleTest.cpp
   WrapperFunctionUtilsTest.cpp
+  JITLinkRedirectionManagerTest.cpp
   )
 
 target_link_libraries(OrcJITTests PRIVATE
diff --git a/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
new file mode 100644
index 000000000000000..0f87c1b7433238b
--- /dev/null
+++ b/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
@@ -0,0 +1,100 @@
+#include "OrcTestCommon.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
+#include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::jitlink;
+
+static int initialTarget() { return 42; }
+static int middleTarget() { return 13; }
+static int finalTarget() { return 53; }
+
+class JITLinkRedirectionManagerTest : public testing::Test {
+public:
+  ~JITLinkRedirectionManagerTest() {
+    if (ES)
+      if (auto Err = ES->endSession())
+        ES->reportError(std::move(Err));
+  }
+
+protected:
+  void SetUp() override {
+    auto JTMB = JITTargetMachineBuilder::detectHost();
+    // Bail out if we can not detect the host.
+    if (!JTMB) {
+      consumeError(JTMB.takeError());
+      GTEST_SKIP();
+    }
+
+    ES = std::make_unique<ExecutionSession>(
+        std::make_unique<UnsupportedExecutorProcessControl>(
+            nullptr, nullptr, JTMB->getTargetTriple().getTriple()));
+    JD = &ES->createBareJITDylib("main");
+    ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
+        *ES, std::make_unique<InProcessMemoryManager>(4096));
+    DL = std::make_unique<DataLayout>(
+        cantFail(JTMB->getDefaultDataLayoutForTarget()));
+  }
+  JITDylib *JD{nullptr};
+  std::unique_ptr<ExecutionSession> ES;
+  std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer;
+  std::unique_ptr<DataLayout> DL;
+};
+
+TEST_F(JITLinkRedirectionManagerTest, BasicRedirectionOperation) {
+  auto RM =
+      JITLinkRedirectableSymbolManager::Create(*ES, *ObjLinkingLayer, *JD);
+  // Bail out if we can not create
+  if (!RM) {
+    consumeError(RM.takeError());
+    GTEST_SKIP();
+  }
+
+  auto DefineTarget = [&](StringRef TargetName, ExecutorAddr Addr) {
+    SymbolStringPtr Target = ES->intern(TargetName);
+    cantFail(JD->define(std::make_unique<SimpleMaterializationUnit>(
+        SymbolFlagsMap({{Target, JITSymbolFlags::Exported}}),
+        [&](std::unique_ptr<MaterializationResponsibility> R) -> void {
+          // No dependencies registered, can't fail.
+          cantFail(
+              R->notifyResolved({{Target, {Addr, JITSymbolFlags::Exported}}}));
+          cantFail(R->notifyEmitted());
+        })));
+    return cantFail(ES->lookup({JD}, TargetName));
+  };
+
+  auto InitialTarget =
+      DefineTarget("InitialTarget", ExecutorAddr::fromPtr(&initialTarget));
+  auto MiddleTarget =
+      DefineTarget("MiddleTarget", ExecutorAddr::fromPtr(&middleTarget));
+  auto FinalTarget =
+      DefineTarget("FinalTarget", ExecutorAddr::fromPtr(&finalTarget));
+
+  auto RedirectableSymbol = ES->intern("RedirectableTarget");
+  EXPECT_THAT_ERROR(
+      (*RM)->createRedirectableSymbols(JD->getDefaultResourceTracker(),
+                                       {{RedirectableSymbol, InitialTarget}}),
+      Succeeded());
+  auto RTDef = cantFail(ES->lookup({JD}, RedirectableSymbol));
+
+  auto RTPtr = RTDef.getAddress().toPtr<int (*)()>();
+  auto Result = RTPtr();
+  EXPECT_EQ(Result, 42) << "Failed to call initial target";
+
+  EXPECT_THAT_ERROR((*RM)->redirect(*JD, {{RedirectableSymbol, MiddleTarget}}),
+                    Succeeded());
+  Result = RTPtr();
+  EXPECT_EQ(Result, 13) << "Failed to call middle redirected target";
+
+  EXPECT_THAT_ERROR((*RM)->redirect(*JD, {{RedirectableSymbol, FinalTarget}}),
+                    Succeeded());
+  Result = RTPtr();
+  EXPECT_EQ(Result, 53) << "Failed to call redirected target";
+}

>From 8e96f8ce2b54fbf404e5df2064da851515920aef Mon Sep 17 00:00:00 2001
From: Zequan Wu <zequanwu at google.com>
Date: Tue, 19 Sep 2023 14:59:34 -0400
Subject: [PATCH 2/5] [Coverage] Skip visiting ctor member initializers with
 invalid source locations.

---
 clang/lib/CodeGen/CoverageMappingGen.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/clang/lib/CodeGen/CoverageMappingGen.cpp b/clang/lib/CodeGen/CoverageMappingGen.cpp
index bb814dfbfd580f1..32c67739915380c 100644
--- a/clang/lib/CodeGen/CoverageMappingGen.cpp
+++ b/clang/lib/CodeGen/CoverageMappingGen.cpp
@@ -1044,7 +1044,8 @@ struct CounterCoverageMappingBuilder
       for (auto *Initializer : Ctor->inits()) {
         if (Initializer->isWritten()) {
           auto *Init = Initializer->getInit();
-          propagateCounts(BodyCounter, Init);
+          if (getStart(Init).isValid() && getEnd(Init).isValid())
+            propagateCounts(BodyCounter, Init);
         }
       }
     }

>From e576d28381104c5cd8243c6e6475ac300be2702d Mon Sep 17 00:00:00 2001
From: Sunho Kim <ksunhokim123 at gmail.com>
Date: Wed, 20 Sep 2023 05:11:15 +0900
Subject: [PATCH 3/5] [ORC] Introduce IRPartitionLayer for common partition
 functionality.

---
 .../BuildingAJIT/Chapter3/KaleidoscopeJIT.h   |   6 +-
 .../SpeculativeJIT/SpeculativeJIT.cpp         |   9 +-
 .../Orc/CompileOnDemandLayer.h                |  36 +--
 .../ExecutionEngine/Orc/IRPartitionLayer.h    |  83 +++++
 llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h |   7 +-
 llvm/lib/ExecutionEngine/Orc/CMakeLists.txt   |   1 +
 .../Orc/CompileOnDemandLayer.cpp              | 282 +---------------
 .../ExecutionEngine/Orc/IRPartitionLayer.cpp  | 303 ++++++++++++++++++
 llvm/lib/ExecutionEngine/Orc/LLJIT.cpp        |   7 +-
 llvm/tools/lli/lli.cpp                        |   3 +-
 10 files changed, 415 insertions(+), 322 deletions(-)
 create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h
 create mode 100644 llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp

diff --git a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h b/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h
index fd0e081ff2b4bd0..fee2d26e5d9254f 100644
--- a/llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h
+++ b/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(
diff --git a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp b/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp
index fdd376d82da5d89..a1722748c6ce069 100644
--- a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp
+++ b/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"
@@ -107,13 +108,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);
     this->ES->setDispatchTask(
@@ -147,6 +149,7 @@ class SpeculativeJIT {
   Speculator S;
   RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr};
   IRSpeculationLayer SpeculateLayer;
+  IRPartitionLayer IPLayer;
   CompileOnDemandLayer CODLayer;
 };
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
index 23a64b5836db85f..ee556d11b708352 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
@@ -54,37 +54,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);
 
@@ -111,22 +89,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;
 };
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h
new file mode 100644
index 000000000000000..0535ce473837474
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/IRPartitionLayer.h
@@ -0,0 +1,83 @@
+//===- 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 {
+
+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
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
index ab54e0c2c288205..896de521f13af80 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
+++ b/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"
@@ -271,9 +272,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.
@@ -293,6 +293,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;
 };
 
diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index 17e611d5795cec3..379c3cc240e8c7b 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
@@ -27,6 +27,7 @@ add_llvm_component_library(LLVMOrcJIT
   IndirectionUtils.cpp
   IRCompileLayer.cpp
   IRTransformLayer.cpp
+  IRPartitionLayer.cpp
   JITTargetMachineBuilder.cpp
   LazyReexports.cpp
   Layer.cpp
diff --git a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
index 6448adaa0ceb36f..9296bc2b389ab24 100644
--- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
@@ -9,6 +9,7 @@
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -17,101 +18,6 @@
 using namespace llvm;
 using namespace llvm::orc;
 
-static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM,
-                                         StringRef Suffix,
-                                         GVPredicate ShouldExtract) {
-
-  auto DeleteExtractedDefs = [](GlobalValue &GV) {
-    // Bump the linkage: this global will be provided by the external module.
-    GV.setLinkage(GlobalValue::ExternalLinkage);
-
-    // Delete the definition in the source module.
-    if (isa<Function>(GV)) {
-      auto &F = cast<Function>(GV);
-      F.deleteBody();
-      F.setPersonalityFn(nullptr);
-    } else if (isa<GlobalVariable>(GV)) {
-      cast<GlobalVariable>(GV).setInitializer(nullptr);
-    } else if (isa<GlobalAlias>(GV)) {
-      // We need to turn deleted aliases into function or variable decls based
-      // on the type of their aliasee.
-      auto &A = cast<GlobalAlias>(GV);
-      Constant *Aliasee = A.getAliasee();
-      assert(A.hasName() && "Anonymous alias?");
-      assert(Aliasee->hasName() && "Anonymous aliasee");
-      std::string AliasName = std::string(A.getName());
-
-      if (isa<Function>(Aliasee)) {
-        auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee));
-        A.replaceAllUsesWith(F);
-        A.eraseFromParent();
-        F->setName(AliasName);
-      } else if (isa<GlobalVariable>(Aliasee)) {
-        auto *G = cloneGlobalVariableDecl(*A.getParent(),
-                                          *cast<GlobalVariable>(Aliasee));
-        A.replaceAllUsesWith(G);
-        A.eraseFromParent();
-        G->setName(AliasName);
-      } else
-        llvm_unreachable("Alias to unsupported type");
-    } else
-      llvm_unreachable("Unsupported global type");
-  };
-
-  auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs);
-  NewTSM.withModuleDo([&](Module &M) {
-    M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
-  });
-
-  return NewTSM;
-}
-
-namespace llvm {
-namespace orc {
-
-class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
-public:
-  PartitioningIRMaterializationUnit(ExecutionSession &ES,
-                                    const IRSymbolMapper::ManglingOptions &MO,
-                                    ThreadSafeModule TSM,
-                                    CompileOnDemandLayer &Parent)
-      : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {}
-
-  PartitioningIRMaterializationUnit(
-      ThreadSafeModule TSM, Interface I,
-      SymbolNameToDefinitionMap SymbolToDefinition,
-      CompileOnDemandLayer &Parent)
-      : IRMaterializationUnit(std::move(TSM), std::move(I),
-                              std::move(SymbolToDefinition)),
-        Parent(Parent) {}
-
-private:
-  void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
-    Parent.emitPartition(std::move(R), std::move(TSM),
-                         std::move(SymbolToDefinition));
-  }
-
-  void discard(const JITDylib &V, const SymbolStringPtr &Name) override {
-    // All original symbols were materialized by the CODLayer and should be
-    // final. The function bodies provided by M should never be overridden.
-    llvm_unreachable("Discard should never be called on an "
-                     "ExtractingIRMaterializationUnit");
-  }
-
-  mutable std::mutex SourceModuleMutex;
-  CompileOnDemandLayer &Parent;
-};
-
-std::optional<CompileOnDemandLayer::GlobalValueSet>
-CompileOnDemandLayer::compileRequested(GlobalValueSet Requested) {
-  return std::move(Requested);
-}
-
-std::optional<CompileOnDemandLayer::GlobalValueSet>
-CompileOnDemandLayer::compileWholeModule(GlobalValueSet Requested) {
-  return std::nullopt;
-}
-
 CompileOnDemandLayer::CompileOnDemandLayer(
     ExecutionSession &ES, IRLayer &BaseLayer, LazyCallThroughManager &LCTMgr,
     IndirectStubsManagerBuilder BuildIndirectStubsManager)
@@ -119,13 +25,10 @@ CompileOnDemandLayer::CompileOnDemandLayer(
       LCTMgr(LCTMgr),
       BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {}
 
-void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) {
-  this->Partition = std::move(Partition);
-}
-
 void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) {
   this->AliaseeImpls = Imp;
 }
+
 void CompileOnDemandLayer::emit(
     std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM) {
   assert(TSM && "Null module");
@@ -138,10 +41,6 @@ void CompileOnDemandLayer::emit(
 
   SymbolAliasMap NonCallables;
   SymbolAliasMap Callables;
-  TSM.withModuleDo([&](Module &M) {
-    // First, do some cleanup on the module:
-    cleanUpModule(M);
-  });
 
   for (auto &KV : R->getSymbols()) {
     auto &Name = KV.first;
@@ -152,11 +51,10 @@ void CompileOnDemandLayer::emit(
       NonCallables[Name] = SymbolAliasMapEntry(Name, Flags);
   }
 
-  // Create a partitioning materialization unit and lodge it with the
-  // implementation dylib.
+  // Lodge symbols with the implementation dylib.
   if (auto Err = PDR.getImplDylib().define(
-          std::make_unique<PartitioningIRMaterializationUnit>(
-              ES, *getManglingOptions(), std::move(TSM), *this))) {
+          std::make_unique<BasicIRLayerMaterializationUnit>(
+              BaseLayer, *getManglingOptions(), std::move(TSM)))) {
     ES.reportError(std::move(Err));
     R->failMaterialization();
     return;
@@ -210,173 +108,3 @@ CompileOnDemandLayer::getPerDylibResources(JITDylib &TargetD) {
 
   return I->second;
 }
-
-void CompileOnDemandLayer::cleanUpModule(Module &M) {
-  for (auto &F : M.functions()) {
-    if (F.isDeclaration())
-      continue;
-
-    if (F.hasAvailableExternallyLinkage()) {
-      F.deleteBody();
-      F.setPersonalityFn(nullptr);
-      continue;
-    }
-  }
-}
-
-void CompileOnDemandLayer::expandPartition(GlobalValueSet &Partition) {
-  // Expands the partition to ensure the following rules hold:
-  // (1) If any alias is in the partition, its aliasee is also in the partition.
-  // (2) If any aliasee is in the partition, its aliases are also in the
-  //     partiton.
-  // (3) If any global variable is in the partition then all global variables
-  //     are in the partition.
-  assert(!Partition.empty() && "Unexpected empty partition");
-
-  const Module &M = *(*Partition.begin())->getParent();
-  bool ContainsGlobalVariables = false;
-  std::vector<const GlobalValue *> GVsToAdd;
-
-  for (const auto *GV : Partition)
-    if (isa<GlobalAlias>(GV))
-      GVsToAdd.push_back(
-          cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));
-    else if (isa<GlobalVariable>(GV))
-      ContainsGlobalVariables = true;
-
-  for (auto &A : M.aliases())
-    if (Partition.count(cast<GlobalValue>(A.getAliasee())))
-      GVsToAdd.push_back(&A);
-
-  if (ContainsGlobalVariables)
-    for (auto &G : M.globals())
-      GVsToAdd.push_back(&G);
-
-  for (const auto *GV : GVsToAdd)
-    Partition.insert(GV);
-}
-
-void CompileOnDemandLayer::emitPartition(
-    std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM,
-    IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {
-
-  // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
-  //        extracted module key, extracted module, and source module key
-  //        together. This could be used, for example, to provide a specific
-  //        memory manager instance to the linking layer.
-
-  auto &ES = getExecutionSession();
-  GlobalValueSet RequestedGVs;
-  for (auto &Name : R->getRequestedSymbols()) {
-    if (Name == R->getInitializerSymbol())
-      TSM.withModuleDo([&](Module &M) {
-        for (auto &GV : getStaticInitGVs(M))
-          RequestedGVs.insert(&GV);
-      });
-    else {
-      assert(Defs.count(Name) && "No definition for symbol");
-      RequestedGVs.insert(Defs[Name]);
-    }
-  }
-
-  /// Perform partitioning with the context lock held, since the partition
-  /// function is allowed to access the globals to compute the partition.
-  auto GVsToExtract =
-      TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); });
-
-  // Take a 'None' partition to mean the whole module (as opposed to an empty
-  // partition, which means "materialize nothing"). Emit the whole module
-  // unmodified to the base layer.
-  if (GVsToExtract == std::nullopt) {
-    Defs.clear();
-    BaseLayer.emit(std::move(R), std::move(TSM));
-    return;
-  }
-
-  // If the partition is empty, return the whole module to the symbol table.
-  if (GVsToExtract->empty()) {
-    if (auto Err =
-            R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
-                std::move(TSM),
-                MaterializationUnit::Interface(R->getSymbols(),
-                                               R->getInitializerSymbol()),
-                std::move(Defs), *this))) {
-      getExecutionSession().reportError(std::move(Err));
-      R->failMaterialization();
-      return;
-    }
-    return;
-  }
-
-  // Ok -- we actually need to partition the symbols. Promote the symbol
-  // linkages/names, expand the partition to include any required symbols
-  // (i.e. symbols that can't be separated from our partition), and
-  // then extract the partition.
-  //
-  // FIXME: We apply this promotion once per partitioning. It's safe, but
-  // overkill.
-  auto ExtractedTSM =
-      TSM.withModuleDo([&](Module &M) -> Expected<ThreadSafeModule> {
-        auto PromotedGlobals = PromoteSymbols(M);
-        if (!PromotedGlobals.empty()) {
-
-          MangleAndInterner Mangle(ES, M.getDataLayout());
-          SymbolFlagsMap SymbolFlags;
-          IRSymbolMapper::add(ES, *getManglingOptions(),
-                              PromotedGlobals, SymbolFlags);
-
-          if (auto Err = R->defineMaterializing(SymbolFlags))
-            return std::move(Err);
-        }
-
-        expandPartition(*GVsToExtract);
-
-        // Submodule name is given by hashing the names of the globals.
-        std::string SubModuleName;
-        {
-          std::vector<const GlobalValue*> HashGVs;
-          HashGVs.reserve(GVsToExtract->size());
-          for (const auto *GV : *GVsToExtract)
-            HashGVs.push_back(GV);
-          llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) {
-              return LHS->getName() < RHS->getName();
-            });
-          hash_code HC(0);
-          for (const auto *GV : HashGVs) {
-            assert(GV->hasName() && "All GVs to extract should be named by now");
-            auto GVName = GV->getName();
-            HC = hash_combine(HC, hash_combine_range(GVName.begin(), GVName.end()));
-          }
-          raw_string_ostream(SubModuleName)
-            << ".submodule."
-            << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}",
-                       static_cast<size_t>(HC))
-            << ".ll";
-        }
-
-        // Extract the requested partiton (plus any necessary aliases) and
-        // put the rest back into the impl dylib.
-        auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
-          return GVsToExtract->count(&GV);
-        };
-
-        return extractSubModule(TSM, SubModuleName , ShouldExtract);
-      });
-
-  if (!ExtractedTSM) {
-    ES.reportError(ExtractedTSM.takeError());
-    R->failMaterialization();
-    return;
-  }
-
-  if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
-          ES, *getManglingOptions(), std::move(TSM), *this))) {
-    ES.reportError(std::move(Err));
-    R->failMaterialization();
-    return;
-  }
-  BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));
-}
-
-} // end namespace orc
-} // end namespace llvm
diff --git a/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp b/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp
new file mode 100644
index 000000000000000..9ad171beac7fe2e
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/IRPartitionLayer.cpp
@@ -0,0 +1,303 @@
+//===----- IRPartitionLayer.cpp - Partition IR module into submodules -----===//
+//
+// 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/IRPartitionLayer.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM,
+                                         StringRef Suffix,
+                                         GVPredicate ShouldExtract) {
+
+  auto DeleteExtractedDefs = [](GlobalValue &GV) {
+    // Bump the linkage: this global will be provided by the external module.
+    GV.setLinkage(GlobalValue::ExternalLinkage);
+
+    // Delete the definition in the source module.
+    if (isa<Function>(GV)) {
+      auto &F = cast<Function>(GV);
+      F.deleteBody();
+      F.setPersonalityFn(nullptr);
+    } else if (isa<GlobalVariable>(GV)) {
+      cast<GlobalVariable>(GV).setInitializer(nullptr);
+    } else if (isa<GlobalAlias>(GV)) {
+      // We need to turn deleted aliases into function or variable decls based
+      // on the type of their aliasee.
+      auto &A = cast<GlobalAlias>(GV);
+      Constant *Aliasee = A.getAliasee();
+      assert(A.hasName() && "Anonymous alias?");
+      assert(Aliasee->hasName() && "Anonymous aliasee");
+      std::string AliasName = std::string(A.getName());
+
+      if (isa<Function>(Aliasee)) {
+        auto *F = cloneFunctionDecl(*A.getParent(), *cast<Function>(Aliasee));
+        A.replaceAllUsesWith(F);
+        A.eraseFromParent();
+        F->setName(AliasName);
+      } else if (isa<GlobalVariable>(Aliasee)) {
+        auto *G = cloneGlobalVariableDecl(*A.getParent(),
+                                          *cast<GlobalVariable>(Aliasee));
+        A.replaceAllUsesWith(G);
+        A.eraseFromParent();
+        G->setName(AliasName);
+      } else
+        llvm_unreachable("Alias to unsupported type");
+    } else
+      llvm_unreachable("Unsupported global type");
+  };
+
+  auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs);
+  NewTSM.withModuleDo([&](Module &M) {
+    M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
+  });
+
+  return NewTSM;
+}
+
+namespace llvm {
+namespace orc {
+
+class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
+public:
+  PartitioningIRMaterializationUnit(ExecutionSession &ES,
+                                    const IRSymbolMapper::ManglingOptions &MO,
+                                    ThreadSafeModule TSM,
+                                    IRPartitionLayer &Parent)
+      : IRMaterializationUnit(ES, MO, std::move(TSM)), Parent(Parent) {}
+
+  PartitioningIRMaterializationUnit(
+      ThreadSafeModule TSM, Interface I,
+      SymbolNameToDefinitionMap SymbolToDefinition, IRPartitionLayer &Parent)
+      : IRMaterializationUnit(std::move(TSM), std::move(I),
+                              std::move(SymbolToDefinition)),
+        Parent(Parent) {}
+
+private:
+  void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
+    Parent.emitPartition(std::move(R), std::move(TSM),
+                         std::move(SymbolToDefinition));
+  }
+
+  void discard(const JITDylib &V, const SymbolStringPtr &Name) override {
+    // All original symbols were materialized by the CODLayer and should be
+    // final. The function bodies provided by M should never be overridden.
+    llvm_unreachable("Discard should never be called on an "
+                     "ExtractingIRMaterializationUnit");
+  }
+
+  IRPartitionLayer &Parent;
+};
+
+} // namespace orc
+} // namespace llvm
+
+IRPartitionLayer::IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer)
+    : IRLayer(ES, BaseLayer.getManglingOptions()), BaseLayer(BaseLayer) {}
+
+void IRPartitionLayer::setPartitionFunction(PartitionFunction Partition) {
+  this->Partition = Partition;
+}
+
+std::optional<IRPartitionLayer::GlobalValueSet>
+IRPartitionLayer::compileRequested(GlobalValueSet Requested) {
+  return std::move(Requested);
+}
+
+std::optional<IRPartitionLayer::GlobalValueSet>
+IRPartitionLayer::compileWholeModule(GlobalValueSet Requested) {
+  return std::nullopt;
+}
+
+void IRPartitionLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
+                            ThreadSafeModule TSM) {
+  assert(TSM && "Null module");
+
+  auto &ES = getExecutionSession();
+  TSM.withModuleDo([&](Module &M) {
+    // First, do some cleanup on the module:
+    cleanUpModule(M);
+  });
+
+  // Create a partitioning materialization unit and pass the responsibility.
+  if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
+          ES, *getManglingOptions(), std::move(TSM), *this))) {
+    ES.reportError(std::move(Err));
+    R->failMaterialization();
+    return;
+  }
+}
+
+void IRPartitionLayer::cleanUpModule(Module &M) {
+  for (auto &F : M.functions()) {
+    if (F.isDeclaration())
+      continue;
+
+    if (F.hasAvailableExternallyLinkage()) {
+      F.deleteBody();
+      F.setPersonalityFn(nullptr);
+      continue;
+    }
+  }
+}
+
+void IRPartitionLayer::expandPartition(GlobalValueSet &Partition) {
+  // Expands the partition to ensure the following rules hold:
+  // (1) If any alias is in the partition, its aliasee is also in the partition.
+  // (2) If any aliasee is in the partition, its aliases are also in the
+  //     partiton.
+  // (3) If any global variable is in the partition then all global variables
+  //     are in the partition.
+  assert(!Partition.empty() && "Unexpected empty partition");
+
+  const Module &M = *(*Partition.begin())->getParent();
+  bool ContainsGlobalVariables = false;
+  std::vector<const GlobalValue *> GVsToAdd;
+
+  for (const auto *GV : Partition)
+    if (isa<GlobalAlias>(GV))
+      GVsToAdd.push_back(
+          cast<GlobalValue>(cast<GlobalAlias>(GV)->getAliasee()));
+    else if (isa<GlobalVariable>(GV))
+      ContainsGlobalVariables = true;
+
+  for (auto &A : M.aliases())
+    if (Partition.count(cast<GlobalValue>(A.getAliasee())))
+      GVsToAdd.push_back(&A);
+
+  if (ContainsGlobalVariables)
+    for (auto &G : M.globals())
+      GVsToAdd.push_back(&G);
+
+  for (const auto *GV : GVsToAdd)
+    Partition.insert(GV);
+}
+
+void IRPartitionLayer::emitPartition(
+    std::unique_ptr<MaterializationResponsibility> R, ThreadSafeModule TSM,
+    IRMaterializationUnit::SymbolNameToDefinitionMap Defs) {
+
+  // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
+  //        extracted module key, extracted module, and source module key
+  //        together. This could be used, for example, to provide a specific
+  //        memory manager instance to the linking layer.
+
+  auto &ES = getExecutionSession();
+  GlobalValueSet RequestedGVs;
+  for (auto &Name : R->getRequestedSymbols()) {
+    if (Name == R->getInitializerSymbol())
+      TSM.withModuleDo([&](Module &M) {
+        for (auto &GV : getStaticInitGVs(M))
+          RequestedGVs.insert(&GV);
+      });
+    else {
+      assert(Defs.count(Name) && "No definition for symbol");
+      RequestedGVs.insert(Defs[Name]);
+    }
+  }
+
+  /// Perform partitioning with the context lock held, since the partition
+  /// function is allowed to access the globals to compute the partition.
+  auto GVsToExtract =
+      TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); });
+
+  // Take a 'None' partition to mean the whole module (as opposed to an empty
+  // partition, which means "materialize nothing"). Emit the whole module
+  // unmodified to the base layer.
+  if (GVsToExtract == std::nullopt) {
+    Defs.clear();
+    BaseLayer.emit(std::move(R), std::move(TSM));
+    return;
+  }
+
+  // If the partition is empty, return the whole module to the symbol table.
+  if (GVsToExtract->empty()) {
+    if (auto Err =
+            R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
+                std::move(TSM),
+                MaterializationUnit::Interface(R->getSymbols(),
+                                               R->getInitializerSymbol()),
+                std::move(Defs), *this))) {
+      getExecutionSession().reportError(std::move(Err));
+      R->failMaterialization();
+      return;
+    }
+    return;
+  }
+
+  // Ok -- we actually need to partition the symbols. Promote the symbol
+  // linkages/names, expand the partition to include any required symbols
+  // (i.e. symbols that can't be separated from our partition), and
+  // then extract the partition.
+  //
+  // FIXME: We apply this promotion once per partitioning. It's safe, but
+  // overkill.
+  auto ExtractedTSM = TSM.withModuleDo([&](Module &M)
+                                           -> Expected<ThreadSafeModule> {
+    auto PromotedGlobals = PromoteSymbols(M);
+    if (!PromotedGlobals.empty()) {
+
+      MangleAndInterner Mangle(ES, M.getDataLayout());
+      SymbolFlagsMap SymbolFlags;
+      IRSymbolMapper::add(ES, *getManglingOptions(), PromotedGlobals,
+                          SymbolFlags);
+
+      if (auto Err = R->defineMaterializing(SymbolFlags))
+        return std::move(Err);
+    }
+
+    expandPartition(*GVsToExtract);
+
+    // Submodule name is given by hashing the names of the globals.
+    std::string SubModuleName;
+    {
+      std::vector<const GlobalValue *> HashGVs;
+      HashGVs.reserve(GVsToExtract->size());
+      for (const auto *GV : *GVsToExtract)
+        HashGVs.push_back(GV);
+      llvm::sort(HashGVs, [](const GlobalValue *LHS, const GlobalValue *RHS) {
+        return LHS->getName() < RHS->getName();
+      });
+      hash_code HC(0);
+      for (const auto *GV : HashGVs) {
+        assert(GV->hasName() && "All GVs to extract should be named by now");
+        auto GVName = GV->getName();
+        HC = hash_combine(HC, hash_combine_range(GVName.begin(), GVName.end()));
+      }
+      raw_string_ostream(SubModuleName)
+          << ".submodule."
+          << formatv(sizeof(size_t) == 8 ? "{0:x16}" : "{0:x8}",
+                     static_cast<size_t>(HC))
+          << ".ll";
+    }
+
+    // Extract the requested partiton (plus any necessary aliases) and
+    // put the rest back into the impl dylib.
+    auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
+      return GVsToExtract->count(&GV);
+    };
+
+    return extractSubModule(TSM, SubModuleName, ShouldExtract);
+  });
+
+  if (!ExtractedTSM) {
+    ES.reportError(ExtractedTSM.takeError());
+    R->failMaterialization();
+    return;
+  }
+
+  if (auto Err = R->replace(std::make_unique<PartitioningIRMaterializationUnit>(
+          ES, *getManglingOptions(), std::move(TSM), *this))) {
+    ES.reportError(std::move(Err));
+    R->failMaterialization();
+    return;
+  }
+  BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));
+}
diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index 18ab9e6b940fefe..15fc50bad94aabf 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -1288,9 +1288,12 @@ LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) {
     return;
   }
 
+  // Create the IP Layer.
+  IPLayer = std::make_unique<IRPartitionLayer>(*ES, *InitHelperTransformLayer);
+
   // Create the COD layer.
-  CODLayer = std::make_unique<CompileOnDemandLayer>(
-      *ES, *InitHelperTransformLayer, *LCTMgr, std::move(ISMBuilder));
+  CODLayer = std::make_unique<CompileOnDemandLayer>(*ES, *IPLayer, *LCTMgr,
+                                                    std::move(ISMBuilder));
 
   if (S.NumCompileThreads > 0)
     CODLayer->setCloneToNewContextOnEmit(true);
diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp
index 360ffa628a47b3f..1378622cb71511d 100644
--- a/llvm/tools/lli/lli.cpp
+++ b/llvm/tools/lli/lli.cpp
@@ -30,6 +30,7 @@
 #include "llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h"
 #include "llvm/ExecutionEngine/Orc/EPCGenericRTDyldMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
@@ -984,7 +985,7 @@ int runOrcJIT(const char *ProgName) {
   }
 
   if (PerModuleLazy)
-    J->setPartitionFunction(orc::CompileOnDemandLayer::compileWholeModule);
+    J->setPartitionFunction(orc::IRPartitionLayer::compileWholeModule);
 
   auto Dump = createDebugDumper();
 

>From 342adf162fc54a20e77b18ad5580622871657f8c Mon Sep 17 00:00:00 2001
From: Sunho Kim <ksunhokim123 at gmail.com>
Date: Wed, 20 Sep 2023 05:22:43 +0900
Subject: [PATCH 4/5] [ORC] Implement basic reoptimization.

---
 compiler-rt/lib/orc/common.h                  |   6 +-
 compiler-rt/lib/orc/elfnix_platform.cpp       |   1 +
 .../ExecutionEngine/Orc/ReOptimizeLayer.h     | 179 +++++++++++
 llvm/lib/ExecutionEngine/Orc/CMakeLists.txt   |   1 +
 .../ExecutionEngine/Orc/ReOptimizeLayer.cpp   | 278 ++++++++++++++++++
 .../ExecutionEngine/Orc/CMakeLists.txt        |   1 +
 .../Orc/ReOptimizeLayerTest.cpp               | 150 ++++++++++
 7 files changed, 613 insertions(+), 3 deletions(-)
 create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
 create mode 100644 llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
 create mode 100644 llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp

diff --git a/compiler-rt/lib/orc/common.h b/compiler-rt/lib/orc/common.h
index 73c5c4a2bd8d474..f34229a615341df 100644
--- a/compiler-rt/lib/orc/common.h
+++ b/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 {};
diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index c087e71038f9504..f76a070240315fe 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/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)
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
new file mode 100644
index 000000000000000..df2f005825b5a7d
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
@@ -0,0 +1,179 @@
+//===- ReOptimizeLayer.h - Re-optimization layer interface ------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Re-optimization layer interface.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_EXECUTIONENGINE_ORC_REOPTLAYER_H
+#define LLVM_EXECUTIONENGINE_ORC_REOPTLAYER_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+
+namespace llvm {
+namespace orc {
+
+using ReOptMaterializationUnitID = uint64_t;
+
+class ReOptimizeLayer : public IRLayer, public ResourceManager {
+public:
+  /// AddProfilerFunc will be called when ReOptimizeLayer emits the first
+  /// version of a materialization unit in order to inject profiling code and
+  /// reoptimization request code.
+  using AddProfilerFunc = unique_function<Error(
+      ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
+      unsigned CurVersion, ThreadSafeModule &TSM)>;
+
+  /// ReOptimizeFunc will be called when ReOptimizeLayer reoptimization of a
+  /// materialization unit was requested in order to reoptimize the IR module
+  /// based on profile data. OldRT is the ResourceTracker that tracks the old
+  /// function definitions. The OldRT must be kept alive until it can be
+  /// guaranteed that every invocation of the old function definitions has been
+  /// terminated.
+  using ReOptimizeFunc = unique_function<Error(
+      ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
+      unsigned CurVersion, ResourceTrackerSP OldRT, ThreadSafeModule &TSM)>;
+
+  ReOptimizeLayer(ExecutionSession &ES, IRLayer &BaseLayer,
+                  RedirectableSymbolManager &RM)
+      : IRLayer(ES, BaseLayer.getManglingOptions()), ES(ES),
+        BaseLayer(BaseLayer), RSManager(RM), ReOptFunc(identity),
+        ProfilerFunc(reoptimizeIfCallFrequent) {}
+
+  void setReoptimizeFunc(ReOptimizeFunc ReOptFunc) {
+    this->ReOptFunc = std::move(ReOptFunc);
+  }
+
+  void setAddProfilerFunc(AddProfilerFunc ProfilerFunc) {
+    this->ProfilerFunc = std::move(ProfilerFunc);
+  }
+
+  /// Registers reoptimize runtime dispatch handlers to given PlatformJD. The
+  /// reoptimization request will not be handled if dispatch handler is not
+  /// registered by using this function.
+  Error reigsterRuntimeFunctions(JITDylib &PlatformJD);
+
+  /// 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;
+
+  static const uint64_t CallCountThreshold = 10;
+
+  /// Basic AddProfilerFunc that reoptimizes the function when the call count
+  /// exceeds CallCountThreshold.
+  static Error reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
+                                        ReOptMaterializationUnitID MUID,
+                                        unsigned CurVersion,
+                                        ThreadSafeModule &TSM);
+
+  static Error identity(ReOptimizeLayer &Parent,
+                        ReOptMaterializationUnitID MUID, unsigned CurVersion,
+                        ResourceTrackerSP OldRT, ThreadSafeModule &TSM) {
+    return Error::success();
+  }
+
+  // Create IR reoptimize request fucntion call.
+  static void createReoptimizeCall(Module &M, Instruction &IP,
+                                   GlobalVariable *ArgBuffer);
+
+  Error handleRemoveResources(JITDylib &JD, ResourceKey K) override;
+  void handleTransferResources(JITDylib &JD, ResourceKey DstK,
+                               ResourceKey SrcK) override;
+
+private:
+  class ReOptMaterializationUnitState {
+  public:
+    ReOptMaterializationUnitState() = default;
+    ReOptMaterializationUnitState(ReOptMaterializationUnitID ID,
+                                  ThreadSafeModule TSM)
+        : ID(ID), TSM(std::move(TSM)) {}
+    ReOptMaterializationUnitState(ReOptMaterializationUnitState &&Other)
+        : ID(Other.ID), TSM(std::move(Other.TSM)), RT(std::move(Other.RT)),
+          Reoptimizing(std::move(Other.Reoptimizing)),
+          CurVersion(Other.CurVersion) {}
+
+    ReOptMaterializationUnitID getID() { return ID; }
+
+    const ThreadSafeModule &getThreadSafeModule() { return TSM; }
+
+    ResourceTrackerSP getResourceTracker() {
+      std::unique_lock<std::mutex> Lock(Mutex);
+      return RT;
+    }
+
+    void setResourceTracker(ResourceTrackerSP RT) {
+      std::unique_lock<std::mutex> Lock(Mutex);
+      this->RT = RT;
+    }
+
+    uint32_t getCurVersion() {
+      std::unique_lock<std::mutex> Lock(Mutex);
+      return CurVersion;
+    }
+
+    bool tryStartReoptimize();
+    void reoptimizeSucceeded();
+    void reoptimizeFailed();
+
+  private:
+    std::mutex Mutex;
+    ReOptMaterializationUnitID ID;
+    ThreadSafeModule TSM;
+    ResourceTrackerSP RT;
+    bool Reoptimizing = false;
+    uint32_t CurVersion = 0;
+  };
+
+  using SPSReoptimizeArgList =
+      shared::SPSArgList<ReOptMaterializationUnitID, uint32_t>;
+  using SendErrorFn = unique_function<void(Error)>;
+
+  Expected<SymbolMap> emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
+                                        uint32_t Version, JITDylib &JD,
+                                        ThreadSafeModule TSM);
+
+  void rt_reoptimize(SendErrorFn SendResult, ReOptMaterializationUnitID MUID,
+                     uint32_t CurVersion);
+
+  static Expected<Constant *>
+  createReoptimizeArgBuffer(Module &M, ReOptMaterializationUnitID MUID,
+                            uint32_t CurVersion);
+
+  ReOptMaterializationUnitState &
+  createMaterializationUnitState(const ThreadSafeModule &TSM);
+
+  void
+  registerMaterializationUnitResource(ResourceKey Key,
+                                      ReOptMaterializationUnitState &State);
+
+  ReOptMaterializationUnitState &
+  getMaterializationUnitState(ReOptMaterializationUnitID MUID);
+
+  ExecutionSession &ES;
+  IRLayer &BaseLayer;
+  RedirectableSymbolManager &RSManager;
+
+  ReOptimizeFunc ReOptFunc;
+  AddProfilerFunc ProfilerFunc;
+
+  std::mutex Mutex;
+  std::map<ReOptMaterializationUnitID, ReOptMaterializationUnitState> MUStates;
+  DenseMap<ResourceKey, DenseSet<ReOptMaterializationUnitID>> MUResources;
+  ReOptMaterializationUnitID NextID = 1;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif
diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index 379c3cc240e8c7b..550a36e7afc4771 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
@@ -52,6 +52,7 @@ add_llvm_component_library(LLVMOrcJIT
   ThreadSafeModule.cpp
   RedirectionManager.cpp
   JITLinkRedirectableSymbolManager.cpp
+  ReOptimizeLayer.cpp
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc
 
diff --git a/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
new file mode 100644
index 000000000000000..2ec82f210e6f625
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
@@ -0,0 +1,278 @@
+#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
+
+using namespace llvm;
+using namespace orc;
+
+bool ReOptimizeLayer::ReOptMaterializationUnitState::tryStartReoptimize() {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  if (Reoptimizing)
+    return false;
+
+  Reoptimizing = true;
+  return true;
+}
+
+void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeSucceeded() {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
+  Reoptimizing = false;
+  CurVersion++;
+}
+
+void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  assert(Reoptimizing && "Tried to mark unstarted reoptimization as done");
+  Reoptimizing = false;
+}
+
+Error ReOptimizeLayer::reigsterRuntimeFunctions(JITDylib &PlatformJD) {
+  ExecutionSession::JITDispatchHandlerAssociationMap WFs;
+  using ReoptimizeSPSSig = shared::SPSError(uint64_t, uint32_t);
+  WFs[ES.intern("__orc_rt_reoptimize_tag")] =
+      ES.wrapAsyncWithSPS<ReoptimizeSPSSig>(this,
+                                            &ReOptimizeLayer::rt_reoptimize);
+  return ES.registerJITDispatchHandlers(PlatformJD, std::move(WFs));
+}
+
+void ReOptimizeLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
+                           ThreadSafeModule TSM) {
+  auto &JD = R->getTargetJITDylib();
+
+  bool HasNonCallable = false;
+  for (auto &KV : R->getSymbols()) {
+    auto &Flags = KV.second;
+    if (!Flags.isCallable())
+      HasNonCallable = true;
+  }
+
+  if (HasNonCallable) {
+    BaseLayer.emit(std::move(R), std::move(TSM));
+    return;
+  }
+
+  auto &MUState = createMaterializationUnitState(TSM);
+
+  if (auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
+        registerMaterializationUnitResource(Key, MUState);
+      })) {
+    ES.reportError(std::move(Err));
+    R->failMaterialization();
+    return;
+  }
+
+  if (auto Err =
+          ProfilerFunc(*this, MUState.getID(), MUState.getCurVersion(), TSM)) {
+    ES.reportError(std::move(Err));
+    R->failMaterialization();
+    return;
+  }
+
+  auto InitialDests =
+      emitMUImplSymbols(MUState, MUState.getCurVersion(), JD, std::move(TSM));
+  if (!InitialDests) {
+    ES.reportError(InitialDests.takeError());
+    R->failMaterialization();
+    return;
+  }
+
+  RSManager.emitRedirectableSymbols(std::move(R), std::move(*InitialDests));
+}
+
+Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
+                                                ReOptMaterializationUnitID MUID,
+                                                unsigned CurVersion,
+                                                ThreadSafeModule &TSM) {
+  return TSM.withModuleDo([&](Module &M) -> Error {
+    Type *I64Ty = Type::getInt64Ty(M.getContext());
+    GlobalVariable *Counter = new GlobalVariable(
+        M, I64Ty, false, GlobalValue::InternalLinkage,
+        Constant::getNullValue(I64Ty), "__orc_reopt_counter");
+    auto ArgBufferConst = createReoptimizeArgBuffer(M, MUID, CurVersion);
+    if (auto Err = ArgBufferConst.takeError())
+      return Err;
+    GlobalVariable *ArgBuffer =
+        new GlobalVariable(M, (*ArgBufferConst)->getType(), true,
+                           GlobalValue::InternalLinkage, (*ArgBufferConst));
+    for (auto &F : M) {
+      if (F.isDeclaration())
+        continue;
+      auto &BB = F.getEntryBlock();
+      auto *IP = &*BB.getFirstInsertionPt();
+      IRBuilder<> IRB(IP);
+      Value *Threshold = ConstantInt::get(I64Ty, CallCountThreshold, true);
+      Value *Cnt = IRB.CreateLoad(I64Ty, Counter);
+      // Use EQ to prevent further reoptimize calls.
+      Value *Cmp = IRB.CreateICmpEQ(Cnt, Threshold);
+      Value *Added = IRB.CreateAdd(Cnt, ConstantInt::get(I64Ty, 1));
+      (void)IRB.CreateStore(Added, Counter);
+      Instruction *SplitTerminator = SplitBlockAndInsertIfThen(Cmp, IP, false);
+      createReoptimizeCall(M, *SplitTerminator, ArgBuffer);
+    }
+    return Error::success();
+  });
+}
+
+Expected<SymbolMap>
+ReOptimizeLayer::emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
+                                   uint32_t Version, JITDylib &JD,
+                                   ThreadSafeModule TSM) {
+  DenseMap<SymbolStringPtr, SymbolStringPtr> RenamedMap;
+  cantFail(TSM.withModuleDo([&](Module &M) -> Error {
+    MangleAndInterner Mangle(ES, M.getDataLayout());
+    for (auto &F : M)
+      if (!F.isDeclaration()) {
+        std::string NewName =
+            (F.getName() + ".__def__." + Twine(Version)).str();
+        RenamedMap[Mangle(F.getName())] = Mangle(NewName);
+        F.setName(NewName);
+      }
+    return Error::success();
+  }));
+
+  auto RT = JD.createResourceTracker();
+  if (auto Err =
+          JD.define(std::make_unique<BasicIRLayerMaterializationUnit>(
+                        BaseLayer, *getManglingOptions(), std::move(TSM)),
+                    RT))
+    return Err;
+  MUState.setResourceTracker(RT);
+
+  SymbolLookupSet LookupSymbols;
+  for (auto [K, V] : RenamedMap)
+    LookupSymbols.add(V);
+
+  auto ImplSymbols =
+      ES.lookup({{&JD, JITDylibLookupFlags::MatchAllSymbols}}, LookupSymbols,
+                LookupKind::Static, SymbolState::Resolved);
+  if (auto Err = ImplSymbols.takeError())
+    return Err;
+
+  SymbolMap Result;
+  for (auto [K, V] : RenamedMap)
+    Result[K] = (*ImplSymbols)[V];
+
+  return Result;
+}
+
+void ReOptimizeLayer::rt_reoptimize(SendErrorFn SendResult,
+                                    ReOptMaterializationUnitID MUID,
+                                    uint32_t CurVersion) {
+  auto &MUState = getMaterializationUnitState(MUID);
+  if (CurVersion < MUState.getCurVersion() || !MUState.tryStartReoptimize()) {
+    SendResult(Error::success());
+    return;
+  }
+
+  ThreadSafeModule TSM = cloneToNewContext(MUState.getThreadSafeModule());
+  auto OldRT = MUState.getResourceTracker();
+  auto &JD = OldRT->getJITDylib();
+
+  if (auto Err = ReOptFunc(*this, MUID, CurVersion + 1, OldRT, TSM)) {
+    ES.reportError(std::move(Err));
+    MUState.reoptimizeFailed();
+    SendResult(Error::success());
+    return;
+  }
+
+  auto SymbolDests =
+      emitMUImplSymbols(MUState, CurVersion + 1, JD, std::move(TSM));
+  if (!SymbolDests) {
+    ES.reportError(SymbolDests.takeError());
+    MUState.reoptimizeFailed();
+    SendResult(Error::success());
+    return;
+  }
+
+  if (auto Err = RSManager.redirect(JD, std::move(*SymbolDests))) {
+    ES.reportError(std::move(Err));
+    MUState.reoptimizeFailed();
+    SendResult(Error::success());
+    return;
+  }
+
+  MUState.reoptimizeSucceeded();
+  SendResult(Error::success());
+}
+
+Expected<Constant *> ReOptimizeLayer::createReoptimizeArgBuffer(
+    Module &M, ReOptMaterializationUnitID MUID, uint32_t CurVersion) {
+  size_t ArgBufferSize = SPSReoptimizeArgList::size(MUID, CurVersion);
+  std::vector<char> ArgBuffer(ArgBufferSize);
+  shared::SPSOutputBuffer OB(ArgBuffer.data(), ArgBuffer.size());
+  if (!SPSReoptimizeArgList::serialize(OB, MUID, CurVersion))
+    return make_error<StringError>("Could not serealize args list",
+                                   inconvertibleErrorCode());
+  return ConstantDataArray::get(M.getContext(), ArrayRef(ArgBuffer));
+}
+
+void ReOptimizeLayer::createReoptimizeCall(Module &M, Instruction &IP,
+                                           GlobalVariable *ArgBuffer) {
+  GlobalVariable *DispatchCtx =
+      M.getGlobalVariable("__orc_rt_jit_dispatch_ctx");
+  if (!DispatchCtx)
+    DispatchCtx = new GlobalVariable(M, Type::getInt8PtrTy(M.getContext()),
+                                     false, GlobalValue::ExternalLinkage,
+                                     nullptr, "__orc_rt_jit_dispatch_ctx");
+  GlobalVariable *ReoptimizeTag =
+      M.getGlobalVariable("__orc_rt_reoptimize_tag");
+  if (!ReoptimizeTag)
+    ReoptimizeTag = new GlobalVariable(M, Type::getInt8PtrTy(M.getContext()),
+                                       false, GlobalValue::ExternalLinkage,
+                                       nullptr, "__orc_rt_reoptimize_tag");
+  Function *DispatchFunc = M.getFunction("__orc_rt_jit_dispatch");
+  if (!DispatchFunc) {
+    std::vector<Type *> Args = {Type::getInt8PtrTy(M.getContext()),
+                                Type::getInt8PtrTy(M.getContext()),
+                                Type::getInt8PtrTy(M.getContext()),
+                                IntegerType::get(M.getContext(), 64)};
+    FunctionType *FuncTy =
+        FunctionType::get(Type::getVoidTy(M.getContext()), Args, false);
+    DispatchFunc = Function::Create(FuncTy, GlobalValue::ExternalLinkage,
+                                    "__orc_rt_jit_dispatch", &M);
+  }
+  size_t ArgBufferSizeConst =
+      SPSReoptimizeArgList::size(ReOptMaterializationUnitID{}, uint32_t{});
+  Constant *ArgBufferSize = ConstantInt::get(
+      IntegerType::get(M.getContext(), 64), ArgBufferSizeConst, false);
+  IRBuilder<> IRB(&IP);
+  (void)IRB.CreateCall(DispatchFunc,
+                       {DispatchCtx, ReoptimizeTag, ArgBuffer, ArgBufferSize});
+}
+
+ReOptimizeLayer::ReOptMaterializationUnitState &
+ReOptimizeLayer::createMaterializationUnitState(const ThreadSafeModule &TSM) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  ReOptMaterializationUnitID MUID = NextID;
+  MUStates.emplace(MUID,
+                   ReOptMaterializationUnitState(MUID, cloneToNewContext(TSM)));
+  ++NextID;
+  return MUStates.at(MUID);
+}
+
+ReOptimizeLayer::ReOptMaterializationUnitState &
+ReOptimizeLayer::getMaterializationUnitState(ReOptMaterializationUnitID MUID) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  return MUStates.at(MUID);
+}
+
+void ReOptimizeLayer::registerMaterializationUnitResource(
+    ResourceKey Key, ReOptMaterializationUnitState &State) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  MUResources[Key].insert(State.getID());
+}
+
+Error ReOptimizeLayer::handleRemoveResources(JITDylib &JD, ResourceKey K) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  for (auto MUID : MUResources[K])
+    MUStates.erase(MUID);
+
+  MUResources.erase(K);
+  return Error::success();
+}
+
+void ReOptimizeLayer::handleTransferResources(JITDylib &JD, ResourceKey DstK,
+                                              ResourceKey SrcK) {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  MUResources[DstK].insert(MUResources[SrcK].begin(), MUResources[SrcK].end());
+  MUResources.erase(SrcK);
+}
diff --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
index e932ff94361a7b7..310cb8276d894a6 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
@@ -41,6 +41,7 @@ add_llvm_unittest(OrcJITTests
   ThreadSafeModuleTest.cpp
   WrapperFunctionUtilsTest.cpp
   JITLinkRedirectionManagerTest.cpp
+  ReOptimizeLayerTest.cpp
   )
 
 target_link_libraries(OrcJITTests PRIVATE
diff --git a/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
new file mode 100644
index 000000000000000..c7e795c749ad000
--- /dev/null
+++ b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
@@ -0,0 +1,150 @@
+#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
+#include "OrcTestCommon.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/CompileUtils.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/JITLinkRedirectableSymbolManager.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/Support/CodeGen.h"
+#include "llvm/TargetParser/Host.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::jitlink;
+
+class ReOptimizeLayerTest : public testing::Test {
+public:
+  ~ReOptimizeLayerTest() {
+    if (ES)
+      if (auto Err = ES->endSession())
+        ES->reportError(std::move(Err));
+  }
+
+protected:
+  void SetUp() override {
+    auto JTMB = JITTargetMachineBuilder::detectHost();
+    // Bail out if we can not detect the host.
+    if (!JTMB) {
+      consumeError(JTMB.takeError());
+      GTEST_SKIP();
+    }
+
+    auto EPC = SelfExecutorProcessControl::Create();
+    if (!EPC) {
+      consumeError(EPC.takeError());
+      GTEST_SKIP();
+    }
+    ES = std::make_unique<ExecutionSession>(std::move(*EPC));
+    JD = &ES->createBareJITDylib("main");
+    ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
+        *ES, std::make_unique<InProcessMemoryManager>(4096));
+    DL = std::make_unique<DataLayout>(
+        cantFail(JTMB->getDefaultDataLayoutForTarget()));
+
+    auto TM = JTMB->createTargetMachine();
+    if (!TM) {
+      consumeError(TM.takeError());
+      GTEST_SKIP();
+    }
+    auto CompileFunction =
+        std::make_unique<TMOwningSimpleCompiler>(std::move(*TM));
+    CompileLayer = std::make_unique<IRCompileLayer>(*ES, *ObjLinkingLayer,
+                                                    std::move(CompileFunction));
+  }
+
+  Error addIRModule(ResourceTrackerSP RT, ThreadSafeModule TSM) {
+    assert(TSM && "Can not add null module");
+
+    TSM.withModuleDo([&](Module &M) { M.setDataLayout(*DL); });
+
+    return ROLayer->add(std::move(RT), std::move(TSM));
+  }
+
+  JITDylib *JD{nullptr};
+  std::unique_ptr<ExecutionSession> ES;
+  std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer;
+  std::unique_ptr<IRCompileLayer> CompileLayer;
+  std::unique_ptr<ReOptimizeLayer> ROLayer;
+  std::unique_ptr<DataLayout> DL;
+};
+
+static Function *createRetFunction(Module *M, StringRef Name,
+                                   uint32_t ReturnCode) {
+  Function *Result = Function::Create(
+      FunctionType::get(Type::getInt32Ty(M->getContext()), {}, false),
+      GlobalValue::ExternalLinkage, Name, M);
+
+  BasicBlock *BB = BasicBlock::Create(M->getContext(), Name, Result);
+  IRBuilder<> Builder(M->getContext());
+  Builder.SetInsertPoint(BB);
+
+  Value *RetValue = ConstantInt::get(M->getContext(), APInt(32, ReturnCode));
+  Builder.CreateRet(RetValue);
+  return Result;
+}
+
+TEST_F(ReOptimizeLayerTest, BasicReOptimization) {
+  auto &EPC = ES->getExecutorProcessControl();
+  EXPECT_THAT_ERROR(JD->define(absoluteSymbols(
+                        {{ES->intern("__orc_rt_jit_dispatch"),
+                          {EPC.getJITDispatchInfo().JITDispatchFunction,
+                           JITSymbolFlags::Exported}},
+                         {ES->intern("__orc_rt_jit_dispatch_ctx"),
+                          {EPC.getJITDispatchInfo().JITDispatchContext,
+                           JITSymbolFlags::Exported}},
+                         {ES->intern("__orc_rt_reoptimize_tag"),
+                          {ExecutorAddr(), JITSymbolFlags::Exported}}})),
+                    Succeeded());
+
+  auto RM =
+      JITLinkRedirectableSymbolManager::Create(*ES, *ObjLinkingLayer, *JD);
+  EXPECT_THAT_ERROR(RM.takeError(), Succeeded());
+
+  ROLayer = std::make_unique<ReOptimizeLayer>(*ES, *CompileLayer, **RM);
+  ROLayer->setReoptimizeFunc(
+      [&](ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
+          unsigned CurVerison, ResourceTrackerSP OldRT, ThreadSafeModule &TSM) {
+        TSM.withModuleDo([&](Module &M) {
+          for (auto &F : M) {
+            if (F.isDeclaration())
+              continue;
+            for (auto &B : F) {
+              for (auto &I : B) {
+                if (ReturnInst *Ret = dyn_cast<ReturnInst>(&I)) {
+                  Value *RetValue =
+                      ConstantInt::get(M.getContext(), APInt(32, 53));
+                  Ret->setOperand(0, RetValue);
+                }
+              }
+            }
+          }
+        });
+        return Error::success();
+      });
+  EXPECT_THAT_ERROR(ROLayer->reigsterRuntimeFunctions(*JD), Succeeded());
+
+  ThreadSafeContext Ctx(std::make_unique<LLVMContext>());
+  auto M = std::make_unique<Module>("<main>", *Ctx.getContext());
+  M->setTargetTriple(sys::getProcessTriple());
+
+  (void)createRetFunction(M.get(), "main", 42);
+
+  EXPECT_THAT_ERROR(addIRModule(JD->getDefaultResourceTracker(),
+                                ThreadSafeModule(std::move(M), std::move(Ctx))),
+                    Succeeded());
+
+  auto Result = cantFail(ES->lookup({JD}, "main"));
+  auto FuncPtr = Result.getAddress().toPtr<int (*)()>();
+  for (size_t I = 0; I <= ReOptimizeLayer::CallCountThreshold; I++)
+    EXPECT_EQ(FuncPtr(), 42);
+  EXPECT_EQ(FuncPtr(), 53);
+}

>From 4e3adab12b509610d81502bb640accbaea39b9f9 Mon Sep 17 00:00:00 2001
From: Sunho Kim <ksunhokim123 at gmail.com>
Date: Sat, 23 Sep 2023 20:50:09 -0700
Subject: [PATCH 5/5] DDD

---
 clang/include/clang/Interpreter/Interpreter.h |   3 +-
 clang/lib/Interpreter/IncrementalExecutor.cpp | 235 +++++++++++++++++-
 clang/lib/Interpreter/IncrementalExecutor.h   |   8 +-
 clang/lib/Interpreter/Interpreter.cpp         |   2 +-
 compiler-rt/lib/orc/common.h                  |   2 +-
 compiler-rt/lib/orc/elfnix_platform.cpp       |  40 ++-
 compiler-rt/lib/orc/elfnix_platform.h         |   3 +
 llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h |   8 +
 .../ExecutionEngine/Orc/ReOptimizeLayer.h     |  22 +-
 .../Orc/CompileOnDemandLayer.cpp              |   2 +-
 llvm/lib/ExecutionEngine/Orc/Core.cpp         |   7 +-
 .../ExecutionEngine/Orc/ELFNixPlatform.cpp    |  12 +-
 llvm/lib/ExecutionEngine/Orc/LLJIT.cpp        |  21 +-
 .../ExecutionEngine/Orc/ReOptimizeLayer.cpp   | 135 +++++++---
 14 files changed, 437 insertions(+), 63 deletions(-)

diff --git a/clang/include/clang/Interpreter/Interpreter.h b/clang/include/clang/Interpreter/Interpreter.h
index 43573fb1a4b8915..5b0b59ace49364a 100644
--- a/clang/include/clang/Interpreter/Interpreter.h
+++ b/clang/include/clang/Interpreter/Interpreter.h
@@ -29,6 +29,7 @@
 namespace llvm {
 namespace orc {
 class LLJIT;
+class LLLazyJIT;
 class ThreadSafeContext;
 } // namespace orc
 } // namespace llvm
@@ -101,7 +102,7 @@ class Interpreter {
   const ASTContext &getASTContext() const;
   ASTContext &getASTContext();
   const CompilerInstance *getCompilerInstance() const;
-  llvm::Expected<llvm::orc::LLJIT &> getExecutionEngine();
+  llvm::Expected<llvm::orc::LLLazyJIT&> getExecutionEngine();
 
   llvm::Expected<PartialTranslationUnit &> Parse(llvm::StringRef Code);
   llvm::Error Execute(PartialTranslationUnit &T);
diff --git a/clang/lib/Interpreter/IncrementalExecutor.cpp b/clang/lib/Interpreter/IncrementalExecutor.cpp
index 2c4dfc9a611e021..9fa76e33d8373b5 100644
--- a/clang/lib/Interpreter/IncrementalExecutor.cpp
+++ b/clang/lib/Interpreter/IncrementalExecutor.cpp
@@ -11,9 +11,25 @@
 //===----------------------------------------------------------------------===//
 
 #include "IncrementalExecutor.h"
+		#include "clang/AST/ASTContext.h"
+#include "clang/Basic/TargetInfo.h"
+#include "clang/CodeGen/ModuleBuilder.h"
+#include "clang/CodeGen/ObjectFilePCHContainerOperations.h"
+#include "clang/Driver/Compilation.h"
+#include "clang/Driver/Driver.h"
+#include "clang/Driver/Job.h"
+#include "clang/Driver/Options.h"
+#include "clang/Driver/Tool.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/TextDiagnosticBuffer.h"
+#include "clang/Lex/PreprocessorOptions.h"
 
+#include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/Bitcode/BitcodeWriter.h"
+#include "llvm/Linker/Linker.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Basic/TargetOptions.h"
+#include "clang/Driver/Driver.h"
 #include "clang/Interpreter/PartialTranslationUnit.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
@@ -26,6 +42,22 @@
 #include "llvm/IR/Module.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/Host.h"
+
+#include "llvm/Analysis/CallGraph.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/StandardInstrumentations.h"
+#include "llvm/Analysis/CallGraphSCCPass.h"
+#include "llvm/Analysis/LoopAnalysisManager.h"
+#include "llvm/Analysis/CGSCCPassManager.h"
+#include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/RegionPass.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
+#include "llvm/IR/PassManager.h"
+
+using namespace llvm;
 
 // Force linking some of the runtimes that helps attaching to a debugger.
 LLVM_ATTRIBUTE_USED void linkComponents() {
@@ -35,6 +67,110 @@ LLVM_ATTRIBUTE_USED void linkComponents() {
 
 namespace clang {
 
+static std::string buildOrcRTBasename(const llvm::Triple &TT, bool AddArch) {
+  bool IsITANMSVCWindows =
+      TT.isWindowsMSVCEnvironment() || TT.isWindowsItaniumEnvironment();
+  IsITANMSVCWindows = false;
+  const char *Prefix = IsITANMSVCWindows ? "" : "lib";
+  const char *Suffix = IsITANMSVCWindows ? ".lib" : ".a";
+  std::string ArchAndEnv;
+  if (AddArch)
+    ArchAndEnv = ("-" + llvm::Triple::getArchTypeName(TT.getArch())).str();
+  return (Prefix + Twine("orc_rt") + ArchAndEnv + Suffix).str();
+}
+
+static std::string findOrcRuntimePath(const std::vector<const char *> &ClangArgv) {
+  IntrusiveRefCntPtr<DiagnosticIDs> DiagID(new DiagnosticIDs());
+  IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts;
+  TextDiagnosticBuffer *DiagsBuffer = new TextDiagnosticBuffer;
+  DiagnosticsEngine Diags(DiagID, &*DiagOpts, DiagsBuffer);
+
+  driver::Driver Driver(/*MainBinaryName=*/ClangArgv[0],
+                        llvm::sys::getProcessTriple(), Diags);
+  Driver.setCheckInputsExist(false);
+  llvm::ArrayRef<const char *> RF = llvm::makeArrayRef(ClangArgv);
+  std::unique_ptr<driver::Compilation> Compilation(Driver.BuildCompilation(RF));
+
+  auto RuntimePaths = *Compilation->getDefaultToolChain().getRuntimePath();
+  auto &TC = Compilation->getDefaultToolChain();
+  auto TT = Compilation->getDefaultToolChain().getTriple();
+  for (const auto &LibPath : TC.getLibraryPaths()) {
+    SmallString<128> P(LibPath);
+    llvm::sys::path::append(P, buildOrcRTBasename(TT, false));
+    if (Driver.getVFS().exists(P))
+      return std::string(P.str());
+  }
+
+  SmallString<256> Path(RuntimePaths);
+  llvm::sys::path::append(Path, buildOrcRTBasename(TT, true));
+  if (Driver.getVFS().exists(Path))
+    return Path.str().str();
+
+  return "";
+}
+
+static void Optimize(TargetMachine* TM, Triple TargetTriple, llvm::Module& M, StringRef PassPipeline) {
+  LoopAnalysisManager LAM;
+  FunctionAnalysisManager FAM;
+  CGSCCAnalysisManager CGAM;
+
+  ModuleAnalysisManager MAM;
+
+  PassInstrumentationCallbacks PIC;
+  PrintPassOptions PrintPassOpts;
+  PrintPassOpts.Verbose = false;
+  PrintPassOpts.SkipAnalyses = false;
+  StandardInstrumentations SI(M.getContext(), false,
+                              false, PrintPassOpts);
+  SI.registerCallbacks(PIC, &MAM);
+
+  PipelineTuningOptions PTO;
+  // LoopUnrolling defaults on to true and DisableLoopUnrolling is initialized
+  // to false above so we shouldn't necessarily need to check whether or not the
+  // option has been enabled.
+  PTO.LoopUnrolling = false;
+  PTO.UnifiedLTO = false;
+  PassBuilder PB;
+
+  // Register all the basic analyses with the managers.
+  PB.registerModuleAnalyses(MAM);
+  PB.registerCGSCCAnalyses(CGAM);
+  PB.registerFunctionAnalyses(FAM);
+  PB.registerLoopAnalyses(LAM);
+  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
+
+
+  auto MPM = PB.buildPerModuleDefaultPipeline(OptimizationLevel::O2);
+  //MPM.printPipeline(dbgs(), [](StringRef x){return x;});
+  MPM.run(M, MAM);
+}
+
+std::unique_ptr<llvm::Module> CloneModuleToContext(llvm::Module& Src, LLVMContext& Ctx) {
+    SmallVector<char, 1> ClonedModuleBuffer;
+
+      std::set<GlobalValue *> ClonedDefsInSrc;
+      ValueToValueMapTy VMap;
+      auto Tmp = CloneModule(Src, VMap);
+
+      BitcodeWriter BCWriter(ClonedModuleBuffer);
+
+      BCWriter.writeModule(*Tmp);
+      BCWriter.writeSymtab();
+      BCWriter.writeStrtab();
+
+    MemoryBufferRef ClonedModuleBufferRef(
+        StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
+        "cloned module buffer");
+
+    auto ClonedModule = cantFail(
+        parseBitcodeFile(ClonedModuleBufferRef, Ctx));
+    ClonedModule->setModuleIdentifier(Src.getName());
+    return ClonedModule;
+}
+
+auto EPC = cantFail(llvm::orc::SelfExecutorProcessControl::Create(
+    std::make_shared<llvm::orc::SymbolStringPool>()));
+
 IncrementalExecutor::IncrementalExecutor(llvm::orc::ThreadSafeContext &TSC,
                                          llvm::Error &Err,
                                          const clang::TargetInfo &TI)
@@ -44,17 +180,112 @@ IncrementalExecutor::IncrementalExecutor(llvm::orc::ThreadSafeContext &TSC,
 
   auto JTMB = JITTargetMachineBuilder(TI.getTriple());
   JTMB.addFeatures(TI.getTargetOpts().Features);
-  LLJITBuilder Builder;
+  LLLazyJITBuilder Builder;
   Builder.setJITTargetMachineBuilder(JTMB);
   // Enable debugging of JIT'd code (only works on JITLink for ELF and MachO).
   Builder.setEnableDebuggerSupport(true);
 
+  Builder.setObjectLinkingLayerCreator([&](llvm::orc::ExecutionSession &ES,
+                                                  const llvm::Triple &TT) {
+    auto L = std::make_unique<llvm::orc::ObjectLinkingLayer>(ES, ES.getExecutorProcessControl().getMemMgr());
+    return L;
+  });
+
+  Builder.setPlatformSetUp(llvm::orc::ExecutorNativePlatform("/home/sunho/dev/llvm-project/build/lib/clang/18/lib/x86_64-unknown-linux-gnu/liborc_rt.a"));
+
   if (auto JitOrErr = Builder.create())
     Jit = std::move(*JitOrErr);
   else {
     Err = JitOrErr.takeError();
     return;
   }
+
+  Jit->getReOptimizeLayer().setReoptimizeFunc(
+      [&](ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
+          unsigned CurVerison, ResourceTrackerSP OldRT, const std::vector<std::pair<uint32_t,uint64_t>>& Profile, ThreadSafeModule &TSM) {
+        TSM.withModuleDo([&](llvm::Module &M) {
+          dbgs() << "Optimizing ---------------" << "\n";
+          dbgs() << "before: " << "\n";
+          dbgs() << M << "\n";
+
+          std::set<uint64_t> ToLink;
+
+          DenseMap<uint32_t, std::vector<StringRef>> ProfileData;
+    
+          for (auto [CID, F] : Profile) {
+            if (Parent.FuncAddrToMU.count(ExecutorAddr(F))) {
+              auto [MUID, Name] = Parent.FuncAddrToMU[ExecutorAddr(F)];
+              ToLink.insert(MUID);
+              ProfileData[CID].push_back(Name);
+            } else {
+              dbgs() << F << "\n";
+              assert(false);
+            }
+          }
+
+          for (auto MUID : ToLink) {
+            auto& State = Parent.getMaterializationUnitState(MUID);
+            State.getThreadSafeModule().withModuleDo([&](llvm::Module& NM) {
+              auto NNM = CloneModuleToContext(NM, M.getContext());
+              for (auto& F : *NNM) {
+                if (F.isDeclaration()) continue;
+                F.setVisibility(GlobalValue::HiddenVisibility);
+              }
+              Linker::linkModules(M, std::move(NNM));
+            });
+          }
+
+          for (auto& F : M) {
+            if (F.isDeclaration()) continue;
+            for (auto& B : F) {
+              std::vector<CallInst*> Insts;
+              for (auto& I : B) {
+                if (auto* Call = dyn_cast<CallInst>(&I)) {
+                  if (Call->isIndirectCall()) {
+                    Insts.push_back(Call);
+                  }
+                }
+              }
+              for (auto* Call : Insts) {
+                IRBuilder<> IRB(Call);
+                auto* a = Call->getMetadata("call_id");
+                if (!a) continue;
+                auto* VAM = cast<ValueAsMetadata>(cast<MDNode>(a)->getOperand(0));
+                int CallID = cast<ConstantInt>(VAM->getValue())->getSExtValue();
+                std::vector<std::pair<BasicBlock*, Value*>> Dones;
+                Instruction* IP = Call;
+                std::vector<Value*> Args(Call->arg_begin(), Call->arg_end());
+                for (auto Name : ProfileData[CallID]) {
+                  Value *Cmp = IRB.CreateICmpEQ(Call->getCalledOperand(), M.getFunction(Name));
+                  Instruction *IfPart, *ElsePart;
+                  SplitBlockAndInsertIfThenElse(Cmp, IP, &IfPart, &ElsePart);
+                  IRBuilder<> Builder(IfPart);
+                  CallInst* Res = Builder.CreateCall(M.getFunction(Name), Args);
+                  InlineFunctionInfo IFI;
+                  InlineFunction(*Res, IFI);
+                  Dones.push_back({IfPart->getParent(), Res});
+                  IP = ElsePart;
+                }
+                IRBuilder<> Builder(IP);
+                Builder.CreateCall(Call->getFunctionType(), Call->getCalledOperand(), Args);
+                if (!Call->getFunctionType()->getReturnType()->isVoidTy()) {
+
+                }
+                Call->eraseFromParent();
+              }
+            }
+          }
+          dbgs() << "inlined: " << "\n";
+          dbgs() << M << "\n";
+
+          Optimize(nullptr, Jit->getTargetTriple(), M, "default<O2>");
+
+          dbgs() << "after: " << "\n";
+          dbgs() << M << "\n";
+
+        });
+        return Error::success();
+      });
 }
 
 IncrementalExecutor::~IncrementalExecutor() {}
@@ -64,7 +295,7 @@ llvm::Error IncrementalExecutor::addModule(PartialTranslationUnit &PTU) {
       Jit->getMainJITDylib().createResourceTracker();
   ResourceTrackers[&PTU] = RT;
 
-  return Jit->addIRModule(RT, {std::move(PTU.TheModule), TSCtx});
+  return Jit->addLazyIRModule(RT, {std::move(PTU.TheModule), TSCtx});
 }
 
 llvm::Error IncrementalExecutor::removeModule(PartialTranslationUnit &PTU) {
diff --git a/clang/lib/Interpreter/IncrementalExecutor.h b/clang/lib/Interpreter/IncrementalExecutor.h
index dd0a210a0614154..9c9e03fbc23144b 100644
--- a/clang/lib/Interpreter/IncrementalExecutor.h
+++ b/clang/lib/Interpreter/IncrementalExecutor.h
@@ -17,6 +17,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/IR/PassManager.h"
 
 #include <memory>
 
@@ -24,6 +25,7 @@ namespace llvm {
 class Error;
 namespace orc {
 class LLJIT;
+class LLLazyJIT;
 class ThreadSafeContext;
 } // namespace orc
 } // namespace llvm
@@ -35,7 +37,7 @@ class TargetInfo;
 
 class IncrementalExecutor {
   using CtorDtorIterator = llvm::orc::CtorDtorIterator;
-  std::unique_ptr<llvm::orc::LLJIT> Jit;
+  std::unique_ptr<llvm::orc::LLLazyJIT> Jit;
   llvm::orc::ThreadSafeContext &TSCtx;
 
   llvm::DenseMap<const PartialTranslationUnit *, llvm::orc::ResourceTrackerSP>
@@ -55,7 +57,9 @@ class IncrementalExecutor {
   llvm::Expected<llvm::orc::ExecutorAddr>
   getSymbolAddress(llvm::StringRef Name, SymbolNameKind NameKind) const;
 
-  llvm::orc::LLJIT &GetExecutionEngine() { return *Jit; }
+  llvm::orc::LLLazyJIT &GetExecutionEngine() { return *Jit; }
+
+  std::unique_ptr<llvm::ModulePassManager> MPM;
 };
 
 } // end namespace clang
diff --git a/clang/lib/Interpreter/Interpreter.cpp b/clang/lib/Interpreter/Interpreter.cpp
index 7968c62cbd3e7b3..aeecf72d703eba6 100644
--- a/clang/lib/Interpreter/Interpreter.cpp
+++ b/clang/lib/Interpreter/Interpreter.cpp
@@ -319,7 +319,7 @@ const CompilerInstance *Interpreter::getCompilerInstance() const {
   return IncrParser->getCI();
 }
 
-llvm::Expected<llvm::orc::LLJIT &> Interpreter::getExecutionEngine() {
+llvm::Expected<llvm::orc::LLLazyJIT &> Interpreter::getExecutionEngine() {
   if (!IncrExecutor) {
     if (auto Err = CreateExecutor())
       return std::move(Err);
diff --git a/compiler-rt/lib/orc/common.h b/compiler-rt/lib/orc/common.h
index f34229a615341df..9d54255b4a14eda 100644
--- a/compiler-rt/lib/orc/common.h
+++ b/compiler-rt/lib/orc/common.h
@@ -20,7 +20,7 @@
 /// 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)                                             \
-  ORC_RT_INTERFACE char X;                                                     \
+ORC_RT_INTERFACE char X;                                                     \
   char X = 0;
 
 /// Opaque struct for external symbols.
diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index f76a070240315fe..cdf65d14bfea9d7 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -16,6 +16,7 @@
 #include "error.h"
 #include "wrapper_function_utils.h"
 
+#include <iostream>
 #include <algorithm>
 #include <map>
 #include <mutex>
@@ -82,6 +83,39 @@ struct TLSDescriptor {
   TLSInfoEntry *InfoEntry;
 };
 
+using CallCountMap = std::unordered_map<uint64_t, size_t>;
+using MUProfileMap = std::unordered_map<uint32_t, CallCountMap>;
+std::unordered_map<uint64_t, MUProfileMap> CallProfiles;
+
+using ReoptimizeParam = SPSSequence<SPSTuple<uint32_t, uint64_t>>;
+
+std::vector<std::pair<uint32_t, uint64_t>> SerealizeProfile(uint64_t MUID) {
+  std::vector<std::pair<uint32_t, uint64_t>> Res;
+  for (auto& [CallID, CountMap] : CallProfiles[MUID]) {
+    for (auto [FuncPtr, _ ] : CountMap) {
+      std::cout << "EECHIII:" << FuncPtr << "\n";
+      Res.push_back({CallID, FuncPtr});
+    }
+  }
+  return Res;
+}
+
+ORC_RT_INTERFACE void __orc_rt_increment_call_count(uint64_t MUID, uint32_t CallID, void* FuncPtr) {
+  CallProfiles[MUID][CallID][(uint64_t)FuncPtr] ++;
+}
+
+ORC_RT_INTERFACE void __orc_rt_reoptimize(uint64_t MUID, uint32_t CurVersion) {
+  auto Profiles = SerealizeProfile(MUID);
+  Error Err2 = Error::success();
+  if (auto Err =
+          WrapperFunction<SPSError(uint64_t, uint32_t, ReoptimizeParam)>::call(&__orc_rt_reoptimize_tag, Err2,
+                                MUID, CurVersion, Profiles))
+    return;
+  if (Err2)
+    return;
+}
+
+
 class ELFNixPlatformRuntimeState {
 private:
   struct AtExitEntry {
@@ -92,9 +126,10 @@ class ELFNixPlatformRuntimeState {
   using AtExitsVector = std::vector<AtExitEntry>;
 
   struct PerJITDylibState {
+    PerJITDylibState() = default;
     void *Header = nullptr;
     size_t RefCount = 0;
-    bool AllowReinitialization = false;
+    bool AllowReinitialization = true;
     AtExitsVector AtExits;
   };
 
@@ -229,6 +264,7 @@ void *ELFNixPlatformRuntimeState::dlopen(std::string_view Path, int Mode) {
     }
   }
 
+
   auto H = dlopenInitialize(Path, Mode);
   if (!H) {
     DLFcnError = toString(H.takeError());
@@ -395,6 +431,8 @@ ELFNixPlatformRuntimeState::dlopenInitialize(std::string_view Path, int Mode) {
     if (auto Err = initializeJITDylib(MOJDIs))
       return std::move(Err);
 
+  printf("ok\n");
+
   // Return the header for the last item in the list.
   auto *JDS = getJITDylibStateByHeaderAddr(
       InitSeq->back().DSOHandleAddress.toPtr<void *>());
diff --git a/compiler-rt/lib/orc/elfnix_platform.h b/compiler-rt/lib/orc/elfnix_platform.h
index e0ee9591dfc622f..6b63b4d7d7a0d7c 100644
--- a/compiler-rt/lib/orc/elfnix_platform.h
+++ b/compiler-rt/lib/orc/elfnix_platform.h
@@ -29,6 +29,9 @@ ORC_RT_INTERFACE int __orc_rt_elfnix_jit_dlclose(void *dso_handle);
 ORC_RT_INTERFACE void *__orc_rt_elfnix_jit_dlsym(void *dso_handle,
                                                  const char *symbol);
 
+ORC_RT_INTERFACE void __orc_rt_increment_call_count(uint64_t MUID, uint32_t CallID, void* FuncPtr);
+ORC_RT_INTERFACE void __orc_rt_reoptimize(uint64_t MUID, uint32_t CurVersion);
+
 namespace __orc_rt {
 namespace elfnix {
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
index 896de521f13af80..ed79d15cfee772b 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
@@ -15,6 +15,7 @@
 
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+#include "llvm/ExecutionEngine/Orc/ReOptimizeLayer.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
@@ -279,6 +280,11 @@ class LLLazyJIT : public LLJIT {
   /// Returns a reference to the on-demand layer.
   CompileOnDemandLayer &getCompileOnDemandLayer() { return *CODLayer; }
 
+  ReOptimizeLayer &getReOptimizeLayer() { return *ROLayer; }
+
+  /// Add a module to be lazily compiled to JITDylib JD.
+  Error addLazyIRModule(ResourceTrackerSP RT, ThreadSafeModule M);
+
   /// Add a module to be lazily compiled to JITDylib JD.
   Error addLazyIRModule(JITDylib &JD, ThreadSafeModule M);
 
@@ -293,8 +299,10 @@ class LLLazyJIT : public LLJIT {
   LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
 
   std::unique_ptr<LazyCallThroughManager> LCTMgr;
+  std::unique_ptr<RedirectableSymbolManager> RSManager;
   std::unique_ptr<IRPartitionLayer> IPLayer;
   std::unique_ptr<CompileOnDemandLayer> CODLayer;
+  std::unique_ptr<ReOptimizeLayer> ROLayer;
 };
 
 class LLJITBuilderState {
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
index df2f005825b5a7d..01a616481f5eb80 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ReOptimizeLayer.h
@@ -42,12 +42,12 @@ class ReOptimizeLayer : public IRLayer, public ResourceManager {
   /// terminated.
   using ReOptimizeFunc = unique_function<Error(
       ReOptimizeLayer &Parent, ReOptMaterializationUnitID MUID,
-      unsigned CurVersion, ResourceTrackerSP OldRT, ThreadSafeModule &TSM)>;
+      unsigned CurVersion, ResourceTrackerSP OldRT, const std::vector<std::pair<uint32_t, uint64_t>>& Profile, ThreadSafeModule &TSM)>;
 
   ReOptimizeLayer(ExecutionSession &ES, IRLayer &BaseLayer,
                   RedirectableSymbolManager &RM)
       : IRLayer(ES, BaseLayer.getManglingOptions()), ES(ES),
-        BaseLayer(BaseLayer), RSManager(RM), ReOptFunc(identity),
+        BaseLayer(BaseLayer), RSManager(RM), ReOptFunc(nullptr),
         ProfilerFunc(reoptimizeIfCallFrequent) {}
 
   void setReoptimizeFunc(ReOptimizeFunc ReOptFunc) {
@@ -85,13 +85,16 @@ class ReOptimizeLayer : public IRLayer, public ResourceManager {
 
   // Create IR reoptimize request fucntion call.
   static void createReoptimizeCall(Module &M, Instruction &IP,
-                                   GlobalVariable *ArgBuffer);
+                                   ReOptMaterializationUnitID MUID, uint32_t CurVersion);
+
+  // Create IR reoptimize request fucntion call.
+  static void createFucnCountCall(Module &M, Instruction &IP, ReOptMaterializationUnitID MUID, uint32_t CallID, Value* FuncPtr);
 
   Error handleRemoveResources(JITDylib &JD, ResourceKey K) override;
   void handleTransferResources(JITDylib &JD, ResourceKey DstK,
                                ResourceKey SrcK) override;
 
-private:
+  DenseMap<ExecutorAddr, std::pair<ReOptMaterializationUnitID,StringRef>> FuncAddrToMU;
   class ReOptMaterializationUnitState {
   public:
     ReOptMaterializationUnitState() = default;
@@ -139,12 +142,17 @@ class ReOptimizeLayer : public IRLayer, public ResourceManager {
       shared::SPSArgList<ReOptMaterializationUnitID, uint32_t>;
   using SendErrorFn = unique_function<void(Error)>;
 
+  ReOptMaterializationUnitState &
+  getMaterializationUnitState(ReOptMaterializationUnitID MUID);
+
   Expected<SymbolMap> emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
                                         uint32_t Version, JITDylib &JD,
                                         ThreadSafeModule TSM);
 
+private:
+
   void rt_reoptimize(SendErrorFn SendResult, ReOptMaterializationUnitID MUID,
-                     uint32_t CurVersion);
+                     uint32_t CurVersion, const std::vector<std::pair<uint32_t,uint64_t>>& Profile);
 
   static Expected<Constant *>
   createReoptimizeArgBuffer(Module &M, ReOptMaterializationUnitID MUID,
@@ -156,10 +164,6 @@ class ReOptimizeLayer : public IRLayer, public ResourceManager {
   void
   registerMaterializationUnitResource(ResourceKey Key,
                                       ReOptMaterializationUnitState &State);
-
-  ReOptMaterializationUnitState &
-  getMaterializationUnitState(ReOptMaterializationUnitID MUID);
-
   ExecutionSession &ES;
   IRLayer &BaseLayer;
   RedirectableSymbolManager &RSManager;
diff --git a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
index 9296bc2b389ab24..75af2c82c6ae513 100644
--- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
@@ -86,7 +86,7 @@ CompileOnDemandLayer::getPerDylibResources(JITDylib &TargetD) {
   auto I = DylibResources.find(&TargetD);
   if (I == DylibResources.end()) {
     auto &ImplD =
-        getExecutionSession().createBareJITDylib(TargetD.getName() + ".impl");
+        cantFail(getExecutionSession().createJITDylib(TargetD.getName() + ".impl"));
     JITDylibSearchOrder NewLinkOrder;
     TargetD.withLinkOrderDo([&](const JITDylibSearchOrder &TargetLinkOrder) {
       NewLinkOrder = TargetLinkOrder;
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index f4c0ecf784cdea4..7fefa0b22f12d0e 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -2232,7 +2232,7 @@ Error ExecutionSession::registerJITDispatchHandlers(
 
   auto TagAddrs = lookup({{&JD, JITDylibLookupFlags::MatchAllSymbols}},
                          SymbolLookupSet::fromMapKeys(
-                             WFs, SymbolLookupFlags::WeaklyReferencedSymbol));
+                             WFs, SymbolLookupFlags::RequiredSymbol));
   if (!TagAddrs)
     return TagAddrs.takeError();
 
@@ -2272,11 +2272,14 @@ void ExecutionSession::runJITDispatchHandler(SendResultFunction SendResult,
 
   if (F)
     (*F)(std::move(SendResult), ArgBuffer.data(), ArgBuffer.size());
-  else
+  else {
+    dbgs() << "dispatch failed" << "\n";
+    assert(false);
     SendResult(shared::WrapperFunctionResult::createOutOfBandError(
         ("No function registered for tag " +
          formatv("{0:x16}", HandlerFnTagAddr))
             .str()));
+  }
 }
 
 void ExecutionSession::dump(raw_ostream &OS) {
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index c08b8b037fa298d..19efdbbb403fe5b 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -332,8 +332,9 @@ void ELFNixPlatform::getInitializersBuildSequencePhase(
       });
       auto ISItr = InitSeqs.find(InitJD.get());
       if (ISItr != InitSeqs.end()) {
-        FullInitSeq.emplace_back(std::move(ISItr->second));
-        InitSeqs.erase(ISItr);
+        FullInitSeq.emplace_back(ISItr->second);
+        //InitSeqs.erase(ISItr);
+        ISItr->second.InitSections.clear();
       }
     }
   }
@@ -549,8 +550,11 @@ Error ELFNixPlatform::registerInitInfo(
 
       Lock.lock();
       I = InitSeqs.find(&JD);
-      assert(I != InitSeqs.end() &&
-             "Entry missing after header symbol lookup?");
+      if (I == InitSeqs.end()) {
+        return Error::success();
+      }
+      /* assert(I != InitSeqs.end() && */
+      /*        "Entry missing after header symbol lookup?"); */
     }
     InitSeq = &I->second;
   }
diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index 15fc50bad94aabf..616b8a40d34597f 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -11,6 +11,7 @@
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/COFFPlatform.h"
 #include "llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h"
+#include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
 #include "llvm/ExecutionEngine/Orc/DebuggerSupportPlugin.h"
 #include "llvm/ExecutionEngine/Orc/ELFNixPlatform.h"
 #include "llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h"
@@ -1241,6 +1242,16 @@ Error LLLazyJITBuilderState::prepareForConstruction() {
   return Error::success();
 }
 
+Error LLLazyJIT::addLazyIRModule(ResourceTrackerSP RT, ThreadSafeModule TSM) {
+  assert(TSM && "Can not add null module");
+
+  if (auto Err = TSM.withModuleDo(
+          [&](Module &M) -> Error { return applyDataLayout(M); }))
+    return Err;
+
+  return IPLayer->add(RT, std::move(TSM));
+}
+
 Error LLLazyJIT::addLazyIRModule(JITDylib &JD, ThreadSafeModule TSM) {
   assert(TSM && "Can not add null module");
 
@@ -1248,7 +1259,7 @@ Error LLLazyJIT::addLazyIRModule(JITDylib &JD, ThreadSafeModule TSM) {
           [&](Module &M) -> Error { return applyDataLayout(M); }))
     return Err;
 
-  return CODLayer->add(JD, std::move(TSM));
+  return IPLayer->add(JD, std::move(TSM));
 }
 
 LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) {
@@ -1288,8 +1299,14 @@ LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) {
     return;
   }
 
+  RSManager = cantFail(JITLinkRedirectableSymbolManager::Create(*ES, *dyn_cast<ObjectLinkingLayer>(&getObjLinkingLayer()), getMainJITDylib()));
+
+  ROLayer = std::make_unique<ReOptimizeLayer>(*ES, *InitHelperTransformLayer, *RSManager);
+
+  cantFail(ROLayer->reigsterRuntimeFunctions(*getPlatformJITDylib()));
+
   // Create the IP Layer.
-  IPLayer = std::make_unique<IRPartitionLayer>(*ES, *InitHelperTransformLayer);
+  IPLayer = std::make_unique<IRPartitionLayer>(*ES, *ROLayer);
 
   // Create the COD layer.
   CODLayer = std::make_unique<CompileOnDemandLayer>(*ES, *IPLayer, *LCTMgr,
diff --git a/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
index 2ec82f210e6f625..08f644a8fa6e8b2 100644
--- a/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp
@@ -27,7 +27,7 @@ void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
 
 Error ReOptimizeLayer::reigsterRuntimeFunctions(JITDylib &PlatformJD) {
   ExecutionSession::JITDispatchHandlerAssociationMap WFs;
-  using ReoptimizeSPSSig = shared::SPSError(uint64_t, uint32_t);
+  using ReoptimizeSPSSig = shared::SPSError(uint64_t, uint32_t, shared::SPSSequence<shared::SPSTuple<uint32_t, uint64_t>>);
   WFs[ES.intern("__orc_rt_reoptimize_tag")] =
       ES.wrapAsyncWithSPS<ReoptimizeSPSSig>(this,
                                             &ReOptimizeLayer::rt_reoptimize);
@@ -46,10 +46,35 @@ void ReOptimizeLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
   }
 
   if (HasNonCallable) {
+    dbgs() << "Skipped" << "\n";
+    TSM.withModuleDo([&](Module& M) {
+      dbgs() << M << "\n";
+    });
     BaseLayer.emit(std::move(R), std::move(TSM));
     return;
   }
 
+  TSM.withModuleDo([&](Module& M){
+      uint32_t ID = 1;
+    for (auto &F : M) {
+      if (F.isDeclaration())
+        continue;
+        for (auto &B : F) {
+          for (auto &I : B) {
+            if (auto* Call = dyn_cast<llvm::CallInst>(&I)) {
+              if (Call->isIndirectCall()) {
+                LLVMContext& C = Call->getContext();
+                Type *I32Ty = Type::getInt32Ty(C);
+                Constant *One = ConstantInt::get(I32Ty, ID);
+                MDNode* N = MDNode::get(C, llvm::ValueAsMetadata::getConstant(One));
+                Call->setMetadata("call_id", N);
+              }
+            }
+          }
+        }
+      }
+    });
+
   auto &MUState = createMaterializationUnitState(TSM);
 
   if (auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
@@ -59,7 +84,7 @@ void ReOptimizeLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
     R->failMaterialization();
     return;
   }
-
+  
   if (auto Err =
           ProfilerFunc(*this, MUState.getID(), MUState.getCurVersion(), TSM)) {
     ES.reportError(std::move(Err));
@@ -75,7 +100,24 @@ void ReOptimizeLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
     return;
   }
 
-  RSManager.emitRedirectableSymbols(std::move(R), std::move(*InitialDests));
+  RSManager.emitRedirectableSymbols(std::move(R), *InitialDests);
+
+  SymbolLookupSet LookupSymbols;
+  for (auto [K, V] : *InitialDests)
+    LookupSymbols.add(K);
+
+  auto FinalSymbols =
+      ES.lookup({{&JD, JITDylibLookupFlags::MatchAllSymbols}}, LookupSymbols,
+                LookupKind::Static, SymbolState::Resolved);
+  if (auto Err = FinalSymbols.takeError()) {
+    ES.reportError(InitialDests.takeError());
+    R->failMaterialization();
+    return;
+  }
+  
+  for (auto [K,V] : (*FinalSymbols)) {
+    FuncAddrToMU[V.getAddress()] = {MUState.getID(), *K};
+  }
 }
 
 Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
@@ -87,12 +129,7 @@ Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
     GlobalVariable *Counter = new GlobalVariable(
         M, I64Ty, false, GlobalValue::InternalLinkage,
         Constant::getNullValue(I64Ty), "__orc_reopt_counter");
-    auto ArgBufferConst = createReoptimizeArgBuffer(M, MUID, CurVersion);
-    if (auto Err = ArgBufferConst.takeError())
-      return Err;
-    GlobalVariable *ArgBuffer =
-        new GlobalVariable(M, (*ArgBufferConst)->getType(), true,
-                           GlobalValue::InternalLinkage, (*ArgBufferConst));
+    dbgs() << "Adding instrumentation" << "\n";
     for (auto &F : M) {
       if (F.isDeclaration())
         continue;
@@ -106,8 +143,21 @@ Error ReOptimizeLayer::reoptimizeIfCallFrequent(ReOptimizeLayer &Parent,
       Value *Added = IRB.CreateAdd(Cnt, ConstantInt::get(I64Ty, 1));
       (void)IRB.CreateStore(Added, Counter);
       Instruction *SplitTerminator = SplitBlockAndInsertIfThen(Cmp, IP, false);
-      createReoptimizeCall(M, *SplitTerminator, ArgBuffer);
+      createReoptimizeCall(M, *SplitTerminator, MUID, CurVersion);
+
+      for (auto &B : F) {
+        for (auto& I : B) {
+          if (auto* Call = dyn_cast<llvm::CallInst>(&I)) {
+            if (Call->isIndirectCall()) {
+                auto* VAM = cast<ValueAsMetadata>(dyn_cast<MDNode>(Call->getMetadata("call_id"))->getOperand(0));
+                int CallID = dyn_cast<ConstantInt>(VAM->getValue())->getSExtValue();
+                createFucnCountCall(M, I, MUID, CallID, Call->getCalledOperand());
+            }
+          }
+        }
+      }
     }
+    dbgs() << M << "\n";
     return Error::success();
   });
 }
@@ -148,15 +198,18 @@ ReOptimizeLayer::emitMUImplSymbols(ReOptMaterializationUnitState &MUState,
     return Err;
 
   SymbolMap Result;
-  for (auto [K, V] : RenamedMap)
+  for (auto [K, V] : RenamedMap) {
     Result[K] = (*ImplSymbols)[V];
+  }
+
 
   return Result;
 }
 
 void ReOptimizeLayer::rt_reoptimize(SendErrorFn SendResult,
                                     ReOptMaterializationUnitID MUID,
-                                    uint32_t CurVersion) {
+                                    uint32_t CurVersion, const std::vector<std::pair<uint32_t,uint64_t>>& Profile) {
+
   auto &MUState = getMaterializationUnitState(MUID);
   if (CurVersion < MUState.getCurVersion() || !MUState.tryStartReoptimize()) {
     SendResult(Error::success());
@@ -164,10 +217,14 @@ void ReOptimizeLayer::rt_reoptimize(SendErrorFn SendResult,
   }
 
   ThreadSafeModule TSM = cloneToNewContext(MUState.getThreadSafeModule());
+  TSM.withModuleDo([&](Module& M) {
+    dbgs() << "Reoptimization requested for" << M.getName() << "\n";
+  });
+
   auto OldRT = MUState.getResourceTracker();
   auto &JD = OldRT->getJITDylib();
 
-  if (auto Err = ReOptFunc(*this, MUID, CurVersion + 1, OldRT, TSM)) {
+  if (auto Err = ReOptFunc(*this, MUID, CurVersion + 1, OldRT, Profile, TSM)) {
     ES.reportError(std::move(Err));
     MUState.reoptimizeFailed();
     SendResult(Error::success());
@@ -206,37 +263,41 @@ Expected<Constant *> ReOptimizeLayer::createReoptimizeArgBuffer(
 }
 
 void ReOptimizeLayer::createReoptimizeCall(Module &M, Instruction &IP,
-                                           GlobalVariable *ArgBuffer) {
-  GlobalVariable *DispatchCtx =
-      M.getGlobalVariable("__orc_rt_jit_dispatch_ctx");
-  if (!DispatchCtx)
-    DispatchCtx = new GlobalVariable(M, Type::getInt8PtrTy(M.getContext()),
-                                     false, GlobalValue::ExternalLinkage,
-                                     nullptr, "__orc_rt_jit_dispatch_ctx");
-  GlobalVariable *ReoptimizeTag =
-      M.getGlobalVariable("__orc_rt_reoptimize_tag");
-  if (!ReoptimizeTag)
-    ReoptimizeTag = new GlobalVariable(M, Type::getInt8PtrTy(M.getContext()),
-                                       false, GlobalValue::ExternalLinkage,
-                                       nullptr, "__orc_rt_reoptimize_tag");
-  Function *DispatchFunc = M.getFunction("__orc_rt_jit_dispatch");
+                                           ReOptMaterializationUnitID MUID, uint32_t CurVersion) {
+  Function *DispatchFunc = M.getFunction("__orc_rt_reoptimize");
+  if (!DispatchFunc) {
+    std::vector<Type *> Args = {IntegerType::get(M.getContext(), 64), IntegerType::get(M.getContext(), 32)};
+    FunctionType *FuncTy =
+        FunctionType::get(Type::getVoidTy(M.getContext()), Args, false);
+    DispatchFunc = Function::Create(FuncTy, GlobalValue::ExternalLinkage,
+                                    "__orc_rt_reoptimize", &M);
+  }
+  Constant *MUIDConst = ConstantInt::get(
+      IntegerType::get(M.getContext(), 64), MUID, false);
+  Constant *VersionConst= ConstantInt::get(
+      IntegerType::get(M.getContext(), 32), CurVersion, false);
+  IRBuilder<> IRB(&IP);
+  (void)IRB.CreateCall(DispatchFunc,
+                       {MUIDConst, VersionConst});
+}
+
+// Create IR reoptimize request fucntion call.
+void ReOptimizeLayer::createFucnCountCall(Module &M, Instruction &IP, ReOptMaterializationUnitID MUID, uint32_t CallID, Value* FuncPtr) {
+  Function *DispatchFunc = M.getFunction("__orc_rt_increment_call_count");
   if (!DispatchFunc) {
-    std::vector<Type *> Args = {Type::getInt8PtrTy(M.getContext()),
-                                Type::getInt8PtrTy(M.getContext()),
-                                Type::getInt8PtrTy(M.getContext()),
-                                IntegerType::get(M.getContext(), 64)};
+    std::vector<Type *> Args = {IntegerType::get(M.getContext(), 64), IntegerType::get(M.getContext(), 32), Type::getInt8PtrTy(M.getContext())};
     FunctionType *FuncTy =
         FunctionType::get(Type::getVoidTy(M.getContext()), Args, false);
     DispatchFunc = Function::Create(FuncTy, GlobalValue::ExternalLinkage,
-                                    "__orc_rt_jit_dispatch", &M);
+                                    "__orc_rt_increment_call_count", &M);
   }
-  size_t ArgBufferSizeConst =
-      SPSReoptimizeArgList::size(ReOptMaterializationUnitID{}, uint32_t{});
-  Constant *ArgBufferSize = ConstantInt::get(
-      IntegerType::get(M.getContext(), 64), ArgBufferSizeConst, false);
+  Constant *MUIDConst = ConstantInt::get(
+      IntegerType::get(M.getContext(), 64), MUID, false);
+  Constant *CallIDConst = ConstantInt::get(
+      IntegerType::get(M.getContext(), 32), CallID, false);
   IRBuilder<> IRB(&IP);
   (void)IRB.CreateCall(DispatchFunc,
-                       {DispatchCtx, ReoptimizeTag, ArgBuffer, ArgBufferSize});
+                       {MUIDConst, CallIDConst, FuncPtr});
 }
 
 ReOptimizeLayer::ReOptMaterializationUnitState &



More information about the cfe-commits mailing list