[ORC] Introduce JITSymbolFlags::HasMaterializeSideEffectsOnly flag.

This flag can be used to mark a symbol as existing only for the purpose of
enabling materialization. Such a symbol can be looked up to trigger
materialization with the lookup returning only once materialization is
complete. Symbols with this flag will never resolve however (to avoid
permanently polluting the symbol table), and should only be looked up using
the SymbolLookupFlags::WeaklyReferencedSymbol flag. The primary use case for
this flag is initialization symbols.

Author: lhames

llvm/include/llvm/ExecutionEngine/JITSymbol.h

@@ -84,7 +84,9 @@ class JITSymbolFlags {
     Absolute = 1U << 3,
     Exported = 1U << 4,
     Callable = 1U << 5,
-    LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ Callable)
+    MaterializationSideEffectsOnly = 1U << 6,
+    LLVM_MARK_AS_BITMASK_ENUM( // LargestValue =
+        MaterializationSideEffectsOnly)
   };
 
   /// Default-construct a JITSymbolFlags instance.
@@ -146,6 +148,21 @@ class JITSymbolFlags {
   /// Returns true if the given symbol is known to be callable.
   bool isCallable() const { return (Flags & Callable) == Callable; }
 
+  /// Returns true if this symbol is a materialization-side-effects-only
+  /// symbol. Such symbols do not have a real address. They exist to trigger
+  /// and support synchronization of materialization side effects, e.g. for
+  /// collecting initialization information. These symbols will vanish from
+  /// the symbol table immediately upon reaching the ready state, and will
+  /// appear to queries as if they were never defined (except that query
+  /// callback execution will be delayed until they reach the ready state).
+  /// MaterializationSideEffectOnly symbols should only be queried using the
+  /// SymbolLookupFlags::WeaklyReferencedSymbol flag (see
+  /// llvm/include/llvm/ExecutionEngine/Orc/Core.h).
+  bool hasMaterializationSideEffectsOnly() const {
+    return (Flags & MaterializationSideEffectsOnly) ==
+           MaterializationSideEffectsOnly;
+  }
+
   /// Get the underlying flags value as an integer.
   UnderlyingType getRawFlagsValue() const {
     return static_cast<UnderlyingType>(Flags);

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -483,6 +483,20 @@ class MaterializationResponsibility {
   /// callbacks, metadata).
   Error defineMaterializing(SymbolFlagsMap SymbolFlags);
 
+  /// Define the given symbols as non-existent, removing it from the symbol
+  /// table and notifying any pending queries. Queries that lookup up the
+  /// symbol using the SymbolLookupFlags::WeaklyReferencedSymbol flag will
+  /// behave as if the symbol had not been matched in the first place. Queries
+  /// that required this symbol will fail with a missing symbol definition
+  /// error.
+  ///
+  /// This method is intended to support cleanup of special symbols like
+  /// initializer symbols: Queries using
+  /// SymbolLookupFlags::WeaklyReferencedSymbol can be used to trigger their
+  /// emission, and this method can be used to remove them from the JITDylib
+  /// once materialization is complete.
+  void defineNonExistent(ArrayRef<SymbolStringPtr> Symbols);
+
   /// Notify all not-yet-emitted covered by this MaterializationResponsibility
   /// instance that an error has occurred.
   /// This will remove all symbols covered by this MaterializationResponsibilty
@@ -721,15 +735,6 @@ class AsynchronousSymbolQuery {
   void notifySymbolMetRequiredState(const SymbolStringPtr &Name,
                                     JITEvaluatedSymbol Sym);
 
-  /// Remove a symbol from the query. This is used to drop weakly referenced
-  /// symbols that are not found.
-  void dropSymbol(const SymbolStringPtr &Name) {
-    assert(ResolvedSymbols.count(Name) &&
-           "Redundant removal of weakly-referenced symbol");
-    ResolvedSymbols.erase(Name);
-    --OutstandingSymbolsCount;
-  }
-
   /// Returns true if all symbols covered by this query have been
   ///        resolved.
   bool isComplete() const { return OutstandingSymbolsCount == 0; }
@@ -747,6 +752,8 @@ class AsynchronousSymbolQuery {
 
   void removeQueryDependence(JITDylib &JD, const SymbolStringPtr &Name);
 
+  void dropSymbol(const SymbolStringPtr &Name);
+
   bool canStillFail();
 
   void handleFailed(Error Err);
@@ -1298,6 +1305,16 @@ auto JITDylib::withSearchOrderDo(Func &&F)
 template <typename MaterializationUnitType>
 Error JITDylib::define(std::unique_ptr<MaterializationUnitType> &&MU) {
   assert(MU && "Can not define with a null MU");
+
+  if (MU->getSymbols().empty()) {
+    // Empty MUs are allowable but pathological, so issue a warning.
+    DEBUG_WITH_TYPE("orc", {
+      dbgs() << "Warning: Discarding empty MU " << MU->getName() << "\n";
+    });
+    return Error::success();
+  } else
+    DEBUG_WITH_TYPE("orc", dbgs() << "Defining MU " << MU->getName() << ":\n");
+
   return ES.runSessionLocked([&, this]() -> Error {
     if (auto Err = defineImpl(*MU))
       return Err;
@@ -1320,6 +1337,15 @@ template <typename MaterializationUnitType>
 Error JITDylib::define(std::unique_ptr<MaterializationUnitType> &MU) {
   assert(MU && "Can not define with a null MU");
 
+  if (MU->getSymbols().empty()) {
+    // Empty MUs are allowable but pathological, so issue a warning.
+    DEBUG_WITH_TYPE("orc", {
+      dbgs() << "Warning: Discarding empty MU " << MU->getName() << "\n";
+    });
+    return Error::success();
+  } else
+    DEBUG_WITH_TYPE("orc", dbgs() << "Defining MU " << MU->getName() << ":\n");
+
   return ES.runSessionLocked([&, this]() -> Error {
     if (auto Err = defineImpl(*MU))
       return Err;

llvm/include/llvm/ExecutionEngine/Orc/DebugUtils.h

@@ -38,6 +38,9 @@ raw_ostream &operator<<(raw_ostream &OS, const SymbolNameSet &Symbols);
 /// Render a SymbolNameVector.
 raw_ostream &operator<<(raw_ostream &OS, const SymbolNameVector &Symbols);
 
+/// Render an array of SymbolStringPtrs.
+raw_ostream &operator<<(raw_ostream &OS, ArrayRef<SymbolStringPtr> Symbols);
+
 /// Render JITSymbolFlags.
 raw_ostream &operator<<(raw_ostream &OS, const JITSymbolFlags &Flags);
 

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -118,7 +118,13 @@ void AsynchronousSymbolQuery::notifySymbolMetRequiredState(
   assert(I != ResolvedSymbols.end() &&
          "Resolving symbol outside the requested set");
   assert(I->second.getAddress() == 0 && "Redundantly resolving symbol Name");
-  I->second = std::move(Sym);
+
+  // If this is a materialization-side-effects-only symbol then drop it,
+  // otherwise update its map entry with its resolved address.
+  if (Sym.getFlags().hasMaterializationSideEffectsOnly())
+    ResolvedSymbols.erase(I);
+  else
+    I->second = std::move(Sym);
   --OutstandingSymbolsCount;
 }
 
@@ -159,6 +165,14 @@ void AsynchronousSymbolQuery::removeQueryDependence(
     QueryRegistrations.erase(QRI);
 }
 
+void AsynchronousSymbolQuery::dropSymbol(const SymbolStringPtr &Name) {
+  auto I = ResolvedSymbols.find(Name);
+  assert(I != ResolvedSymbols.end() &&
+         "Redundant removal of weakly-referenced symbol");
+  ResolvedSymbols.erase(I);
+  --OutstandingSymbolsCount;
+}
+
 void AsynchronousSymbolQuery::detach() {
   ResolvedSymbols.clear();
   OutstandingSymbolsCount = 0;
@@ -186,6 +200,8 @@ Error MaterializationResponsibility::notifyResolved(const SymbolMap &Symbols) {
     auto I = SymbolFlags.find(KV.first);
     assert(I != SymbolFlags.end() &&
            "Resolving symbol outside this responsibility set");
+    assert(!I->second.hasMaterializationSideEffectsOnly() &&
+           "Can't resolve materialization-side-effects-only symbol");
     assert((KV.second.getFlags() & ~WeakFlags) == (I->second & ~WeakFlags) &&
            "Resolving symbol with incorrect flags");
   }
@@ -398,11 +414,13 @@ void ReExportsMaterializationUnit::materialize(
     SymbolAliasMap Aliases;
   };
 
-  // Build a list of queries to issue. In each round we build the largest set of
-  // aliases that we can resolve without encountering a chain definition of the
-  // form Foo -> Bar, Bar -> Baz. Such a form would deadlock as the query would
-  // be waitin on a symbol that it itself had to resolve. Usually this will just
-  // involve one round and a single query.
+  // Build a list of queries to issue. In each round we build a query for the
+  // largest set of aliases that we can resolve without encountering a chain of
+  // aliases (e.g. Foo -> Bar, Bar -> Baz). Such a chain would deadlock as the
+  // query would be waiting on a symbol that it itself had to resolve. Creating
+  // a new query for each link in such a chain eliminates the possibility of
+  // deadlock. In practice chains are likely to be rare, and this algorithm will
+  // usually result in a single query to issue.
 
   std::vector<std::pair<SymbolLookupSet, std::shared_ptr<OnResolveInfo>>>
       QueryInfos;
@@ -419,7 +437,10 @@ void ReExportsMaterializationUnit::materialize(
         continue;
 
       ResponsibilitySymbols.insert(KV.first);
-      QuerySymbols.add(KV.second.Aliasee);
+      QuerySymbols.add(KV.second.Aliasee,
+                       KV.second.AliasFlags.hasMaterializationSideEffectsOnly()
+                           ? SymbolLookupFlags::WeaklyReferencedSymbol
+                           : SymbolLookupFlags::RequiredSymbol);
       QueryAliases[KV.first] = std::move(KV.second);
     }
 
@@ -468,8 +489,13 @@ void ReExportsMaterializationUnit::materialize(
       if (Result) {
         SymbolMap ResolutionMap;
         for (auto &KV : QueryInfo->Aliases) {
-          assert(Result->count(KV.second.Aliasee) &&
+          assert((KV.second.AliasFlags.hasMaterializationSideEffectsOnly() ||
+                  Result->count(KV.second.Aliasee)) &&
                  "Result map missing entry?");
+          // Don't try to resolve materialization-side-effects-only symbols.
+          if (KV.second.AliasFlags.hasMaterializationSideEffectsOnly())
+            continue;
+
           ResolutionMap[KV.first] = JITEvaluatedSymbol(
               (*Result)[KV.second.Aliasee].getAddress(), KV.second.AliasFlags);
         }
@@ -906,7 +932,9 @@ Error JITDylib::emit(const SymbolFlagsMap &Emitted) {
       auto &SymEntry = SymI->second;
 
       // Move symbol to the emitted state.
-      assert(SymEntry.getState() == SymbolState::Resolved &&
+      assert(((SymEntry.getFlags().hasMaterializationSideEffectsOnly() &&
+               SymEntry.getState() == SymbolState::Materializing) ||
+              SymEntry.getState() == SymbolState::Resolved) &&
              "Emitting from state other than Resolved");
       SymEntry.setState(SymbolState::Emitted);
 
@@ -1311,6 +1339,12 @@ Error JITDylib::lodgeQueryImpl(MaterializationUnitList &MUs,
         if (SymI == Symbols.end())
           return false;
 
+        // If we match against a materialization-side-effects only symbol then
+        // make sure it is weakly-referenced. Otherwise bail out with an error.
+        if (SymI->second.getFlags().hasMaterializationSideEffectsOnly() &&
+            SymLookupFlags != SymbolLookupFlags::WeaklyReferencedSymbol)
+          return make_error<SymbolsNotFound>(SymbolNameVector({Name}));
+
         // If this is a non exported symbol and we're matching exported symbols
         // only then skip this symbol without removal.
         if (!SymI->second.getFlags().isExported() &&
@@ -1580,6 +1614,9 @@ JITDylib::JITDylib(ExecutionSession &ES, std::string Name)
 }
 
 Error JITDylib::defineImpl(MaterializationUnit &MU) {
+
+  LLVM_DEBUG({ dbgs() << "  " << MU.getSymbols() << "\n"; });
+
   SymbolNameSet Duplicates;
   std::vector<SymbolStringPtr> ExistingDefsOverridden;
   std::vector<SymbolStringPtr> MUDefsOverridden;
@@ -1604,14 +1641,26 @@ Error JITDylib::defineImpl(MaterializationUnit &MU) {
   }
 
   // If there were any duplicate definitions then bail out.
-  if (!Duplicates.empty())
+  if (!Duplicates.empty()) {
+    LLVM_DEBUG(
+        { dbgs() << "  Error: Duplicate symbols " << Duplicates << "\n"; });
     return make_error<DuplicateDefinition>(std::string(**Duplicates.begin()));
+  }
 
   // Discard any overridden defs in this MU.
+  LLVM_DEBUG({
+    if (!MUDefsOverridden.empty())
+      dbgs() << "  Defs in this MU overridden: " << MUDefsOverridden << "\n";
+  });
   for (auto &S : MUDefsOverridden)
     MU.doDiscard(*this, S);
 
   // Discard existing overridden defs.
+  LLVM_DEBUG({
+    if (!ExistingDefsOverridden.empty())
+      dbgs() << "  Existing defs overridden by this MU: " << MUDefsOverridden
+             << "\n";
+  });
   for (auto &S : ExistingDefsOverridden) {
 
     auto UMII = UnmaterializedInfos.find(S);

llvm/lib/ExecutionEngine/Orc/DebugUtils.cpp

@@ -150,6 +150,10 @@ raw_ostream &operator<<(raw_ostream &OS, const SymbolNameVector &Symbols) {
   return OS << printSequence(Symbols, '[', ']', PrintAll<SymbolStringPtr>());
 }
 
+raw_ostream &operator<<(raw_ostream &OS, ArrayRef<SymbolStringPtr> Symbols) {
+  return OS << printSequence(Symbols, '[', ']', PrintAll<SymbolStringPtr>());
+}
+
 raw_ostream &operator<<(raw_ostream &OS, const JITSymbolFlags &Flags) {
   if (Flags.hasError())
     OS << "[*ERROR*]";

llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -176,7 +176,7 @@ class GenericLLVMIRPlatformSupport : public LLJIT::PlatformSupport {
 
   Error notifyAdding(JITDylib &JD, const MaterializationUnit &MU) {
     if (auto &InitSym = MU.getInitializerSymbol())
-      InitSymbols[&JD].add(InitSym);
+      InitSymbols[&JD].add(InitSym, SymbolLookupFlags::WeaklyReferencedSymbol);
     else {
       // If there's no identified init symbol attached, but there is a symbol
       // with the GenericIRPlatform::InitFunctionPrefix, then treat that as
@@ -185,7 +185,8 @@ class GenericLLVMIRPlatformSupport : public LLJIT::PlatformSupport {
       // map (which holds the names of the symbols to execute).
       for (auto &KV : MU.getSymbols())
         if ((*KV.first).startswith(*InitFunctionPrefix)) {
-          InitSymbols[&JD].add(KV.first);
+          InitSymbols[&JD].add(KV.first,
+                               SymbolLookupFlags::WeaklyReferencedSymbol);
           InitFunctions[&JD].add(KV.first);
         }
     }

llvm/lib/ExecutionEngine/Orc/Layer.cpp

@@ -81,11 +81,19 @@ IRMaterializationUnit::IRMaterializationUnit(
 
     // If we need an init symbol for this module then create one.
     if (!llvm::empty(getStaticInitGVs(M))) {
-      std::string InitSymbolName;
-      raw_string_ostream(InitSymbolName)
-          << "$." << M.getModuleIdentifier() << ".__inits";
-      InitSymbol = ES.intern(InitSymbolName);
-      SymbolFlags[InitSymbol] = JITSymbolFlags();
+      size_t Counter = 0;
+
+      while (true) {
+        std::string InitSymbolName;
+        raw_string_ostream(InitSymbolName)
+            << "$." << M.getModuleIdentifier() << ".__inits." << Counter++;
+        InitSymbol = ES.intern(InitSymbolName);
+        if (SymbolFlags.count(InitSymbol))
+          continue;
+        SymbolFlags[InitSymbol] =
+            JITSymbolFlags::MaterializationSideEffectsOnly;
+        break;
+      }
     }
   });
 }

llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp

@@ -164,7 +164,8 @@ Error MachOPlatform::notifyAdding(JITDylib &JD, const MaterializationUnit &MU) {
   if (!InitSym)
     return Error::success();
 
-  RegisteredInitSymbols[&JD].add(InitSym);
+  RegisteredInitSymbols[&JD].add(InitSym,
+                                 SymbolLookupFlags::WeaklyReferencedSymbol);
   LLVM_DEBUG({
     dbgs() << "MachOPlatform: Registered init symbol " << *InitSym << " for MU "
            << MU.getName() << "\n";

llvm/lib/ExecutionEngine/Orc/Mangling.cpp

@@ -130,11 +130,19 @@ getObjectSymbolInfo(ExecutionSession &ES, MemoryBufferRef ObjBuffer) {
     for (auto &Sec : MachOObj.sections()) {
       auto SecType = MachOObj.getSectionType(Sec);
       if ((SecType & MachO::SECTION_TYPE) == MachO::S_MOD_INIT_FUNC_POINTERS) {
-        std::string InitSymString;
-        raw_string_ostream(InitSymString)
-            << "$." << ObjBuffer.getBufferIdentifier() << ".__inits";
-        InitSymbol = ES.intern(InitSymString);
-        SymbolFlags[InitSymbol] = JITSymbolFlags();
+        size_t Counter = 0;
+        while (true) {
+          std::string InitSymString;
+          raw_string_ostream(InitSymString)
+              << "$." << ObjBuffer.getBufferIdentifier() << ".__inits."
+              << Counter++;
+          InitSymbol = ES.intern(InitSymString);
+          if (SymbolFlags.count(InitSymbol))
+            continue;
+          SymbolFlags[InitSymbol] =
+              JITSymbolFlags::MaterializationSideEffectsOnly;
+          break;
+        }
         break;
       }
     }