[JITLink][PowerPC][WIP] Handle case when local branch target is out of range

[lkail]

Summary

R_PPC64_REL24 relocation is able to branch directly if the displacement of the branching target is within int26 range. If not, we have to create a stub to perform indirect branching.

This patch assumes all local branching need a stub and optimize it via a PreFixupPass.

Since rtdyld for ppc64 has handled such case, now we implement it for jitlink for ppc64.

llvm/include/llvm/ExecutionEngine/JITLink/ppc64.h

enum EdgeKind_ppc64 : Edge::Kind {
   TOC,
   TOCDelta16,
   TOCDelta16DS,
   TOCDelta16HA,
   TOCDelta16HI,
   TOCDelta16LO,
   TOCDelta16LODS,
   CallBranchDelta,
+  CallBranchDeltaRelaxable,
   // Need to restore r2 after the bl, suggesting the bl is followed by a nop.
   CallBranchDeltaRestoreTOC,
   // Request calling function with TOC.
   RequestCall,
   // Request calling function without TOC.
   RequestCallNoTOC,
   RequestTLSDescInGOTAndTransformToTOCDelta16HA,
   RequestTLSDescInGOTAndTransformToTOCDelta16LO,
inline Symbol &createAnonymousPointerJumpStub(LinkGraph &G,
   PLTCallStubInfo StubInfo = pickStub<Endianness>(StubKind);
   Block &B = G.createContentBlock(StubSection, StubInfo.Content,
                                   orc::ExecutorAddr(), 4, 0);
   for (auto const &Reloc : StubInfo.Relocs)
     B.addEdge(Reloc.K, Reloc.Offset, PointerSymbol, Reloc.A);
   return G.addAnonymousSymbol(B, 0, StubInfo.Content.size(), true, false);
 }
 
+inline Edge &findEdgeToFinalTargetFromStub(Block &StubEntryBlock) {
+  Symbol *GOTTarget = nullptr;
+  for (auto &E : StubEntryBlock.edges()) {
+    if (GOTTarget == nullptr) {
+      GOTTarget = &E.getTarget();
+      continue;
+    }
+    assert(&E.getTarget() == GOTTarget &&
+           "The stub entry is expected pointing to exact one target");
+  }
+  assert(StubEntryBlock.edges_size() > 0 &&
+         "The stub entry is expected to have at least one edge");
+  Symbol &GOTEntry = StubEntryBlock.edges().begin()->getTarget();
+  Block &GOTEntryBlock = GOTEntry.getBlock();
+  assert(GOTEntryBlock.edges_size() == 1 &&
+         "The GOT entry is expected pointing to exact one target");
+  return *GOTEntryBlock.edges().begin();
+}
+
+inline Symbol &findFinalTargetFromStub(Block &StubEntryBlock) {
+  return findEdgeToFinalTargetFromStub(StubEntryBlock).getTarget();
+}
+
 template <support::endianness Endianness>
 class TOCTableManager : public TableManager<TOCTableManager<Endianness>> {
 public:
   // FIXME: `llvm-jitlink -check` relies this name to be $__GOT.
   static StringRef getSectionName() { return "$__GOT"; }
 
   bool visitEdge(LinkGraph &G, Block *B, Edge &E) {
     Edge::Kind K = E.getKind();
if (K == ppc64::RequestCall) {
         E.setKind(ppc64::CallBranchDeltaRestoreTOC);
         this->StubKind = LongBranchSaveR2;
         // FIXME: We assume the addend to the external target is zero. It's
         // quite unusual that the addend of an external target to be non-zero as
         // if we have known the layout of the external object.
         E.setTarget(this->getEntryForTarget(G, E.getTarget()));
         // Addend to the stub is zero.
         E.setAddend(0);
-      } else
+      } else {
         // TODO: There are cases a local function call need a call stub.
         // 1. Caller uses TOC, the callee doesn't, need a r2 save stub.
         // 2. Caller doesn't use TOC, the callee does, need a r12 setup stub.
-        // 3. Branching target is out of range.
-        E.setKind(ppc64::CallBranchDelta);
+        // We always create a stub to perform local branching. If the branch
+        // target is within branching range, we can optimize that in a
+        // PreFixupPass.
+        this->StubKind = LongBranch;
+        E.setTarget(this->getEntryForTarget(G, E.getTarget()));
+        E.setKind(ppc64::CallBranchDeltaRelaxable);
+        Edge &TargetEdge =
+            findEdgeToFinalTargetFromStub(E.getTarget().getBlock());
+        TargetEdge.setAddend(E.getAddend());
+        E.setAddend(0);
+      }
       return true;
     }
     if (K == ppc64::RequestCallNoTOC) {
       E.setKind(ppc64::CallBranchDelta);
       this->StubKind = LongBranchNoTOC;
       E.setTarget(this->getEntryForTarget(G, E.getTarget()));
       return true;
     }
inline Error applyFixup(LinkGraph &G, Block &B, const Edge &E,
   default:
     return make_error<JITLinkError>(
         "In graph " + G.getName() + ", section " + B.getSection().getName() +
         " unsupported edge kind " + getEdgeKindName(E.getKind()));
   }
   return Error::success();
 }
 
+/// Optimize the GOT and Stub relocations if the edge target address is in
+/// range.
+Error optimizeGOTAndStubAccesses(LinkGraph &G);
+
 } // end namespace llvm::jitlink::ppc64
 
 #endif // LLVM_EXECUTIONENGINE_JITLINK_PPC64_H

llvm/lib/ExecutionEngine/JITLink/ELF_ppc64.cpp

if (Ctx->shouldAddDefaultTargetPasses(G->getTargetTriple())) {
     if (auto MarkLive = Ctx->getMarkLivePass(G->getTargetTriple()))
       Config.PrePrunePasses.push_back(std::move(MarkLive));
     else
       Config.PrePrunePasses.push_back(markAllSymbolsLive);
   }
 
   Config.PostPrunePasses.push_back(buildTables_ELF_ppc64<Endianness>);
 
+  // Add GOT/Stubs optimizer pass.
+  Config.PreFixupPasses.push_back(ppc64::optimizeGOTAndStubAccesses);
+
   if (auto Err = Ctx->modifyPassConfig(*G, Config))
     return Ctx->notifyFailed(std::move(Err));
 
   ELFJITLinker_ppc64<Endianness>::link(std::move(Ctx), std::move(G),
                                        std::move(Config));
 }
 
 Expected<std::unique_ptr<LinkGraph>>

llvm/lib/ExecutionEngine/JITLink/ppc64.cpp

const char *getEdgeKindName(Edge::Kind K) {
   case TOCDelta16HI:
     return "TOCDelta16HI";
   case TOCDelta16LO:
     return "TOCDelta16LO";
   case TOCDelta16LODS:
     return "TOCDelta16LODS";
   case CallBranchDelta:
     return "CallBranchDelta";
+  case CallBranchDeltaRelaxable:
+    return "CallBranchDeltaRelaxable";
   case CallBranchDeltaRestoreTOC:
     return "CallBranchDeltaRestoreTOC";
   case RequestCall:
     return "RequestCall";
   case RequestCallNoTOC:
     return "RequestCallNoTOC";
   case RequestTLSDescInGOTAndTransformToTOCDelta16HA:
     return "RequestTLSDescInGOTAndTransformToTOCDelta16HA";
   case RequestTLSDescInGOTAndTransformToTOCDelta16LO:
     return "RequestTLSDescInGOTAndTransformToTOCDelta16LO";
   default:
     return getGenericEdgeKindName(static_cast<Edge::Kind>(K));
   }
 }
 
+Error optimizeGOTAndStubAccesses(LinkGraph &G) {
+  LLVM_DEBUG(dbgs() << "Optimizing GOT entries and stubs:\n");
+
+  for (auto *B : G.blocks()) {
+    for (auto &E : B->edges()) {
+      Edge::Kind K = E.getKind();
+      if (K == ppc64::CallBranchDeltaRelaxable) {
+        Symbol &StubEntry = E.getTarget();
+        Block &StubEntryBlock = StubEntry.getBlock();
+        Edge &TargetEdge = findEdgeToFinalTargetFromStub(StubEntryBlock);
+        Symbol &Target = TargetEdge.getTarget();
+        orc::ExecutorAddr TargetAddr = Target.getAddress();
+        orc::ExecutorAddr EdgeAddr = B->getFixupAddress(E);
+        int64_t Displacement = TargetAddr - EdgeAddr;
+        if (LLVM_LIKELY(isInt<26>(Displacement))) {
+          E.setKind(CallBranchDelta);
+          E.setTarget(Target);
+          E.setAddend(TargetEdge.getAddend());
+          LLVM_DEBUG({
+            dbgs() << "  Replaced branching to stub with direct branching:\n";
+            printEdge(dbgs(), *B, E, getEdgeKindName(K));
+            dbgs() << "\n";
+          });
+        } else {
+          E.setKind(CallBranchDelta);
+          LLVM_DEBUG({
+            dbgs() << "  Keep branching to stub:\n";
+            printEdge(dbgs(), *B, E, getEdgeKindName(K));
+            dbgs() << "\n";
+          });
+        }
+      }
+    }
+  }
+  return Error::success();
+}
+
 } // end namespace llvm::jitlink::ppc64