From 18fc39590978949fb75969e4bd63f8d2f13288ad Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko@google.com>
Date: Fri, 13 May 2022 18:03:08 +0200
Subject: [PATCH] [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
---
 .../mlir/Dialect/Transform/IR/TransformDialect.td    | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td b/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
index 2802d65..055cd78e6 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>{
-- 
2.7.4