[RFC] Introduce the concept of IR listeners in MLIR

I worked on a solution to be able to track all IR mutations by attaching “listeners” to an Operation directly. At the moment we built some infrastructure for listeners around the Rewriter (and relevant subclasses). This allows to keep track of most IR mutations, but require the clients to ensure that all mutation are going through the rewriter API: any direct mutation doing through APIs on the Operation class or updating an OpOperand in place won’t be caught.

Here is the first draft implementation: ⚙ D142291 Introduce the concept of IR listeners in MLIR

The Listener interface is expected to grow with more hooks (for example modifying attributes…), but at the moment I implemented support for these:

/// This class attaches to an operation and provides a mechanism to listen to IR
/// modifications. The listeners are notified when an operation is inserted,
/// detached, destroyed, moved, or operands are updated.
class IRListener : public llvm::RefCountedBase<IRListener> {
public:
  virtual ~IRListener() = default;
  // This method is called when the listener is attached to an operation.
  virtual void attachToOperation(Operation *op) {}
  // This method is called when an operation is inserted into a block. The oldBlock is nullptr is the operation wasn't previously in a block.
  virtual void notifyOpInserted(Operation *op, Block *oldBlock,
                                Block *newBlock) {}
                                // This method is called when the operation is detached from a block.
  virtual void notifyOpDetached(Operation *op) {}
  // This method is called when the operation is about to be destroyed.
  virtual void notifyOpDestroyed(Operation *op) {}
  // This method is called when the operation is moved.
  virtual void notifyOpMoved(Operation *op) {}
  // This method is called when an operand is updated.
  virtual void notifyOpOperandChanged(OpOperand &operand, Value newValue) {}
  // This method is called when a block operand is updated.
  virtual void notifyBlockOperandChanged(BlockOperand &operand,
                                         Block *newBlock) {}
};

To keep track of changes, a client would derive this interface and implement the desired hooks.
To showcase this with an example, a -test-ir-listeners is included in the patch, it installs the following listener on the current IR unit (and all nested operations):

class Listener final : public IRListener {
public:
  ~Listener() final {
    llvm::errs() << "===== IRListener Trace =====\n";
    llvm::errs() << traceStream.str();
    llvm::errs() << "===== End IRListener Trace =====\n";
  }
  void printOp(Operation *op) {
    op->print(traceStream, OpPrintingFlags()
                               .elideLargeElementsAttrs()
                               .printGenericOpForm()
                               .assumeVerified()
                               .useLocalScope()
                               .elideRegions());
  }
  void notifyOpInserted(Operation *op, Block *oldBlock, Block *newBlock) final {
    traceStream << "Op inserted: ";
    printOp(op);
    traceStream << "\n";
  }
  void notifyOpDetached(Operation *op) final {
    traceStream << "Op detached: ";
    printOp(op);
    traceStream << "\n";
  }
  void notifyOpDestroyed(Operation *op) final {
    traceStream << "Op destroyed: ";
    printOp(op);
    traceStream << "\n";
  }
  void notifyOpMoved(Operation *op) final {
    traceStream << "Op moved: ";
    printOp(op);
    traceStream << "\n";
  }
  void notifyOpOperandChanged(OpOperand &operand, Value newValue) final {
    traceStream << "OpOperand #" << operand.getOperandNumber()
                << " changed on Operation ";
    printOp(operand.getOwner());
    traceStream << " to " << newValue << "\n";
  }
  void notifyBlockOperandChanged(BlockOperand &operand, Block *newBlock) final {
  }

private:
  std::string trace;
  llvm::raw_string_ostream traceStream{trace};
};

This listener will be destroyed only when the IR is destroyed and print a trace of every event notified, see the test case listener.mlir:

$ ./bin/mlir-opt  -test-ir-listeners -canonicalize   ../mlir/test/IR/listener.mlir
module {
  func.func @andOfExtSI(%arg0: i8, %arg1: i8) -> i64 {
    %0 = arith.andi %arg0, %arg1 : i8
    %1 = arith.extsi %0 : i8 to i64
    return %1 : i64
  }
}

===== IRListener Trace =====
Op inserted: <<UNKNOWN SSA VALUE>> = "arith.andi"(%arg0, %arg1) : (i8, i8) -> i8
Op inserted: <<UNKNOWN SSA VALUE>> = "arith.extsi"(%2) : (i8) -> i64
OpOperand #0 changed on Operation "func.return"(%4) : (i64) -> () to %3 = arith.extsi %2 : i8 to i64
Op detached: %4 = "arith.andi"(%0, %1) : (i64, i64) -> i64
Op destroyed: %0 = "arith.andi"(<<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>) : (i64, i64) -> i64
Op detached: %1 = "arith.extsi"(%arg1) : (i8) -> i64
Op destroyed: %0 = "arith.extsi"(<<UNKNOWN SSA VALUE>>) : (i8) -> i64
Op detached: %0 = "arith.extsi"(%arg0) : (i8) -> i64
Op destroyed: %0 = "arith.extsi"(<<UNKNOWN SSA VALUE>>) : (i8) -> i64
Op destroyed: "builtin.module"() (1 elided regions...) : () -> ()
Op detached: "func.func"() (1 elided regions...) {function_type = (i8, i8) -> i64, sym_name = "andOfExtSI"} : () -> ()
Op destroyed: "func.func"() (1 elided regions...) {function_type = (i8, i8) -> i64, sym_name = "andOfExtSI"} : () -> ()
Op detached: "func.return"(<<NULL VALUE>>) : (<<NULL TYPE>>) -> ()
Op destroyed: "func.return"(<<NULL VALUE>>) : (<<NULL TYPE>>) -> ()
Op detached: %0 = "arith.extsi"(<<NULL VALUE>>) : (<<NULL TYPE>>) -> i64
Op destroyed: %0 = "arith.extsi"(<<NULL VALUE>>) : (<<NULL TYPE>>) -> i64
Op detached: %0 = "arith.andi"(<<NULL VALUE>>, <<NULL VALUE>>) : (<<NULL TYPE>>, <<NULL TYPE>>) -> i8
Op destroyed: %0 = "arith.andi"(<<NULL VALUE>>, <<NULL VALUE>>) : (<<NULL TYPE>>, <<NULL TYPE>>) -> i8
===== End IRListener Trace =====

When we run some passes after -test-ir-listeners installed the listener, we collect IR modification events in a trace, and when the module is destroyed the listener prints a trace of the modifications that were recorded.

Implementation

This will cost a new pointer-size member on each operation to be able to attach listener. When no listener is set, we will also pay the cost a branch for checking this member on every mutation.
When one or multiple listeners are attached to an operation, we pay the price of iterating a vector and one virtual dispatch per listener.

A listener instance is refcounted, it’ll be kept alive for as long at there is an operation referring to it. This is what allows the -test-ir-listeners pass to install a listener and have it outlive the IR.

To be able to track mutation on Operation that didn’t exist at the time the listener was set, we added some custom logic so that when an Operation is inserted in a block, it’ll inherit the listeners from the parentOp (if any).
This is what enables the listener to track mutation to operations created during canonicalization in the code above.
In the extreme, if you start from an empty ModuleOp and install a listener on it, every Operation added somewhere nested in this module will inherit this listener: we will be able to track every single IR mutation and trace it.

6 Likes

Could the listener add support to trace the replaceUsers operation? In a pattern-rewrite transformation, the replaceUsers operations are very common, if the listener could trace these operations, it will be very convenient to track the transformation process.

Let’s talk about this some time before you go further. This is directly related to the llvm/IR/ValueHandle.h stuff that we built in LLVM IR and that turns out to be a really bad thing IMO. I added it originally for alias analyses, with the goal of making it so the analysis would automatically stay up to date as transformations mutated the IR.

There are two problems with this:

  1. It bloats the IR with the extra pointer as you mention. This can be reduced to a single bit + on the side hash table though.

  2. The bigger issue is that this is effectively impossible to use. You don’t get context to update the analyses, and the granularity of updates is tiny. Building anything on top of this sort of thing is really slow because of this, and impossible to make “actually general” because there are transformations that do all sorts of things. Reflecting them would require tons of high level transformation methods to reflect them into the hooks.

In LLVM, the ValueHandle stuff was originally adopted very widely and then backed out because of these reasons. It is now effectively only used for very local and simple things that know how the transformations are being applied. At that point though, the overhead of cost and complexity defeats the purpose of using something like this: if you have tight coupling between the code doing the IR transforms and the things that need updating, you might as well directly update them.

7 Likes

I just want to add that GlobalIsel uses GISelChangeObserver for the same idea.

Think this could be interesting for debugging passes as it allows to hook in at a low granularity and observe what happens.

This sounds interesting, but might have to fine of a granularity as Chris mentions above. In some of the cases for rewriter listeners that I have, it is interesting to see even higher level mutations like “replace an op with another op” (rather than a list of values) or “combine these two ops into one”. That being said, I would really love this to be implemented as a sort of sanitizer to intercept all direct IR mutations bypassing the rewriter when it is supposed to be used.

This was also on my mind: a mechanism for adding safety features. I haven’t fully thought this through, but if using more for those kind of use cases, I might have considered a global context-level IR mutation listener instead of fine grained registration per op. Then you could gate all of the logic and overhead on one branch on a context-owned bool.

I’m also not sure that the granularity of the registration matches use cases I have seen.

It seems like this feature might be useful for a “better” MLIR diff experience. (Maybe that was the goal?)

Could be useful for diffing, but I think it’s not needed (you could just give each op a unique location and achieve ~the same thing, I think?).

As one more point in this problem space, there’s the CFG update interface (llvm/IR/CFGUpdate.h / llvm::cfg::Update) which can be used to update (post-)dominator trees in bulk. A similar bulk update mechanism would be interesting for mutating IR, especially from a safety perspective.

1 Like

Kind of, but it is a bit of a clutch: you rely on location propagation (which is imperfect) and there is also the fundamental issue of many-to-many rewrites which can make it hard to trace things back.
It kind of work for many debugging use-cases though (where the imperfection does not prevent to get far enough, especially in combination with other techniques like -debug traces).