[llvm] [llvm-jitlink] Add optional index argument in jitlink-check stub_addr() expressions (PR #78369)

Stefan Gränitz via llvm-commits llvm-commits at lists.llvm.org
Tue Jan 16 16:16:02 PST 2024


https://github.com/weliveindetail created https://github.com/llvm/llvm-project/pull/78369

We use `jitlink-check` lines in LIT tests as the primary tool for testing JITLink backends. Parsing and evaluation of the expressions is implemented in `RuntimeDyldChecker`. The `stub_addr(obj, name)` expression allows to obtain the linker-generated stub for the external symbol `name` in object file `obj`. This is limiting JITLink backends to a single stub for an external symbol.

This patch adds support for an index parameter: `stub_addr(obj, name, index=0)`. It allows us to inspect more than a single stub in our tests and, respectively, it enables JITLink backends to emit multiple stubs. This is necessary for the AArch32 backend, which must be able to emit two different kinds of stubs depending on the instruction set state (Arm/Thumb) of the relocation site. Since the new parameter is optional, we don't have to update existing tests.

>From 6e3b317fcf8a34c4fe36df51105172cda08281c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefan=20Gr=C3=A4nitz?= <stefan.graenitz at gmail.com>
Date: Tue, 16 Jan 2024 22:33:02 +0100
Subject: [PATCH 1/2] [llvm-jitlink] Refactor GOT and stubs registration (NFC)

---
 llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp | 35 +++-----
 llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp  | 79 ++++++++-----------
 .../tools/llvm-jitlink/llvm-jitlink-macho.cpp | 30 ++-----
 llvm/tools/llvm-jitlink/llvm-jitlink.cpp      | 30 +++++++
 llvm/tools/llvm-jitlink/llvm-jitlink.h        | 10 +++
 5 files changed, 89 insertions(+), 95 deletions(-)

diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp
index 283e655205d780..5271fdb5565904 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp
@@ -108,34 +108,17 @@ Error registerCOFFGraphInfo(Session &S, LinkGraph &G) {
       if (Sym->getAddress() > LastSym->getAddress())
         LastSym = Sym;
 
-      if (isGOTSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in GOT section",
-                                         inconvertibleErrorCode());
-
-        // If this is a GOT symbol with size (i.e. not the GOT start symbol)
-        // then add it to the GOT entry info table.
-        if (Sym->getSize() != 0) {
-          if (auto TS = getCOFFGOTTarget(G, Sym->getBlock()))
-            FileInfo.GOTEntryInfos[TS->getName()] = {
-                Sym->getSymbolContent(), Sym->getAddress().getValue(),
-                Sym->getTargetFlags()};
-          else
-            return TS.takeError();
+      if (isGOTSection || isStubsSection) {
+        if (isGOTSection) {
+          // Skip the GOT start symbol
+          if (Sym->getSize() != 0)
+            if (Error E = FileInfo.registerGOTEntry(G, *Sym, getCOFFGOTTarget))
+              return E;
+        } else {
+          if (Error E = FileInfo.registerStubEntry(G, *Sym, getCOFFStubTarget))
+            return E;
         }
         SectionContainsContent = true;
-      } else if (isStubsSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in Stub section",
-                                         inconvertibleErrorCode());
-
-        if (auto TS = getCOFFStubTarget(G, Sym->getBlock()))
-          FileInfo.StubInfos[TS->getName()] = {Sym->getSymbolContent(),
-                                               Sym->getAddress().getValue(),
-                                               Sym->getTargetFlags()};
-        else
-          return TS.takeError();
-        SectionContainsContent = true;
       }
 
       if (Sym->hasName()) {
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
index a02468758b33f6..c6b4218aad7af8 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
@@ -72,13 +72,29 @@ static Expected<Symbol &> getELFStubTarget(LinkGraph &G, Block &B) {
   return getELFGOTTarget(G, GOTSym.getBlock());
 }
 
-static Expected<std::string> getELFAArch32StubTargetName(LinkGraph &G,
-                                                         Block &B) {
+static Expected<Symbol &> getELFAArch32StubTarget(LinkGraph &G, Block &B) {
   auto E = getFirstRelocationEdge(G, B);
   if (!E)
     return E.takeError();
-  Symbol &StubTarget = E->getTarget();
-  return StubTarget.getName().str();
+  return E->getTarget();
+}
+
+enum SectionType { GOT, Stubs, AArch32Stubs, Other };
+
+static Error registerSymbol(LinkGraph &G, Symbol &Sym, Session::FileInfo &FI,
+                            SectionType SecType) {
+  switch (SecType) {
+  case GOT:
+    if (Sym.getSize() == 0)
+      return Error::success(); // Skip the GOT start symbol
+    return FI.registerGOTEntry(G, Sym, getELFGOTTarget);
+  case Stubs:
+    return FI.registerStubEntry(G, Sym, getELFStubTarget);
+  case AArch32Stubs:
+    return FI.registerStubEntry(G, Sym, getELFAArch32StubTarget);
+  case Other:
+    return Error::success();
+  }
 }
 
 namespace llvm {
@@ -113,9 +129,16 @@ Error registerELFGraphInfo(Session &S, LinkGraph &G) {
                                          "\"",
                                      inconvertibleErrorCode());
 
-    bool isGOTSection = isELFGOTSection(Sec);
-    bool isStubsSection = isELFStubsSection(Sec);
-    bool isAArch32StubsSection = isELFAArch32StubsSection(Sec);
+    SectionType SecType;
+    if (isELFGOTSection(Sec)) {
+      SecType = GOT;
+    } else if (isELFStubsSection(Sec)) {
+      SecType = Stubs;
+    } else if (isELFAArch32StubsSection(Sec)) {
+      SecType = AArch32Stubs;
+    } else {
+      SecType = Other;
+    }
 
     bool SectionContainsContent = false;
     bool SectionContainsZeroFill = false;
@@ -128,45 +151,9 @@ Error registerELFGraphInfo(Session &S, LinkGraph &G) {
       if (Sym->getAddress() > LastSym->getAddress())
         LastSym = Sym;
 
-      if (isGOTSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in GOT section",
-                                         inconvertibleErrorCode());
-
-        // If this is a GOT symbol with size (i.e. not the GOT start symbol)
-        // then add it to the GOT entry info table.
-        if (Sym->getSize() != 0) {
-          if (auto TS = getELFGOTTarget(G, Sym->getBlock()))
-            FileInfo.GOTEntryInfos[TS->getName()] = {
-                Sym->getSymbolContent(), Sym->getAddress().getValue(),
-                Sym->getTargetFlags()};
-          else
-            return TS.takeError();
-        }
-        SectionContainsContent = true;
-      } else if (isStubsSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in Stub section",
-                                         inconvertibleErrorCode());
-
-        if (auto TS = getELFStubTarget(G, Sym->getBlock()))
-          FileInfo.StubInfos[TS->getName()] = {Sym->getSymbolContent(),
-                                               Sym->getAddress().getValue(),
-                                               Sym->getTargetFlags()};
-        else
-          return TS.takeError();
-        SectionContainsContent = true;
-      } else if (isAArch32StubsSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in Stub section",
-                                         inconvertibleErrorCode());
-
-        if (auto Name = getELFAArch32StubTargetName(G, Sym->getBlock()))
-          FileInfo.StubInfos[*Name] = {Sym->getSymbolContent(),
-                                       Sym->getAddress().getValue(),
-                                       Sym->getTargetFlags()};
-        else
-          return Name.takeError();
+      if (SecType != Other) {
+        if (Error Err = registerSymbol(G, *Sym, FileInfo, SecType))
+          return Err;
         SectionContainsContent = true;
       }
 
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp
index 7dcadd94c2365a..2c60c802293a1e 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp
@@ -111,29 +111,13 @@ Error registerMachOGraphInfo(Session &S, LinkGraph &G) {
         FirstSym = Sym;
       if (Sym->getAddress() > LastSym->getAddress())
         LastSym = Sym;
-      if (isGOTSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in GOT section",
-                                         inconvertibleErrorCode());
-
-        if (auto TS = getMachOGOTTarget(G, Sym->getBlock()))
-          FileInfo.GOTEntryInfos[TS->getName()] = {Sym->getSymbolContent(),
-                                                   Sym->getAddress().getValue(),
-                                                   Sym->getTargetFlags()};
-        else
-          return TS.takeError();
-        SectionContainsContent = true;
-      } else if (isStubsSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in Stub section",
-                                         inconvertibleErrorCode());
-
-        if (auto TS = getMachOStubTarget(G, Sym->getBlock()))
-          FileInfo.StubInfos[TS->getName()] = {Sym->getSymbolContent(),
-                                               Sym->getAddress().getValue(),
-                                               Sym->getTargetFlags()};
-        else
-          return TS.takeError();
+      if (isGOTSection || isStubsSection) {
+        Error Err =
+            isGOTSection
+                ? FileInfo.registerGOTEntry(G, *Sym, getMachOGOTTarget)
+                : FileInfo.registerStubEntry(G, *Sym, getMachOStubTarget);
+        if (Err)
+          return Err;
         SectionContainsContent = true;
       } else if (Sym->hasName()) {
         if (Sym->isSymbolZeroFill()) {
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index aa032c97485fbf..8c18610313ce8f 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -1183,6 +1183,36 @@ Error Session::loadAndLinkDynamicLibrary(JITDylib &JD, StringRef LibPath) {
   return Error::success();
 }
 
+Error Session::FileInfo::registerGOTEntry(
+    LinkGraph &G, Symbol &Sym, GetSymbolTargetFunction GetSymbolTarget) {
+  if (Sym.isSymbolZeroFill())
+    return make_error<StringError>("Unexpected zero-fill symbol in section " +
+                                       Sym.getBlock().getSection().getName(),
+                                   inconvertibleErrorCode());
+  auto TS = GetSymbolTarget(G, Sym.getBlock());
+  if (!TS)
+    return TS.takeError();
+  GOTEntryInfos[TS->getName()] = {Sym.getSymbolContent(),
+                                  Sym.getAddress().getValue(),
+                                  Sym.getTargetFlags()};
+  return Error::success();
+}
+
+Error Session::FileInfo::registerStubEntry(
+    LinkGraph &G, Symbol &Sym, GetSymbolTargetFunction GetSymbolTarget) {
+  if (Sym.isSymbolZeroFill())
+    return make_error<StringError>("Unexpected zero-fill symbol in section " +
+                                       Sym.getBlock().getSection().getName(),
+                                   inconvertibleErrorCode());
+  auto TS = GetSymbolTarget(G, Sym.getBlock());
+  if (!TS)
+    return TS.takeError();
+  StubInfos[TS->getName()] = {Sym.getSymbolContent(),
+                              Sym.getAddress().getValue(),
+                              Sym.getTargetFlags()};
+  return Error::success();
+}
+
 Expected<Session::FileInfo &> Session::findFileInfo(StringRef FileName) {
   auto FileInfoItr = FileInfos.find(FileName);
   if (FileInfoItr == FileInfos.end())
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
index 3ff406b7b82dfd..93a00266b15043 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.h
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.h
@@ -51,6 +51,16 @@ struct Session {
     StringMap<MemoryRegionInfo> SectionInfos;
     StringMap<MemoryRegionInfo> StubInfos;
     StringMap<MemoryRegionInfo> GOTEntryInfos;
+
+    using Symbol = jitlink::Symbol;
+    using LinkGraph = jitlink::LinkGraph;
+    using GetSymbolTargetFunction =
+        unique_function<Expected<Symbol &>(LinkGraph &G, jitlink::Block &)>;
+
+    Error registerGOTEntry(LinkGraph &G, Symbol &Sym,
+                           GetSymbolTargetFunction GetSymbolTarget);
+    Error registerStubEntry(LinkGraph &G, Symbol &Sym,
+                            GetSymbolTargetFunction GetSymbolTarget);
   };
 
   using DynLibJDMap = std::map<std::string, orc::JITDylib *>;

>From 3104039c376190ff78edc84b0de5a9cc4ef69230 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefan=20Gr=C3=A4nitz?= <stefan.graenitz at gmail.com>
Date: Sat, 13 Jan 2024 23:35:54 +0100
Subject: [PATCH 2/2] [llvm-jitlink] Allow optional index argument in
 jitlink-check stub_addr() expressions

---
 .../llvm/ExecutionEngine/RuntimeDyldChecker.h |  2 +-
 .../RuntimeDyld/RuntimeDyldChecker.cpp        | 21 +++++--
 .../RuntimeDyld/RuntimeDyldCheckerImpl.h      |  3 +-
 llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp  |  2 +-
 llvm/tools/llvm-jitlink/llvm-jitlink.cpp      | 55 ++++++++++++++++---
 llvm/tools/llvm-jitlink/llvm-jitlink.h        |  8 ++-
 llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp        | 15 +++--
 7 files changed, 81 insertions(+), 25 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/RuntimeDyldChecker.h b/llvm/include/llvm/ExecutionEngine/RuntimeDyldChecker.h
index 80e4bbf494339c..55cc16834976d0 100644
--- a/llvm/include/llvm/ExecutionEngine/RuntimeDyldChecker.h
+++ b/llvm/include/llvm/ExecutionEngine/RuntimeDyldChecker.h
@@ -154,7 +154,7 @@ class RuntimeDyldChecker {
   using GetSectionInfoFunction = std::function<Expected<MemoryRegionInfo>(
       StringRef FileName, StringRef SectionName)>;
   using GetStubInfoFunction = std::function<Expected<MemoryRegionInfo>(
-      StringRef StubContainer, StringRef TargetName)>;
+      StringRef StubContainer, StringRef TargetName, int64_t StubIndex)>;
   using GetGOTInfoFunction = std::function<Expected<MemoryRegionInfo>(
       StringRef GOTContainer, StringRef TargetName)>;
 
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldChecker.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldChecker.cpp
index 7fadbdd6a1fff2..801496b8b91035 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldChecker.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldChecker.cpp
@@ -400,6 +400,15 @@ class RuntimeDyldCheckerExprEval {
     StringRef Symbol;
     std::tie(Symbol, RemainingExpr) = parseSymbol(RemainingExpr);
 
+    // Parse optional stub index parameter
+    int64_t StubIndex = 0;
+    if (RemainingExpr.starts_with(",")) {
+      RemainingExpr = RemainingExpr.substr(1).ltrim();
+      EvalResult Number;
+      std::tie(Number, RemainingExpr) = evalNumberExpr(RemainingExpr);
+      StubIndex = Number.getValue();
+    }
+
     if (!RemainingExpr.starts_with(")"))
       return std::make_pair(
           unexpectedToken(RemainingExpr, Expr, "expected ')'"), "");
@@ -408,7 +417,7 @@ class RuntimeDyldCheckerExprEval {
     uint64_t StubAddr;
     std::string ErrorMsg;
     std::tie(StubAddr, ErrorMsg) = Checker.getStubOrGOTAddrFor(
-        StubContainerName, Symbol, PCtx.IsInsideLoad, IsStubAddr);
+        StubContainerName, Symbol, StubIndex, PCtx.IsInsideLoad, IsStubAddr);
 
     if (ErrorMsg != "")
       return std::make_pair(EvalResult(ErrorMsg), "");
@@ -985,11 +994,13 @@ std::pair<uint64_t, std::string> RuntimeDyldCheckerImpl::getSectionAddr(
 }
 
 std::pair<uint64_t, std::string> RuntimeDyldCheckerImpl::getStubOrGOTAddrFor(
-    StringRef StubContainerName, StringRef SymbolName, bool IsInsideLoad,
-    bool IsStubAddr) const {
+    StringRef StubContainerName, StringRef SymbolName, int64_t StubIndex,
+    bool IsInsideLoad, bool IsStubAddr) const {
 
-  auto StubInfo = IsStubAddr ? GetStubInfo(StubContainerName, SymbolName)
-                             : GetGOTInfo(StubContainerName, SymbolName);
+  assert((StubIndex == 0 || IsStubAddr) && "Indexing only supported for stubs");
+  auto StubInfo = IsStubAddr
+                      ? GetStubInfo(StubContainerName, SymbolName, StubIndex)
+                      : GetGOTInfo(StubContainerName, SymbolName);
 
   if (!StubInfo) {
     std::string ErrMsg;
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCheckerImpl.h b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCheckerImpl.h
index 9f44a9389f4734..0f43e877318895 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCheckerImpl.h
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCheckerImpl.h
@@ -64,7 +64,8 @@ class RuntimeDyldCheckerImpl {
 
   std::pair<uint64_t, std::string>
   getStubOrGOTAddrFor(StringRef StubContainerName, StringRef Symbol,
-                      bool IsInsideLoad, bool IsStubAddr) const;
+                      int64_t StubIndex, bool IsInsideLoad,
+                      bool IsStubAddr) const;
 
   std::optional<uint64_t> getSectionLoadAddress(void *LocalAddr) const;
 
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
index c6b4218aad7af8..1f0fca2202a0ef 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
@@ -91,7 +91,7 @@ static Error registerSymbol(LinkGraph &G, Symbol &Sym, Session::FileInfo &FI,
   case Stubs:
     return FI.registerStubEntry(G, Sym, getELFStubTarget);
   case AArch32Stubs:
-    return FI.registerStubEntry(G, Sym, getELFAArch32StubTarget);
+    return FI.registerMultiStubEntry(G, Sym, getELFAArch32StubTarget);
   case Other:
     return Error::success();
   }
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 8c18610313ce8f..68f39b880c69c9 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -331,8 +331,12 @@ operator<<(raw_ostream &OS, const Session::FileInfo &FI) {
     OS << "  Section \"" << SIKV.first() << "\": " << SIKV.second << "\n";
   for (auto &GOTKV : FI.GOTEntryInfos)
     OS << "  GOT \"" << GOTKV.first() << "\": " << GOTKV.second << "\n";
-  for (auto &StubKV : FI.StubInfos)
-    OS << "  Stub \"" << StubKV.first() << "\": " << StubKV.second << "\n";
+  for (auto &StubKVs : FI.StubInfos) {
+    OS << "  Stubs \"" << StubKVs.first() << "\":";
+    for (auto MemRegion : StubKVs.second)
+      OS << " " << MemRegion;
+    OS << "\n";
+  }
   return OS;
 }
 
@@ -1207,9 +1211,35 @@ Error Session::FileInfo::registerStubEntry(
   auto TS = GetSymbolTarget(G, Sym.getBlock());
   if (!TS)
     return TS.takeError();
-  StubInfos[TS->getName()] = {Sym.getSymbolContent(),
-                              Sym.getAddress().getValue(),
-                              Sym.getTargetFlags()};
+
+  SmallVector<MemoryRegionInfo> &Entry = StubInfos[TS->getName()];
+  Entry.insert(Entry.begin(),
+               {Sym.getSymbolContent(), Sym.getAddress().getValue(),
+                Sym.getTargetFlags()});
+  return Error::success();
+}
+
+Error Session::FileInfo::registerMultiStubEntry(
+    LinkGraph &G, Symbol &Sym, GetSymbolTargetFunction GetSymbolTarget) {
+  if (Sym.isSymbolZeroFill())
+    return make_error<StringError>("Unexpected zero-fill symbol in section " +
+                                       Sym.getBlock().getSection().getName(),
+                                   inconvertibleErrorCode());
+
+  auto Target = GetSymbolTarget(G, Sym.getBlock());
+  if (!Target)
+    return Target.takeError();
+
+  SmallVector<MemoryRegionInfo> &Entry = StubInfos[Target->getName()];
+  Entry.emplace_back(Sym.getSymbolContent(), Sym.getAddress().getValue(),
+                     Sym.getTargetFlags());
+
+  // Let's keep stubs ordered by ascending address.
+  std::sort(Entry.begin(), Entry.end(),
+            [](const MemoryRegionInfo &L, const MemoryRegionInfo &R) {
+              return L.getTargetAddress() < R.getTargetAddress();
+            });
+
   return Error::success();
 }
 
@@ -1236,7 +1266,8 @@ Session::findSectionInfo(StringRef FileName, StringRef SectionName) {
 }
 
 Expected<Session::MemoryRegionInfo &>
-Session::findStubInfo(StringRef FileName, StringRef TargetName) {
+Session::findStubInfo(StringRef FileName, StringRef TargetName,
+                      uint64_t StubIndex) {
   auto FI = findFileInfo(FileName);
   if (!FI)
     return FI.takeError();
@@ -1246,7 +1277,12 @@ Session::findStubInfo(StringRef FileName, StringRef TargetName) {
                                        "\" registered for file \"" + FileName +
                                        "\"",
                                    inconvertibleErrorCode());
-  return StubInfoItr->second;
+  if (StubIndex >= StubInfoItr->second.size())
+    return make_error<StringError>(
+        "no stub for \"" + TargetName + " with index " + Twine(StubIndex) +
+            "\" registered for file \"" + FileName + "\"",
+        inconvertibleErrorCode());
+  return StubInfoItr->second[StubIndex];
 }
 
 Expected<Session::MemoryRegionInfo &>
@@ -2015,8 +2051,9 @@ static Error runChecks(Session &S, Triple TT, SubtargetFeatures Features) {
     return S.findSectionInfo(FileName, SectionName);
   };
 
-  auto GetStubInfo = [&S](StringRef FileName, StringRef SectionName) {
-    return S.findStubInfo(FileName, SectionName);
+  auto GetStubInfo = [&S](StringRef FileName, StringRef SectionName,
+                          int64_t StubIndex) {
+    return S.findStubInfo(FileName, SectionName, StubIndex);
   };
 
   auto GetGOTInfo = [&S](StringRef FileName, StringRef SectionName) {
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
index 93a00266b15043..46346bfde5e957 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.h
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.h
@@ -49,7 +49,7 @@ struct Session {
 
   struct FileInfo {
     StringMap<MemoryRegionInfo> SectionInfos;
-    StringMap<MemoryRegionInfo> StubInfos;
+    StringMap<SmallVector<MemoryRegionInfo, 1>> StubInfos;
     StringMap<MemoryRegionInfo> GOTEntryInfos;
 
     using Symbol = jitlink::Symbol;
@@ -61,6 +61,8 @@ struct Session {
                            GetSymbolTargetFunction GetSymbolTarget);
     Error registerStubEntry(LinkGraph &G, Symbol &Sym,
                             GetSymbolTargetFunction GetSymbolTarget);
+    Error registerMultiStubEntry(LinkGraph &G, Symbol &Sym,
+                                 GetSymbolTargetFunction GetSymbolTarget);
   };
 
   using DynLibJDMap = std::map<std::string, orc::JITDylib *>;
@@ -73,8 +75,8 @@ struct Session {
   Expected<FileInfo &> findFileInfo(StringRef FileName);
   Expected<MemoryRegionInfo &> findSectionInfo(StringRef FileName,
                                                StringRef SectionName);
-  Expected<MemoryRegionInfo &> findStubInfo(StringRef FileName,
-                                            StringRef TargetName);
+  Expected<MemoryRegionInfo &>
+  findStubInfo(StringRef FileName, StringRef TargetName, uint64_t StubIndex);
   Expected<MemoryRegionInfo &> findGOTEntryInfo(StringRef FileName,
                                                 StringRef TargetName);
 
diff --git a/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp b/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
index 107b555a99faa4..266e24e067aa4f 100644
--- a/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
+++ b/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
@@ -925,9 +925,10 @@ static int linkAndVerify() {
     return SecInfo;
   };
 
-  auto GetStubInfo = [&Dyld, &StubMap](StringRef StubContainer,
-                                       StringRef SymbolName)
-      -> Expected<RuntimeDyldChecker::MemoryRegionInfo> {
+  auto GetStubInfo =
+      [&Dyld, &StubMap](
+          StringRef StubContainer, StringRef SymbolName,
+          int64_t StubIndex) -> Expected<RuntimeDyldChecker::MemoryRegionInfo> {
     if (!StubMap.count(StubContainer))
       return make_error<StringError>("Stub container not found: " +
                                          StubContainer,
@@ -947,6 +948,11 @@ static int linkAndVerify() {
     return StubMemInfo;
   };
 
+  auto GetGOTInfo = [&GetStubInfo](StringRef StubContainer,
+                                   StringRef SymbolName) {
+    return GetStubInfo(StubContainer, SymbolName, 0);
+  };
+
   // We will initialize this below once we have the first object file and can
   // know the endianness.
   std::unique_ptr<RuntimeDyldChecker> Checker;
@@ -977,8 +983,7 @@ static int linkAndVerify() {
 
     if (!Checker)
       Checker = std::make_unique<RuntimeDyldChecker>(
-          IsSymbolValid, GetSymbolInfo, GetSectionInfo, GetStubInfo,
-          GetStubInfo,
+          IsSymbolValid, GetSymbolInfo, GetSectionInfo, GetStubInfo, GetGOTInfo,
           Obj.isLittleEndian() ? llvm::endianness::little
                                : llvm::endianness::big,
           TheTriple, MCPU, SubtargetFeatures(), dbgs());



More information about the llvm-commits mailing list