[Mlir-commits] [mlir] 18fc395 - [mlir] allow for re-registering extension ops

[Alex Zinenko]

Author: Alex Zinenko
Date: 2022-05-17T00:03:40+02:00
New Revision: 18fc39590978949fb75969e4bd63f8d2f13288ad

URL: https://github.com/llvm/llvm-project/commit/18fc39590978949fb75969e4bd63f8d2f13288ad
DIFF: https://github.com/llvm/llvm-project/commit/18fc39590978949fb75969e4bd63f8d2f13288ad.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td

Removed: 
    


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