[MLIR] Register JIT event listeners with RTDyldObjectLinkingLayer

Use a new API to register JIT event listeners.

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

diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
index 958a8b1..87515ac 100644 (file)
--- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
+++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
@@ -60,7 +60,8 @@ private:
 /// be used to invoke the JIT-compiled function.
 class ExecutionEngine {
 public:
-  ExecutionEngine(bool enableObjectCache, bool enableGDBNotificationListener);
+  ExecutionEngine(bool enableObjectCache, bool enableGDBNotificationListener,
+                  bool enablePerfNotificationListener);
 
   /// Creates an execution engine for the given module.  If `transformer` is
   /// provided, it will be called on the LLVM module during JIT-compilation and
@@ -70,14 +71,17 @@ public:
   /// JIT-compilation will open and link the shared libraries for symbol
   /// resolution. If `enableObjectCache` is set, the JIT compiler will create
   /// one to store the object generated for the given module. If enable
-  // `enableGDBNotificationListener` is set, the JIT compiler will notify
-  /// the llvm's global GDB notification listener.
+  /// `enableGDBNotificationListener` is set, the JIT compiler will notify
+  /// the llvm's global GDB notification listener. If
+  /// `enablePerfNotificationListener` is set, the JIT compiler will notify
+  /// the llvm's global Perf notification listener.
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
   create(ModuleOp m,
          std::function<llvm::Error(llvm::Module *)> transformer = {},
          Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel = llvm::None,
          ArrayRef<StringRef> sharedLibPaths = {}, bool enableObjectCache = true,
-         bool enableGDBNotificationListener = true);
+         bool enableGDBNotificationListener = true,
+         bool enablePerfNotificationListener = true);
 
   /// Looks up a packed-argument function with the given name and returns a
   /// pointer to it.  Propagates errors in case of failure.
@@ -114,6 +118,9 @@ private:
 
   /// GDB notification listener.
   llvm::JITEventListener *gdbListener;
+
+  /// Perf notification listener.
+  llvm::JITEventListener *perfListener;
 };
 
 template <typename... Args>

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 25bd45f..4d5cfc0 100644 (file)
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -183,19 +183,24 @@ static void packFunctionArguments(Module *module) {
 }
 
 ExecutionEngine::ExecutionEngine(bool enableObjectCache,
-                                 bool enableGDBNotificationListener)
+                                 bool enableGDBNotificationListener,
+                                 bool enablePerfNotificationListener)
     : cache(enableObjectCache ? new SimpleObjectCache() : nullptr),
       gdbListener(enableGDBNotificationListener
                       ? llvm::JITEventListener::createGDBRegistrationListener()
-                      : nullptr) {}
+                      : nullptr),
+      perfListener(enablePerfNotificationListener
+                       ? llvm::JITEventListener::createPerfJITEventListener()
+                       : nullptr) {}
 
 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,
-    bool enableGDBNotificationListener) {
+    bool enableGDBNotificationListener, bool enablePerfNotificationListener) {
   auto engine = std::make_unique<ExecutionEngine>(
-      enableObjectCache, enableGDBNotificationListener);
+      enableObjectCache, enableGDBNotificationListener,
+      enablePerfNotificationListener);
 
   std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
   auto llvmModule = translateModuleToLLVMIR(m);
@@ -220,16 +225,12 @@ 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) {
-          if (engine->gdbListener) {
-            uint64_t key = static_cast<uint64_t>(
-                reinterpret_cast<uintptr_t>(object.getData().data()));
-            engine->gdbListener->notifyObjectLoaded(key, object, objectInfo);
-          }
-        });
+
+    // Register JIT event listeners if they are enabled.
+    if (engine->gdbListener)
+      objectLayer->registerJITEventListener(*engine->gdbListener);
+    if (engine->perfListener)
+      objectLayer->registerJITEventListener(*engine->perfListener);
 
     // Resolve symbols from shared libraries.
     for (auto libPath : sharedLibPaths) {