Add the ability to have "Ops" defined as small C++ classes, with some nice
properties:
- They allow type checked dynamic casting from their base Operation.
- They allow nice accessors for C++ clients, e.g. a "getIndex()" method on
'dim' that returns an unsigned.
- They work with both OperationInst/OperationStmt (once OperationStmt is
implemented).
- They get custom printing logic. They will eventually get custom parsing,
verifier, and builder logic as well.
- Out of tree clients can register their own operation set without having to
change MLIR core, e.g. for TensorFlow or custom target instructions.
This registers addf and dim as examples.
PiperOrigin-RevId: 203382993
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 1cefc0a..61d1698 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -26,6 +26,7 @@
#include "mlir/IR/CFGFunction.h"
#include "mlir/IR/MLFunction.h"
#include "mlir/IR/Module.h"
+#include "mlir/IR/OperationSet.h"
#include "mlir/IR/Types.h"
#include "mlir/Support/STLExtras.h"
#include "llvm/ADT/DenseMap.h"
@@ -103,12 +104,14 @@
private:
const CFGFunction *function;
raw_ostream &os;
+ const OperationSet &operationSet;
DenseMap<const BasicBlock*, unsigned> basicBlockIDs;
};
} // end anonymous namespace
CFGFunctionState::CFGFunctionState(const CFGFunction *function, raw_ostream &os)
- : function(function), os(os) {
+ : function(function), os(os),
+ operationSet(OperationSet::get(function->getContext())) {
// Each basic block gets a unique ID per function.
unsigned blockID = 0;
@@ -148,6 +151,14 @@
}
void CFGFunctionState::print(const OperationInst *inst) {
+ // Check to see if this is a known operation. If so, use the registered
+ // custom printer hook.
+ if (auto opInfo = operationSet.lookup(inst->getName().str())) {
+ os << " ";
+ opInfo->printAssembly(inst, os);
+ return;
+ }
+
// TODO: escape name if necessary.
os << " \"" << inst->getName().str() << "\"()";