[ORC] Enable JITEventListeners in the RTDyldObjectLinkingLayer.

[lhames]

Summary

Enable use of ExecutionEngine JITEventListeners in RTDyldObjectLinkingLayer.
This allows existing MCJIT clients to more easily migrate to LLJIT / ORCv2.

Still needs some polish, but posting now so that potential clients can road
test it early.

[wenju]

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

87

&MemMgr should be MemMgr.get()

[lhames]

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

87

You're right. Thanks Wenju!

[lhames]

I am going to fix the issue Wenju identified and commit this: further development can happen in-tree.

[thakis]

This broke tests in exactly the way the precommit test predicted (http://45.33.8.238/linux/12615/step_12.txt). Please fix or revert. (And please look at precommit test results going forward.)

[lhames]

In D75838#1923564, @thakis wrote:

This broke tests in exactly the way the precommit test predicted (http://45.33.8.238/linux/12615/step_12.txt). Please fix or revert. (And please look at precommit test results going forward.)

Already reverted. I didn't realize the pre-commit tests were actually reliable: previous Harbormaster messages seemed to just be noise. That's very handy!

[lhames]

Re-enabled in 39253a50f0f with fixes for the use-after-free bug.

llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_EXECUTIONENGINE_ORC_RTDYLDOBJECTLINKINGLAYER_H
 #define LLVM_EXECUTIONENGINE_ORC_RTDYLDOBJECTLINKINGLAYER_H
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ExecutionEngine/JITEventListener.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/Layer.h"
 #include "llvm/ExecutionEngine/Orc/Legacy.h"
 #include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Error.h"
 #include <algorithm>
public:
   /// symbols may not be detected until late, and detection may depend on
   /// the flow of control through JIT'd code. Use with care.
   RTDyldObjectLinkingLayer &
   setAutoClaimResponsibilityForObjectSymbols(bool AutoClaimObjectSymbols) {
     this->AutoClaimObjectSymbols = AutoClaimObjectSymbols;
     return *this;
   }
 
+  /// Register a JITEventListener.
+  void registerJITEventListener(JITEventListener &L);
+
+  /// Unregister a JITEventListener.
+  void unregisterJITEventListener(JITEventListener &L);
+
 private:
   Error onObjLoad(VModuleKey K, MaterializationResponsibility &R,
-                  object::ObjectFile &Obj,
+                  object::ObjectFile &Obj, RuntimeDyld::MemoryManager *MemMgr,
                   std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
                   std::map<StringRef, JITEvaluatedSymbol> Resolved,
                   std::set<StringRef> &InternalSymbols);
 
-  void onObjEmit(VModuleKey K, std::unique_ptr<MemoryBuffer> ObjBuffer,
-                 MaterializationResponsibility &R, Error Err);
+  void onObjEmit(VModuleKey K, MaterializationResponsibility &R,
+                 object::ObjectFile &Obj,
+                 std::unique_ptr<MemoryBuffer> ObjBuffer,
+                 RuntimeDyld::MemoryManager *MemMgr, Error Err);
 
   mutable std::mutex RTDyldLayerMutex;
   GetMemoryManagerFunction GetMemoryManager;
   NotifyLoadedFunction NotifyLoaded;
   NotifyEmittedFunction NotifyEmitted;
   bool ProcessAllSections = false;
   bool OverrideObjectFlags = false;
   bool AutoClaimObjectSymbols = false;
   std::vector<std::unique_ptr<RuntimeDyld::MemoryManager>> MemMgrs;
+  std::vector<JITEventListener *> EventListeners;
+  DenseMap<RuntimeDyld::MemoryManager *,
+           std::unique_ptr<RuntimeDyld::LoadedObjectInfo>>
+      LoadedObjInfos;
 };
 
 class LegacyRTDyldObjectLinkingLayerBase {
 public:
   using ObjectPtr = std::unique_ptr<MemoryBuffer>;
 
 protected:
 

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

 namespace orc {
 
 RTDyldObjectLinkingLayer::RTDyldObjectLinkingLayer(
     ExecutionSession &ES, GetMemoryManagerFunction GetMemoryManager)
     : ObjectLayer(ES), GetMemoryManager(GetMemoryManager) {}
 
 RTDyldObjectLinkingLayer::~RTDyldObjectLinkingLayer() {
   std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
-  for (auto &MemMgr : MemMgrs)
+  for (auto &MemMgr : MemMgrs) {
+    for (auto *L : EventListeners)
+      L->notifyFreeingObject(
+          static_cast<uint64_t>(reinterpret_cast<uintptr_t>(&MemMgr)));
     MemMgr->deregisterEHFrames();
   }
+}
 
 void RTDyldObjectLinkingLayer::emit(MaterializationResponsibility R,
                                     std::unique_ptr<MemoryBuffer> O) {
   assert(O && "Object must not be null");
 
   // This method launches an asynchronous link step that will fulfill our
   // materialization responsibility. We need to switch R to be heap
   // allocated before that happens so it can live as long as the asynchronous
// Create a record a memory manager for this object.
     MemMgrs.push_back(std::move(Tmp));
     MemMgr = MemMgrs.back().get();
   }
 
   JITDylibSearchOrderResolver Resolver(*SharedR);
 
   jitLinkForORC(
       **Obj, std::move(O), *MemMgr, Resolver, ProcessAllSections,
-      [this, K, SharedR, &Obj, InternalSymbols](
+      [this, K, SharedR, &Obj, MemMgr, InternalSymbols](
           std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
           std::map<StringRef, JITEvaluatedSymbol> ResolvedSymbols) {
-        return onObjLoad(K, *SharedR, **Obj, std::move(LoadedObjInfo),
+        return onObjLoad(K, *SharedR, **Obj, MemMgr, std::move(LoadedObjInfo),
                          ResolvedSymbols, *InternalSymbols);
       },
-      [this, K, SharedR, O = std::move(O)](Error Err) mutable {
-        onObjEmit(K, std::move(O), *SharedR, std::move(Err));
+      [this, K, SharedR, &Obj, O = std::move(O), MemMgr](Error Err) mutable {
+        onObjEmit(K, *SharedR, **Obj, std::move(O), MemMgr, std::move(Err));
       });
 }
 
+void RTDyldObjectLinkingLayer::registerJITEventListener(JITEventListener &L) {
+  std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  assert(llvm::none_of(EventListeners,
+                       [&](JITEventListener *O) { return O == &L; }) &&
+         "Listener has already been registered");
+  EventListeners.push_back(&L);
+}
+
+void RTDyldObjectLinkingLayer::unregisterJITEventListener(JITEventListener &L) {
+  std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  auto I = llvm::find(EventListeners, &L);
+  assert(I != EventListeners.end() && "Listener not registered");
+  EventListeners.erase(I);
+}
+
 Error RTDyldObjectLinkingLayer::onObjLoad(
     VModuleKey K, MaterializationResponsibility &R, object::ObjectFile &Obj,
+    RuntimeDyld::MemoryManager *MemMgr,
     std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
     std::map<StringRef, JITEvaluatedSymbol> Resolved,
     std::set<StringRef> &InternalSymbols) {
   SymbolFlagsMap ExtraSymbolsToClaim;
   SymbolMap Symbols;
 
   // Hack to support COFF constant pool comdats introduced during compilation:
   // (See http://llvm.org/PR40074)
Error RTDyldObjectLinkingLayer::onObjLoad(
   if (auto Err = R.notifyResolved(Symbols)) {
     R.failMaterialization();
     return Err;
   }
 
   if (NotifyLoaded)
     NotifyLoaded(K, Obj, *LoadedObjInfo);
 
+  std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  assert(!LoadedObjInfos.count(MemMgr) && "Duplicate loaded info for MemMgr");
+  LoadedObjInfos[MemMgr] = std::move(LoadedObjInfo);
+
   return Error::success();
 }
 
 void RTDyldObjectLinkingLayer::onObjEmit(
-    VModuleKey K, std::unique_ptr<MemoryBuffer> ObjBuffer,
-    MaterializationResponsibility &R, Error Err) {
+    VModuleKey K, MaterializationResponsibility &R, object::ObjectFile &Obj,
+    std::unique_ptr<MemoryBuffer> ObjBuffer, RuntimeDyld::MemoryManager *MemMgr,
+    Error Err) {
   if (Err) {
     getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();
     return;
   }
 
   if (auto Err = R.notifyEmitted()) {
     getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();
     return;
   }
 
   if (NotifyEmitted)
     NotifyEmitted(K, std::move(ObjBuffer));
+
+  // Run EventListener notifyLoaded callbacks.
+  std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  auto LOIItr = LoadedObjInfos.find(MemMgr);
+  assert(LOIItr != LoadedObjInfos.end() && "LoadedObjInfo missing");
+  for (auto *L : EventListeners)
+    L->notifyObjectLoaded(
+        static_cast<uint64_t>(reinterpret_cast<uintptr_t>(MemMgr)), Obj,
+        *LOIItr->second);
+  LoadedObjInfos.erase(MemMgr);
 }
 
 LegacyRTDyldObjectLinkingLayer::LegacyRTDyldObjectLinkingLayer(
     ExecutionSession &ES, ResourcesGetter GetResources,
     NotifyLoadedFtor NotifyLoaded, NotifyFinalizedFtor NotifyFinalized,
     NotifyFreedFtor NotifyFreed)
     : ES(ES), GetResources(std::move(GetResources)),
       NotifyLoaded(std::move(NotifyLoaded)),
       NotifyFinalized(std::move(NotifyFinalized)),
       NotifyFreed(std::move(NotifyFreed)), ProcessAllSections(false) {}
 
 } // End namespace orc.
 } // End namespace llvm.

llvm/tools/lli/lli.cpp

 #include "llvm/ExecutionEngine/MCJIT.h"
 #include "llvm/ExecutionEngine/ObjectCache.h"
 #include "llvm/ExecutionEngine/Orc/DebugUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/ExecutionEngine/Orc/MachOPlatform.h"
 #include "llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/OrcMCJITReplacement.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/IRReader/IRReader.h"
case LLJITPlatform::MachO:
       break;
     default:
       llvm_unreachable("Unrecognized platform value");
     }
   }
 
   auto J = ExitOnErr(Builder.create());
 
+  if (!TT->isOSBinFormatMachO())
+    static_cast<llvm::orc::RTDyldObjectLinkingLayer &>(J->getObjLinkingLayer())
+        .registerJITEventListener(
+            *JITEventListener::createGDBRegistrationListener());
+
   if (PerModuleLazy)
     J->setPartitionFunction(orc::CompileOnDemandLayer::compileWholeModule);
 
   auto Dump = createDebugDumper();
 
   J->getIRTransformLayer().setTransform(
       [&](orc::ThreadSafeModule TSM,
           const orc::MaterializationResponsibility &R) {