[MLIR] Add an Op util which returns its name with the dialect stripped.

[rriddle]

mlir/include/mlir/IR/Operation.h

59 ↗

(On Diff #269368)

nit: ///

61 ↗

(On Diff #269368)

This won't work for 'func' and 'module' as they don't have a prefix. You need to check if .first is empty.

[rriddle]

mlir/include/mlir/IR/Operation.h

61 ↗

(On Diff #269368)

Oops, meant that you need to check if .second is empty.

[lucyrfox]

Added a check for the case where there is no dialect prefix.

Also moved this util to be in OperationSupport alongside getDialect() because it seems to
fit better there.

[mehdi_amini]

Are there use cases in-tree that could be updated?

[lucyrfox]

Add one usage of getting dialect directly.

[lucyrfox]

In D81435#2081398, @mehdi_amini wrote:

Are there use cases in-tree that could be updated?

Most similar cases in MLIR core are operating on a StringRef instead of an OperationName. There is probably more usage downstream which can be updated incrementally. (I'll take a look separately at what can be updated within TF.)

mlir/include/mlir/IR/OperationSupport.h

using RepresentationUnion =
       PointerUnion<Identifier, const AbstractOperation *>;
 
   OperationName(AbstractOperation *op) : representation(op) {}
   OperationName(StringRef name, MLIRContext *context);
 
   /// Return the name of the dialect this operation is registered to.
   StringRef getDialect() const;
 
+  /// Return the operation name with dialect name stripped, if it has one.
+  StringRef stripDialect() const;
+
   /// Return the name of this operation.  This always succeeds.
   StringRef getStringRef() const;
 
   /// If this operation has a registered operation description, return it.
   /// Otherwise return null.
   const AbstractOperation *getAbstractOperation() const;
 
   void print(raw_ostream &os) const;

mlir/lib/IR/MLIRContext.cpp

 bool MLIRContext::isOperationRegistered(StringRef name) {
   // Lock access to the context registry.
   ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled);
 
   return impl->registeredOperations.count(name);
 }
 
 void Dialect::addOperation(AbstractOperation opInfo) {
-  assert((getNamespace().empty() ||
-          opInfo.name.split('.').first == getNamespace()) &&
+  assert((getNamespace().empty() || opInfo.dialect.name == getNamespace()) &&
          "op name doesn't start with dialect namespace");
   assert(&opInfo.dialect == this && "Dialect object mismatch");
   auto &impl = context->getImpl();
 
   // Lock access to the context registry.
   ScopedWriterLock registryLock(impl.contextMutex, impl.threadingIsEnabled);
   if (!impl.registeredOperations.insert({opInfo.name, opInfo}).second) {
     llvm::errs() << "error: operation named '" << opInfo.name

mlir/lib/IR/Operation.cpp

else
     representation = Identifier::get(name, context);
 }
 
 /// Return the name of the dialect this operation is registered to.
 StringRef OperationName::getDialect() const {
   return getStringRef().split('.').first;
 }
 
+/// Return the operation name with dialect name stripped, if it has one.
+StringRef OperationName::stripDialect() const {
+  auto splitName = getStringRef().split(".");
+  return splitName.second.empty() ? splitName.first : splitName.second;
+}
+
 /// Return the name of this operation.  This always succeeds.
 StringRef OperationName::getStringRef() const {
   if (auto *op = representation.dyn_cast<const AbstractOperation *>())
     return op->name;
   return representation.get<Identifier>().strref();
 }
 
 const AbstractOperation *OperationName::getAbstractOperation() const {