[JITLink][PowerPC] Change method to check if a symbol is external to current object

[lkail]

Summary

After PrePrunePass claimOrExternalizeWeakAndCommonSymbols, a defined symbol might become external. So determine a function call is external or not when building the linkgraph is not accurate. This largely affects updating TOC pointer on PowerPC. TOC pointer is supposed to be the same in one object file(if no mulitple TOC appears) and is updated when control flow transferred to another object file.

This patch defers checking a function call is external or not, in buildTables_ELF_ppc64 which is a PostPrunePass.

This patch fixes failures when jitlink -orc-runtime=/path/to/libort_rt.a is used.

[lkail]

Pre-commit test is in https://github.com/llvm/llvm-project/commit/cee4494c14d821104ad839c44c047bacd652b8dd.

[lhames]

LGTM. Thanks @lkail!

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

enum EdgeKind_ppc64 : Edge::Kind {
   Delta16LO,
   TOCDelta16HA,
   TOCDelta16LO,
   TOCDelta16DS,
   TOCDelta16LODS,
   CallBranchDelta,
   // Need to restore r2 after the bl, suggesting the bl is followed by a nop.
   CallBranchDeltaRestoreTOC,
-  // Need PLT call stub using TOC, TOC pointer is not saved before branching.
-  RequestPLTCallStub,
-  // Need PLT call stub using TOC, TOC pointer is saved before branching.
-  RequestPLTCallStubSaveTOC,
-  // Need PLT call stub without using TOC.
-  RequestPLTCallStubNoTOC,
+  // Request calling function with TOC.
+  RequestCall,
+  // Request calling function without TOC.
+  RequestCallNoTOC,
 };
 
 enum PLTCallStubKind {
+  // Setup function entry(r12) and long branch to target using TOC.
   LongBranch,
+  // Save TOC pointer, setup function entry and long branch to target using TOC.
   LongBranchSaveR2,
+  // Setup function entry(r12) and long branch to target without using TOC.
   LongBranchNoTOC,
 };
 
 extern const char NullPointerContent[8];
 extern const char PointerJumpStubContent_big[20];
 extern const char PointerJumpStubContent_little[20];
 extern const char PointerJumpStubNoTOCContent_big[32];
 extern const char PointerJumpStubNoTOCContent_little[32];
 
 struct PLTCallStubReloc {
   Edge::Kind K;
   size_t Offset;
   Edge::AddendT A;
 };
 
 struct PLTCallStubInfo {
   ArrayRef<char> Content;
   SmallVector<PLTCallStubReloc, 2> Relocs;
 };
 
+inline bool isExternalSymbol(LinkGraph &G, Symbol &S) {
+  return llvm::count_if(G.external_symbols(), [&](const Symbol *ES) {
+           return ES->getName() == S.getName();
+         }) > 0;
+}
+
 template <support::endianness Endianness>
 inline PLTCallStubInfo pickStub(PLTCallStubKind StubKind) {
   constexpr bool isLE = Endianness == support::endianness::little;
   switch (StubKind) {
   case LongBranch: {
     ArrayRef<char> Content =
         isLE ? PointerJumpStubContent_little : PointerJumpStubContent_big;
     // Skip save r2.
public:
   bool visitEdge(LinkGraph &G, Block *B, Edge &E) {
     Edge::Kind K = E.getKind();
     switch (K) {
     case TOCDelta16HA:
     case TOCDelta16LO:
     case TOCDelta16DS:
     case TOCDelta16LODS:
     case CallBranchDeltaRestoreTOC:
-    case RequestPLTCallStub:
-    case RequestPLTCallStubSaveTOC:
+    case RequestCall:
       // Create TOC section if TOC relocation, PLT or GOT is used.
       getOrCreateTOCSection(G);
       return false;
     default:
       return false;
     }
   }
 

 class PLTTableManager : public TableManager<PLTTableManager<Endianness>> {
 public:
   PLTTableManager(TOCTableManager<Endianness> &TOC) : TOC(TOC) {}
 
   static StringRef getSectionName() { return "$__STUBS"; }
 
   bool visitEdge(LinkGraph &G, Block *B, Edge &E) {
     Edge::Kind K = E.getKind();
-    if (K == ppc64::RequestPLTCallStubSaveTOC && E.getTarget().isExternal()) {
+    bool isExternal = isExternalSymbol(G, E.getTarget());
+    if (K == ppc64::RequestCall) {
+      if (isExternal) {
         E.setKind(ppc64::CallBranchDeltaRestoreTOC);
         this->StubKind = LongBranchSaveR2;
         E.setTarget(this->getEntryForTarget(G, E.getTarget()));
+        // We previously set branching to local entry. Now reverse that
+        // operation.
+        E.setAddend(0);
+      } 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);
       return true;
     }
-    if (K == ppc64::RequestPLTCallStubNoTOC && E.getTarget().isExternal()) {
+    if (K == ppc64::RequestCallNoTOC) {
       E.setKind(ppc64::CallBranchDelta);
       this->StubKind = LongBranchNoTOC;
       E.setTarget(this->getEntryForTarget(G, E.getTarget()));
       return true;
     }
     return false;
   }
 

llvm/lib/ExecutionEngine/JITLink/ELF_ppc64.cpp

Symbol &createELFGOTHeader(LinkGraph &G,
   return TOC.getEntryForTarget(G, *TOCSymbol);
 }
 
 // Register preexisting GOT entries with TOC table manager.
 template <support::endianness Endianness>
 inline void
 registerExistingGOTEntries(LinkGraph &G,
                            ppc64::TOCTableManager<Endianness> &TOC) {
-  auto isGOTEntry = [](const Edge &E) {
-    return E.getKind() == ppc64::Pointer64 && E.getTarget().isExternal();
+  auto isGOTEntry = [&](const Edge &E) {
+    return E.getKind() == ppc64::Pointer64 &&
+           ppc64::isExternalSymbol(G, E.getTarget());
   };
   if (Section *dotTOCSection = G.findSectionByName(".toc")) {
     for (Block *B : dotTOCSection->blocks())
       for (Edge &E : B->edges())
         if (isGOTEntry(E))
           TOC.registerPreExistingEntry(E.getTarget(),
                                        G.addAnonymousSymbol(*B, E.getOffset(),
                                                             G.getPointerSize(),
case ELF::R_PPC64_REL16_HA:
       break;
     case ELF::R_PPC64_REL16_LO:
       Kind = ppc64::Delta16LO;
       break;
     case ELF::R_PPC64_REL32:
       Kind = ppc64::Delta32;
       break;
     case ELF::R_PPC64_REL24_NOTOC:
-    case ELF::R_PPC64_REL24: {
-      bool isLocal = !GraphSymbol->isExternal();
-      if (isLocal) {
-        // 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.
-        // FIXME: For a local call, we might need a thunk if branch target is
-        // out of range.
-        Kind = ppc64::CallBranchDelta;
-        // Branch to local entry.
-        Addend += ELF::decodePPC64LocalEntryOffset((*ObjSymbol)->st_other);
-      } else {
-        Kind = ELFReloc == ELF::R_PPC64_REL24 ? ppc64::RequestPLTCallStubSaveTOC
-                                              : ppc64::RequestPLTCallStubNoTOC;
-      }
+      Kind = ppc64::RequestCallNoTOC;
+      break;
+    case ELF::R_PPC64_REL24:
+      Kind = ppc64::RequestCall;
+      Addend = ELF::decodePPC64LocalEntryOffset((*ObjSymbol)->st_other);
       break;
-    }
     case ELF::R_PPC64_REL64:
       Kind = ppc64::Delta64;
       break;
     }
 
     Edge GE(Kind, Offset, *GraphSymbol, Addend);
     BlockToFix.addEdge(std::move(GE));
     return Error::success();

llvm/lib/ExecutionEngine/JITLink/ppc64.cpp

const char *getEdgeKindName(Edge::Kind K) {
   case TOCDelta16DS:
     return "TOCDelta16DS";
   case TOCDelta16LODS:
     return "TOCDelta16LODS";
   case CallBranchDelta:
     return "CallBranchDelta";
   case CallBranchDeltaRestoreTOC:
     return "CallBranchDeltaRestoreTOC";
-  case RequestPLTCallStub:
-    return "RequestPLTCallStub";
-  case RequestPLTCallStubSaveTOC:
-    return "RequestPLTCallStubSaveTOC";
-  case RequestPLTCallStubNoTOC:
-    return "RequestPLTCallStubNoTOC";
+  case RequestCall:
+    return "RequestCall";
+  case RequestCallNoTOC:
+    return "RequestCallNoTOC";
   default:
     return getGenericEdgeKindName(static_cast<Edge::Kind>(K));
   }
 }
 
 } // end namespace llvm::jitlink::ppc64