Add a utility Instruction::getDialect method to return the dialect an operation is associated with, or nullptr if the associated dialect has not been registered.
PiperOrigin-RevId: 240402300
diff --git a/include/mlir/IR/Instruction.h b/include/mlir/IR/Instruction.h
index 53d2220..8b4d2e6 100644
--- a/include/mlir/IR/Instruction.h
+++ b/include/mlir/IR/Instruction.h
@@ -92,6 +92,10 @@
/// Return the context this operation is associated with.
MLIRContext *getContext();
+ /// Return the dialact this operation is associated with, or nullptr if the
+ /// associated dialect is not registered.
+ Dialect *getDialect();
+
/// The source location the operation was defined or derived from.
Location getLoc() { return location; }
diff --git a/lib/IR/Instruction.cpp b/lib/IR/Instruction.cpp
index 1560e31..f8ecd29 100644
--- a/lib/IR/Instruction.cpp
+++ b/lib/IR/Instruction.cpp
@@ -300,6 +300,19 @@
return getFunction()->getContext();
}
+/// Return the dialact this operation is associated with, or nullptr if the
+/// associated dialect is not registered.
+Dialect *Instruction::getDialect() {
+ if (auto *abstractOp = getAbstractOperation())
+ return &abstractOp->dialect;
+
+ // If this operation hasn't been registered or doesn't have abstract
+ // operation, fall back to a dialect which matches the prefix.
+ auto opName = getName().getStringRef();
+ auto dialectPrefix = opName.split('.').first;
+ return getContext()->getRegisteredDialect(dialectPrefix);
+}
+
Instruction *Instruction::getParentInst() {
return block ? block->getContainingInst() : nullptr;
}
