[mlir] Support DialectRegistry extension comparison

mlir/examples/transform/Ch2/lib/MyExtension.cpp

@@ -29,6 +29,9 @@
 class MyExtension
     : public ::mlir::transform::TransformDialectExtension<MyExtension> {
 public:
+  // The TypeID of this extension.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MyExtension)
+
   // The extension must derive the base constructor.
   using Base::Base;
 

mlir/examples/transform/Ch3/lib/MyExtension.cpp

@@ -35,6 +35,9 @@
 class MyExtension
     : public ::mlir::transform::TransformDialectExtension<MyExtension> {
 public:
+  // The TypeID of this extension.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MyExtension)
+
   // The extension must derive the base constructor.
   using Base::Base;
 

mlir/examples/transform/Ch4/lib/MyExtension.cpp

@@ -31,6 +31,9 @@
 class MyExtension
     : public ::mlir::transform::TransformDialectExtension<MyExtension> {
 public:
+  // The TypeID of this extension.
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MyExtension)
+
   // The extension must derive the base constructor.
   using Base::Base;
 

mlir/include/mlir/IR/DialectRegistry.h

@@ -14,9 +14,9 @@
 #define MLIR_IR_DIALECTREGISTRY_H
 
 #include "mlir/IR/MLIRContext.h"
+#include "mlir/Support/TypeID.h"
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/MapVector.h"
 
 #include <map>
 #include <tuple>
@@ -187,7 +187,8 @@ class DialectRegistry {
                          nameAndRegistrationIt.second.second);
     // Merge the extensions.
     for (const auto &extension : extensions)
-      destination.extensions.push_back(extension->clone());
+      destination.extensions.try_emplace(extension.first,
+                                         extension.second->clone());
   }
 
   /// Return the names of dialects known to this registry.
@@ -206,47 +207,47 @@ class DialectRegistry {
   void applyExtensions(MLIRContext *ctx) const;
 
   /// Add the given extension to the registry.
-  void addExtension(std::unique_ptr<DialectExtensionBase> extension) {
-    extensions.push_back(std::move(extension));
+  bool addExtension(TypeID extensionID,
+                    std::unique_ptr<DialectExtensionBase> extension) {
+    return extensions.try_emplace(extensionID, std::move(extension)).second;
   }
 
   /// Add the given extensions to the registry.
   template <typename... ExtensionsT>
   void addExtensions() {
-    (addExtension(std::make_unique<ExtensionsT>()), ...);
+    (addExtension(TypeID::get<ExtensionsT>(), std::make_unique<ExtensionsT>()),
+     ...);
   }
 
   /// Add an extension function that requires the given dialects.
   /// Note: This bare functor overload is provided in addition to the
   /// std::function variant to enable dialect type deduction, e.g.:
-  ///  registry.addExtension(+[](MLIRContext *ctx, MyDialect *dialect) { ... })
+  ///  registry.addExtension(+[](MLIRContext *ctx, MyDialect *dialect) {
+  ///  ... })
   ///
   /// is equivalent to:
   ///  registry.addExtension<MyDialect>(
   ///     [](MLIRContext *ctx, MyDialect *dialect){ ... }
   ///  )
   template <typename... DialectsT>
-  void addExtension(void (*extensionFn)(MLIRContext *, DialectsT *...)) {
-    addExtension<DialectsT...>(
-        std::function<void(MLIRContext *, DialectsT * ...)>(extensionFn));
-  }
-  template <typename... DialectsT>
-  void
-  addExtension(std::function<void(MLIRContext *, DialectsT *...)> extensionFn) {
-    using ExtensionFnT = std::function<void(MLIRContext *, DialectsT * ...)>;
+  bool addExtension(void (*extensionFn)(MLIRContext *, DialectsT *...)) {
+    using ExtensionFnT = void (*)(MLIRContext *, DialectsT *...);
 
     struct Extension : public DialectExtension<Extension, DialectsT...> {
       Extension(const Extension &) = default;
       Extension(ExtensionFnT extensionFn)
-          : extensionFn(std::move(extensionFn)) {}
+          : DialectExtension<Extension, DialectsT...>(),
+            extensionFn(extensionFn) {}
       ~Extension() override = default;
 
       void apply(MLIRContext *context, DialectsT *...dialects) const final {
         extensionFn(context, dialects...);
       }
       ExtensionFnT extensionFn;
     };
-    addExtension(std::make_unique<Extension>(std::move(extensionFn)));
+    return addExtension(TypeID::getFromOpaquePointer(
+                            reinterpret_cast<const void *>(extensionFn)),
+                        std::make_unique<Extension>(extensionFn));
   }
 
   /// Returns true if the current registry is a subset of 'rhs', i.e. if 'rhs'
@@ -255,7 +256,7 @@ class DialectRegistry {
 
 private:
   MapTy registry;
-  std::vector<std::unique_ptr<DialectExtensionBase>> extensions;
+  llvm::MapVector<TypeID, std::unique_ptr<DialectExtensionBase>> extensions;
 };
 
 } // namespace mlir

mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp

@@ -35,6 +35,8 @@ namespace {
 /// starting a pass pipeline that involves dialect conversion to LLVM.
 class LoadDependentDialectExtension : public DialectExtensionBase {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LoadDependentDialectExtension)
+
   LoadDependentDialectExtension() : DialectExtensionBase(/*dialectNames=*/{}) {}
 
   void apply(MLIRContext *context,

mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp

@@ -157,6 +157,8 @@ class AffineTransformDialectExtension
     : public transform::TransformDialectExtension<
           AffineTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AffineTransformDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp

@@ -150,6 +150,9 @@ class BufferizationTransformDialectExtension
     : public transform::TransformDialectExtension<
           BufferizationTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      BufferizationTransformDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp

@@ -236,6 +236,8 @@ class FuncTransformDialectExtension
     : public transform::TransformDialectExtension<
           FuncTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FuncTransformDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

@@ -924,6 +924,8 @@ class GPUTransformDialectExtension
     : public transform::TransformDialectExtension<
           GPUTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GPUTransformDialectExtension)
+
   GPUTransformDialectExtension() {
     declareGeneratedDialect<scf::SCFDialect>();
     declareGeneratedDialect<arith::ArithDialect>();

mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp

@@ -30,6 +30,8 @@ class LinalgTransformDialectExtension
     : public transform::TransformDialectExtension<
           LinalgTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LinalgTransformDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp

@@ -309,6 +309,8 @@ class MemRefTransformDialectExtension
     : public transform::TransformDialectExtension<
           MemRefTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MemRefTransformDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp

@@ -1135,6 +1135,8 @@ class NVGPUTransformDialectExtension
     : public transform::TransformDialectExtension<
           NVGPUTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(NVGPUTransformDialectExtension)
+
   NVGPUTransformDialectExtension() {
     declareGeneratedDialect<arith::ArithDialect>();
     declareGeneratedDialect<affine::AffineDialect>();

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

@@ -613,6 +613,8 @@ class SCFTransformDialectExtension
     : public transform::TransformDialectExtension<
           SCFTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SCFTransformDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp

@@ -38,6 +38,9 @@ class SparseTensorTransformDialectExtension
     : public transform::TransformDialectExtension<
           SparseTensorTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      SparseTensorTransformDialectExtension)
+
   SparseTensorTransformDialectExtension() {
     declareGeneratedDialect<sparse_tensor::SparseTensorDialect>();
     registerTransformOps<

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

@@ -236,6 +236,8 @@ class TensorTransformDialectExtension
     : public transform::TransformDialectExtension<
           TensorTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TensorTransformDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp

@@ -20,6 +20,8 @@ namespace {
 class DebugExtension
     : public transform::TransformDialectExtension<DebugExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DebugExtension)
+
   void init() {
     registerTransformOps<
 #define GET_OP_LIST

mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp

@@ -18,6 +18,8 @@ namespace {
 class IRDLExtension
     : public transform::TransformDialectExtension<IRDLExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IRDLExtension)
+
   void init() {
     registerTransformOps<
 #define GET_OP_LIST

mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp

@@ -20,6 +20,8 @@ namespace {
 class LoopExtension
     : public transform::TransformDialectExtension<LoopExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LoopExtension)
+
   void init() {
     registerTransformOps<
 #define GET_OP_LIST

mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp

@@ -38,6 +38,8 @@ namespace {
 /// with Transform dialect operations.
 class PDLExtension : public transform::TransformDialectExtension<PDLExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PDLExtension)
+
   void init() {
     registerTransformOps<
 #define GET_OP_LIST

mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp

@@ -212,6 +212,8 @@ class VectorTransformDialectExtension
     : public transform::TransformDialectExtension<
           VectorTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VectorTransformDialectExtension)
+
   VectorTransformDialectExtension() {
     declareGeneratedDialect<vector::VectorDialect>();
     declareGeneratedDialect<LLVM::LLVMDialect>();

mlir/lib/IR/Dialect.cpp

@@ -11,14 +11,20 @@
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/DialectInterface.h"
+#include "mlir/IR/DialectRegistry.h"
 #include "mlir/IR/ExtensibleDialect.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Support/TypeID.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SetOperations.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Regex.h"
+#include <memory>
 
 #define DEBUG_TYPE "dialect"
 
@@ -173,6 +179,40 @@ bool dialect_extension_detail::hasPromisedInterface(Dialect &dialect,
 // DialectRegistry
 //===----------------------------------------------------------------------===//
 
+namespace {
+template <typename Fn>
+void applyExtensionsFn(
+    Fn &&applyExtension,
+    const llvm::MapVector<TypeID, std::unique_ptr<DialectExtensionBase>>
+        &extensions) {
+  // Note: Additional extensions may be added while applying an extension.
+  // The iterators will be invalidated if extensions are added so we'll keep
+  // a copy of the extensions for ourselves.
+
+  const auto extractExtension =
+      [](const auto &entry) -> DialectExtensionBase * {
+    return entry.second.get();
+  };
+
+  auto startIt = extensions.begin(), endIt = extensions.end();
+  size_t count = 0;
+  while (startIt != endIt) {
+    count += endIt - startIt;
+
+    // Grab the subset of extensions we'll apply in this iteration.
+    const auto subset =
+        llvm::map_to_vector(llvm::make_range(startIt, endIt), extractExtension);
+
+    for (const auto *ext : subset)
+      applyExtension(*ext);
+
+    // Book-keep for the next iteration.
+    startIt = extensions.begin() + count;
+    endIt = extensions.end();
+  }
+}
+} // namespace
+
 DialectRegistry::DialectRegistry() { insert<BuiltinDialect>(); }
 
 DialectAllocatorFunctionRef
@@ -258,9 +298,7 @@ void DialectRegistry::applyExtensions(Dialect *dialect) const {
     extension.apply(ctx, requiredDialects);
   };
 
-  // Note: Additional extensions may be added while applying an extension.
-  for (int i = 0; i < static_cast<int>(extensions.size()); ++i)
-    applyExtension(*extensions[i]);
+  applyExtensionsFn(applyExtension, extensions);
 }
 
 void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
@@ -285,15 +323,17 @@ void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
     extension.apply(ctx, requiredDialects);
   };
 
-  // Note: Additional extensions may be added while applying an extension.
-  for (int i = 0; i < static_cast<int>(extensions.size()); ++i)
-    applyExtension(*extensions[i]);
+  applyExtensionsFn(applyExtension, extensions);
 }
 
 bool DialectRegistry::isSubsetOf(const DialectRegistry &rhs) const {
-  // Treat any extensions conservatively.
-  if (!extensions.empty())
+  // Check that all extension keys are present in 'rhs'.
+  const auto hasExtension = [&](const auto &key) {
+    return rhs.extensions.contains(key);
+  };
+  if (!llvm::all_of(make_first_range(extensions), hasExtension))
     return false;
+
   // Check that the current dialects fully overlap with the dialects in 'rhs'.
   return llvm::all_of(
       registry, [&](const auto &it) { return rhs.registry.count(it.first); });

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

@@ -874,6 +874,8 @@ class TestTransformDialectExtension
     : public transform::TransformDialectExtension<
           TestTransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTransformDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp

@@ -382,6 +382,9 @@ class TestTilingInterfaceDialectExtension
     : public transform::TransformDialectExtension<
           TestTilingInterfaceDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestTilingInterfaceDialectExtension)
+
   using Base::Base;
 
   void init() {

mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp

@@ -18,6 +18,8 @@ using namespace mlir::transform;
 namespace {
 class Extension : public TransformDialectExtension<Extension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(Extension)
+
   using Base::Base;
   void init() { declareGeneratedDialect<func::FuncDialect>(); }
 };