[Mlir-commits] [mlir] 2e312a7 - WIP listeners

[Mehdi Amini]

Author: Mehdi Amini
Date: 2023-01-21T01:03:43+01:00
New Revision: 2e312a7baec5e8d8f841ee819966ff3f53f274c6

URL: https://github.com/llvm/llvm-project/commit/2e312a7baec5e8d8f841ee819966ff3f53f274c6
DIFF: https://github.com/llvm/llvm-project/commit/2e312a7baec5e8d8f841ee819966ff3f53f274c6.diff

LOG: WIP listeners

Added: 
    

Modified: 
    mlir/include/mlir/IR/Operation.h
    mlir/lib/IR/Operation.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 59d450ea97bb8..3fbcfe871a430 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -22,6 +22,30 @@
 #include <optional>
 
 namespace mlir {
+class IRListeners {
+public:
+  class ListenerInterface {
+   public:
+     ~ListenerInterface() = default;
+     virtual void notifyOpInserted(Operation *op, Block *oldBlock, Block *newBlock);
+     virtual void notifyOpDestroyed(Operation *op);
+  };
+  void notifyOpInserted(Operation *op, Block *oldBlock, Block *newBlock) { 
+    for (auto &listener : listeners)
+      listener.notifyOpInserted(op, oldBlock, newBlock);
+  }
+  void notifyOpDestroyed(Operation *op) {
+    for (auto &listener : listeners)
+      listener.notifyOpDestroyed(op);
+  }
+  void addListener(std::shared_ptr<ListenerInterface> listener) {
+    listeners.push_back(std::move(listener));
+  }
+  private:
+  std::vector<std::shared_ptr<ListenerInterface>> listeners;
+};
+
+
 /// Operation is the basic unit of execution within MLIR.
 ///
 /// The following documentation are recommended to understand this class:
@@ -735,6 +759,12 @@ class alignas(8) Operation final
   /// Returns true if this operation has a valid order.
   bool hasValidOrder() { return orderIndex != kInvalidOrderIdx; }
 
+  /// Attach a listener to this operation.
+  void addListener(std::shared_ptr<ListenerInterface> listener) {
+    if (!listeners) listeners = std::make_unique<IRListeners>();
+    listeners->addListener(std::move(listener));
+  }
+
 private:
   Operation(Location location, OperationName name, unsigned numResults,
             unsigned numSuccessors, unsigned numRegions,
@@ -743,7 +773,7 @@ class alignas(8) Operation final
   // Operations are deleted through the destroy() member because they are
   // allocated with malloc.
   ~Operation();
-
+g
   /// Returns the additional size necessary for allocating the given objects
   /// before an Operation in-memory.
   static size_t prefixAllocSize(unsigned numOutOfLineResults,
@@ -836,6 +866,9 @@ class alignas(8) Operation final
   /// This holds general named attributes for the operation.
   DictionaryAttr attrs;
 
+  /// Optionally defined listeners that register here to capture IR modification events.
+  std::unique_ptr<IRListeners> listeners;
+
   // allow ilist_traits access to 'block' field.
   friend struct llvm::ilist_traits<Operation>;
 

diff  --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 50815b3738bfe..a700a54a784aa 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -162,6 +162,7 @@ Operation::~Operation() {
 
 /// Destroy this operation or one of its subclasses.
 void Operation::destroy() {
+  if (listeners) listeners->notifyOpDestroyed(this);
   // Operations may have additional prefixed allocation, which needs to be
   // accounted for here when computing the address to free.
   char *rawMem = reinterpret_cast<char *>(this) -
@@ -380,8 +381,9 @@ Block *llvm::ilist_traits<::mlir::Operation>::getContainingBlock() {
 /// keep the block pointer up to date.
 void llvm::ilist_traits<::mlir::Operation>::addNodeToList(Operation *op) {
   assert(!op->getBlock() && "already in an operation block!");
+  auto oldBlock = op->block;
   op->block = getContainingBlock();
-
+  if (listeners) listeners->notifyOpInserted(op, oldBlock, newBlock);
   // Invalidate the order on the operation.
   op->orderIndex = Operation::kInvalidOrderIdx;
 }
@@ -390,6 +392,7 @@ void llvm::ilist_traits<::mlir::Operation>::addNodeToList(Operation *op) {
 /// We keep the block pointer up to date.
 void llvm::ilist_traits<::mlir::Operation>::removeNodeFromList(Operation *op) {
   assert(op->block && "not already in an operation block!");
+  if (listeners) listeners->notifyOpRemoved(op);
   op->block = nullptr;
 }