[mlir] allow for re-registering extension ops

Op registration mechanism does not allow for ops with the same name to be
re-registered. This is okay to avoid name conflicts and debug
double-registration, but may be problematic for dialect extensions that may get
registered several times (unlike dialects that are deduplicated in the
registry). When registering ops through the Transform dialect extension
mechanism, check first if the ops are already registered and only complain in
the case of repeated registration with the same name but different TypeID.

Differential Revision: https://reviews.llvm.org/D125554

1 files changed, 19 insertions, 1 deletions

diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
index 2802d659312e..055cd78e6130 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
@@ -242,11 +242,29 @@ def Transform_Dialect : Dialect {
       getPDLConstraintHooks() const;
 
     private:
+      template <typename OpTy>
+      void addOperationIfNotRegistered() {
+        Optional<RegisteredOperationName> opName =
+            RegisteredOperationName::lookup(OpTy::getOperationName(),
+                                            getContext());
+        if (!opName)
+          return addOperations<OpTy>();
+
+        if (opName->getTypeID() == TypeID::get<OpTy>())
+          return;
+
+        llvm::errs() << "error: extensible dialect operation '"
+                     << OpTy::getOperationName()
+                     << "' is already registered with a mismatching TypeID";
+        abort();
+      }
+
       /// Registers operations specified as template parameters with this
       /// dialect. Checks that they implement the required interfaces.
       template <typename... OpTys>
       void addOperationsChecked() {
-        addOperations<OpTys...>();
+        (void)std::initializer_list<int>{(addOperationIfNotRegistered<OpTys>(),
+                                          0)...};
 
         #ifndef NDEBUG
         (void)std::initializer_list<int>{