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.