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

[merge_guards_bot]

Unit tests: pass. 62405 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[nicolasvasilache]

Great, thanks @rriddle !

[stephenneuendorffer]

This broke BUILD_SHARED_LIBS=on, and I have submitted a fix.

[dblaikie]

@lhames just for some examples of debugging and orc jit

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

 
 #include <functional>
 #include <memory>
 
 namespace llvm {
 template <typename T> class Expected;
 class Module;
 class ExecutionEngine;
+class JITEventListener;
 class MemoryBuffer;
 } // namespace llvm
 
 namespace mlir {
 
 class ModuleOp;
 
 /// A simple object cache following Lang's LLJITWithObjectCache example.
public:
   /// Set the target triple on the module. This is implicitly done when creating
   /// the engine.
   static bool setupTargetTriple(llvm::Module *llvmModule);
 
   /// Dump object code to output file `filename`.
   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>
 llvm::Error ExecutionEngine::invoke(StringRef name, Args &... args) {
   auto expectedFPtr = lookup(name);
   if (!expectedFPtr)
     return expectedFPtr.takeError();
   auto fptr = *expectedFPtr;

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

 #include "mlir/ExecutionEngine/ExecutionEngine.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Target/LLVMIR.h"
 
 #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"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
for (auto &func : module->getFunctionList()) {
     }
 
     // The interface function returns void.
     builder.CreateRetVoid();
   }
 }
 
 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,
     Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel,
     ArrayRef<StringRef> sharedLibPaths, bool enableObjectCache) {
   auto engine = std::make_unique<ExecutionEngine>(enableObjectCache);
 
   std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
   std::unique_ptr<Module> deserModule = std::move(*expectedModule);
 
   // Callback to create the object layer with symbol resolution to current
   // process and dynamically linked libraries.
   auto objectLinkingLayerCreator = [&](ExecutionSession &session,
                                        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)
       mainJD = &session.createJITDylib("<main>");
 
     // Resolve symbols that are statically linked in the current process.
     mainJD->addGenerator(
         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(