[PATCH] D78435: [MLIR] Register JIT event listeners with RTDyldObjectLinkingLayer

[Eugene Zhulenev via Phabricator via llvm-commits]

ezhulenev created this revision.
Herald added subscribers: llvm-commits, frgossen, grosul1, Joonsoo, liufengdb, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, burmako, jpienaar, rriddle, mehdi_amini.
Herald added a project: LLVM.
ezhulenev added a reviewer: rriddle.
rriddle accepted this revision.
rriddle added a comment.
This revision is now accepted and ready to land.

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



================
Comment at: mlir/include/mlir/ExecutionEngine/ExecutionEngine.h:122
+
+  // Perf notification listener.
+  llvm::JITEventListener *perfListener;
----------------
nit: Use /// for comments.


Use a new API to register JIT event listeners.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D78435

Files:
  mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
  mlir/lib/ExecutionEngine/ExecutionEngine.cpp


Index: mlir/lib/ExecutionEngine/ExecutionEngine.cpp
===================================================================
--- mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ 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<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 @@
                                        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) {
Index: mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
===================================================================
--- mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
+++ 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
@@ -71,13 +72,16 @@
   /// 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.
@@ -114,6 +118,9 @@
 
   /// GDB notification listener.
   llvm::JITEventListener *gdbListener;
+
+  // Perf notification listener.
+  llvm::JITEventListener *perfListener;
 };
 
 template <typename... Args>


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D78435.258554.patch
Type: text/x-patch
Size: 4742 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200418/21e63b50/attachment.bin>