[ORC] Add Auto-Loading DyLib Feature with Symbol Resolution

llvm/include/llvm/ExecutionEngine/Orc/DylibManager.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_DYLIBMANAGER_H
 #define LLVM_EXECUTIONENGINE_ORC_DYLIBMANAGER_H
 
+#include "llvm/ExecutionEngine/Orc/Shared/AutoLoadDylibUtils.h"
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MSVCErrorWorkarounds.h"
@@ -60,6 +61,9 @@ class DylibManager {
   using SymbolLookupCompleteFn =
       unique_function<void(Expected<std::vector<tpctypes::LookupResult>>)>;
 
+  using ResolveSymbolsCompleteFn =
+      unique_function<void(Expected<std::vector<ResolveResult>>)>;
+
   /// Search for symbols in the target process.
   ///
   /// The result of the lookup is a 2-dimensional array of target addresses
@@ -68,6 +72,9 @@ class DylibManager {
   /// symbol is not found then it be assigned a '0' value.
   virtual void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
                                   SymbolLookupCompleteFn F) = 0;
+
+  virtual void resolveSymbolsAsync(ArrayRef<SymbolLookupSet> Request,
+                                   ResolveSymbolsCompleteFn F) = 0;
 };
 
 } // end namespace llvm::orc

llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h

@@ -15,6 +15,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_EPCDYNAMICLIBRARYSEARCHGENERATOR_H
 
 #include "llvm/ADT/FunctionExtras.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 
 namespace llvm {
@@ -71,6 +72,46 @@ class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
   AddAbsoluteSymbolsFn AddAbsoluteSymbols;
 };
 
+class AutoLoadDynamicLibrarySearchGenerator : public DefinitionGenerator {
+public:
+  using AddAbsoluteSymbolsFn = unique_function<Error(JITDylib &, SymbolMap)>;
+
+  /// Creates an AutoLoadDynamicLibrarySearchGenerator that searches for symbols
+  /// across all currently loaded libraries. If a symbol is not found, it scans
+  /// all potential dynamic libraries (dylibs), and if the symbol is located,
+  /// the corresponding library is loaded, and the symbol's definition is
+  /// returned.
+  ///
+  /// If \p AddAbsoluteSymbols is provided, it is used to add the symbols to the
+  /// \c JITDylib; otherwise it uses JD.define(absoluteSymbols(...)).
+  AutoLoadDynamicLibrarySearchGenerator(
+      ExecutionSession &ES, AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr)
+      : EPC(ES.getExecutorProcessControl()),
+        AddAbsoluteSymbols(std::move(AddAbsoluteSymbols)) {}
+
+  /// Creates a AutoLoadDynamicLibrarySearchGenerator that searches for symbols
+  /// in the target process.
+  static Expected<std::unique_ptr<AutoLoadDynamicLibrarySearchGenerator>>
+  GetForTargetProcess(ExecutionSession &ES,
+                      AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr) {
+    return std::make_unique<AutoLoadDynamicLibrarySearchGenerator>(
+        ES, std::move(AddAbsoluteSymbols));
+  }
+
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
+                      JITDylibLookupFlags JDLookupFlags,
+                      const SymbolLookupSet &Symbols) override;
+
+  Error tryToResolve(SymbolNameSet CandidateSyms,
+                     DylibManager::ResolveSymbolsCompleteFn OnCompleteFn);
+
+private:
+  ExecutorProcessControl &EPC;
+  BloomFilter GlobalFilter;
+  StringSet<> ExcludedSymbols;
+  AddAbsoluteSymbolsFn AddAbsoluteSymbols;
+};
+
 } // end namespace orc
 } // end namespace llvm
 

llvm/include/llvm/ExecutionEngine/Orc/EPCGenericDylibManager.h

@@ -34,6 +34,7 @@ class EPCGenericDylibManager {
     ExecutorAddr Instance;
     ExecutorAddr Open;
     ExecutorAddr Lookup;
+    ExecutorAddr Resolve;
   };
 
   /// Create an EPCGenericMemoryAccess instance from a given set of
@@ -70,6 +71,9 @@ class EPCGenericDylibManager {
   using SymbolLookupCompleteFn =
       unique_function<void(Expected<std::vector<ExecutorSymbolDef>>)>;
 
+  using ResolveSymbolsCompleteFn =
+      unique_function<void(Expected<ResolveResult>)>;
+
   /// Looks up symbols within the given dylib.
   void lookupAsync(tpctypes::DylibHandle H, const SymbolLookupSet &Lookup,
                    SymbolLookupCompleteFn Complete);
@@ -78,6 +82,14 @@ class EPCGenericDylibManager {
   void lookupAsync(tpctypes::DylibHandle H, const RemoteSymbolLookupSet &Lookup,
                    SymbolLookupCompleteFn Complete);
 
+  /// Look up and resolve symbols across all available dynamic libraries.
+  void resolveAsync(const SymbolLookupSet &Lookup,
+                    ResolveSymbolsCompleteFn Complete);
+
+  /// Look up and resolve symbols across all available dynamic libraries.
+  void resolveAsync(const RemoteSymbolLookupSet &Lookup,
+                    ResolveSymbolsCompleteFn Complete);
+
 private:
   ExecutorProcessControl &EPC;
   SymbolAddrs SAs;

llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h

@@ -506,6 +506,9 @@ class SelfExecutorProcessControl : public ExecutorProcessControl,
   void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
                           SymbolLookupCompleteFn F) override;
 
+  void resolveSymbolsAsync(ArrayRef<SymbolLookupSet> Request,
+                           ResolveSymbolsCompleteFn F) override;
+
   std::unique_ptr<jitlink::JITLinkMemoryManager> OwnedMemMgr;
   char GlobalManglingPrefix = 0;
 };

llvm/include/llvm/ExecutionEngine/Orc/Shared/AutoLoadDylibUtils.h

@@ -0,0 +1,232 @@
+//===------ AutoLoadDylibUtils.h - Auto-Loading Dynamic Library -------*- C++
+//-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_SHARED_AUTOLOADDYLIBUTILS_H
+#define LLVM_EXECUTIONENGINE_ORC_SHARED_AUTOLOADDYLIBUTILS_H
+
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
+#include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
+
+#include <math.h>
+#include <type_traits>
+#include <vector>
+
+namespace llvm {
+namespace orc {
+
+namespace shared {
+using SPSBloomFilter =
+    SPSTuple<bool, uint32_t, uint32_t, uint32_t, SPSSequence<uint64_t>>;
+}
+
+constexpr uint32_t log2u(std::uint32_t n) {
+  return (n > 1) ? 1 + log2u(n >> 1) : 0;
+}
+
+extern const char *const kEnvDelim;
+
+enum class SplitMode {
+  PruneNonExistant, ///< Don't add non-existant paths into output
+  FailNonExistant,  ///< Fail on any non-existant paths
+  AllowNonExistant  ///< Add all paths whether they exist or not
+};
+
+/// Collect the constituant paths from a PATH string.
+/// /bin:/usr/bin:/usr/local/bin -> {/bin, /usr/bin, /usr/local/bin}
+bool SplitPaths(llvm::StringRef PathStr,
+                llvm::SmallVectorImpl<llvm::StringRef> &Paths,
+                SplitMode Mode = SplitMode::PruneNonExistant,
+                llvm::StringRef Delim = kEnvDelim, bool Verbose = false);
+
+///
+bool GetSystemLibraryPaths(llvm::SmallVectorImpl<std::string> &Paths);
+
+/// Returns a normalized version of the given Path
+std::string NormalizePath(const std::string &Path);
+
+class BloomFilter {
+private:
+  static constexpr int Bits = 8 * sizeof(uint64_t);
+  static constexpr float P = 0.02f;
+
+  bool Initialized = false;
+  uint32_t SymbolsCount = 0;
+  uint32_t BloomSize = 0;
+  uint32_t BloomShift = 0;
+  std::vector<uint64_t> BloomTable;
+
+  // This is a GNU implementation of hash used in bloom filter!
+  static uint32_t GNUHash(StringRef S) {
+    uint32_t H = 5381;
+    for (uint8_t C : S)
+      H = (H << 5) + H + C;
+    return H;
+  }
+  // Helper method for hash testing
+  bool TestHash(uint32_t hash) const {
+    assert(Initialized && "Bloom filter is not initialized!");
+    uint32_t hash2 = hash >> BloomShift;
+    uint32_t n = (hash >> log2u(Bits)) % BloomSize;
+    uint64_t mask = ((1ULL << (hash % Bits)) | (1ULL << (hash2 % Bits)));
+    return (mask & BloomTable[n]) == mask;
+  }
+
+  // Helper method to add a hash
+  void AddHash(uint32_t hash) {
+    assert(Initialized && "Bloom filter is not initialized!");
+    uint32_t hash2 = hash >> BloomShift;
+    uint32_t n = (hash >> log2u(Bits)) % BloomSize;
+    uint64_t mask = ((1ULL << (hash % Bits)) | (1ULL << (hash2 % Bits)));
+    BloomTable[n] |= mask;
+  }
+
+  // Resizes the Bloom filter table based on symbol count
+  void ResizeTable(uint32_t newSymbolsCount) {
+    assert(SymbolsCount == 0 && "Resize not supported after initialization!");
+    SymbolsCount = newSymbolsCount;
+    BloomSize =
+        static_cast<uint32_t>(ceil((-1.44f * SymbolsCount * log2f(P)) / Bits));
+    BloomShift = std::min(6u, log2u(SymbolsCount));
+    BloomTable.resize(BloomSize, 0);
+  }
+
+  friend class shared::SPSSerializationTraits<shared::SPSBloomFilter,
+                                              BloomFilter>;
+
+public:
+  BloomFilter() = default;
+  BloomFilter(const BloomFilter &other) noexcept
+      : Initialized(other.Initialized), SymbolsCount(other.SymbolsCount),
+        BloomSize(other.BloomSize), BloomShift(other.BloomShift),
+        BloomTable(other.BloomTable) {}
+  BloomFilter &operator=(const BloomFilter &other) = delete;
+
+  BloomFilter(BloomFilter &&other) noexcept
+      : Initialized(other.Initialized), SymbolsCount(other.SymbolsCount),
+        BloomSize(other.BloomSize), BloomShift(other.BloomShift),
+        BloomTable(std::move(other.BloomTable)) {
+    other.Initialized = false;
+    other.SymbolsCount = 0;
+    other.BloomSize = 0;
+    other.BloomShift = 0;
+  }
+
+  BloomFilter &operator=(BloomFilter &&other) noexcept {
+    if (this != &other) {
+      Initialized = other.Initialized;
+      SymbolsCount = other.SymbolsCount;
+      BloomSize = other.BloomSize;
+      BloomShift = other.BloomShift;
+      BloomTable = std::move(other.BloomTable);
+
+      other.Initialized = false;
+      other.SymbolsCount = 0;
+      other.BloomSize = 0;
+      other.BloomShift = 0;
+    }
+    return *this;
+  }
+
+  void swap(BloomFilter &other) noexcept {
+    std::swap(Initialized, other.Initialized);
+    std::swap(SymbolsCount, other.SymbolsCount);
+    std::swap(BloomSize, other.BloomSize);
+    std::swap(BloomShift, other.BloomShift);
+    std::swap(BloomTable, other.BloomTable);
+  }
+
+  void Initialize(uint32_t newSymbolsCount) {
+    assert(!Initialized && "Cannot reinitialize the Bloom filter!");
+    Initialized = true;
+    ResizeTable(newSymbolsCount);
+  }
+
+  bool IsEmpty() const { return SymbolsCount == 0; }
+
+  uint32_t getSymCount() const { return SymbolsCount; }
+
+  bool IsInitialized() const { return Initialized; }
+
+  bool MayContain(uint32_t hash) const {
+    if (IsEmpty())
+      return false;
+    return TestHash(hash);
+  }
+
+  bool MayContain(StringRef symbol) const {
+    return MayContain(GNUHash(symbol));
+  }
+
+  void AddSymbol(StringRef symbol) { AddHash(GNUHash(symbol)); }
+};
+
+struct ResolveResult {
+  std::optional<BloomFilter> Filter;
+  std::vector<ExecutorSymbolDef> SymbolDef;
+};
+
+namespace shared {
+
+template <> class SPSSerializationTraits<SPSBloomFilter, BloomFilter> {
+public:
+  static size_t size(const BloomFilter &Filter) {
+    return SPSBloomFilter::AsArgList::size(
+        Filter.Initialized, Filter.SymbolsCount, Filter.BloomSize,
+        Filter.BloomShift, Filter.BloomTable);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB, const BloomFilter &Filter) {
+    return SPSBloomFilter::AsArgList::serialize(
+        OB, Filter.Initialized, Filter.SymbolsCount, Filter.BloomSize,
+        Filter.BloomShift, Filter.BloomTable);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB, BloomFilter &Filter) {
+    bool IsInitialized;
+    uint32_t SymbolsCount = 0, BloomSize = 0, BloomShift = 0;
+    std::vector<uint64_t> BloomTable;
+
+    if (!SPSBloomFilter::AsArgList::deserialize(
+            IB, IsInitialized, SymbolsCount, BloomSize, BloomShift, BloomTable))
+      return false;
+
+    Filter.Initialized = IsInitialized;
+    Filter.SymbolsCount = SymbolsCount;
+    Filter.BloomSize = BloomSize;
+    Filter.BloomShift = BloomShift;
+    Filter.BloomTable = std::move(BloomTable);
+
+    return true;
+  }
+};
+
+using SPSResolveResult =
+    SPSTuple<SPSOptional<SPSBloomFilter>, SPSSequence<SPSExecutorSymbolDef>>;
+template <> class SPSSerializationTraits<SPSResolveResult, ResolveResult> {
+public:
+  static size_t size(const ResolveResult &Result) {
+    return SPSResolveResult::AsArgList::size(Result.Filter, Result.SymbolDef);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB, const ResolveResult &Result) {
+    return SPSResolveResult::AsArgList::serialize(OB, Result.Filter,
+                                                  Result.SymbolDef);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB, ResolveResult &Result) {
+    return SPSResolveResult::AsArgList::deserialize(IB, Result.Filter,
+                                                    Result.SymbolDef);
+  }
+};
+
+} // end namespace shared
+} // end namespace orc
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_SHARED_AUTOLOADDYLIBUTILS_H

llvm/include/llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_SHARED_ORCRTBRIDGE_H
 #define LLVM_EXECUTIONENGINE_ORC_SHARED_ORCRTBRIDGE_H
 
+#include "llvm/ExecutionEngine/Orc/Shared/AutoLoadDylibUtils.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
@@ -25,6 +26,7 @@ namespace rt {
 extern const char *SimpleExecutorDylibManagerInstanceName;
 extern const char *SimpleExecutorDylibManagerOpenWrapperName;
 extern const char *SimpleExecutorDylibManagerLookupWrapperName;
+extern const char *SimpleExecutorDylibManagerResolveWrapperName;
 
 extern const char *SimpleExecutorMemoryManagerInstanceName;
 extern const char *SimpleExecutorMemoryManagerReserveWrapperName;
@@ -60,6 +62,11 @@ using SPSSimpleExecutorDylibManagerLookupSignature =
         shared::SPSExecutorAddr, shared::SPSExecutorAddr,
         shared::SPSRemoteSymbolLookupSet);
 
+using SPSSimpleExecutorDylibManagerResolveSignature = shared::SPSExpected<
+    shared::SPSTuple<shared::SPSOptional<shared::SPSBloomFilter>,
+                     shared::SPSSequence<shared::SPSExecutorSymbolDef>>>(
+    shared::SPSExecutorAddr, shared::SPSRemoteSymbolLookupSet);
+
 using SPSSimpleExecutorMemoryManagerReserveSignature =
     shared::SPSExpected<shared::SPSExecutorAddr>(shared::SPSExecutorAddr,
                                                  uint64_t);

llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h

@@ -122,6 +122,9 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
   void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
                           SymbolLookupCompleteFn F) override;
 
+  void resolveSymbolsAsync(ArrayRef<SymbolLookupSet> Request,
+                           ResolveSymbolsCompleteFn F) override;
+
   using PendingCallWrapperResultsMap =
     DenseMap<uint64_t, IncomingWFRHandler>;