[llvm] bb41fc6 - [ORC-RT][ORC][MachO] Add executor-side symbol tables to MachO platform support.

Lang Hames via llvm-commits llvm-commits at lists.llvm.org
Sat Dec 2 15:32:46 PST 2023


Author: Lang Hames
Date: 2023-12-02T15:25:25-08:00
New Revision: bb41fc682ee779c775185ada1aeab6f4be5c282f

URL: https://github.com/llvm/llvm-project/commit/bb41fc682ee779c775185ada1aeab6f4be5c282f
DIFF: https://github.com/llvm/llvm-project/commit/bb41fc682ee779c775185ada1aeab6f4be5c282f.diff

LOG: [ORC-RT][ORC][MachO] Add executor-side symbol tables to MachO platform support.

Adds symbol tables to the JITDylibState struct in the ORC runtime
MachOPlatformRuntimeState class. This table will hold the addresses of
materialized symbols (registered by a new JITLink pass in MachOPlatform),
allowing these to be looked up in the executor without an IPC request to the
controller.

The old lookup-symbols callback (made by the runtime in response to dlsym
lookups) is replaced with a push-symbols callback that can trigger
materialization of requested symbols.

Holding a symbol table on the executor side should make repeat calls to dlsym
(and other symbol lookup operations) cheaper since the IPC to trigger
materialization happens at most once per symbol. It should also enable us (at
some point in the future) to symbolicate backtraces in JIT'd code even if the
controller process is gone (e.g. detached or crashed). The trade-off for this
is increased memory consumption in the executor and larger JIT'd data transfers
(since symbol names are now transferred to the executor unconditionally, even
though they may never be used).

Added: 
    

Modified: 
    compiler-rt/lib/orc/macho_platform.cpp
    compiler-rt/lib/orc/simple_packed_serialization.h
    llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
    llvm/include/llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h
    llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
    llvm/lib/ExecutionEngine/Orc/Shared/ObjectFormats.cpp

Removed: 
    


################################################################################
diff  --git a/compiler-rt/lib/orc/macho_platform.cpp b/compiler-rt/lib/orc/macho_platform.cpp
index 679163140e700..73b17a0799c4c 100644
--- a/compiler-rt/lib/orc/macho_platform.cpp
+++ b/compiler-rt/lib/orc/macho_platform.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "macho_platform.h"
+#include "bitmask_enum.h"
 #include "common.h"
 #include "debug.h"
 #include "error.h"
@@ -34,7 +35,7 @@ using namespace __orc_rt::macho;
 
 // Declare function tags for functions in the JIT process.
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_macho_push_initializers_tag)
-ORC_RT_JIT_DISPATCH_TAG(__orc_rt_macho_symbol_lookup_tag)
+ORC_RT_JIT_DISPATCH_TAG(__orc_rt_macho_push_symbols_tag)
 
 struct objc_image_info;
 struct mach_header;
@@ -148,6 +149,16 @@ struct TLVDescriptor {
 };
 
 class MachOPlatformRuntimeState {
+public:
+  // Used internally by MachOPlatformRuntimeState, but made public to enable
+  // serialization.
+  enum class MachOExecutorSymbolFlags : uint8_t {
+    None = 0,
+    Weak = 1U << 0,
+    Callable = 1U << 1,
+    ORC_RT_MARK_AS_BITMASK_ENUM(/* LargestValue = */ Callable)
+  };
+
 private:
   struct AtExitEntry {
     void (*Func)(void *);
@@ -256,11 +267,17 @@ class MachOPlatformRuntimeState {
       IntervalMap<char *, UnwindSections, IntervalCoalescing::Disabled>;
 
   struct JITDylibState {
+
+    using SymbolTableMap =
+        std::unordered_map<std::string_view,
+                           std::pair<ExecutorAddr, MachOExecutorSymbolFlags>>;
+
     std::string Name;
     void *Header = nullptr;
     bool Sealed = false;
     size_t LinkedAgainstRefCount = 0;
     size_t DlRefCount = 0;
+    SymbolTableMap SymbolTable;
     std::vector<JITDylibState *> Deps;
     AtExitsVector AtExits;
     const objc_image_info *ObjCImageInfo = nullptr;
@@ -296,6 +313,14 @@ class MachOPlatformRuntimeState {
   Error deregisterJITDylib(void *Header);
   Error registerThreadDataSection(span<const char> ThreadDataSection);
   Error deregisterThreadDataSection(span<const char> ThreadDataSection);
+  Error registerObjectSymbolTable(
+      ExecutorAddr HeaderAddr,
+      const std::vector<std::tuple<ExecutorAddr, ExecutorAddr,
+                                   MachOExecutorSymbolFlags>> &Entries);
+  Error deregisterObjectSymbolTable(
+      ExecutorAddr HeaderAddr,
+      const std::vector<std::tuple<ExecutorAddr, ExecutorAddr,
+                                   MachOExecutorSymbolFlags>> &Entries);
   Error registerObjectPlatformSections(
       ExecutorAddr HeaderAddr, std::optional<UnwindSectionInfo> UnwindSections,
       std::vector<std::pair<std::string_view, ExecutorAddrRange>> Secs);
@@ -306,7 +331,7 @@ class MachOPlatformRuntimeState {
   const char *dlerror();
   void *dlopen(std::string_view Name, int Mode);
   int dlclose(void *DSOHandle);
-  void *dlsym(void *DSOHandle, std::string_view Symbol);
+  void *dlsym(void *DSOHandle, const char *Symbol);
 
   int registerAtExit(void (*F)(void *), void *Arg, void *DSOHandle);
   void runAtExits(std::unique_lock<std::mutex> &JDStatesLock,
@@ -321,8 +346,42 @@ class MachOPlatformRuntimeState {
   JITDylibState *getJITDylibStateByHeader(void *DSOHandle);
   JITDylibState *getJITDylibStateByName(std::string_view Path);
 
-  Expected<ExecutorAddr> lookupSymbolInJITDylib(void *DSOHandle,
-                                                std::string_view Symbol);
+  /// Requests materialization of the given symbols. For each pair, the bool
+  /// element indicates whether the symbol is required (true) or weakly
+  /// referenced (false).
+  Error requestPushSymbols(JITDylibState &JDS,
+                           span<std::pair<std::string_view, bool>> Symbols);
+
+  /// Visits the symbol table for the JITDylib associated with DSOHandle.
+  /// Visitor should be callable as
+  ///
+  ///   void (size_t,
+  ///         std::optional<std::pair<ExecutorAddr, MachOExecutorSymbolFlags>>)
+  ///
+  /// The visitor function will be called for each element of the Symbols, but
+  /// in an arbitrary order. The first argument of the callback will indicate
+  /// the index of the result. The second argument will be std::nullopt (if the
+  /// symbol at the given index was not present in the symbol table), or a
+  /// pair containing the symbol's address and flags.
+  ///
+  /// This function will remove all elements of Symbols that are found, leaving
+  /// only the symbols that were not. This allows it to dovetail with
+  /// requestPushSymbols, enabling the following idiom:
+  ///
+  /// ...
+  /// visitSymbolAddrs(DSO, Symbols);
+  /// if (!Symbols.empty()) {
+  ///   requestPushSymbols(DSO, Symbols);
+  ///   visitSymbolAddrs(DSO, Symbols);
+  ///   for (auto &Sym : Symbols) {
+  ///     -- handle symbols that were not found --
+  ///   }
+  /// }
+  ///
+  template <typename VisitorFn>
+  void visitSymbolAddrs(JITDylibState &JDS,
+                        std::vector<std::pair<std::string_view, bool>> &Symbols,
+                        VisitorFn &&Visit);
 
   bool lookupUnwindSections(void *Addr, unw_dynamic_unwind_sections &Info);
 
@@ -366,6 +425,47 @@ class MachOPlatformRuntimeState {
   std::map<const char *, size_t> ThreadDataSections;
 };
 
+} // anonymous namespace
+
+namespace __orc_rt {
+
+class SPSMachOExecutorSymbolFlags;
+
+template <>
+class SPSSerializationTraits<
+    SPSMachOExecutorSymbolFlags,
+    MachOPlatformRuntimeState::MachOExecutorSymbolFlags> {
+private:
+  using UT = std::underlying_type_t<
+      MachOPlatformRuntimeState::MachOExecutorSymbolFlags>;
+
+public:
+  static size_t
+  size(const MachOPlatformRuntimeState::MachOExecutorSymbolFlags &SF) {
+    return sizeof(UT);
+  }
+
+  static bool
+  serialize(SPSOutputBuffer &OB,
+            const MachOPlatformRuntimeState::MachOExecutorSymbolFlags &SF) {
+    return SPSArgList<UT>::serialize(OB, static_cast<UT>(SF));
+  }
+
+  static bool
+  deserialize(SPSInputBuffer &IB,
+              MachOPlatformRuntimeState::MachOExecutorSymbolFlags &SF) {
+    UT Tmp;
+    if (!SPSArgList<UT>::deserialize(IB, Tmp))
+      return false;
+    SF = static_cast<MachOPlatformRuntimeState::MachOExecutorSymbolFlags>(Tmp);
+    return true;
+  }
+};
+
+} // namespace __orc_rt
+
+namespace {
+
 MachOPlatformRuntimeState *MachOPlatformRuntimeState::MOPS = nullptr;
 
 Error MachOPlatformRuntimeState::create() {
@@ -492,6 +592,48 @@ Error MachOPlatformRuntimeState::deregisterThreadDataSection(
   return Error::success();
 }
 
+Error MachOPlatformRuntimeState::registerObjectSymbolTable(
+    ExecutorAddr HeaderAddr,
+    const std::vector<std::tuple<ExecutorAddr, ExecutorAddr,
+                                 MachOExecutorSymbolFlags>> &Entries) {
+
+  std::lock_guard<std::mutex> Lock(JDStatesMutex);
+  auto *JDS = getJITDylibStateByHeader(HeaderAddr.toPtr<void *>());
+  if (!JDS) {
+    std::ostringstream ErrStream;
+    ErrStream << "Could not register object platform sections for "
+                 "unrecognized header "
+              << HeaderAddr.toPtr<void *>();
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  for (auto &[NameAddr, SymAddr, Flags] : Entries)
+    JDS->SymbolTable[NameAddr.toPtr<const char *>()] = {SymAddr, Flags};
+
+  return Error::success();
+}
+
+Error MachOPlatformRuntimeState::deregisterObjectSymbolTable(
+    ExecutorAddr HeaderAddr,
+    const std::vector<std::tuple<ExecutorAddr, ExecutorAddr,
+                                 MachOExecutorSymbolFlags>> &Entries) {
+
+  std::lock_guard<std::mutex> Lock(JDStatesMutex);
+  auto *JDS = getJITDylibStateByHeader(HeaderAddr.toPtr<void *>());
+  if (!JDS) {
+    std::ostringstream ErrStream;
+    ErrStream << "Could not register object platform sections for "
+                 "unrecognized header "
+              << HeaderAddr.toPtr<void *>();
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  for (auto &[NameAddr, SymAddr, Flags] : Entries)
+    JDS->SymbolTable.erase(NameAddr.toPtr<const char *>());
+
+  return Error::success();
+}
+
 Error MachOPlatformRuntimeState::registerObjectPlatformSections(
     ExecutorAddr HeaderAddr, std::optional<UnwindSectionInfo> UnwindInfo,
     std::vector<std::pair<std::string_view, ExecutorAddrRange>> Secs) {
@@ -687,15 +829,51 @@ int MachOPlatformRuntimeState::dlclose(void *DSOHandle) {
   return 0;
 }
 
-void *MachOPlatformRuntimeState::dlsym(void *DSOHandle,
-                                       std::string_view Symbol) {
-  auto Addr = lookupSymbolInJITDylib(DSOHandle, Symbol);
-  if (!Addr) {
-    DLFcnError = toString(Addr.takeError());
-    return 0;
+void *MachOPlatformRuntimeState::dlsym(void *DSOHandle, const char *Symbol) {
+  std::lock_guard<std::mutex> Lock(JDStatesMutex);
+  auto *JDS = getJITDylibStateByHeader(DSOHandle);
+  if (!JDS) {
+    std::ostringstream ErrStream;
+    ErrStream << "In call to dlsym, unrecognized header address " << DSOHandle;
+    DLFcnError = ErrStream.str();
+    return nullptr;
+  }
+
+  std::string MangledName("_");
+  MangledName += Symbol;
+  std::vector<std::pair<std::string_view, bool>> Symbols;
+  Symbols.push_back({MangledName, false});
+
+  ExecutorAddr Result;
+  using ElemResult =
+      std::optional<std::pair<ExecutorAddr, MachOExecutorSymbolFlags>>;
+
+  // Try to resolve the symbol in the local symbol tables.
+  visitSymbolAddrs(*JDS, Symbols, [&](size_t Idx, ElemResult E) {
+    if (E)
+      Result = E->first;
+  });
+
+  // Return early if we found it.
+  if (Symbols.empty())
+    return Result.toPtr<void *>();
+
+  // Otherwise call back to the controller to try to request that the symbol
+  // be materialized.
+  if (auto Err = requestPushSymbols(*JDS, {Symbols.data(), Symbols.size()})) {
+    DLFcnError = toString(std::move(Err));
+    return nullptr;
   }
 
-  return Addr->toPtr<void *>();
+  // Try another local resolution.
+  visitSymbolAddrs(*JDS, Symbols, [&](size_t Idx, ElemResult E) {
+    if (E)
+      Result = E->first;
+  });
+
+  // At this point Result has either been set (if we found the symbol) or is
+  // still null (if we didn't). Either way it's the right value.
+  return Result.toPtr<void *>();
 }
 
 int MachOPlatformRuntimeState::registerAtExit(void (*F)(void *), void *Arg,
@@ -774,17 +952,35 @@ MachOPlatformRuntimeState::getJITDylibStateByName(std::string_view Name) {
   return nullptr;
 }
 
-Expected<ExecutorAddr>
-MachOPlatformRuntimeState::lookupSymbolInJITDylib(void *DSOHandle,
-                                                  std::string_view Sym) {
-  Expected<ExecutorAddr> Result((ExecutorAddr()));
-  if (auto Err = WrapperFunction<SPSExpected<SPSExecutorAddr>(
-          SPSExecutorAddr, SPSString)>::call(&__orc_rt_macho_symbol_lookup_tag,
-                                             Result,
-                                             ExecutorAddr::fromPtr(DSOHandle),
-                                             Sym))
+Error MachOPlatformRuntimeState::requestPushSymbols(
+    JITDylibState &JDS, span<std::pair<std::string_view, bool>> Symbols) {
+  Error OpErr = Error::success();
+  if (auto Err = WrapperFunction<SPSError(
+          SPSExecutorAddr, SPSSequence<SPSTuple<SPSString, bool>>)>::
+          call(&__orc_rt_macho_push_symbols_tag, OpErr,
+               ExecutorAddr::fromPtr(JDS.Header), Symbols)) {
+    cantFail(std::move(OpErr));
     return std::move(Err);
-  return Result;
+  }
+  return OpErr;
+}
+
+template <typename VisitorFn>
+void MachOPlatformRuntimeState::visitSymbolAddrs(
+    JITDylibState &JDS, std::vector<std::pair<std::string_view, bool>> &Symbols,
+    VisitorFn &&Visit) {
+
+  std::vector<std::pair<std::string_view, bool>> RemainingSymbols;
+
+  for (size_t Idx = 0; Idx != Symbols.size(); ++Idx) {
+    auto I = JDS.SymbolTable.find(Symbols[Idx].first);
+    if (I != JDS.SymbolTable.end())
+      Visit(Idx, I->second);
+    else
+      RemainingSymbols.push_back(Symbols[Idx]);
+  }
+
+  Symbols = std::move(RemainingSymbols);
 }
 
 // eh-frame registration functions.
@@ -1193,6 +1389,38 @@ __orc_rt_macho_register_object_platform_sections(char *ArgData,
           .release();
 }
 
+ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
+__orc_rt_macho_register_object_symbol_table(char *ArgData, size_t ArgSize) {
+  using SymtabContainer = std::vector<
+      std::tuple<ExecutorAddr, ExecutorAddr,
+                 MachOPlatformRuntimeState::MachOExecutorSymbolFlags>>;
+  return WrapperFunction<SPSError(
+      SPSExecutorAddr, SPSSequence<SPSTuple<SPSExecutorAddr, SPSExecutorAddr,
+                                            SPSMachOExecutorSymbolFlags>>)>::
+      handle(ArgData, ArgSize,
+             [](ExecutorAddr HeaderAddr, SymtabContainer &Symbols) {
+               return MachOPlatformRuntimeState::get()
+                   .registerObjectSymbolTable(HeaderAddr, Symbols);
+             })
+          .release();
+}
+
+ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
+__orc_rt_macho_deregister_object_symbol_table(char *ArgData, size_t ArgSize) {
+  using SymtabContainer = std::vector<
+      std::tuple<ExecutorAddr, ExecutorAddr,
+                 MachOPlatformRuntimeState::MachOExecutorSymbolFlags>>;
+  return WrapperFunction<SPSError(
+      SPSExecutorAddr, SPSSequence<SPSTuple<SPSExecutorAddr, SPSExecutorAddr,
+                                            SPSMachOExecutorSymbolFlags>>)>::
+      handle(ArgData, ArgSize,
+             [](ExecutorAddr HeaderAddr, SymtabContainer &Symbols) {
+               return MachOPlatformRuntimeState::get()
+                   .deregisterObjectSymbolTable(HeaderAddr, Symbols);
+             })
+          .release();
+}
+
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_macho_deregister_object_platform_sections(char *ArgData,
                                                    size_t ArgSize) {

diff  --git a/compiler-rt/lib/orc/simple_packed_serialization.h b/compiler-rt/lib/orc/simple_packed_serialization.h
index f24ebae6cbee4..488d2407ddd43 100644
--- a/compiler-rt/lib/orc/simple_packed_serialization.h
+++ b/compiler-rt/lib/orc/simple_packed_serialization.h
@@ -281,6 +281,13 @@ class TrivialSPSSequenceSerialization<SPSElementTagT, std::vector<T>> {
   static constexpr bool available = true;
 };
 
+/// Trivial span<T> -> SPSSequence<SPSElementTagT> serialization.
+template <typename SPSElementTagT, typename T>
+class TrivialSPSSequenceSerialization<SPSElementTagT, span<T>> {
+public:
+  static constexpr bool available = true;
+};
+
 /// Trivial SPSSequence<SPSElementTagT> -> std::vector<T> deserialization.
 template <typename SPSElementTagT, typename T>
 class TrivialSPSSequenceDeserialization<SPSElementTagT, std::vector<T>> {

diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
index ac6ff3cb0861d..37c044a7415d6 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
@@ -39,6 +39,14 @@ class MachOPlatform : public Platform {
   using MachOJITDylibDepInfoMap =
       std::vector<std::pair<ExecutorAddr, MachOJITDylibDepInfo>>;
 
+  // Used internally by MachOPlatform, but made public to enable serialization.
+  enum class MachOExecutorSymbolFlags : uint8_t {
+    None = 0,
+    Weak = 1U << 0,
+    Callable = 1U << 1,
+    LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ Callable)
+  };
+
   /// Try to create a MachOPlatform instance, adding the ORC runtime to the
   /// given JITDylib.
   ///
@@ -164,6 +172,12 @@ class MachOPlatform : public Platform {
       bool Finalized = false;
     };
 
+    struct SymbolTablePair {
+      jitlink::Symbol *OriginalSym = nullptr;
+      jitlink::Symbol *NameSym = nullptr;
+    };
+    using JITSymTabVector = SmallVector<SymbolTablePair>;
+
     Error bootstrapPipelineStart(jitlink::LinkGraph &G);
     Error bootstrapPipelineRecordRuntimeFunctions(jitlink::LinkGraph &G);
     Error bootstrapPipelineEnd(jitlink::LinkGraph &G);
@@ -183,7 +197,6 @@ class MachOPlatform : public Platform {
     Error fixTLVSectionsAndEdges(jitlink::LinkGraph &G, JITDylib &JD);
 
     std::optional<UnwindSections> findUnwindSectionInfo(jitlink::LinkGraph &G);
-
     Error registerObjectPlatformSections(jitlink::LinkGraph &G, JITDylib &JD,
                                          bool InBootstrapPhase);
 
@@ -191,8 +204,15 @@ class MachOPlatform : public Platform {
     Error populateObjCRuntimeObject(jitlink::LinkGraph &G,
                                     MaterializationResponsibility &MR);
 
+    Error prepareSymbolTableRegistration(jitlink::LinkGraph &G,
+                                         JITSymTabVector &JITSymTabInfo);
+    Error addSymbolTableRegistration(jitlink::LinkGraph &G,
+                                     JITSymTabVector &JITSymTabInfo,
+                                     bool InBootstrapPhase);
+
     std::mutex PluginMutex;
     MachOPlatform &MP;
+    ExecutorAddr HeaderAddr;
 
     // FIXME: ObjCImageInfos and HeaderAddrs need to be cleared when
     // JITDylibs are removed.
@@ -208,9 +228,14 @@ class MachOPlatform : public Platform {
   using PushInitializersSendResultFn =
       unique_function<void(Expected<MachOJITDylibDepInfoMap>)>;
   using SendSymbolAddressFn = unique_function<void(Expected<ExecutorAddr>)>;
+  using PushSymbolsInSendResultFn = unique_function<void(Error)>;
 
   static bool supportedTarget(const Triple &TT);
 
+  static jitlink::Edge::Kind getPointerEdgeKind(jitlink::LinkGraph &G);
+
+  static MachOExecutorSymbolFlags flagsForSymbol(jitlink::Symbol &Sym);
+
   MachOPlatform(ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
                 JITDylib &PlatformJD,
                 std::unique_ptr<DefinitionGenerator> OrcRuntimeGenerator,
@@ -229,9 +254,12 @@ class MachOPlatform : public Platform {
   void rt_pushInitializers(PushInitializersSendResultFn SendResult,
                            ExecutorAddr JDHeaderAddr);
 
-  // Handle requests for symbol addresses from the ORC runtime.
-  void rt_lookupSymbol(SendSymbolAddressFn SendResult, ExecutorAddr Handle,
-                       StringRef SymbolName);
+  // Request that that the given symbols be materialized. The bool element of
+  // each pair indicates whether the symbol must be initialized, or whether it
+  // is optional. If any required symbol is not found then the pushSymbols
+  // function will return an error.
+  void rt_pushSymbols(PushSymbolsInSendResultFn SendResult, ExecutorAddr Handle,
+                      const std::vector<std::pair<StringRef, bool>> &Symbols);
 
   // Call the ORC runtime to create a pthread key.
   Expected<uint64_t> createPThreadKey();
@@ -260,6 +288,10 @@ class MachOPlatform : public Platform {
       ES.intern("___orc_rt_macho_register_jitdylib")};
   RuntimeFunction DeregisterJITDylib{
       ES.intern("___orc_rt_macho_deregister_jitdylib")};
+  RuntimeFunction RegisterObjectSymbolTable{
+      ES.intern("___orc_rt_macho_register_object_symbol_table")};
+  RuntimeFunction DeregisterObjectSymbolTable{
+      ES.intern("___orc_rt_macho_deregister_object_symbol_table")};
   RuntimeFunction RegisterObjectPlatformSections{
       ES.intern("___orc_rt_macho_register_object_platform_sections")};
   RuntimeFunction DeregisterObjectPlatformSections{

diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h b/llvm/include/llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h
index e5b95858013a8..0e75d16e39510 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h
@@ -24,6 +24,7 @@ extern StringRef MachODataCommonSectionName;
 extern StringRef MachODataDataSectionName;
 extern StringRef MachOEHFrameSectionName;
 extern StringRef MachOCompactUnwindInfoSectionName;
+extern StringRef MachOCStringSectionName;
 extern StringRef MachOModInitFuncSectionName;
 extern StringRef MachOObjCCatListSectionName;
 extern StringRef MachOObjCCatList2SectionName;

diff  --git a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
index 164adce444a15..850e29cf9952d 100644
--- a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
@@ -34,6 +34,8 @@ using SPSMachOJITDylibDepInfo = SPSTuple<bool, SPSSequence<SPSExecutorAddr>>;
 using SPSMachOJITDylibDepInfoMap =
     SPSSequence<SPSTuple<SPSExecutorAddr, SPSMachOJITDylibDepInfo>>;
 
+class SPSMachOExecutorSymbolFlags;
+
 template <>
 class SPSSerializationTraits<SPSMachOJITDylibDepInfo,
                              MachOPlatform::MachOJITDylibDepInfo> {
@@ -55,6 +57,32 @@ class SPSSerializationTraits<SPSMachOJITDylibDepInfo,
   }
 };
 
+template <>
+class SPSSerializationTraits<SPSMachOExecutorSymbolFlags,
+                             MachOPlatform::MachOExecutorSymbolFlags> {
+private:
+  using UT = std::underlying_type_t<MachOPlatform::MachOExecutorSymbolFlags>;
+
+public:
+  static size_t size(const MachOPlatform::MachOExecutorSymbolFlags &SF) {
+    return sizeof(UT);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB,
+                        const MachOPlatform::MachOExecutorSymbolFlags &SF) {
+    return SPSArgList<UT>::serialize(OB, static_cast<UT>(SF));
+  }
+
+  static bool deserialize(SPSInputBuffer &IB,
+                          MachOPlatform::MachOExecutorSymbolFlags &SF) {
+    UT Tmp;
+    if (!SPSArgList<UT>::deserialize(IB, Tmp))
+      return false;
+    SF = static_cast<MachOPlatform::MachOExecutorSymbolFlags>(Tmp);
+    return true;
+  }
+};
+
 } // namespace shared
 } // namespace orc
 } // namespace llvm
@@ -185,7 +213,8 @@ class MachOPlatformCompleteBootstrapMaterializationUnit
       SymbolStringPtr CompleteBootstrapSymbol, shared::AllocActions DeferredAAs,
       ExecutorAddr PlatformBootstrap, ExecutorAddr PlatformShutdown,
       ExecutorAddr RegisterJITDylib, ExecutorAddr DeregisterJITDylib,
-      ExecutorAddr MachOHeaderAddr)
+      ExecutorAddr RegisterObjectSymbolTable,
+      ExecutorAddr DeregisterObjectSymbolTable, ExecutorAddr MachOHeaderAddr)
       : MaterializationUnit(
             {{{CompleteBootstrapSymbol, JITSymbolFlags::None}}, nullptr}),
         MOP(MOP), PlatformJDName(PlatformJDName),
@@ -194,6 +223,8 @@ class MachOPlatformCompleteBootstrapMaterializationUnit
         PlatformBootstrap(PlatformBootstrap),
         PlatformShutdown(PlatformShutdown), RegisterJITDylib(RegisterJITDylib),
         DeregisterJITDylib(DeregisterJITDylib),
+        RegisterObjectSymbolTable(RegisterObjectSymbolTable),
+        DeregisterObjectSymbolTable(DeregisterObjectSymbolTable),
         MachOHeaderAddr(MachOHeaderAddr) {}
 
   StringRef getName() const override {
@@ -245,6 +276,8 @@ class MachOPlatformCompleteBootstrapMaterializationUnit
   ExecutorAddr PlatformShutdown;
   ExecutorAddr RegisterJITDylib;
   ExecutorAddr DeregisterJITDylib;
+  ExecutorAddr RegisterObjectSymbolTable;
+  ExecutorAddr DeregisterObjectSymbolTable;
   ExecutorAddr MachOHeaderAddr;
 };
 
@@ -446,6 +479,29 @@ bool MachOPlatform::supportedTarget(const Triple &TT) {
   }
 }
 
+jitlink::Edge::Kind MachOPlatform::getPointerEdgeKind(jitlink::LinkGraph &G) {
+  switch (G.getTargetTriple().getArch()) {
+  case Triple::aarch64:
+    return jitlink::aarch64::Pointer64;
+  case Triple::x86_64:
+    return jitlink::x86_64::Pointer64;
+  default:
+    llvm_unreachable("Unsupported architecture");
+  }
+}
+
+MachOPlatform::MachOExecutorSymbolFlags
+MachOPlatform::flagsForSymbol(jitlink::Symbol &Sym) {
+  MachOPlatform::MachOExecutorSymbolFlags Flags{};
+  if (Sym.getLinkage() == jitlink::Linkage::Weak)
+    Flags |= MachOExecutorSymbolFlags::Weak;
+
+  if (Sym.isCallable())
+    Flags |= MachOExecutorSymbolFlags::Callable;
+
+  return Flags;
+}
+
 MachOPlatform::MachOPlatform(
     ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
     JITDylib &PlatformJD,
@@ -525,6 +581,8 @@ MachOPlatform::MachOPlatform(
                        SymbolLookupSet(
                            {PlatformBootstrap.Name, PlatformShutdown.Name,
                             RegisterJITDylib.Name, DeregisterJITDylib.Name,
+                            RegisterObjectSymbolTable.Name,
+                            DeregisterObjectSymbolTable.Name,
                             RegisterObjectPlatformSections.Name,
                             DeregisterObjectPlatformSections.Name,
                             CreatePThreadKey.Name}))
@@ -545,7 +603,8 @@ MachOPlatform::MachOPlatform(
                *this, PlatformJD.getName(), BootstrapCompleteSymbol,
                std::move(BI.DeferredAAs), PlatformBootstrap.Addr,
                PlatformShutdown.Addr, RegisterJITDylib.Addr,
-               DeregisterJITDylib.Addr, BI.MachOHeaderAddr))))
+               DeregisterJITDylib.Addr, RegisterObjectSymbolTable.Addr,
+               DeregisterObjectSymbolTable.Addr, BI.MachOHeaderAddr))))
     return;
   if ((Err = ES.lookup(makeJITDylibSearchOrder(
                            &PlatformJD, JITDylibLookupFlags::MatchAllSymbols),
@@ -567,11 +626,11 @@ Error MachOPlatform::associateRuntimeSupportFunctions() {
       ES.wrapAsyncWithSPS<PushInitializersSPSSig>(
           this, &MachOPlatform::rt_pushInitializers);
 
-  using LookupSymbolSPSSig =
-      SPSExpected<SPSExecutorAddr>(SPSExecutorAddr, SPSString);
-  WFs[ES.intern("___orc_rt_macho_symbol_lookup_tag")] =
-      ES.wrapAsyncWithSPS<LookupSymbolSPSSig>(this,
-                                              &MachOPlatform::rt_lookupSymbol);
+  using PushSymbolsSPSSig =
+      SPSError(SPSExecutorAddr, SPSSequence<SPSTuple<SPSString, bool>>);
+  WFs[ES.intern("___orc_rt_macho_push_symbols_tag")] =
+      ES.wrapAsyncWithSPS<PushSymbolsSPSSig>(this,
+                                             &MachOPlatform::rt_pushSymbols);
 
   return ES.registerJITDispatchHandlers(PlatformJD, std::move(WFs));
 }
@@ -692,10 +751,14 @@ void MachOPlatform::rt_pushInitializers(PushInitializersSendResultFn SendResult,
   pushInitializersLoop(std::move(SendResult), JD);
 }
 
-void MachOPlatform::rt_lookupSymbol(SendSymbolAddressFn SendResult,
-                                    ExecutorAddr Handle, StringRef SymbolName) {
+void MachOPlatform::rt_pushSymbols(
+    PushSymbolsInSendResultFn SendResult, ExecutorAddr Handle,
+    const std::vector<std::pair<StringRef, bool>> &SymbolNames) {
   LLVM_DEBUG({
-    dbgs() << "MachOPlatform::rt_lookupSymbol(\"" << Handle << "\")\n";
+    dbgs() << "MachOPlatform::rt_pushSymbols(" << Handle << ", [ ";
+    for (auto &Name : SymbolNames)
+      dbgs() << "\"" << Name.first << "\" ";
+    dbgs() << "])\n";
   });
 
   JITDylib *JD = nullptr;
@@ -708,37 +771,26 @@ void MachOPlatform::rt_lookupSymbol(SendSymbolAddressFn SendResult,
   }
 
   if (!JD) {
-    LLVM_DEBUG(dbgs() << "  No JITDylib for handle " << Handle << "\n");
     SendResult(make_error<StringError>("No JITDylib associated with handle " +
                                            formatv("{0:x}", Handle),
                                        inconvertibleErrorCode()));
     return;
   }
 
-  // Use functor class to work around XL build compiler issue on AIX.
-  class RtLookupNotifyComplete {
-  public:
-    RtLookupNotifyComplete(SendSymbolAddressFn &&SendResult)
-        : SendResult(std::move(SendResult)) {}
-    void operator()(Expected<SymbolMap> Result) {
-      if (Result) {
-        assert(Result->size() == 1 && "Unexpected result map count");
-        SendResult(Result->begin()->second.getAddress());
-      } else {
-        SendResult(Result.takeError());
-      }
-    }
-
-  private:
-    SendSymbolAddressFn SendResult;
-  };
+  SymbolLookupSet LS;
+  for (auto &[Name, Required] : SymbolNames)
+    LS.add(ES.intern(Name), Required
+                                ? SymbolLookupFlags::RequiredSymbol
+                                : SymbolLookupFlags::WeaklyReferencedSymbol);
 
-  // FIXME: Proper mangling.
-  auto MangledName = ("_" + SymbolName).str();
   ES.lookup(
       LookupKind::DLSym, {{JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
-      SymbolLookupSet(ES.intern(MangledName)), SymbolState::Ready,
-      RtLookupNotifyComplete(std::move(SendResult)), NoDependenciesToRegister);
+      std::move(LS), SymbolState::Ready,
+      [SendResult = std::move(SendResult)](Expected<SymbolMap> Result) mutable {
+        dbgs() << "Sending result pushSymbols result...\n";
+        SendResult(Result.takeError());
+      },
+      NoDependenciesToRegister);
 }
 
 Expected<uint64_t> MachOPlatform::createPThreadKey() {
@@ -761,6 +813,14 @@ void MachOPlatform::MachOPlatformPlugin::modifyPassConfig(
 
   using namespace jitlink;
 
+  // Check for a header address.
+  {
+    std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
+    auto I = MP.JITDylibToHeaderAddr.find(&MR.getTargetJITDylib());
+    if (I != MP.JITDylibToHeaderAddr.end())
+      HeaderAddr = I->second;
+  }
+
   bool InBootstrapPhase =
       &MR.getTargetJITDylib() == &MP.PlatformJD && MP.Bootstrap;
 
@@ -808,6 +868,18 @@ void MachOPlatform::MachOPlatformPlugin::modifyPassConfig(
         return fixTLVSectionsAndEdges(G, JD);
       });
 
+  // Add symbol table prepare and register passes: These will add strings for
+  // all symbols to the c-strings section, and build a symbol table registration
+  // call.
+  auto JITSymTabInfo = std::make_shared<JITSymTabVector>();
+  Config.PostPrunePasses.push_back([this, JITSymTabInfo](LinkGraph &G) {
+    return prepareSymbolTableRegistration(G, *JITSymTabInfo);
+  });
+  Config.PostFixupPasses.push_back(
+      [this, JITSymTabInfo, InBootstrapPhase](LinkGraph &G) {
+        return addSymbolTableRegistration(G, *JITSymTabInfo, InBootstrapPhase);
+      });
+
   // Add a pass to register the final addresses of any special sections in the
   // object with the runtime.
   Config.PostAllocationPasses.push_back(
@@ -853,6 +925,9 @@ Error MachOPlatform::MachOPlatformPlugin::
       {*MP.PlatformShutdown.Name, &MP.PlatformShutdown.Addr},
       {*MP.RegisterJITDylib.Name, &MP.RegisterJITDylib.Addr},
       {*MP.DeregisterJITDylib.Name, &MP.DeregisterJITDylib.Addr},
+      {*MP.RegisterObjectSymbolTable.Name, &MP.RegisterObjectSymbolTable.Addr},
+      {*MP.DeregisterObjectSymbolTable.Name,
+       &MP.DeregisterObjectSymbolTable.Addr},
       {*MP.RegisterObjectPlatformSections.Name,
        &MP.RegisterObjectPlatformSections.Addr},
       {*MP.DeregisterObjectPlatformSections.Name,
@@ -1335,15 +1410,6 @@ Error MachOPlatform::MachOPlatformPlugin::registerObjectPlatformSections(
                                  UI->CompactUnwindSection);
 
   if (!MachOPlatformSecs.empty() || UnwindInfo) {
-    ExecutorAddr HeaderAddr;
-    {
-      std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
-      auto I = MP.JITDylibToHeaderAddr.find(&JD);
-      assert(I != MP.JITDylibToHeaderAddr.end() &&
-             "Missing header for JITDylib");
-      HeaderAddr = I->second;
-    }
-
     // Dump the scraped inits.
     LLVM_DEBUG({
       dbgs() << "MachOPlatform: Scraped " << G.getName() << " init sections:\n";
@@ -1361,6 +1427,7 @@ Error MachOPlatform::MachOPlatformPlugin::registerObjectPlatformSections(
                                              ? G.allocActions()
                                              : MP.Bootstrap.load()->DeferredAAs;
 
+    assert(HeaderAddr && "No HeaderAddr for JITDylib");
     allocActions.push_back(
         {cantFail(
              WrapperFunctionCall::Create<SPSRegisterObjectPlatformSectionsArgs>(
@@ -1459,17 +1526,7 @@ Error MachOPlatform::MachOPlatformPlugin::populateObjCRuntimeObject(
     strcpy(SD.Sec.segname, "__DATA");
     SD.Sec.size = 8;
     SD.AddFixups = [&](size_t RecordOffset) {
-      jitlink::Edge::Kind PointerEdge = jitlink::Edge::Invalid;
-      switch (G.getTargetTriple().getArch()) {
-      case Triple::aarch64:
-        PointerEdge = jitlink::aarch64::Pointer64;
-        break;
-      case Triple::x86_64:
-        PointerEdge = jitlink::x86_64::Pointer64;
-        break;
-      default:
-        llvm_unreachable("Unsupported architecture");
-      }
+      auto PointerEdge = getPointerEdgeKind(G);
 
       // Look for an existing __objc_imageinfo symbol.
       jitlink::Symbol *ObjCImageInfoSym = nullptr;
@@ -1595,5 +1652,79 @@ Error MachOPlatform::MachOPlatformPlugin::populateObjCRuntimeObject(
   return Error::success();
 }
 
+Error MachOPlatform::MachOPlatformPlugin::prepareSymbolTableRegistration(
+    jitlink::LinkGraph &G, JITSymTabVector &JITSymTabInfo) {
+
+  auto *CStringSec = G.findSectionByName(MachOCStringSectionName);
+  if (!CStringSec)
+    CStringSec = &G.createSection(MachOCStringSectionName,
+                                  MemProt::Read | MemProt::Exec);
+
+  // Make a map of existing strings so that we can re-use them:
+  DenseMap<StringRef, jitlink::Symbol *> ExistingStrings;
+  for (auto *Sym : CStringSec->symbols()) {
+
+    // The LinkGraph builder should have created single strings blocks, and all
+    // plugins should have maintained this invariant.
+    auto Content = Sym->getBlock().getContent();
+    ExistingStrings.insert(
+        std::make_pair(StringRef(Content.data(), Content.size()), Sym));
+  }
+
+  // Add all symbol names to the string section, and record the symbols for
+  // those names.
+  {
+    SmallVector<jitlink::Symbol *> SymsToProcess;
+    for (auto *Sym : G.defined_symbols())
+      SymsToProcess.push_back(Sym);
+
+    for (auto *Sym : SymsToProcess) {
+      if (!Sym->hasName())
+        continue;
+
+      auto I = ExistingStrings.find(Sym->getName());
+      if (I == ExistingStrings.end()) {
+        auto &NameBlock = G.createMutableContentBlock(
+            *CStringSec, G.allocateCString(Sym->getName()), orc::ExecutorAddr(),
+            1, 0);
+        auto &SymbolNameSym = G.addAnonymousSymbol(
+            NameBlock, 0, NameBlock.getSize(), false, true);
+        JITSymTabInfo.push_back({Sym, &SymbolNameSym});
+      } else
+        JITSymTabInfo.push_back({Sym, I->second});
+    }
+  }
+
+  return Error::success();
+}
+
+Error MachOPlatform::MachOPlatformPlugin::addSymbolTableRegistration(
+    jitlink::LinkGraph &G, JITSymTabVector &JITSymTabInfo,
+    bool InBootstrapPhase) {
+
+  SmallVector<std::tuple<ExecutorAddr, ExecutorAddr, MachOExecutorSymbolFlags>>
+      SymTab;
+  for (auto &[OriginalSymbol, NameSym] : JITSymTabInfo) {
+    // dbgs() << "Original symbol: \"" << OriginalSymbol->getName() << "\"\n";
+    SymTab.push_back({NameSym->getAddress(), OriginalSymbol->getAddress(),
+                      flagsForSymbol(*OriginalSymbol)});
+  }
+
+  using SPSRegisterSymbolsArgs =
+      SPSArgList<SPSExecutorAddr,
+                 SPSSequence<SPSTuple<SPSExecutorAddr, SPSExecutorAddr,
+                                      SPSMachOExecutorSymbolFlags>>>;
+
+  shared::AllocActions &allocActions = LLVM_LIKELY(!InBootstrapPhase)
+                                           ? G.allocActions()
+                                           : MP.Bootstrap.load()->DeferredAAs;
+  allocActions.push_back(
+      {cantFail(WrapperFunctionCall::Create<SPSRegisterSymbolsArgs>(
+           MP.RegisterObjectSymbolTable.Addr, HeaderAddr, SymTab)),
+       cantFail(WrapperFunctionCall::Create<SPSRegisterSymbolsArgs>(
+           MP.DeregisterObjectSymbolTable.Addr, HeaderAddr, SymTab))});
+
+  return Error::success();
+}
 } // End namespace orc.
 } // End namespace llvm.

diff  --git a/llvm/lib/ExecutionEngine/Orc/Shared/ObjectFormats.cpp b/llvm/lib/ExecutionEngine/Orc/Shared/ObjectFormats.cpp
index 9955130fc3b19..1d9d23e641587 100644
--- a/llvm/lib/ExecutionEngine/Orc/Shared/ObjectFormats.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Shared/ObjectFormats.cpp
@@ -19,6 +19,7 @@ StringRef MachODataCommonSectionName = "__DATA,__common";
 StringRef MachODataDataSectionName = "__DATA,__data";
 StringRef MachOEHFrameSectionName = "__TEXT,__eh_frame";
 StringRef MachOCompactUnwindInfoSectionName = "__TEXT,__unwind_info";
+StringRef MachOCStringSectionName = "__TEXT,__cstring";
 StringRef MachOModInitFuncSectionName = "__DATA,__mod_init_func";
 StringRef MachOObjCCatListSectionName = "__DATA,__objc_catlist";
 StringRef MachOObjCCatList2SectionName = "__DATA,__objc_catlist2";


        


More information about the llvm-commits mailing list