[mlir] Register the GDB listener with ExecutionEngine to enable debugging JIT'd code

Differential Revision: https://reviews.llvm.org/D73932

Author: River707

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

@@ -26,6 +26,7 @@ namespace llvm {
 template <typename T> class Expected;
 class Module;
 class ExecutionEngine;
+class JITEventListener;
 class MemoryBuffer;
 } // namespace llvm
 
@@ -97,15 +98,18 @@ class ExecutionEngine {
   void dumpToObjectFile(StringRef filename);
 
 private:
-  // Ordering of llvmContext and jit is important for destruction purposes: the
-  // jit must be destroyed before the context.
+  /// Ordering of llvmContext and jit is important for destruction purposes: the
+  /// jit must be destroyed before the context.
   llvm::LLVMContext llvmContext;
 
-  // Underlying LLJIT.
+  /// Underlying LLJIT.
   std::unique_ptr<llvm::orc::LLJIT> jit;
 
-  // Underlying cache.
+  /// Underlying cache.
   std::unique_ptr<SimpleObjectCache> cache;
+
+  /// GDB notification listener.
+  llvm::JITEventListener *gdbListener;
 };
 
 template <typename... Args>

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

@@ -18,6 +18,7 @@
 
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
+#include "llvm/ExecutionEngine/JITEventListener.h"
 #include "llvm/ExecutionEngine/ObjectCache.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
@@ -182,7 +183,8 @@ static void packFunctionArguments(Module *module) {
 }
 
 ExecutionEngine::ExecutionEngine(bool enableObjectCache)
-    : cache(enableObjectCache ? nullptr : new SimpleObjectCache()) {}
+    : cache(enableObjectCache ? nullptr : new SimpleObjectCache()),
+      gdbListener(llvm::JITEventListener::createGDBRegistrationListener()) {}
 
 Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
     ModuleOp m, std::function<Error(llvm::Module *)> transformer,
@@ -221,6 +223,14 @@ Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
                                        const Triple &TT) {
     auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>(
         session, []() { return std::make_unique<SectionMemoryManager>(); });
+    objectLayer->setNotifyLoaded(
+        [engine = engine.get()](
+            llvm::orc::VModuleKey, const llvm::object::ObjectFile &object,
+            const llvm::RuntimeDyld::LoadedObjectInfo &objectInfo) {
+          uint64_t key = static_cast<uint64_t>(
+              reinterpret_cast<uintptr_t>(object.getData().data()));
+          engine->gdbListener->notifyObjectLoaded(key, object, objectInfo);
+        });
     auto dataLayout = deserModule->getDataLayout();
     llvm::orc::JITDylib *mainJD = session.getJITDylibByName("<main>");
     if (!mainJD)