diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h --- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h +++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h @@ -60,7 +60,8 @@ /// 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 @@ /// 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> create(ModuleOp m, std::function transformer = {}, Optional jitCodeGenOptLevel = llvm::None, ArrayRef 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 @@ /// GDB notification listener. llvm::JITEventListener *gdbListener; + + /// Perf notification listener. + llvm::JITEventListener *perfListener; }; template diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -183,19 +183,24 @@ } 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> ExecutionEngine::create( ModuleOp m, std::function transformer, Optional jitCodeGenOptLevel, ArrayRef sharedLibPaths, bool enableObjectCache, - bool enableGDBNotificationListener) { + bool enableGDBNotificationListener, bool enablePerfNotificationListener) { auto engine = std::make_unique( - enableObjectCache, enableGDBNotificationListener); + enableObjectCache, enableGDBNotificationListener, + enablePerfNotificationListener); std::unique_ptr ctx(new llvm::LLVMContext); auto llvmModule = translateModuleToLLVMIR(m); @@ -220,16 +225,12 @@ const Triple &TT) { auto objectLayer = std::make_unique( session, []() { return std::make_unique(); }); - 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( - reinterpret_cast(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) {