[MLIR] Register JIT event listeners with RTDyldObjectLinkingLayer

[ezhulenev]

Summary

Use a new API to register JIT event listeners.

[rriddle]

Thanks for fixing this! I've been meaning to come back and clean this up after listener support was added.

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

122

nit: Use /// for comments.

[ezhulenev]

Fix comments.

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

 ///     void _mlir_funcName(void **)
 ///
 /// where the only argument is interpreted as a list of pointers to the actual
 /// arguments of the function, followed by a pointer to the result.  This allows
 /// the engine to provide the caller with a generic function pointer that can
 /// 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
   /// can be used, e.g., for reporting or optimization. `jitCodeGenOptLevel`,
   /// when provided, is used as the optimization level for target code
   /// generation. If `sharedLibPaths` are provided, the underlying
   /// 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.
+  /// 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.
   llvm::Expected<void (*)(void **)> lookup(StringRef name) const;
 
   /// Invokes the function with the given name passing it the list of arguments.
   /// The arguments are accepted by lvalue-reference since the packed function
   /// interface expects a list of non-null pointers.
private:
   /// Underlying LLJIT.
   std::unique_ptr<llvm::orc::LLJIT> jit;
 
   /// Underlying cache.
   std::unique_ptr<SimpleObjectCache> cache;
 
   /// GDB notification listener.
   llvm::JITEventListener *gdbListener;
+
+  // Perf notification listener.
+  llvm::JITEventListener *perfListener;
 };
 
 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

for (auto &func : module->getFunctionList()) {
     }
 
     // The interface function returns void.
     builder.CreateRetVoid();
   }
 }
 
 ExecutionEngine::ExecutionEngine(bool enableObjectCache,
-                                 bool enableGDBNotificationListener)
+                                 bool enableGDBNotificationListener,
+                                 bool enablePerfNotificationListener)
     : cache(enableObjectCache ? new SimpleObjectCache() : nullptr),
       gdbListener(enableGDBNotificationListener
                       ? llvm::JITEventListener::createGDBRegistrationListener()
+                      : 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);
   if (!llvmModule)
     return make_string_error("could not convert to LLVM IR");
   // FIXME: the triple should be passed to the translation or dialect conversion
   // instead of this.  Currently, the LLVM module created above has no triple
   // associated with it.
   setupTargetTriple(llvmModule.get());
   packFunctionArguments(llvmModule.get());
 
   // Clone module in a new LLVMContext since translateModuleToLLVMIR buries
   // ownership too deeply.
   // TODO(zinenko): Reevaluate model of ownership of LLVMContext in LLVMDialect.
   std::unique_ptr<Module> deserModule =
       LLVM::cloneModuleIntoNewContext(ctx.get(), llvmModule.get());
   auto dataLayout = deserModule->getDataLayout();
 
   // 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) {
-          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) {
       auto mb = llvm::MemoryBuffer::getFile(libPath);
       if (!mb) {
         errs() << "Fail to create MemoryBuffer for: " << libPath << "\n";
         continue;
       }