From 97db10d41307f1133a6efb9dba257b2206f28333 Mon Sep 17 00:00:00 2001 From: River Riddle Date: Tue, 26 Mar 2019 12:22:52 -0700 Subject: [PATCH] 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 --- mlir/include/mlir/IR/Instruction.h | 4 ++++ mlir/lib/IR/Instruction.cpp | 13 +++++++++++++ 2 files changed, 17 insertions(+) diff --git a/mlir/include/mlir/IR/Instruction.h b/mlir/include/mlir/IR/Instruction.h index 53d222042ffb..8b4d2e65f174 100644 --- a/mlir/include/mlir/IR/Instruction.h +++ b/mlir/include/mlir/IR/Instruction.h @@ -92,6 +92,10 @@ public: /// 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/mlir/lib/IR/Instruction.cpp b/mlir/lib/IR/Instruction.cpp index 1560e31feb6a..f8ecd292f9da 100644 --- a/mlir/lib/IR/Instruction.cpp +++ b/mlir/lib/IR/Instruction.cpp @@ -300,6 +300,19 @@ MLIRContext *Instruction::getContext() { 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; }