[clang] [ORC] Implement basic reoptimization. (PR #67050)

[via cfe-commits]

================
@@ -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];
+  }
----------------
lhames wrote:

We only need one anchor symbol to trigger emission of the graph. Eventually I think this could be sped up by emitting:
```asm
.section <JumpStubTableName>
__<JumpStubTableName>.<++Idx>:
  ... stubs ...
.section <StubPtrTableName>
__<StubPtrTableName>.<Idx>:
  ... ptrs ...
```

Then looking up `__JumpStubTableName.<Idx>` and `__StubPtrTableName.<Idx>` and using the stub and pointer sizes to stride through the resulting addresses.

I don't think this change needs to happen before the commit though -- you could add a FIXME and this can be handled in the future.

https://github.com/llvm/llvm-project/pull/67050