[mlir] Support DialectRegistry extension comparison (#101119)

`PassManager::run` loads the dependent dialects for each pass into the
current context prior to invoking the individual passes. If the
dependent dialect is already loaded into the context, this should be a
no-op. However, if there are extensions registered in the
`DialectRegistry`, the dependent dialects are unconditionally registered
into the context.

This poses a problem for dynamic pass pipelines, however, because they
will likely be executing while the context is in an immutable state
(because of the parent pass pipeline being run).

To solve this, we'll update the extension registration API on
`DialectRegistry` to require a type ID for each extension that is
registered. Then, instead of unconditionally registered dialects into a
context if extensions are present, we'll check against the extension
type IDs already present in the context's internal `DialectRegistry`.
The context will only be marked as dirty if there are net-new extension
types present in the `DialectRegistry` populated by
`PassManager::getDependentDialects`.

Note: this PR removes the `addExtension` overload that utilizes
`std::function` as the parameter. This is because `std::function` is
copyable and potentially allocates memory for the contained function so
we can't use the function pointer as the unique type ID for the
extension.

Downstream changes required:
- Existing `DialectExtension` subclasses will need a type ID to be
registered for each subclass. More details on how to register a type ID
can be found here:
https://github.com/llvm/llvm-project/blob/8b68e06731e0033ed3f8d6fe6292ae671611cfa1/mlir/include/mlir/Support/TypeID.h#L30
- Existing uses of the `std::function` overload of `addExtension` will
need to be refactored into dedicated `DialectExtension` classes with
associated type IDs. The attached `std::function` can either be inlined
into or called directly from `DialectExtension::apply`.

---------

Co-authored-by: Mehdi Amini <[email protected]>

Author: nikalra

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() {