[JITLink] Unify x86-64 MachO and ELF 's optimize GOT/Stub function

[StephenFan]

Summary

This patch unify optimizeELF_x86_64_GOTAndStubs and optimizeMachO_x86_64_GOTAndStubs into a pure optimize_x86_64_GOTAndStubs

[lhames]

Hi Stephen,

Please address the inline comments before committing. Otherwise LGTM. Thank you very much for this work!

• Lang.

llvm/lib/ExecutionEngine/JITLink/x86_64.cpp

80–83

This should assert, rather than silently continuing -- a graph with an invalid edge is malformed.

96–97

This can use the x86_64::isInRangeForImmS32 function now (the original optimize functions predate the utility, which is why it wasn't used before).

109

This should use x86_64::getEdgeKindName(E.getKind()).

131–132

This can also use isInRangeForImmS32.

137

This should also use x86_64::getEdgeKindName(E.getKind()).

[lhames]

Stephen and I discussed this further on discord and realized that some generic x86-64 edge kinds needed to be renamed to avoid confusion (ELF allows several relaxations that aren't possible in MachO).

The necessary renames have been made in 632135acae2 and 27ea3f16072.

Stephen -- You'll need to rebase to account for the renames, but otherwise this should be good to land now.

[StephenFan]

Rebase

[lhames]

LGTM. Thanks Stephen!

llvm/include/llvm/ExecutionEngine/JITLink/ELF_x86_64.h

 namespace jitlink {
 
 namespace ELF_x86_64_Edges {
 enum ELFX86RelocationKind : Edge::Kind {
   Branch32 = Edge::FirstRelocation,
   Pointer64,
   PCRel32,
   PCRel32GOTLoad,
+  PCRel32REXGOTLoad,
   PCRel64GOT,
   GOTOFF64,
   GOT64,
   Delta64,
 };
 
 } // end namespace ELF_x86_64_Edges
 

llvm/include/llvm/ExecutionEngine/JITLink/x86_64.h

enum EdgeKind_x86_64 : Edge::Kind {
   /// Fixup expression:
   ///   NONE
   ///
   /// Errors:
   ///   - *ASSERTION* Failure to handle edges of this kind prior to the fixup
   ///     phase will result in an assert/unreachable during the fixup phase
   RequestGOTAndTransformToDelta64FromGOT,
 
+  /// A PC-relative load of a GOT entry, relaxable if GOT entry target is
+  /// in-range of the fixup
+  ///
+  /// TODO: Explain the optimization
+  ///
+  /// Fixup expression
+  ///   Fixup <- Target - (Fixup + 4) + Addend : int32
+  ///
+  /// Errors:
+  ///   - The result of the fixup expression must fit into an int32, otherwise
+  ///     an out-of-range error will be returned.
+  //
+  PCRel32GOTLoadRelaxable,
+
   /// A PC-relative REX load of a GOT entry, relaxable if GOT entry target
   /// is in-range of the fixup.
   ///
   /// If the GOT entry target is in-range of the fixup then the load from the
   /// GOT may be replaced with a direct memory address calculation.
   ///
   /// Fixup expression:
   ///   Fixup <- Target - (Fixup + 4) + Addend : int32
enum EdgeKind_x86_64 : Edge::Kind {
   ///   NONE
   ///
   /// Errors:
   ///   - *ASSERTION* Failure to handle edges of this kind prior to the fixup
   ///     phase will result in an assert/unreachable during the fixup phase.
   ///
   RequestGOTAndTransformToPCRel32GOTLoadREXRelaxable,
 
+  /// A GOT entry getter/constructor, transformed to
+  /// PCRel32ToGOTLoadRelaxable pointing at the GOT entry for the original
+  /// target.
+  ///
+  /// Indicates that this edge should be lowered to a PC32ToGOTLoadRelaxable
+  /// targeting the GOT entry for the edge's current target, maintaining the
+  /// same addend. A GOT entry for the target should be created if one does not
+  /// already exist.
+  ///
+  /// Edges of this kind are usually lowered by a GOT builder pass inserted by
+  /// default.
+  ///
+  /// Fixup expression:
+  ///   NONE
+  ///
+  /// Errors:
+  ///   - *ASSERTION* Failure to handle edges of this kind prior to the fixup
+  ///     phase will result in an assert/unreachable during the fixup phase.
+  ///
+  RequestGOTAndTransformToPCRel32GOTLoadRelaxable,
+
   /// A PC-relative REX load of a Thread Local Variable Pointer (TLVP) entry,
   /// relaxable if the TLVP entry target is in-range of the fixup.
   ///
   /// If the TLVP entry target is in-range of the fixup then the load from the
   /// TLVP may be replaced with a direct memory address calculation.
   ///
   /// The target of this edge must be a thread local variable entry of the form
   ///   .quad <tlv getter thunk>
enum EdgeKind_x86_64 : Edge::Kind {
   ///
   RequestTLVPAndTransformToPCRel32TLVPLoadREXRelaxable
 };
 
 /// Returns a string name for the given x86-64 edge. For debugging purposes
 /// only.
 const char *getEdgeKindName(Edge::Kind K);
 
+/// Optimize the GOT and Stub relocations if the edge target address is in range
+/// 1. PCRel32GOTLoadRelaxable. For this edge kind, if the target is in range,
+/// then replace GOT load with lea
+/// 2. BranchPCRel32ToPtrJumpStubRelaxable. For this edge kind, if the target is
+/// in range, replace a indirect jump by plt stub with a direct jump to the
+/// target
+Error optimize_x86_64_GOTAndStubs(LinkGraph &G);
+
 /// Returns true if the given uint64_t value is in range for a uint32_t.
 inline bool isInRangeForImmU32(uint64_t Value) {
   return Value <= std::numeric_limits<uint32_t>::max();
 }
 
 /// Returns true if the given int64_t value is in range for an int32_t.
 inline bool isInRangeForImmS32(int64_t Value) {
   return (Value >= std::numeric_limits<int32_t>::min() &&
case Pointer32: {
     else
       return makeTargetOutOfRangeError(G, B, E);
     break;
   }
 
   case BranchPCRel32:
   case BranchPCRel32ToPtrJumpStub:
   case BranchPCRel32ToPtrJumpStubBypassable:
+  case PCRel32GOTLoadRelaxable:
   case PCRel32GOTLoadREXRelaxable:
   case PCRel32TLVPLoadREXRelaxable: {
     int64_t Value =
         E.getTarget().getAddress() - (FixupAddress + 4) + E.getAddend();
     if (LLVM_LIKELY(isInRangeForImmS32(Value)))
       *(little32_t *)FixupPtr = Value;
     else
       return makeTargetOutOfRangeError(G, B, E);

llvm/lib/ExecutionEngine/JITLink/ELF_x86_64.cpp

if (E.getKind() == x86_64::Delta64FromGOT) {
       // need to fix up this edge.
       getGOTSection();
       return false;
     }
     return E.getKind() == x86_64::RequestGOTAndTransformToDelta32 ||
            E.getKind() == x86_64::RequestGOTAndTransformToDelta64 ||
            E.getKind() ==
                x86_64::RequestGOTAndTransformToPCRel32GOTLoadREXRelaxable ||
-           E.getKind() == x86_64::RequestGOTAndTransformToDelta64FromGOT;
+           E.getKind() == x86_64::RequestGOTAndTransformToDelta64FromGOT ||
+           E.getKind() ==
+               x86_64::RequestGOTAndTransformToPCRel32GOTLoadRelaxable;
   }
 
   Symbol &createGOTEntry(Symbol &Target) {
     auto &GOTEntryBlock = G.createContentBlock(
         getGOTSection(), getGOTEntryBlockContent(), 0, 8, 0);
     GOTEntryBlock.addEdge(x86_64::Pointer64, 0, Target, 0);
     return G.addAnonymousSymbol(GOTEntryBlock, 0, 8, false, false);
   }
 
   void fixGOTEdge(Edge &E, Symbol &GOTEntry) {
     // If this is a PCRel32GOT/PCRel64GOT then change it to an ordinary
     // PCRel32/PCRel64. If it is a PCRel32GOTLoad then leave it as-is for now:
     // We will use the kind to check for GOT optimization opportunities in the
     // optimizeMachO_x86_64_GOTAndStubs pass below.
     // If it's a GOT64 leave it as is.
     switch (E.getKind()) {
     case x86_64::RequestGOTAndTransformToPCRel32GOTLoadREXRelaxable:
       E.setKind(x86_64::PCRel32GOTLoadREXRelaxable);
       break;
+    case x86_64::RequestGOTAndTransformToPCRel32GOTLoadRelaxable:
+      E.setKind(x86_64::PCRel32GOTLoadRelaxable);
+      break;
     case x86_64::RequestGOTAndTransformToDelta64:
       E.setKind(x86_64::Delta64);
       break;
     case x86_64::RequestGOTAndTransformToDelta64FromGOT:
       E.setKind(x86_64::Delta64FromGOT);
       break;
     case x86_64::RequestGOTAndTransformToDelta32:
       E.setKind(x86_64::Delta32);

 
 } // namespace
 
 const uint8_t PerGraphGOTAndPLTStubsBuilder_ELF_x86_64::NullGOTEntryContent[8] =
     {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
 const uint8_t PerGraphGOTAndPLTStubsBuilder_ELF_x86_64::StubContent[6] = {
     0xFF, 0x25, 0x00, 0x00, 0x00, 0x00};
 
-static Error optimizeELF_x86_64_GOTAndStubs(LinkGraph &G) {
-  LLVM_DEBUG(dbgs() << "Optimizing GOT entries and stubs:\n");
-
-  for (auto *B : G.blocks())
-    for (auto &E : B->edges())
-      if (E.getKind() == x86_64::PCRel32GOTLoadREXRelaxable) {
-        // Replace GOT load with LEA only for MOVQ instructions.
-        constexpr uint8_t MOVQRIPRel[] = {0x48, 0x8b};
-        if (E.getOffset() < 3 ||
-            strncmp(B->getContent().data() + E.getOffset() - 3,
-                    reinterpret_cast<const char *>(MOVQRIPRel), 2) != 0)
-          continue;
-
-        auto &GOTBlock = E.getTarget().getBlock();
-        assert(GOTBlock.getSize() == G.getPointerSize() &&
-               "GOT entry block should be pointer sized");
-        assert(GOTBlock.edges_size() == 1 &&
-               "GOT entry should only have one outgoing edge");
-
-        auto &GOTTarget = GOTBlock.edges().begin()->getTarget();
-        JITTargetAddress EdgeAddr = B->getAddress() + E.getOffset();
-        JITTargetAddress TargetAddr = GOTTarget.getAddress();
-
-        int64_t Displacement = TargetAddr - EdgeAddr + 4;
-        if (Displacement >= std::numeric_limits<int32_t>::min() &&
-            Displacement <= std::numeric_limits<int32_t>::max()) {
-          // Change the edge kind as we don't go through GOT anymore. This is
-          // for formal correctness only. Technically, the two relocation kinds
-          // are resolved the same way.
-          E.setKind(x86_64::Delta32);
-          E.setTarget(GOTTarget);
-          E.setAddend(E.getAddend() - 4);
-          auto *BlockData = reinterpret_cast<uint8_t *>(
-              const_cast<char *>(B->getContent().data()));
-          BlockData[E.getOffset() - 2] = 0x8d;
-          LLVM_DEBUG({
-            dbgs() << "  Replaced GOT load wih LEA:\n    ";
-            printEdge(dbgs(), *B, E, getELFX86RelocationKindName(E.getKind()));
-            dbgs() << "\n";
-          });
-        }
-      } else if (E.getKind() == x86_64::BranchPCRel32ToPtrJumpStubBypassable) {
-        auto &StubBlock = E.getTarget().getBlock();
-        assert(
-            StubBlock.getSize() ==
-                sizeof(PerGraphGOTAndPLTStubsBuilder_ELF_x86_64::StubContent) &&
-            "Stub block should be stub sized");
-        assert(StubBlock.edges_size() == 1 &&
-               "Stub block should only have one outgoing edge");
-
-        auto &GOTBlock = StubBlock.edges().begin()->getTarget().getBlock();
-        assert(GOTBlock.getSize() == G.getPointerSize() &&
-               "GOT block should be pointer sized");
-        assert(GOTBlock.edges_size() == 1 &&
-               "GOT block should only have one outgoing edge");
-
-        auto &GOTTarget = GOTBlock.edges().begin()->getTarget();
-        JITTargetAddress EdgeAddr = B->getAddress() + E.getOffset();
-        JITTargetAddress TargetAddr = GOTTarget.getAddress();
-
-        int64_t Displacement = TargetAddr - EdgeAddr + 4;
-        if (Displacement >= std::numeric_limits<int32_t>::min() &&
-            Displacement <= std::numeric_limits<int32_t>::max()) {
-          E.setKind(x86_64::BranchPCRel32);
-          E.setTarget(GOTTarget);
-          LLVM_DEBUG({
-            dbgs() << "  Replaced stub branch with direct branch:\n    ";
-            printEdge(dbgs(), *B, E, getELFX86RelocationKindName(E.getKind()));
-            dbgs() << "\n";
-          });
-        }
-      }
-
-  return Error::success();
-}
-
 static const char *getELFX86_64RelocName(uint32_t Type) {
   switch (Type) {
 #define ELF_RELOC(Name, Number)                                                \
   case Number:                                                                 \
     return #Name;
 #include "llvm/BinaryFormat/ELFRelocs/x86_64.def"
 #undef ELF_RELOC
   }
case ELF::R_X86_64_PC32:
       return ELF_x86_64_Edges::ELFX86RelocationKind::PCRel32;
     case ELF::R_X86_64_PC64:
     case ELF::R_X86_64_GOTPC64:
       return ELF_x86_64_Edges::ELFX86RelocationKind::Delta64;
     case ELF::R_X86_64_64:
       return ELF_x86_64_Edges::ELFX86RelocationKind::Pointer64;
     case ELF::R_X86_64_GOTPCREL:
     case ELF::R_X86_64_GOTPCRELX:
-    case ELF::R_X86_64_REX_GOTPCRELX:
       return ELF_x86_64_Edges::ELFX86RelocationKind::PCRel32GOTLoad;
+    case ELF::R_X86_64_REX_GOTPCRELX:
+      return ELF_x86_64_Edges::ELFX86RelocationKind::PCRel32REXGOTLoad;
     case ELF::R_X86_64_GOTPCREL64:
       return ELF_x86_64_Edges::ELFX86RelocationKind::PCRel64GOT;
     case ELF::R_X86_64_GOT64:
       return ELF_x86_64_Edges::ELFX86RelocationKind::GOT64;
     case ELF::R_X86_64_GOTOFF64:
       return ELF_x86_64_Edges::ELFX86RelocationKind::GOTOFF64;
     case ELF::R_X86_64_PLT32:
       return ELF_x86_64_Edges::ELFX86RelocationKind::Branch32;
for (auto &SecRef : Sections) {
           break;
         case Delta64:
           Kind = x86_64::Delta64;
           break;
         case Pointer64:
           Kind = x86_64::Pointer64;
           break;
         case PCRel32GOTLoad: {
+          Kind = x86_64::RequestGOTAndTransformToPCRel32GOTLoadRelaxable;
+          Addend = 0;
+          break;
+        }
+        case PCRel32REXGOTLoad: {
           Kind = x86_64::RequestGOTAndTransformToPCRel32GOTLoadREXRelaxable;
           Addend = 0;
           break;
         }
         case PCRel64GOT: {
           Kind = x86_64::RequestGOTAndTransformToDelta64;
           break;
         }
Config.PostPrunePasses.push_back(
         PerGraphGOTAndPLTStubsBuilder_ELF_x86_64::asPass);
 
     // Resolve any external section start / end symbols.
     Config.PostAllocationPasses.push_back(
         createDefineExternalSectionStartAndEndSymbolsPass(
             identifyELFSectionStartAndEndSymbols));
 
     // Add GOT/Stubs optimizer pass.
-    Config.PreFixupPasses.push_back(optimizeELF_x86_64_GOTAndStubs);
+    Config.PreFixupPasses.push_back(x86_64::optimize_x86_64_GOTAndStubs);
   }
 
   if (auto Err = Ctx->modifyPassConfig(*G, Config))
     return Ctx->notifyFailed(std::move(Err));
 
   ELFJITLinker_x86_64::link(std::move(Ctx), std::move(G), std::move(Config));
 }
 const char *getELFX86RelocationKindName(Edge::Kind R) {
   switch (R) {
   case Branch32:
     return "Branch32";
   case Pointer64:
     return "Pointer64";
   case PCRel32:
     return "PCRel32";
   case PCRel32GOTLoad:
     return "PCRel32GOTLoad";
+  case PCRel32REXGOTLoad:
+    return "PCRel32REXGOTLoad";
   case PCRel64GOT:
     return "PCRel64GOT";
   case Delta64:
     return "Delta64";
   case GOT64:
     return "GOT64";
   case GOTOFF64:
     return "GOTOFF64";
   }
   return getGenericEdgeKindName(static_cast<Edge::Kind>(R));
 }
 } // end namespace jitlink
 } // end namespace llvm

llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp

private:
   }
 
   Section *GOTSection = nullptr;
   Section *StubsSection = nullptr;
 };
 
 } // namespace
 
-static Error optimizeMachO_x86_64_GOTAndStubs(LinkGraph &G) {
-  LLVM_DEBUG(dbgs() << "Optimizing GOT entries and stubs:\n");
-
-  for (auto *B : G.blocks())
-    for (auto &E : B->edges())
-      if (E.getKind() == x86_64::PCRel32GOTLoadREXRelaxable) {
-        assert(E.getOffset() >= 3 && "GOT edge occurs too early in block");
-
-        // Optimize GOT references.
-        auto &GOTBlock = E.getTarget().getBlock();
-        assert(GOTBlock.getSize() == G.getPointerSize() &&
-               "GOT entry block should be pointer sized");
-        assert(GOTBlock.edges_size() == 1 &&
-               "GOT entry should only have one outgoing edge");
-
-        auto &GOTTarget = GOTBlock.edges().begin()->getTarget();
-        JITTargetAddress EdgeAddr = B->getAddress() + E.getOffset();
-        JITTargetAddress TargetAddr = GOTTarget.getAddress();
-
-        // Check that this is a recognized MOV instruction.
-        // FIXME: Can we assume this?
-        constexpr uint8_t MOVQRIPRel[] = {0x48, 0x8b};
-        if (strncmp(B->getContent().data() + E.getOffset() - 3,
-                    reinterpret_cast<const char *>(MOVQRIPRel), 2) != 0)
-          continue;
-
-        int64_t Displacement = TargetAddr - EdgeAddr + 4;
-        if (Displacement >= std::numeric_limits<int32_t>::min() &&
-            Displacement <= std::numeric_limits<int32_t>::max()) {
-          E.setTarget(GOTTarget);
-          E.setKind(x86_64::Delta32);
-          E.setAddend(E.getAddend() - 4);
-          char *BlockData = B->getMutableContent(G).data();
-          BlockData[E.getOffset() - 2] = (char)0x8d;
-          LLVM_DEBUG({
-            dbgs() << "  Replaced GOT load wih LEA:\n    ";
-            printEdge(dbgs(), *B, E, x86_64::getEdgeKindName(E.getKind()));
-            dbgs() << "\n";
-          });
-        }
-      } else if (E.getKind() == x86_64::BranchPCRel32ToPtrJumpStubBypassable) {
-        auto &StubBlock = E.getTarget().getBlock();
-        assert(StubBlock.getSize() == sizeof(x86_64::PointerJumpStubContent) &&
-               "Stub block should be stub sized");
-        assert(StubBlock.edges_size() == 1 &&
-               "Stub block should only have one outgoing edge");
-
-        auto &GOTBlock = StubBlock.edges().begin()->getTarget().getBlock();
-        assert(GOTBlock.getSize() == G.getPointerSize() &&
-               "GOT block should be pointer sized");
-        assert(GOTBlock.edges_size() == 1 &&
-               "GOT block should only have one outgoing edge");
-
-        auto &GOTTarget = GOTBlock.edges().begin()->getTarget();
-        JITTargetAddress EdgeAddr = B->getAddress() + E.getOffset();
-        JITTargetAddress TargetAddr = GOTTarget.getAddress();
-
-        int64_t Displacement = TargetAddr - EdgeAddr + 4;
-        if (Displacement >= std::numeric_limits<int32_t>::min() &&
-            Displacement <= std::numeric_limits<int32_t>::max()) {
-          E.setKind(x86_64::BranchPCRel32);
-          E.setTarget(GOTTarget);
-          LLVM_DEBUG({
-            dbgs() << "  Replaced stub branch with direct branch:\n    ";
-            printEdge(dbgs(), *B, E, x86_64::getEdgeKindName(E.getKind()));
-            dbgs() << "\n";
-          });
-        }
-      }
-
-  return Error::success();
-}
-
 namespace llvm {
 namespace jitlink {
 
 class MachOJITLinker_x86_64 : public JITLinker<MachOJITLinker_x86_64> {
   friend class JITLinker<MachOJITLinker_x86_64>;
 
 public:
   MachOJITLinker_x86_64(std::unique_ptr<JITLinkContext> Ctx,
if (Ctx->shouldAddDefaultTargetPasses(G->getTargetTriple())) {
     else
       Config.PrePrunePasses.push_back(markAllSymbolsLive);
 
     // Add an in-place GOT/Stubs pass.
     Config.PostPrunePasses.push_back(
         PerGraphGOTAndPLTStubsBuilder_MachO_x86_64::asPass);
 
     // Add GOT/Stubs optimizer pass.
-    Config.PreFixupPasses.push_back(optimizeMachO_x86_64_GOTAndStubs);
+    Config.PreFixupPasses.push_back(x86_64::optimize_x86_64_GOTAndStubs);
   }
 
   if (auto Err = Ctx->modifyPassConfig(*G, Config))
     return Ctx->notifyFailed(std::move(Err));
 
   // Construct a JITLinker and run the link function.
   MachOJITLinker_x86_64::link(std::move(Ctx), std::move(G), std::move(Config));
 }

llvm/lib/ExecutionEngine/JITLink/x86_64.cpp

const char *getEdgeKindName(Edge::Kind K) {
   case RequestGOTAndTransformToDelta64:
     return "RequestGOTAndTransformToDelta64";
   case RequestGOTAndTransformToDelta64FromGOT:
     return "RequestGOTAndTransformToDelta64FromGOT";
   case PCRel32GOTLoadREXRelaxable:
     return "PCRel32GOTLoadREXRelaxable";
   case RequestGOTAndTransformToPCRel32GOTLoadREXRelaxable:
     return "RequestGOTAndTransformToPCRel32GOTLoadREXRelaxable";
+  case PCRel32GOTLoadRelaxable:
+    return "PCRel32GOTLoadRelaxable";
+  case RequestGOTAndTransformToPCRel32GOTLoadRelaxable:
+    return "RequestGOTAndTransformToPCRel32GOTLoadRelaxable";
   case PCRel32TLVPLoadREXRelaxable:
     return "PCRel32TLVPLoadREXRelaxable";
   case RequestTLVPAndTransformToPCRel32TLVPLoadREXRelaxable:
     return "RequestTLVPAndTransformToPCRel32TLVPLoadREXRelaxable";
   default:
     return getGenericEdgeKindName(static_cast<Edge::Kind>(K));
   }
 }
 
 const char NullPointerContent[PointerSize] = {0x00, 0x00, 0x00, 0x00,
                                               0x00, 0x00, 0x00, 0x00};
 
 const char PointerJumpStubContent[6] = {
     static_cast<char>(0xFFu), 0x25, 0x00, 0x00, 0x00, 0x00};
 
+Error optimize_x86_64_GOTAndStubs(LinkGraph &G) {
+  LLVM_DEBUG(dbgs() << "Optimizing GOT entries and stubs:\n");
+
+  for (auto *B : G.blocks())
+    for (auto &E : B->edges())
+      if (E.getKind() == x86_64::PCRel32GOTLoadREXRelaxable) {
+        // Replace GOT load with LEA only for MOVQ instructions.
+        assert(E.getOffset() >= 3 && "GOT edge occurs too early in block");
+
+        constexpr uint8_t MOVQRIPRel[] = {0x48, 0x8b};
+        if (strncmp(B->getContent().data() + E.getOffset() - 3,
+                    reinterpret_cast<const char *>(MOVQRIPRel), 2) != 0)
+          continue;
+
+        auto &GOTBlock = E.getTarget().getBlock();
+        assert(GOTBlock.getSize() == G.getPointerSize() &&
+               "GOT entry block should be pointer sized");
+        assert(GOTBlock.edges_size() == 1 &&
+               "GOT entry should only have one outgoing edge");
+
+        auto &GOTTarget = GOTBlock.edges().begin()->getTarget();
+        JITTargetAddress EdgeAddr = B->getAddress() + E.getOffset();
+        JITTargetAddress TargetAddr = GOTTarget.getAddress();
+
+        int64_t Displacement = TargetAddr - EdgeAddr + 4;
+        if (isInRangeForImmS32(Displacement)) {
+          // Change the edge kind as we don't go through GOT anymore. This is
+          // for formal correctness only. Technically, the two relocation kinds
+          // are resolved the same way.
+          E.setKind(x86_64::Delta32);
+          E.setTarget(GOTTarget);
+          E.setAddend(E.getAddend() - 4);
+          auto *BlockData = reinterpret_cast<uint8_t *>(
+              const_cast<char *>(B->getContent().data()));
+          BlockData[E.getOffset() - 2] = 0x8d;
+          LLVM_DEBUG({
+            dbgs() << "  Replaced GOT load wih LEA:\n    ";
+            printEdge(dbgs(), *B, E, getEdgeKindName(E.getKind()));
+            dbgs() << "\n";
+          });
+        }
+      } else if (E.getKind() == x86_64::BranchPCRel32ToPtrJumpStubBypassable) {
+        auto &StubBlock = E.getTarget().getBlock();
+        assert(StubBlock.getSize() == sizeof(PointerJumpStubContent) &&
+               "Stub block should be stub sized");
+        assert(StubBlock.edges_size() == 1 &&
+               "Stub block should only have one outgoing edge");
+
+        auto &GOTBlock = StubBlock.edges().begin()->getTarget().getBlock();
+        assert(GOTBlock.getSize() == G.getPointerSize() &&
+               "GOT block should be pointer sized");
+        assert(GOTBlock.edges_size() == 1 &&
+               "GOT block should only have one outgoing edge");
+
+        auto &GOTTarget = GOTBlock.edges().begin()->getTarget();
+        JITTargetAddress EdgeAddr = B->getAddress() + E.getOffset();
+        JITTargetAddress TargetAddr = GOTTarget.getAddress();
+
+        int64_t Displacement = TargetAddr - EdgeAddr + 4;
+        if (isInRangeForImmS32(Displacement)) {
+          E.setKind(x86_64::BranchPCRel32);
+          E.setTarget(GOTTarget);
+          LLVM_DEBUG({
+            dbgs() << "  Replaced stub branch with direct branch:\n    ";
+            printEdge(dbgs(), *B, E, getEdgeKindName(E.getKind()));
+            dbgs() << "\n";
+          });
+        }
+      }
+
+  return Error::success();
+}
+
 } // end namespace x86_64
 } // end namespace jitlink
 } // end namespace llvm