[Mlir-commits] [flang] [mlir] [mlir] Allow dialect registry updates in a single-threaded context (PR #101119)

Nikhil Kalra llvmlistbot at llvm.org
Wed Jul 31 19:37:42 PDT 2024


https://github.com/nikalra updated https://github.com/llvm/llvm-project/pull/101119

>From c6532793c01b129ff7cbf6cf0db86d8ca7325069 Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Wed, 31 Jul 2024 19:24:14 -0700
Subject: [PATCH 1/3] DialectRegistry: add extension ID

---
 mlir/include/mlir/IR/DialectRegistry.h | 32 ++++++++++++++++++--------
 mlir/lib/IR/Dialect.cpp                |  4 ++--
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir/IR/DialectRegistry.h b/mlir/include/mlir/IR/DialectRegistry.h
index c13a1a1858eb1..ce5a766843228 100644
--- a/mlir/include/mlir/IR/DialectRegistry.h
+++ b/mlir/include/mlir/IR/DialectRegistry.h
@@ -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.emplace_back(extension.first,
+                                          extension.second->clone());
   }
 
   /// Return the names of dialects known to this registry.
@@ -206,39 +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));
+  void addExtension(StringRef extensionID,
+                    std::unique_ptr<DialectExtensionBase> extension) {
+    extensions.emplace_back(extensionID, std::move(extension));
   }
 
   /// Add the given extensions to the registry.
   template <typename... ExtensionsT>
   void addExtensions() {
-    (addExtension(std::make_unique<ExtensionsT>()), ...);
+    (addExtension(ExtensionsT::extensionID, 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("ID", +[](MLIRContext *ctx, MyDialect *dialect) {
+  ///  ... })
   ///
   /// is equivalent to:
   ///  registry.addExtension<MyDialect>(
+  ///     "ID",
   ///     [](MLIRContext *ctx, MyDialect *dialect){ ... }
   ///  )
   template <typename... DialectsT>
-  void addExtension(void (*extensionFn)(MLIRContext *, DialectsT *...)) {
+  void addExtension(StringRef extensionID,
+                    void (*extensionFn)(MLIRContext *, DialectsT *...)) {
     addExtension<DialectsT...>(
+        extensionID,
         std::function<void(MLIRContext *, DialectsT * ...)>(extensionFn));
   }
   template <typename... DialectsT>
   void
-  addExtension(std::function<void(MLIRContext *, DialectsT *...)> extensionFn) {
+  addExtension(StringRef extensionID,
+               std::function<void(MLIRContext *, DialectsT *...)> extensionFn) {
     using ExtensionFnT = std::function<void(MLIRContext *, DialectsT * ...)>;
 
     struct Extension : public DialectExtension<Extension, DialectsT...> {
       Extension(const Extension &) = default;
       Extension(ExtensionFnT extensionFn)
-          : extensionFn(std::move(extensionFn)) {}
+          : DialectExtension<Extension, DialectsT...>(),
+            extensionFn(std::move(extensionFn)) {}
       ~Extension() override = default;
 
       void apply(MLIRContext *context, DialectsT *...dialects) const final {
@@ -246,7 +255,8 @@ class DialectRegistry {
       }
       ExtensionFnT extensionFn;
     };
-    addExtension(std::make_unique<Extension>(std::move(extensionFn)));
+    addExtension(extensionID,
+                 std::make_unique<Extension>(std::move(extensionFn)));
   }
 
   /// Returns true if the current registry is a subset of 'rhs', i.e. if 'rhs'
@@ -255,7 +265,9 @@ class DialectRegistry {
 
 private:
   MapTy registry;
-  std::vector<std::unique_ptr<DialectExtensionBase>> extensions;
+  using KeyExtensionPair =
+      std::pair<llvm::StringRef, std::unique_ptr<DialectExtensionBase>>;
+  llvm::SmallVector<KeyExtensionPair> extensions;
 };
 
 } // namespace mlir
diff --git a/mlir/lib/IR/Dialect.cpp b/mlir/lib/IR/Dialect.cpp
index 965386681f270..0bdc34bbc5ae8 100644
--- a/mlir/lib/IR/Dialect.cpp
+++ b/mlir/lib/IR/Dialect.cpp
@@ -260,7 +260,7 @@ void DialectRegistry::applyExtensions(Dialect *dialect) const {
 
   // Note: Additional extensions may be added while applying an extension.
   for (int i = 0; i < static_cast<int>(extensions.size()); ++i)
-    applyExtension(*extensions[i]);
+    applyExtension(*extensions[i].second);
 }
 
 void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
@@ -287,7 +287,7 @@ void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
 
   // Note: Additional extensions may be added while applying an extension.
   for (int i = 0; i < static_cast<int>(extensions.size()); ++i)
-    applyExtension(*extensions[i]);
+    applyExtension(*extensions[i].second);
 }
 
 bool DialectRegistry::isSubsetOf(const DialectRegistry &rhs) const {

>From 26669bd8a53297dfe545b407d438c85d0ce32bfd Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Wed, 31 Jul 2024 19:36:25 -0700
Subject: [PATCH 2/3] attach extension IDs

---
 flang/lib/Optimizer/Dialect/FIRDialect.cpp    |  3 +-
 .../transform/Ch2/lib/MyExtension.cpp         |  3 +
 .../transform/Ch3/lib/MyExtension.cpp         |  3 +
 .../transform/Ch4/lib/MyExtension.cpp         |  3 +
 .../Conversion/ArithToLLVM/ArithToLLVM.cpp    |  7 +-
 .../ComplexToLLVM/ComplexToLLVM.cpp           |  1 +
 .../ControlFlowToLLVM/ControlFlowToLLVM.cpp   |  7 +-
 .../ConvertToLLVM/ConvertToLLVMPass.cpp       |  5 ++
 mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp |  7 +-
 .../Conversion/IndexToLLVM/IndexToLLVM.cpp    |  7 +-
 mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp |  7 +-
 .../Conversion/MemRefToLLVM/MemRefToLLVM.cpp  |  7 +-
 mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp |  7 +-
 mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp     |  7 +-
 .../Affine/IR/ValueBoundsOpInterfaceImpl.cpp  | 11 +--
 .../TransformOps/AffineTransformOps.cpp       |  2 +
 .../Arith/IR/ValueBoundsOpInterfaceImpl.cpp   | 16 +++--
 .../BufferDeallocationOpInterfaceImpl.cpp     |  7 +-
 .../BufferViewFlowOpInterfaceImpl.cpp         |  7 +-
 .../BufferizableOpInterfaceImpl.cpp           | 11 +--
 .../BufferizationTransformOps.cpp             |  2 +
 .../FuncBufferizableOpInterfaceImpl.cpp       | 14 ++--
 .../BufferDeallocationOpInterfaceImpl.cpp     |  7 +-
 .../BufferizableOpInterfaceImpl.cpp           | 10 +--
 .../Func/Extensions/InlinerExtension.cpp      | 11 +--
 .../Extensions/MeshShardingExtensions.cpp     | 11 +--
 .../Func/TransformOps/FuncTransformOps.cpp    |  2 +
 .../GPU/TransformOps/GPUTransformOps.cpp      |  2 +
 .../BufferDeallocationOpInterfaceImpl.cpp     |  7 +-
 .../Linalg/IR/ValueBoundsOpInterfaceImpl.cpp  | 13 ++--
 .../Linalg/TransformOps/DialectExtension.cpp  |  2 +
 .../BufferizableOpInterfaceImpl.cpp           | 19 ++---
 .../Transforms/MeshShardingInterfaceImpl.cpp  | 25 +++----
 .../Transforms/RuntimeOpVerification.cpp      | 17 ++---
 .../SubsetInsertionOpInterfaceImpl.cpp        |  9 +--
 .../Linalg/Transforms/TilingInterfaceImpl.cpp | 11 +--
 .../BufferizableOpInterfaceImpl.cpp           | 11 +--
 .../Dialect/MemRef/IR/MemRefMemorySlot.cpp    |  9 ++-
 .../MemRef/IR/ValueBoundsOpInterfaceImpl.cpp  | 25 ++++---
 .../TransformOps/MemRefTransformOps.cpp       |  2 +
 .../Transforms/AllocationOpInterfaceImpl.cpp  | 13 ++--
 .../BufferViewFlowOpInterfaceImpl.cpp         |  8 ++-
 .../Transforms/RuntimeOpVerification.cpp      | 26 +++----
 .../NVGPU/TransformOps/NVGPUTransformOps.cpp  |  2 +
 .../SCF/IR/ValueBoundsOpInterfaceImpl.cpp     |  9 +--
 .../SCF/TransformOps/SCFTransformOps.cpp      |  2 +
 .../BufferDeallocationOpInterfaceImpl.cpp     |  9 +--
 .../BufferizableOpInterfaceImpl.cpp           | 23 +++---
 .../BufferizableOpInterfaceImpl.cpp           | 10 +--
 .../TransformOps/SparseTensorTransformOps.cpp |  2 +
 .../BufferizableOpInterfaceImpl.cpp           | 40 ++++++-----
 .../IR/TensorInferTypeOpInterfaceImpl.cpp     | 15 ++--
 .../Tensor/IR/TensorTilingInterfaceImpl.cpp   | 21 +++---
 .../Tensor/IR/ValueBoundsOpInterfaceImpl.cpp  | 25 ++++---
 .../TransformOps/TensorTransformOps.cpp       | 23 +++---
 .../BufferizableOpInterfaceImpl.cpp           | 46 ++++++------
 .../SubsetInsertionOpInterfaceImpl.cpp        | 33 ++++-----
 .../Dialect/Tosa/IR/ShardingInterfaceImpl.cpp | 22 +++---
 .../DebugExtension/DebugExtension.cpp         |  2 +
 .../Transform/IRDLExtension/IRDLExtension.cpp |  2 +
 .../Transform/LoopExtension/LoopExtension.cpp |  2 +
 .../Transform/PDLExtension/PDLExtension.cpp   |  2 +
 .../Vector/IR/ValueBoundsOpInterfaceImpl.cpp  | 10 +--
 .../TransformOps/VectorTransformOps.cpp       |  2 +
 .../BufferizableOpInterfaceImpl.cpp           | 16 +++--
 .../Transforms/SubsetOpInterfaceImpl.cpp      | 21 +++---
 mlir/lib/Interfaces/CastInterfaces.cpp        |  9 +--
 mlir/lib/Target/LLVM/NVVM/Target.cpp          |  7 +-
 mlir/lib/Target/LLVM/ROCDL/Target.cpp         |  7 +-
 .../Dialect/AMX/AMXToLLVMIRTranslation.cpp    |  7 +-
 .../ArmNeon/ArmNeonToLLVMIRTranslation.cpp    |  1 +
 .../ArmSME/ArmSMEToLLVMIRTranslation.cpp      |  8 ++-
 .../ArmSVE/ArmSVEToLLVMIRTranslation.cpp      |  8 ++-
 .../Builtin/BuiltinToLLVMIRTranslation.cpp    |  7 +-
 .../Dialect/GPU/GPUToLLVMIRTranslation.cpp    |  7 +-
 .../LLVMIR/Dialect/GPU/SelectObjectAttr.cpp   |  7 +-
 .../LLVMIR/LLVMIRToLLVMTranslation.cpp        |  7 +-
 .../LLVMIR/LLVMToLLVMIRTranslation.cpp        |  7 +-
 .../Dialect/NVVM/LLVMIRToNVVMTranslation.cpp  |  7 +-
 .../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp  |  7 +-
 .../OpenACC/OpenACCToLLVMIRTranslation.cpp    |  7 +-
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      |  7 +-
 .../ROCDL/ROCDLToLLVMIRTranslation.cpp        |  7 +-
 .../Dialect/VCIX/VCIXToLLVMIRTranslation.cpp  |  7 +-
 .../X86VectorToLLVMIRTranslation.cpp          |  1 +
 mlir/lib/Target/SPIRV/Target.cpp              |  7 +-
 .../Test/TestFromLLVMIRTranslation.cpp        |  1 +
 .../Dialect/Test/TestToLLVMIRTranslation.cpp  |  1 +
 .../TestTransformDialectExtension.cpp         |  2 +
 .../TestTilingInterfaceTransformOps.cpp       |  2 +
 .../Transform/BuildOnlyExtensionTest.cpp      |  2 +
 mlir/unittests/IR/DialectTest.cpp             | 25 +++++--
 mlir/unittests/IR/InterfaceAttachmentTest.cpp | 72 +++++++++++--------
 93 files changed, 567 insertions(+), 393 deletions(-)

diff --git a/flang/lib/Optimizer/Dialect/FIRDialect.cpp b/flang/lib/Optimizer/Dialect/FIRDialect.cpp
index 4b1dadaac6728..3bf8b43c3572e 100644
--- a/flang/lib/Optimizer/Dialect/FIRDialect.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRDialect.cpp
@@ -74,7 +74,7 @@ void fir::FIROpsDialect::initialize() {
 // Register the FIRInlinerInterface to FIROpsDialect
 void fir::addFIRInlinerExtension(mlir::DialectRegistry &registry) {
   registry.addExtension(
-      +[](mlir::MLIRContext *ctx, fir::FIROpsDialect *dialect) {
+      "FIR_INLINER", +[](mlir::MLIRContext *ctx, fir::FIROpsDialect *dialect) {
         dialect->addInterface<FIRInlinerInterface>();
       });
 }
@@ -90,6 +90,7 @@ void fir::addFIRInlinerExtension(mlir::DialectRegistry &registry) {
 // when more sophisticated translation is required.
 void fir::addFIRToLLVMIRExtension(mlir::DialectRegistry &registry) {
   registry.addExtension(
+      "FIR_LLVM_TRANSLATION",
       +[](mlir::MLIRContext *ctx, fir::FIROpsDialect *dialect) {
         dialect->addInterface<mlir::LLVMTranslationDialectInterface>();
       });
diff --git a/mlir/examples/transform/Ch2/lib/MyExtension.cpp b/mlir/examples/transform/Ch2/lib/MyExtension.cpp
index 68d538a098018..a3f4f9aeb8106 100644
--- a/mlir/examples/transform/Ch2/lib/MyExtension.cpp
+++ b/mlir/examples/transform/Ch2/lib/MyExtension.cpp
@@ -36,6 +36,9 @@ class MyExtension
   // dialect definitions. List individual operations and dependent dialects
   // here.
   void init();
+
+  // Declare a unique ID for this extension.
+  static constexpr llvm::StringRef extensionID = "CH2_TRANSFORM";
 };
 
 void MyExtension::init() {
diff --git a/mlir/examples/transform/Ch3/lib/MyExtension.cpp b/mlir/examples/transform/Ch3/lib/MyExtension.cpp
index f7a99423a52ee..0eeeb3b4071d3 100644
--- a/mlir/examples/transform/Ch3/lib/MyExtension.cpp
+++ b/mlir/examples/transform/Ch3/lib/MyExtension.cpp
@@ -42,6 +42,9 @@ class MyExtension
   // dialect definitions. List individual operations and dependent dialects
   // here.
   void init();
+
+  // Declare a unique ID for this extension.
+  static constexpr llvm::StringRef extensionID = "CH3_TRANSFORM";
 };
 
 void MyExtension::init() {
diff --git a/mlir/examples/transform/Ch4/lib/MyExtension.cpp b/mlir/examples/transform/Ch4/lib/MyExtension.cpp
index 38c8ca1125a24..1d440f588e60b 100644
--- a/mlir/examples/transform/Ch4/lib/MyExtension.cpp
+++ b/mlir/examples/transform/Ch4/lib/MyExtension.cpp
@@ -38,6 +38,9 @@ class MyExtension
   // dialect definitions. List individual operations and dependent dialects
   // here.
   void init();
+
+  // Declare a unique ID for this extension.
+  static constexpr llvm::StringRef extensionID = "CH4_TRANSFORM";
 };
 
 void MyExtension::init() {
diff --git a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
index d882f1184f457..d2794a24cd05a 100644
--- a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
+++ b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
@@ -510,9 +510,10 @@ struct ArithToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
 
 void mlir::arith::registerConvertArithToLLVMInterface(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
-    dialect->addInterfaces<ArithToLLVMDialectInterface>();
-  });
+  registry.addExtension(
+      "ARITH_TO_LLVM", +[](MLIRContext *ctx, arith::ArithDialect *dialect) {
+        dialect->addInterfaces<ArithToLLVMDialectInterface>();
+      });
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp b/mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp
index 0a3c3a330ff69..3701d62efba50 100644
--- a/mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp
+++ b/mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp
@@ -385,6 +385,7 @@ struct ComplexToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
 
 void mlir::registerConvertComplexToLLVMInterface(DialectRegistry &registry) {
   registry.addExtension(
+      "COMPLEX_TO_LLVM",
       +[](MLIRContext *ctx, complex::ComplexDialect *dialect) {
         dialect->addInterfaces<ComplexToLLVMDialectInterface>();
       });
diff --git a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
index b8e5aec25286d..4a154554c30bc 100644
--- a/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
+++ b/mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp
@@ -273,7 +273,8 @@ struct ControlFlowToLLVMDialectInterface
 
 void mlir::cf::registerConvertControlFlowToLLVMInterface(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
-    dialect->addInterfaces<ControlFlowToLLVMDialectInterface>();
-  });
+  registry.addExtension(
+      "CF_TO_LLVM", +[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
+        dialect->addInterfaces<ControlFlowToLLVMDialectInterface>();
+      });
 }
diff --git a/mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp b/mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp
index 6135117348a5b..c55607f5e7c60 100644
--- a/mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp
+++ b/mlir/lib/Conversion/ConvertToLLVM/ConvertToLLVMPass.cpp
@@ -15,6 +15,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/StringRef.h"
 #include <memory>
 
 #define DEBUG_TYPE "convert-to-llvm"
@@ -54,6 +55,10 @@ class LoadDependentDialectExtension : public DialectExtensionBase {
   std::unique_ptr<DialectExtensionBase> clone() const final {
     return std::make_unique<LoadDependentDialectExtension>(*this);
   }
+
+  /// Unique ID
+  constexpr static llvm::StringRef extensionID =
+      "CONVERT_TO_LLVM_LOAD_DEPENDENT";
 };
 
 /// This is a generic pass to convert to LLVM, it uses the
diff --git a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
index c1f6d8bc5b361..5994a4c5fdf0f 100644
--- a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
+++ b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
@@ -786,7 +786,8 @@ struct FuncToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
 } // namespace
 
 void mlir::registerConvertFuncToLLVMInterface(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
-    dialect->addInterfaces<FuncToLLVMDialectInterface>();
-  });
+  registry.addExtension(
+      "FUNC_TO_LLVM", +[](MLIRContext *ctx, func::FuncDialect *dialect) {
+        dialect->addInterfaces<FuncToLLVMDialectInterface>();
+      });
 }
diff --git a/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp b/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
index 9d8a5d8a0e1c0..e045605475b78 100644
--- a/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
+++ b/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
@@ -392,7 +392,8 @@ struct IndexToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
 
 void mlir::index::registerConvertIndexToLLVMInterface(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, index::IndexDialect *dialect) {
-    dialect->addInterfaces<IndexToLLVMDialectInterface>();
-  });
+  registry.addExtension(
+      "INDEX_TO_LLVM", +[](MLIRContext *ctx, index::IndexDialect *dialect) {
+        dialect->addInterfaces<IndexToLLVMDialectInterface>();
+      });
 }
diff --git a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
index 23e957288eb95..f43ce2800e310 100644
--- a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
+++ b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
@@ -356,7 +356,8 @@ struct MathToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
 } // namespace
 
 void mlir::registerConvertMathToLLVMInterface(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, math::MathDialect *dialect) {
-    dialect->addInterfaces<MathToLLVMDialectInterface>();
-  });
+  registry.addExtension(
+      "MATH_TO_LLVM", +[](MLIRContext *ctx, math::MathDialect *dialect) {
+        dialect->addInterfaces<MathToLLVMDialectInterface>();
+      });
 }
diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 054827d40f0f3..d0d4b159ddd57 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -1753,7 +1753,8 @@ struct MemRefToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
 } // namespace
 
 void mlir::registerConvertMemRefToLLVMInterface(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
-    dialect->addInterfaces<MemRefToLLVMDialectInterface>();
-  });
+  registry.addExtension(
+      "MEMREF_TO_LLVM", +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+        dialect->addInterfaces<MemRefToLLVMDialectInterface>();
+      });
 }
diff --git a/mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp b/mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp
index 662ee9e483bc5..c6b856b84b10b 100644
--- a/mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp
+++ b/mlir/lib/Conversion/NVVMToLLVM/NVVMToLLVM.cpp
@@ -113,7 +113,8 @@ void mlir::populateNVVMToLLVMConversionPatterns(RewritePatternSet &patterns) {
 }
 
 void mlir::registerConvertNVVMToLLVMInterface(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, NVVMDialect *dialect) {
-    dialect->addInterfaces<NVVMToLLVMDialectInterface>();
-  });
+  registry.addExtension(
+      "NVVM_TO_LLVM", +[](MLIRContext *ctx, NVVMDialect *dialect) {
+        dialect->addInterfaces<NVVMToLLVMDialectInterface>();
+      });
 }
diff --git a/mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp b/mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp
index 0051333a35dcd..12150d2e1977b 100644
--- a/mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp
+++ b/mlir/lib/Conversion/UBToLLVM/UBToLLVM.cpp
@@ -119,7 +119,8 @@ struct UBToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
 } // namespace
 
 void mlir::ub::registerConvertUBToLLVMInterface(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, ub::UBDialect *dialect) {
-    dialect->addInterfaces<UBToLLVMDialectInterface>();
-  });
+  registry.addExtension(
+      "UB_TO_LLVM", +[](MLIRContext *ctx, ub::UBDialect *dialect) {
+        dialect->addInterfaces<UBToLLVMDialectInterface>();
+      });
 }
diff --git a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
index 82a9fb0d49088..2fe39c757fcd3 100644
--- a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -96,11 +96,12 @@ struct AffineMaxOpInterface
 
 void mlir::affine::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, AffineDialect *dialect) {
-    AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
-    AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
-    AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "AFFINE_VALUE_BOUNDS", +[](MLIRContext *ctx, AffineDialect *dialect) {
+        AffineApplyOp::attachInterface<AffineApplyOpInterface>(*ctx);
+        AffineMaxOp::attachInterface<AffineMaxOpInterface>(*ctx);
+        AffineMinOp::attachInterface<AffineMinOpInterface>(*ctx);
+      });
 }
 
 FailureOr<int64_t>
diff --git a/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp b/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
index 6457655cfe416..c6f5cfd47b046 100644
--- a/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
+++ b/mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp
@@ -167,6 +167,8 @@ class AffineTransformDialectExtension
 #include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "AFFINE_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
index 7cfcc4180539c..981a52090ad66 100644
--- a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -150,11 +150,13 @@ struct SelectOpInterface
 
 void mlir::arith::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
-    arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);
-    arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
-    arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
-    arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
-    arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "ARITH_VALUE_BOUNDS",
+      +[](MLIRContext *ctx, arith::ArithDialect *dialect) {
+        arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);
+        arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
+        arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
+        arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
+        arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index f2e7732e8ea4a..c17f6aacd88f8 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -79,7 +79,8 @@ struct SelectOpInterface
 
 void mlir::arith::registerBufferDeallocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, ArithDialect *dialect) {
-    SelectOp::attachInterface<SelectOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "ARITH_BUFFER_DEALLOC", +[](MLIRContext *ctx, ArithDialect *dialect) {
+        SelectOp::attachInterface<SelectOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.cpp
index 9df9df86b64fb..9f03fa947085f 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.cpp
@@ -38,7 +38,8 @@ struct SelectOpInterface
 
 void arith::registerBufferViewFlowOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
-    SelectOp::attachInterface<SelectOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "ARITH_BUFFER_FLOW", +[](MLIRContext *ctx, arith::ArithDialect *dialect) {
+        SelectOp::attachInterface<SelectOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
index 5e69a98db8f1e..74b55c124e480 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -206,9 +206,10 @@ struct SelectOpInterface
 
 void mlir::arith::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, ArithDialect *dialect) {
-    ConstantOp::attachInterface<ConstantOpInterface>(*ctx);
-    IndexCastOp::attachInterface<IndexCastOpInterface>(*ctx);
-    SelectOp::attachInterface<SelectOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "ARITH_BUFFERIZATION", +[](MLIRContext *ctx, ArithDialect *dialect) {
+        ConstantOp::attachInterface<ConstantOpInterface>(*ctx);
+        IndexCastOp::attachInterface<IndexCastOpInterface>(*ctx);
+        SelectOp::attachInterface<SelectOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
index e10c7bd914e35..d0993a5b85e8d 100644
--- a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
@@ -161,6 +161,8 @@ class BufferizationTransformDialectExtension
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "BUFFERIZATION_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 053ea7935260a..6f2ec90bc866e 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -82,7 +82,8 @@ getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
 
 /// Return the FuncOp called by `callOp`.
 static FuncOp getCalledFunction(CallOpInterface callOp) {
-  SymbolRefAttr sym = llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
+  SymbolRefAttr sym =
+      llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
   if (!sym)
     return nullptr;
   return dyn_cast_or_null<FuncOp>(
@@ -489,9 +490,10 @@ struct FuncOpInterface
 
 void mlir::bufferization::func_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
-    func::CallOp::attachInterface<func_ext::CallOpInterface>(*ctx);
-    func::FuncOp::attachInterface<func_ext::FuncOpInterface>(*ctx);
-    func::ReturnOp::attachInterface<func_ext::ReturnOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "FUNC_BUFFERIZATION", +[](MLIRContext *ctx, func::FuncDialect *dialect) {
+        func::CallOp::attachInterface<func_ext::CallOpInterface>(*ctx);
+        func::FuncOp::attachInterface<func_ext::FuncOpInterface>(*ctx);
+        func::ReturnOp::attachInterface<func_ext::ReturnOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index 89546da428fa2..f58f21b35547d 100644
--- a/mlir/lib/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -157,7 +157,8 @@ struct CondBranchOpInterface
 
 void mlir::cf::registerBufferDeallocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, ControlFlowDialect *dialect) {
-    CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "CF_BUFFER_DEALLOC", +[](MLIRContext *ctx, ControlFlowDialect *dialect) {
+        CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp
index 72f4a1a4f4c66..0db95f92dd95b 100644
--- a/mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -63,8 +63,10 @@ struct CondBranchOpInterface
 
 void mlir::cf::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
-    cf::BranchOp::attachInterface<BranchOpInterface>(*ctx);
-    cf::CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "CF_BUFFERIZATION",
+      +[](MLIRContext *ctx, cf::ControlFlowDialect *dialect) {
+        cf::BranchOp::attachInterface<BranchOpInterface>(*ctx);
+        cf::CondBranchOp::attachInterface<CondBranchOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
index 719a74a29a622..43e998fbc04f8 100644
--- a/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
+++ b/mlir/lib/Dialect/Func/Extensions/InlinerExtension.cpp
@@ -80,10 +80,11 @@ struct FuncInlinerInterface : public DialectInlinerInterface {
 //===----------------------------------------------------------------------===//
 
 void mlir::func::registerInlinerExtension(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
-    dialect->addInterfaces<FuncInlinerInterface>();
+  registry.addExtension(
+      "FUNC_INLINER", +[](MLIRContext *ctx, func::FuncDialect *dialect) {
+        dialect->addInterfaces<FuncInlinerInterface>();
 
-    // The inliner extension relies on the ControlFlow dialect.
-    ctx->getOrLoadDialect<cf::ControlFlowDialect>();
-  });
+        // The inliner extension relies on the ControlFlow dialect.
+        ctx->getOrLoadDialect<cf::ControlFlowDialect>();
+      });
 }
diff --git a/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp b/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp
index da508cc95bfe1..833362041848b 100644
--- a/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp
+++ b/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp
@@ -14,11 +14,12 @@
 namespace mlir::func {
 
 void registerShardingInterfaceExternalModels(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, FuncDialect *dialect) {
-    ReturnOp::attachInterface<
-        mesh::IndependentParallelIteratorDomainShardingInterface<ReturnOp>>(
-        *ctx);
-  });
+  registry.addExtension(
+      "FUNC_SHARDING", +[](MLIRContext *ctx, FuncDialect *dialect) {
+        ReturnOp::attachInterface<
+            mesh::IndependentParallelIteratorDomainShardingInterface<ReturnOp>>(
+            *ctx);
+      });
 }
 
 } // namespace mlir::func
diff --git a/mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp b/mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp
index b632b25d0cc67..711a2d9685f6d 100644
--- a/mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp
+++ b/mlir/lib/Dialect/Func/TransformOps/FuncTransformOps.cpp
@@ -246,6 +246,8 @@ class FuncTransformDialectExtension
 #include "mlir/Dialect/Func/TransformOps/FuncTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "FUNC_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
index 3661c5dea4525..d6a044698746f 100644
--- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
+++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
@@ -933,6 +933,8 @@ class GPUTransformDialectExtension
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "GPU_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index 6ccc0a26426c1..ddad79f173628 100644
--- a/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -31,7 +31,8 @@ struct GPUTerminatorOpInterface
 
 void mlir::gpu::registerBufferDeallocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, GPUDialect *dialect) {
-    gpu::TerminatorOp::attachInterface<GPUTerminatorOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "GPU_BUFFER_DEALLOC", +[](MLIRContext *ctx, GPUDialect *dialect) {
+        gpu::TerminatorOp::attachInterface<GPUTerminatorOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
index f56ef485069f8..390d6582c487e 100644
--- a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -53,9 +53,12 @@ struct IndexOpInterface
 
 void mlir::linalg::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
-    IndexOp::attachInterface<IndexOpInterface>(*ctx);
-    // Note: ValueBoundsOpInterface implementation is not required for ops that
-    // implement `DestinationStyleOpInterface` (for querying shaped OpResults).
-  });
+  registry.addExtension(
+      "LINALG_VALUE_BOUNDS",
+      +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+        IndexOp::attachInterface<IndexOpInterface>(*ctx);
+        // Note: ValueBoundsOpInterface implementation is not required for ops
+        // that implement `DestinationStyleOpInterface` (for querying shaped
+        // OpResults).
+      });
 }
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp b/mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp
index f4244ca962232..01c6814bc2082 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/DialectExtension.cpp
@@ -52,6 +52,8 @@ class LinalgTransformDialectExtension
 #include "mlir/Dialect/Linalg/TransformOps/LinalgMatchOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "LINALG_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
index be158af09d398..dc3bea56da8de 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -195,15 +195,18 @@ struct SoftmaxOpInterface
 
 void mlir::linalg::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
-    // Register all Linalg structured ops. `LinalgOp` is an interface and it is
-    // not possible to attach an external interface to an existing interface.
-    // Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
-    LinalgOpInterfaceHelper<
+  registry.addExtension(
+      "LINALG_BUFFERIZATION",
+      +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+        // Register all Linalg structured ops. `LinalgOp` is an interface and it
+        // is not possible to attach an external interface to an existing
+        // interface. Therefore, attach the `BufferizableOpInterface` to all ops
+        // one-by-one.
+        LinalgOpInterfaceHelper<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
-        >::registerOpInterface(ctx);
+            >::registerOpInterface(ctx);
 
-    SoftmaxOp::attachInterface<SoftmaxOpInterface>(*ctx);
-  });
+        SoftmaxOp::attachInterface<SoftmaxOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp
index 36b6088b83cc2..afb5ffcf31613 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.cpp
@@ -347,20 +347,21 @@ static void registerAll(MLIRContext *ctx) {
 }
 
 void registerMeshShardingInterfaceExternalModels(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, LinalgDialect *dialect) {
-    DialectRegistry registry;
-    registry.insert<affine::AffineDialect, arith::ArithDialect, scf::SCFDialect,
-                    tensor::TensorDialect>();
-    ctx->appendDialectRegistry(registry);
-    for (StringRef name : registry.getDialectNames())
-      ctx->getOrLoadDialect(name);
-
-    registerOne<linalg::GenericOp>(ctx);
-    registerAll<
+  registry.addExtension(
+      "LINALG_SHARDING", +[](MLIRContext *ctx, LinalgDialect *dialect) {
+        DialectRegistry registry;
+        registry.insert<affine::AffineDialect, arith::ArithDialect,
+                        scf::SCFDialect, tensor::TensorDialect>();
+        ctx->appendDialectRegistry(registry);
+        for (StringRef name : registry.getDialectNames())
+          ctx->getOrLoadDialect(name);
+
+        registerOne<linalg::GenericOp>(ctx);
+        registerAll<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
-        >(ctx);
-  });
+            >(ctx);
+      });
 }
 
 } // namespace mlir::linalg
diff --git a/mlir/lib/Dialect/Linalg/Transforms/RuntimeOpVerification.cpp b/mlir/lib/Dialect/Linalg/Transforms/RuntimeOpVerification.cpp
index b30182dc84079..6b3ab1a42bd99 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/RuntimeOpVerification.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/RuntimeOpVerification.cpp
@@ -121,15 +121,16 @@ void attachInterface(MLIRContext *ctx) {
 
 void mlir::linalg::registerRuntimeVerifiableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, LinalgDialect *) {
-    attachInterface<
+  registry.addExtension(
+      "LINALG_RUNTIME_VERIFICATION", +[](MLIRContext *ctx, LinalgDialect *) {
+        attachInterface<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
-        >(ctx);
+            >(ctx);
 
-    // Load additional dialects of which ops may get created.
-    ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,
-                     cf::ControlFlowDialect, index::IndexDialect,
-                     tensor::TensorDialect>();
-  });
+        // Load additional dialects of which ops may get created.
+        ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,
+                         cf::ControlFlowDialect, index::IndexDialect,
+                         tensor::TensorDialect>();
+      });
 }
diff --git a/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
index 6fcfa05468eea..04e4b86f8a855 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -74,8 +74,9 @@ struct LinalgCopyOpInterface
 
 void mlir::linalg::registerSubsetOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
-    linalg::CopyOp::attachInterface<LinalgCopyOpSubsetOpInterface>(*ctx);
-    linalg::CopyOp::attachInterface<LinalgCopyOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "LINALG_SUBSET", +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+        linalg::CopyOp::attachInterface<LinalgCopyOpSubsetOpInterface>(*ctx);
+        linalg::CopyOp::attachInterface<LinalgCopyOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index 2133458efe74c..9cd4b4c6adbd9 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -506,10 +506,11 @@ static void registerAll(MLIRContext *ctx) {
 
 void mlir::linalg::registerTilingInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
-    registerOne<linalg::GenericOp>(ctx);
-    registerAll<
+  registry.addExtension(
+      "LINALG_TILING", +[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+        registerOne<linalg::GenericOp>(ctx);
+        registerAll<
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
-        >(ctx);
-  });
+            >(ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
index 926d580ac7852..a07a412f11e85 100644
--- a/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -149,11 +149,12 @@ struct GlobalStoreOpInterface
 } // namespace
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, MLProgramDialect *) {
-    GlobalOp::attachInterface<GlobalOpInterface>(*ctx);
-    GlobalLoadOp::attachInterface<GlobalLoadOpInterface>(*ctx);
-    GlobalStoreOp::attachInterface<GlobalStoreOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "MLPROGRAM_BUFFERIZATION", +[](MLIRContext *ctx, MLProgramDialect *) {
+        GlobalOp::attachInterface<GlobalOpInterface>(*ctx);
+        GlobalLoadOp::attachInterface<GlobalLoadOpInterface>(*ctx);
+        GlobalStoreOp::attachInterface<GlobalStoreOpInterface>(*ctx);
+      });
 }
 } // namespace ml_program
 } // namespace mlir
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
index f630c48cdcaa1..7c21858adc5fb 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefMemorySlot.cpp
@@ -347,7 +347,10 @@ struct MemRefDestructurableTypeExternalModel
 //===----------------------------------------------------------------------===//
 
 void mlir::memref::registerMemorySlotExternalModels(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
-    MemRefType::attachInterface<MemRefDestructurableTypeExternalModel>(*ctx);
-  });
+  registry.addExtension(
+      "MEMREF_MEMORY_SLOT_MODELS",
+      +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+        MemRefType::attachInterface<MemRefDestructurableTypeExternalModel>(
+            *ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
index daec22cf6ebdc..fde96046b3ca2 100644
--- a/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -115,15 +115,18 @@ struct SubViewOpInterface
 
 void mlir::memref::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
-    memref::AllocOp::attachInterface<memref::AllocOpInterface<memref::AllocOp>>(
-        *ctx);
-    memref::AllocaOp::attachInterface<
-        memref::AllocOpInterface<memref::AllocaOp>>(*ctx);
-    memref::CastOp::attachInterface<memref::CastOpInterface>(*ctx);
-    memref::DimOp::attachInterface<memref::DimOpInterface>(*ctx);
-    memref::GetGlobalOp::attachInterface<memref::GetGlobalOpInterface>(*ctx);
-    memref::RankOp::attachInterface<memref::RankOpInterface>(*ctx);
-    memref::SubViewOp::attachInterface<memref::SubViewOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "MEMREF_VALUE_BOUNDS",
+      +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+        memref::AllocOp::attachInterface<
+            memref::AllocOpInterface<memref::AllocOp>>(*ctx);
+        memref::AllocaOp::attachInterface<
+            memref::AllocOpInterface<memref::AllocaOp>>(*ctx);
+        memref::CastOp::attachInterface<memref::CastOpInterface>(*ctx);
+        memref::DimOp::attachInterface<memref::DimOpInterface>(*ctx);
+        memref::GetGlobalOp::attachInterface<memref::GetGlobalOpInterface>(
+            *ctx);
+        memref::RankOp::attachInterface<memref::RankOpInterface>(*ctx);
+        memref::SubViewOp::attachInterface<memref::SubViewOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp b/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
index 8469e84c668cb..20f96a018d330 100644
--- a/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
+++ b/mlir/lib/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp
@@ -323,6 +323,8 @@ class MemRefTransformDialectExtension
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "MEMREF_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.cpp
index c433415944323..2cde261180e9e 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.cpp
@@ -60,10 +60,11 @@ struct DefaultReallocationInterface
 
 void mlir::memref::registerAllocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
-    memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
-    memref::AllocaOp::attachInterface<
-        DefaultAutomaticAllocationHoistingInterface>(*ctx);
-    memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
-  });
+  registry.addExtension(
+      "MEMREF_ALLOC", +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+        memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
+        memref::AllocaOp::attachInterface<
+            DefaultAutomaticAllocationHoistingInterface>(*ctx);
+        memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.cpp
index bbb269bd00161..5076c0be916c9 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.cpp
@@ -42,7 +42,9 @@ struct ReallocOpInterface
 
 void memref::registerBufferViewFlowOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
-    ReallocOp::attachInterface<ReallocOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "MEMREF_BUFFER_FLOW",
+      +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+        ReallocOp::attachInterface<ReallocOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp b/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
index 450bfa0cec0c7..80fc16d058b41 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/RuntimeOpVerification.cpp
@@ -333,16 +333,18 @@ struct ExpandShapeOpInterface
 
 void mlir::memref::registerRuntimeVerifiableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
-    CastOp::attachInterface<CastOpInterface>(*ctx);
-    ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
-    LoadOp::attachInterface<LoadStoreOpInterface<LoadOp>>(*ctx);
-    ReinterpretCastOp::attachInterface<ReinterpretCastOpInterface>(*ctx);
-    StoreOp::attachInterface<LoadStoreOpInterface<StoreOp>>(*ctx);
-    SubViewOp::attachInterface<SubViewOpInterface>(*ctx);
-
-    // Load additional dialects of which ops may get created.
-    ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,
-                     cf::ControlFlowDialect>();
-  });
+  registry.addExtension(
+      "MEMREF_RUNTIME_VERIFICATION",
+      +[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+        CastOp::attachInterface<CastOpInterface>(*ctx);
+        ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
+        LoadOp::attachInterface<LoadStoreOpInterface<LoadOp>>(*ctx);
+        ReinterpretCastOp::attachInterface<ReinterpretCastOpInterface>(*ctx);
+        StoreOp::attachInterface<LoadStoreOpInterface<StoreOp>>(*ctx);
+        SubViewOp::attachInterface<SubViewOpInterface>(*ctx);
+
+        // Load additional dialects of which ops may get created.
+        ctx->loadDialect<affine::AffineDialect, arith::ArithDialect,
+                         cf::ControlFlowDialect>();
+      });
 }
diff --git a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
index 7d3d868b326c6..8315ad2f5892b 100644
--- a/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp
@@ -1146,6 +1146,8 @@ class NVGPUTransformDialectExtension
 #include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "NVGPU_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 17a1c016ea16d..c6ede24e605df 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -159,8 +159,9 @@ struct IfOpInterface
 
 void mlir::scf::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) {
-    scf::ForOp::attachInterface<scf::ForOpInterface>(*ctx);
-    scf::IfOp::attachInterface<scf::IfOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "SCF_VALUE_BOUNDS", +[](MLIRContext *ctx, scf::SCFDialect *dialect) {
+        scf::ForOp::attachInterface<scf::ForOpInterface>(*ctx);
+        scf::IfOp::attachInterface<scf::IfOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index c4a55c302d0a3..c56d1e46bc0fd 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -624,6 +624,8 @@ class SCFTransformDialectExtension
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "SCF_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
index 24fbc1dca8361..bea0728e16987 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.cpp
@@ -75,8 +75,9 @@ struct ReduceReturnOpInterface
 
 void mlir::scf::registerBufferDeallocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, SCFDialect *dialect) {
-    InParallelOp::attachInterface<InParallelOpInterface>(*ctx);
-    ReduceReturnOp::attachInterface<ReduceReturnOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "SCF_BUFFER_DEALLOC", +[](MLIRContext *ctx, SCFDialect *dialect) {
+        InParallelOp::attachInterface<InParallelOpInterface>(*ctx);
+        ReduceReturnOp::attachInterface<ReduceReturnOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index cf40443ff3839..0aa9fd4385bb5 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1352,15 +1352,16 @@ struct InParallelOpInterface
 
 void mlir::scf::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) {
-    ConditionOp::attachInterface<ConditionOpInterface>(*ctx);
-    ExecuteRegionOp::attachInterface<ExecuteRegionOpInterface>(*ctx);
-    ForOp::attachInterface<ForOpInterface>(*ctx);
-    IfOp::attachInterface<IfOpInterface>(*ctx);
-    IndexSwitchOp::attachInterface<IndexSwitchOpInterface>(*ctx);
-    ForallOp::attachInterface<ForallOpInterface>(*ctx);
-    InParallelOp::attachInterface<InParallelOpInterface>(*ctx);
-    WhileOp::attachInterface<WhileOpInterface>(*ctx);
-    YieldOp::attachInterface<YieldOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "SCF_BUFFERIZATION", +[](MLIRContext *ctx, scf::SCFDialect *dialect) {
+        ConditionOp::attachInterface<ConditionOpInterface>(*ctx);
+        ExecuteRegionOp::attachInterface<ExecuteRegionOpInterface>(*ctx);
+        ForOp::attachInterface<ForOpInterface>(*ctx);
+        IfOp::attachInterface<IfOpInterface>(*ctx);
+        IndexSwitchOp::attachInterface<IndexSwitchOpInterface>(*ctx);
+        ForallOp::attachInterface<ForallOpInterface>(*ctx);
+        InParallelOp::attachInterface<InParallelOpInterface>(*ctx);
+        WhileOp::attachInterface<WhileOpInterface>(*ctx);
+        YieldOp::attachInterface<YieldOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp
index 66a2e45001781..b82c6624d0c32 100644
--- a/mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -137,8 +137,10 @@ struct AssumingYieldOpInterface
 
 void mlir::shape::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, shape::ShapeDialect *dialect) {
-    shape::AssumingOp::attachInterface<AssumingOpInterface>(*ctx);
-    shape::AssumingYieldOp::attachInterface<AssumingYieldOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "SHAPE_BUFFERIZATION",
+      +[](MLIRContext *ctx, shape::ShapeDialect *dialect) {
+        shape::AssumingOp::attachInterface<AssumingOpInterface>(*ctx);
+        shape::AssumingYieldOp::attachInterface<AssumingYieldOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp b/mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp
index ca19259ebffa6..e2b0c77f5f5a0 100644
--- a/mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp
+++ b/mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp
@@ -45,6 +45,8 @@ class SparseTensorTransformDialectExtension
 #include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "SPARSE_TENSOR_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 7734d1d258453..aca1a13a8b913 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -326,22 +326,26 @@ struct ToValuesOpInterface
 
 void mlir::sparse_tensor::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx,
-                            sparse_tensor::SparseTensorDialect *dialect) {
-    sparse_tensor::ConcatenateOp::attachInterface<ConcatenateOpInterface>(*ctx);
-    sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
-    sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
-    sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
-    sparse_tensor::NumberOfEntriesOp::attachInterface<
-        NumberOfEntriesOpInterface>(*ctx);
-    sparse_tensor::AssembleOp::attachInterface<AssembleOpInterface>(*ctx);
-    sparse_tensor::DisassembleOp::attachInterface<DisassembleOpInterface>(*ctx);
-    sparse_tensor::ForeachOp::attachInterface<ForeachOpInterface>(*ctx);
-    sparse_tensor::ToCoordinatesBufferOp::attachInterface<
-        ToCoordinatesBufferOpInterface>(*ctx);
-    sparse_tensor::ToCoordinatesOp::attachInterface<ToCoordinatesOpInterface>(
-        *ctx);
-    sparse_tensor::ToPositionsOp::attachInterface<ToPositionsOpInterface>(*ctx);
-    sparse_tensor::ToValuesOp::attachInterface<ToValuesOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "SPARSE_TENSOR_BUFFERIZATION",
+      +[](MLIRContext *ctx, sparse_tensor::SparseTensorDialect *dialect) {
+        sparse_tensor::ConcatenateOp::attachInterface<ConcatenateOpInterface>(
+            *ctx);
+        sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
+        sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
+        sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
+        sparse_tensor::NumberOfEntriesOp::attachInterface<
+            NumberOfEntriesOpInterface>(*ctx);
+        sparse_tensor::AssembleOp::attachInterface<AssembleOpInterface>(*ctx);
+        sparse_tensor::DisassembleOp::attachInterface<DisassembleOpInterface>(
+            *ctx);
+        sparse_tensor::ForeachOp::attachInterface<ForeachOpInterface>(*ctx);
+        sparse_tensor::ToCoordinatesBufferOp::attachInterface<
+            ToCoordinatesBufferOpInterface>(*ctx);
+        sparse_tensor::ToCoordinatesOp::attachInterface<
+            ToCoordinatesOpInterface>(*ctx);
+        sparse_tensor::ToPositionsOp::attachInterface<ToPositionsOpInterface>(
+            *ctx);
+        sparse_tensor::ToValuesOp::attachInterface<ToValuesOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
index 7ff435a033985..48b1dc2bdaa6c 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp
@@ -201,11 +201,12 @@ struct ReifyPadOp
 
 void mlir::tensor::registerInferTypeOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
-    ExpandShapeOp::attachInterface<
-        ReifyExpandOrCollapseShapeOp<tensor::ExpandShapeOp>>(*ctx);
-    CollapseShapeOp::attachInterface<
-        ReifyExpandOrCollapseShapeOp<tensor::CollapseShapeOp>>(*ctx);
-    PadOp::attachInterface<ReifyPadOp>(*ctx);
-  });
+  registry.addExtension(
+      "TENSOR_INFER_TYPE", +[](MLIRContext *ctx, TensorDialect *dialect) {
+        ExpandShapeOp::attachInterface<
+            ReifyExpandOrCollapseShapeOp<tensor::ExpandShapeOp>>(*ctx);
+        CollapseShapeOp::attachInterface<
+            ReifyExpandOrCollapseShapeOp<tensor::CollapseShapeOp>>(*ctx);
+        PadOp::attachInterface<ReifyPadOp>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
index 9b2a97eb2b006..88801be553148 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
@@ -785,17 +785,20 @@ FailureOr<TilingResult> tensor::bubbleUpPadSlice(OpBuilder &b,
 
 void mlir::tensor::registerTilingInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
-    tensor::PadOp::attachInterface<PadOpTiling>(*ctx);
-    tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
-    tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
-  });
+  registry.addExtension(
+      "TENSOR_TILING", +[](MLIRContext *ctx, TensorDialect *dialect) {
+        tensor::PadOp::attachInterface<PadOpTiling>(*ctx);
+        tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
+        tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
+      });
 }
 
 void mlir::tensor::registerTilingInterfaceExternalModelsForPackUnPackOps(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
-    tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
-    tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
-  });
+  registry.addExtension(
+      "TENOR_TILING_PACK_UNPACK",
+      +[](MLIRContext *ctx, TensorDialect *dialect) {
+        tensor::PackOp::attachInterface<PackOpTiling>(*ctx);
+        tensor::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
index 06f2c16406d3c..3b68822a4813d 100644
--- a/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -114,15 +114,18 @@ struct RankOpInterface
 
 void mlir::tensor::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
-    tensor::CastOp::attachInterface<tensor::CastOpInterface>(*ctx);
-    tensor::DimOp::attachInterface<tensor::DimOpInterface>(*ctx);
-    tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
-    tensor::ExtractSliceOp::attachInterface<tensor::ExtractSliceOpInterface>(
-        *ctx);
-    tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
-    tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
-    // Note: ValueBoundsOpInterface implementation is not required for ops that
-    // implement `DestinationStyleOpInterface` (for querying shaped OpResults).
-  });
+  registry.addExtension(
+      "TENSOR_VALUE_BOUNDS",
+      +[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+        tensor::CastOp::attachInterface<tensor::CastOpInterface>(*ctx);
+        tensor::DimOp::attachInterface<tensor::DimOpInterface>(*ctx);
+        tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
+        tensor::ExtractSliceOp::attachInterface<
+            tensor::ExtractSliceOpInterface>(*ctx);
+        tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
+        tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
+        // Note: ValueBoundsOpInterface implementation is not required for ops
+        // that implement `DestinationStyleOpInterface` (for querying shaped
+        // OpResults).
+      });
 }
diff --git a/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp b/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
index 33016f84056e9..74f795bf02ba4 100644
--- a/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
+++ b/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
@@ -70,15 +70,18 @@ struct ReassociativeReshapeOpReplacementInterface
 
 void tensor::registerFindPayloadReplacementOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
-    CollapseShapeOp::attachInterface<
-        ReassociativeReshapeOpReplacementInterface<CollapseShapeOp>>(*ctx);
-    ExpandShapeOp::attachInterface<
-        ReassociativeReshapeOpReplacementInterface<ExpandShapeOp>>(*ctx);
-    ExtractSliceOp::attachInterface<ExtractSliceOpReplacementInterface>(*ctx);
-    InsertSliceOp::attachInterface<InsertSliceOpReplacementInterface>(*ctx);
-    ReshapeOp::attachInterface<ReshapeOpReplacementInterface>(*ctx);
-  });
+  registry.addExtension(
+      "TENSOR_PAYLOAD_REPLACEMENT",
+      +[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+        CollapseShapeOp::attachInterface<
+            ReassociativeReshapeOpReplacementInterface<CollapseShapeOp>>(*ctx);
+        ExpandShapeOp::attachInterface<
+            ReassociativeReshapeOpReplacementInterface<ExpandShapeOp>>(*ctx);
+        ExtractSliceOp::attachInterface<ExtractSliceOpReplacementInterface>(
+            *ctx);
+        InsertSliceOp::attachInterface<InsertSliceOpReplacementInterface>(*ctx);
+        ReshapeOp::attachInterface<ReshapeOpReplacementInterface>(*ctx);
+      });
 }
 
 //===----------------------------------------------------------------------===//
@@ -247,6 +250,8 @@ class TensorTransformDialectExtension
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "TENSOR_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 87464ccb71720..b7f5fe676a7d2 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1055,28 +1055,30 @@ struct SplatOpInterface
 
 void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
-    CastOp::attachInterface<CastOpInterface>(*ctx);
-    CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
-    DimOp::attachInterface<DimOpInterface>(*ctx);
-    EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
-    ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
-    ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
-    ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
-    FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
-    GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
-    InsertOp::attachInterface<InsertOpInterface>(*ctx);
-    InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
-    PadOp::attachInterface<PadOpInterface>(*ctx);
-    ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
-        *ctx);
-    RankOp::attachInterface<RankOpInterface>(*ctx);
-    ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
-    SplatOp::attachInterface<SplatOpInterface>(*ctx);
-
-    // Load additional dialects of which ops may get created.
-    ctx->loadDialect<arith::ArithDialect, linalg::LinalgDialect>();
-  });
+  registry.addExtension(
+      "TENSOR_BUFFERIZATION",
+      +[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+        CastOp::attachInterface<CastOpInterface>(*ctx);
+        CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
+        DimOp::attachInterface<DimOpInterface>(*ctx);
+        EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
+        ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
+        ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
+        ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
+        FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
+        GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
+        InsertOp::attachInterface<InsertOpInterface>(*ctx);
+        InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
+        PadOp::attachInterface<PadOpInterface>(*ctx);
+        ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
+            *ctx);
+        RankOp::attachInterface<RankOpInterface>(*ctx);
+        ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
+        SplatOp::attachInterface<SplatOpInterface>(*ctx);
+
+        // Load additional dialects of which ops may get created.
+        ctx->loadDialect<arith::ArithDialect, linalg::LinalgDialect>();
+      });
 
   // Bufferization requires SubsetInsertionOpInterface models. Make sure that
   // they are registered.
diff --git a/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
index d50d7c62b789c..2b052cbb762e1 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -86,20 +86,21 @@ struct InsertSliceLikeOpSubsetInsertionOpInterface
 
 void mlir::tensor::registerSubsetOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
-    // Note: `SubsetExtractionOpInterface` and `SubsetInsertionOpInterface`
-    // require `SubsetOpInterface`.
-    ExtractSliceOp::attachInterface<ExtractSliceOpSubsetOpInterface>(*ctx);
-    ExtractSliceOp::attachInterface<ExtractSliceOpSubsetExtractionOpInterface>(
-        *ctx);
-    InsertSliceOp::attachInterface<
-        InsertSliceLikeOpSubsetOpInterface<InsertSliceOp>>(*ctx);
-    InsertSliceOp::attachInterface<
-        InsertSliceLikeOpSubsetInsertionOpInterface<InsertSliceOp>>(*ctx);
-    ParallelInsertSliceOp::attachInterface<
-        InsertSliceLikeOpSubsetOpInterface<ParallelInsertSliceOp>>(*ctx);
-    ParallelInsertSliceOp::attachInterface<
-        InsertSliceLikeOpSubsetInsertionOpInterface<ParallelInsertSliceOp>>(
-        *ctx);
-  });
+  registry.addExtension(
+      "TENSOR_SUBSET", +[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+        // Note: `SubsetExtractionOpInterface` and `SubsetInsertionOpInterface`
+        // require `SubsetOpInterface`.
+        ExtractSliceOp::attachInterface<ExtractSliceOpSubsetOpInterface>(*ctx);
+        ExtractSliceOp::attachInterface<
+            ExtractSliceOpSubsetExtractionOpInterface>(*ctx);
+        InsertSliceOp::attachInterface<
+            InsertSliceLikeOpSubsetOpInterface<InsertSliceOp>>(*ctx);
+        InsertSliceOp::attachInterface<
+            InsertSliceLikeOpSubsetInsertionOpInterface<InsertSliceOp>>(*ctx);
+        ParallelInsertSliceOp::attachInterface<
+            InsertSliceLikeOpSubsetOpInterface<ParallelInsertSliceOp>>(*ctx);
+        ParallelInsertSliceOp::attachInterface<
+            InsertSliceLikeOpSubsetInsertionOpInterface<ParallelInsertSliceOp>>(
+            *ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
index ffbb707344b8c..6161d5ae7e3ee 100644
--- a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
@@ -76,15 +76,17 @@ static void registerElemwiseAll(MLIRContext *ctx) {
 void mlir::tosa::registerShardingInterfaceExternalModels(
     DialectRegistry &registry) {
 
-  registry.addExtension(+[](MLIRContext *ctx, TosaDialect *dialect) {
-    registerElemwiseAll<
-        ClampOp, SigmoidOp, TanhOp, AddOp, ArithmeticRightShiftOp, BitwiseAndOp,
-        BitwiseOrOp, BitwiseXorOp, IntDivOp, LogicalAndOp, LogicalLeftShiftOp,
-        LogicalRightShiftOp, LogicalOrOp, LogicalXorOp, MaximumOp, MinimumOp,
-        MulOp, PowOp, SubOp, AbsOp, BitwiseNotOp, CeilOp, ClzOp, ExpOp, FloorOp,
-        LogOp, LogicalNotOp, NegateOp, ReciprocalOp, RsqrtOp, SelectOp, EqualOp,
-        GreaterOp, GreaterEqualOp>(ctx);
+  registry.addExtension(
+      "TOSA_SHARDING", +[](MLIRContext *ctx, TosaDialect *dialect) {
+        registerElemwiseAll<
+            ClampOp, SigmoidOp, TanhOp, AddOp, ArithmeticRightShiftOp,
+            BitwiseAndOp, BitwiseOrOp, BitwiseXorOp, IntDivOp, LogicalAndOp,
+            LogicalLeftShiftOp, LogicalRightShiftOp, LogicalOrOp, LogicalXorOp,
+            MaximumOp, MinimumOp, MulOp, PowOp, SubOp, AbsOp, BitwiseNotOp,
+            CeilOp, ClzOp, ExpOp, FloorOp, LogOp, LogicalNotOp, NegateOp,
+            ReciprocalOp, RsqrtOp, SelectOp, EqualOp, GreaterOp,
+            GreaterEqualOp>(ctx);
 
-    MatMulOp::attachInterface<MatMulOpSharding>(*ctx);
-  });
+        MatMulOp::attachInterface<MatMulOpSharding>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp b/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp
index e369daddb00cb..1e3c426162462 100644
--- a/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp
+++ b/mlir/lib/Dialect/Transform/DebugExtension/DebugExtension.cpp
@@ -26,6 +26,8 @@ class DebugExtension
 #include "mlir/Dialect/Transform/DebugExtension/DebugExtensionOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "TRANSFORM_DEBUG";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp b/mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp
index 94004365b8a1a..8cae64ff8de0a 100644
--- a/mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp
+++ b/mlir/lib/Dialect/Transform/IRDLExtension/IRDLExtension.cpp
@@ -26,6 +26,8 @@ class IRDLExtension
 
     declareDependentDialect<irdl::IRDLDialect>();
   }
+
+  static constexpr llvm::StringRef extensionID = "IRDL_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp b/mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp
index b33288fd7b991..032cd65e1065f 100644
--- a/mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp
+++ b/mlir/lib/Dialect/Transform/LoopExtension/LoopExtension.cpp
@@ -26,6 +26,8 @@ class LoopExtension
 #include "mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "TRANSFORM_LOOP";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp b/mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp
index 2c770abd56d52..c7f1e63e3d921 100644
--- a/mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp
+++ b/mlir/lib/Dialect/Transform/PDLExtension/PDLExtension.cpp
@@ -61,6 +61,8 @@ class PDLExtension : public transform::TransformDialectExtension<PDLExtension> {
           PDLOperationTypeTransformHandleTypeInterfaceImpl>(*context);
     });
   }
+
+  static constexpr llvm::StringRef extensionID = "TRANSFORM_PDL";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp
index ca95072d9bb0f..91bf4ef633a6d 100644
--- a/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -44,8 +44,10 @@ struct VectorScaleOpInterface
 
 void mlir::vector::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
-    vector::VectorScaleOp::attachInterface<vector::VectorScaleOpInterface>(
-        *ctx);
-  });
+  registry.addExtension(
+      "VECTOR_VALUE_BOUNDS",
+      +[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+        vector::VectorScaleOp::attachInterface<vector::VectorScaleOpInterface>(
+            *ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
index 2e9aa88011825..0b42040a4d63b 100644
--- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
@@ -220,6 +220,8 @@ class VectorTransformDialectExtension
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "VECTOR_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
index 1caec5bb8644f..d3e01e33b3f1d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -317,11 +317,13 @@ struct YieldOpInterface
 
 void mlir::vector::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
-    TransferReadOp::attachInterface<TransferReadOpInterface>(*ctx);
-    TransferWriteOp::attachInterface<TransferWriteOpInterface>(*ctx);
-    GatherOp::attachInterface<GatherOpInterface>(*ctx);
-    MaskOp::attachInterface<MaskOpInterface>(*ctx);
-    YieldOp::attachInterface<YieldOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "VECTOR_BUFFERIZATION",
+      +[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+        TransferReadOp::attachInterface<TransferReadOpInterface>(*ctx);
+        TransferWriteOp::attachInterface<TransferWriteOpInterface>(*ctx);
+        GatherOp::attachInterface<GatherOpInterface>(*ctx);
+        MaskOp::attachInterface<MaskOpInterface>(*ctx);
+        YieldOp::attachInterface<YieldOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
index b450d5b78a466..91a014a7ab24b 100644
--- a/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
@@ -69,14 +69,15 @@ struct TransferWriteOpSubsetInsertionOpInterface
 
 void mlir::vector::registerSubsetOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
-    TransferReadOp::attachInterface<XferOpSubsetOpInterface<TransferReadOp>>(
-        *ctx);
-    TransferReadOp::attachInterface<TransferReadOpSubsetExtractionOpInterface>(
-        *ctx);
-    TransferWriteOp::attachInterface<XferOpSubsetOpInterface<TransferWriteOp>>(
-        *ctx);
-    TransferWriteOp::attachInterface<TransferWriteOpSubsetInsertionOpInterface>(
-        *ctx);
-  });
+  registry.addExtension(
+      "VECTOR_SUBSET", +[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+        TransferReadOp::attachInterface<
+            XferOpSubsetOpInterface<TransferReadOp>>(*ctx);
+        TransferReadOp::attachInterface<
+            TransferReadOpSubsetExtractionOpInterface>(*ctx);
+        TransferWriteOp::attachInterface<
+            XferOpSubsetOpInterface<TransferWriteOp>>(*ctx);
+        TransferWriteOp::attachInterface<
+            TransferWriteOpSubsetInsertionOpInterface>(*ctx);
+      });
 }
diff --git a/mlir/lib/Interfaces/CastInterfaces.cpp b/mlir/lib/Interfaces/CastInterfaces.cpp
index 05c872daf5dab..11e581a5fefa8 100644
--- a/mlir/lib/Interfaces/CastInterfaces.cpp
+++ b/mlir/lib/Interfaces/CastInterfaces.cpp
@@ -79,10 +79,11 @@ struct UnrealizedConversionCastOpInterface
 
 void mlir::builtin::registerCastOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
-    UnrealizedConversionCastOp::attachInterface<
-        UnrealizedConversionCastOpInterface>(*ctx);
-  });
+  registry.addExtension(
+      "BUILTIN_CAST", +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+        UnrealizedConversionCastOp::attachInterface<
+            UnrealizedConversionCastOpInterface>(*ctx);
+      });
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index e608d26e8d2ec..f5d1babd01180 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -58,9 +58,10 @@ class NVVMTargetAttrImpl
 // Register the NVVM dialect, the NVVM translation & the target interface.
 void mlir::NVVM::registerNVVMTargetInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
-    NVVMTargetAttr::attachInterface<NVVMTargetAttrImpl>(*ctx);
-  });
+  registry.addExtension(
+      "NVVM_TARGET", +[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+        NVVMTargetAttr::attachInterface<NVVMTargetAttrImpl>(*ctx);
+      });
 }
 
 void mlir::NVVM::registerNVVMTargetInterfaceExternalModels(
diff --git a/mlir/lib/Target/LLVM/ROCDL/Target.cpp b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
index 4d23f987eb05e..3a56688117fda 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -68,9 +68,10 @@ class ROCDLTargetAttrImpl
 // Register the ROCDL dialect, the ROCDL translation and the target interface.
 void mlir::ROCDL::registerROCDLTargetInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
-    ROCDLTargetAttr::attachInterface<ROCDLTargetAttrImpl>(*ctx);
-  });
+  registry.addExtension(
+      "ROCDL_TARGET", +[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
+        ROCDLTargetAttr::attachInterface<ROCDLTargetAttrImpl>(*ctx);
+      });
 }
 
 void mlir::ROCDL::registerROCDLTargetInterfaceExternalModels(
diff --git a/mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp
index 044462d33cfd1..4ebcfa5941d0e 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp
@@ -44,9 +44,10 @@ class AMXDialectLLVMIRTranslationInterface
 
 void mlir::registerAMXDialectTranslation(DialectRegistry &registry) {
   registry.insert<amx::AMXDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, amx::AMXDialect *dialect) {
-    dialect->addInterfaces<AMXDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "AMX_TO_LLVMIR", +[](MLIRContext *ctx, amx::AMXDialect *dialect) {
+        dialect->addInterfaces<AMXDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerAMXDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp
index 7098592d506e0..d317f68731eb8 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp
@@ -46,6 +46,7 @@ class ArmNeonDialectLLVMIRTranslationInterface
 void mlir::registerArmNeonDialectTranslation(DialectRegistry &registry) {
   registry.insert<arm_neon::ArmNeonDialect>();
   registry.addExtension(
+      "ARM_NEON_TO_LLVMIR",
       +[](MLIRContext *ctx, arm_neon::ArmNeonDialect *dialect) {
         dialect->addInterfaces<ArmNeonDialectLLVMIRTranslationInterface>();
       });
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp
index e6ee41188d594..a1f63b9f07bdd 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp
@@ -45,9 +45,11 @@ class ArmSMEDialectLLVMIRTranslationInterface
 
 void mlir::registerArmSMEDialectTranslation(DialectRegistry &registry) {
   registry.insert<arm_sme::ArmSMEDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, arm_sme::ArmSMEDialect *dialect) {
-    dialect->addInterfaces<ArmSMEDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "ARM_SME_TO_LLVMIR",
+      +[](MLIRContext *ctx, arm_sme::ArmSMEDialect *dialect) {
+        dialect->addInterfaces<ArmSMEDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerArmSMEDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp
index cd10811b68f02..3411f5d147a34 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp
@@ -44,9 +44,11 @@ class ArmSVEDialectLLVMIRTranslationInterface
 
 void mlir::registerArmSVEDialectTranslation(DialectRegistry &registry) {
   registry.insert<arm_sve::ArmSVEDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, arm_sve::ArmSVEDialect *dialect) {
-    dialect->addInterfaces<ArmSVEDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "ARM_SVE_TO_LLVMIR",
+      +[](MLIRContext *ctx, arm_sve::ArmSVEDialect *dialect) {
+        dialect->addInterfaces<ArmSVEDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerArmSVEDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.cpp
index 51c304cfbb8e5..c120eb243b783 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.cpp
@@ -34,9 +34,10 @@ class BuiltinDialectLLVMIRTranslationInterface
 } // namespace
 
 void mlir::registerBuiltinDialectTranslation(DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
-    dialect->addInterfaces<BuiltinDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "BUILTIN_TO_LLVMIR", +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+        dialect->addInterfaces<BuiltinDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerBuiltinDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
index eecc8f1001ca4..c4fff8cfee7d1 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp
@@ -65,9 +65,10 @@ class GPUDialectLLVMIRTranslationInterface
 
 void mlir::registerGPUDialectTranslation(DialectRegistry &registry) {
   registry.insert<gpu::GPUDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
-    dialect->addInterfaces<GPUDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "GPU_TO_LLVMIR", +[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
+        dialect->addInterfaces<GPUDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerGPUDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
index b023c4c126da3..c982e5f67900f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
@@ -56,9 +56,10 @@ std::string getBinaryIdentifier(StringRef binaryName) {
 
 void mlir::gpu::registerOffloadingLLVMTranslationInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
-    SelectObjectAttr::attachInterface<SelectObjectAttrImpl>(*ctx);
-  });
+  registry.addExtension(
+      "GPU_OFFLOADING_LLVM", +[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
+        SelectObjectAttr::attachInterface<SelectObjectAttrImpl>(*ctx);
+      });
 }
 
 gpu::ObjectAttr
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
index 06673965245c0..14485c54ffdde 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
@@ -281,9 +281,10 @@ class LLVMDialectLLVMIRImportInterface : public LLVMImportDialectInterface {
 
 void mlir::registerLLVMDialectImport(DialectRegistry &registry) {
   registry.insert<LLVM::LLVMDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
-    dialect->addInterfaces<LLVMDialectLLVMIRImportInterface>();
-  });
+  registry.addExtension(
+      "LLVMIR_TO_LLVM", +[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
+        dialect->addInterfaces<LLVMDialectLLVMIRImportInterface>();
+      });
 }
 
 void mlir::registerLLVMDialectImport(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index bdb15a290209b..d5bc600fd7e13 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -439,9 +439,10 @@ class LLVMDialectLLVMIRTranslationInterface
 
 void mlir::registerLLVMDialectTranslation(DialectRegistry &registry) {
   registry.insert<LLVM::LLVMDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
-    dialect->addInterfaces<LLVMDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "LLVM_TO_LLVMIR", +[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
+        dialect->addInterfaces<LLVMDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerLLVMDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
index 855abc12a909e..9978513e59ec0 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/LLVMIRToNVVMTranslation.cpp
@@ -81,9 +81,10 @@ class NVVMDialectLLVMIRImportInterface : public LLVMImportDialectInterface {
 
 void mlir::registerNVVMDialectImport(DialectRegistry &registry) {
   registry.insert<NVVM::NVVMDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
-    dialect->addInterfaces<NVVMDialectLLVMIRImportInterface>();
-  });
+  registry.addExtension(
+      "LLVMIR_TO_NVVM", +[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+        dialect->addInterfaces<NVVMDialectLLVMIRImportInterface>();
+      });
 }
 
 void mlir::registerNVVMDialectImport(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index a09c24dda82af..e12364b1122f0 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -261,9 +261,10 @@ class NVVMDialectLLVMIRTranslationInterface
 
 void mlir::registerNVVMDialectTranslation(DialectRegistry &registry) {
   registry.insert<NVVM::NVVMDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
-    dialect->addInterfaces<NVVMDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "NVVM_TO_LLVMIR", +[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+        dialect->addInterfaces<NVVMDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerNVVMDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp
index d9cf85e4aecab..67f0f8fc4f374 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp
@@ -532,9 +532,10 @@ LogicalResult OpenACCDialectLLVMIRTranslationInterface::convertOperation(
 
 void mlir::registerOpenACCDialectTranslation(DialectRegistry &registry) {
   registry.insert<acc::OpenACCDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, acc::OpenACCDialect *dialect) {
-    dialect->addInterfaces<OpenACCDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "OPENACC_TO_LLVMIR", +[](MLIRContext *ctx, acc::OpenACCDialect *dialect) {
+        dialect->addInterfaces<OpenACCDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerOpenACCDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index ddee117838697..a43761460fa31 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -3741,9 +3741,10 @@ LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
 
 void mlir::registerOpenMPDialectTranslation(DialectRegistry &registry) {
   registry.insert<omp::OpenMPDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, omp::OpenMPDialect *dialect) {
-    dialect->addInterfaces<OpenMPDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "OPENMV_TO_LLVMIR", +[](MLIRContext *ctx, omp::OpenMPDialect *dialect) {
+        dialect->addInterfaces<OpenMPDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerOpenMPDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
index 2a146f5efed30..7fe58e5d161a8 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
@@ -190,9 +190,10 @@ class ROCDLDialectLLVMIRTranslationInterface
 
 void mlir::registerROCDLDialectTranslation(DialectRegistry &registry) {
   registry.insert<ROCDL::ROCDLDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
-    dialect->addInterfaces<ROCDLDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "ROCDL_TO_LLVMIR", +[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
+        dialect->addInterfaces<ROCDLDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerROCDLDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.cpp
index b78b002d32292..17d63863d9ddc 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/VCIX/VCIXToLLVMIRTranslation.cpp
@@ -76,9 +76,10 @@ class VCIXDialectLLVMIRTranslationInterface
 
 void mlir::registerVCIXDialectTranslation(DialectRegistry &registry) {
   registry.insert<vcix::VCIXDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, vcix::VCIXDialect *dialect) {
-    dialect->addInterfaces<VCIXDialectLLVMIRTranslationInterface>();
-  });
+  registry.addExtension(
+      "VCIX_TO_LLVMIR", +[](MLIRContext *ctx, vcix::VCIXDialect *dialect) {
+        dialect->addInterfaces<VCIXDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerVCIXDialectTranslation(MLIRContext &context) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp
index fa5f61420ee8a..5a6104162123b 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp
@@ -46,6 +46,7 @@ class X86VectorDialectLLVMIRTranslationInterface
 void mlir::registerX86VectorDialectTranslation(DialectRegistry &registry) {
   registry.insert<x86vector::X86VectorDialect>();
   registry.addExtension(
+      "X86_TO_LLVMIR",
       +[](MLIRContext *ctx, x86vector::X86VectorDialect *dialect) {
         dialect->addInterfaces<X86VectorDialectLLVMIRTranslationInterface>();
       });
diff --git a/mlir/lib/Target/SPIRV/Target.cpp b/mlir/lib/Target/SPIRV/Target.cpp
index 4c416abe71cac..15832a43b6620 100644
--- a/mlir/lib/Target/SPIRV/Target.cpp
+++ b/mlir/lib/Target/SPIRV/Target.cpp
@@ -43,9 +43,10 @@ class SPIRVTargetAttrImpl
 // Register the SPIR-V dialect, the SPIR-V translation & the target interface.
 void mlir::spirv::registerSPIRVTargetInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addExtension(+[](MLIRContext *ctx, spirv::SPIRVDialect *dialect) {
-    spirv::TargetEnvAttr::attachInterface<SPIRVTargetAttrImpl>(*ctx);
-  });
+  registry.addExtension(
+      "SPIRV_TARGET", +[](MLIRContext *ctx, spirv::SPIRVDialect *dialect) {
+        spirv::TargetEnvAttr::attachInterface<SPIRVTargetAttrImpl>(*ctx);
+      });
 }
 
 void mlir::spirv::registerSPIRVTargetInterfaceExternalModels(
diff --git a/mlir/test/lib/Dialect/Test/TestFromLLVMIRTranslation.cpp b/mlir/test/lib/Dialect/Test/TestFromLLVMIRTranslation.cpp
index dc6413b25707e..0422705663b81 100644
--- a/mlir/test/lib/Dialect/Test/TestFromLLVMIRTranslation.cpp
+++ b/mlir/test/lib/Dialect/Test/TestFromLLVMIRTranslation.cpp
@@ -104,6 +104,7 @@ void registerTestFromLLVMIR() {
         registry.insert<test::TestDialect>();
         registerLLVMDialectImport(registry);
         registry.addExtension(
+            "TEST_FROM_LLVMIR",
             +[](MLIRContext *ctx, test::TestDialect *dialect) {
               dialect->addInterfaces<TestDialectLLVMImportDialectInterface>();
             });
diff --git a/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp b/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp
index 157c6265be834..3d4b8dc6b1495 100644
--- a/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp
+++ b/mlir/test/lib/Dialect/Test/TestToLLVMIRTranslation.cpp
@@ -140,6 +140,7 @@ void registerTestToLLVMIR() {
         registerBuiltinDialectTranslation(registry);
         registerLLVMDialectTranslation(registry);
         registry.addExtension(
+            "TEST_TO_LLVMIR",
             +[](MLIRContext *ctx, test::TestDialect *dialect) {
               dialect->addInterfaces<TestDialectLLVMIRTranslationInterface>();
             });
diff --git a/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp b/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
index b8a4b9470d736..cc8f0e8988ea6 100644
--- a/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
+++ b/mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp
@@ -905,6 +905,8 @@ class TestTransformDialectExtension
           hooks.mergeInPDLMatchHooks(std::move(constraints));
         });
   }
+
+  static constexpr llvm::StringRef extensionID = "TEST_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
index 8f206d9077272..172001123d4da 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
@@ -399,6 +399,8 @@ class TestTilingInterfaceDialectExtension
 #include "TestTilingInterfaceTransformOps.cpp.inc"
         >();
   }
+
+  static constexpr llvm::StringRef extensionID = "TEST_TILING_TRANSFORM";
 };
 } // namespace
 
diff --git a/mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp b/mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp
index 40fb752ffd6eb..2d8fbf078618f 100644
--- a/mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp
+++ b/mlir/unittests/Dialect/Transform/BuildOnlyExtensionTest.cpp
@@ -20,6 +20,8 @@ class Extension : public TransformDialectExtension<Extension> {
 public:
   using Base::Base;
   void init() { declareGeneratedDialect<func::FuncDialect>(); }
+
+  static constexpr llvm::StringRef extensionID = "TRANSFORM_BUILD_ONLY";
 };
 } // end namespace
 
diff --git a/mlir/unittests/IR/DialectTest.cpp b/mlir/unittests/IR/DialectTest.cpp
index e99d46e6d2643..b79d7941a12a0 100644
--- a/mlir/unittests/IR/DialectTest.cpp
+++ b/mlir/unittests/IR/DialectTest.cpp
@@ -75,7 +75,8 @@ TEST(Dialect, DelayedInterfaceRegistration) {
   registry.insert<TestDialect, SecondTestDialect>();
 
   // Delayed registration of an interface for TestDialect.
-  registry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+  registry.addExtension(
+      "TEST_DIALECT_DELAYED", +[](MLIRContext *ctx, TestDialect *dialect) {
     dialect->addInterfaces<TestDialectInterface>();
   });
 
@@ -100,7 +101,7 @@ TEST(Dialect, DelayedInterfaceRegistration) {
   DialectRegistry secondRegistry;
   secondRegistry.insert<SecondTestDialect>();
   secondRegistry.addExtension(
-      +[](MLIRContext *ctx, SecondTestDialect *dialect) {
+      "SECOND_TEST", +[](MLIRContext *ctx, SecondTestDialect *dialect) {
         dialect->addInterfaces<SecondTestDialectInterface>();
       });
   context.appendDialectRegistry(secondRegistry);
@@ -113,7 +114,8 @@ TEST(Dialect, RepeatedDelayedRegistration) {
   // Set up the delayed registration.
   DialectRegistry registry;
   registry.insert<TestDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+  registry.addExtension(
+      "TEST_DIALECT", +[](MLIRContext *ctx, TestDialect *dialect) {
     dialect->addInterfaces<TestDialectInterface>();
   });
   MLIRContext context(registry);
@@ -128,7 +130,8 @@ TEST(Dialect, RepeatedDelayedRegistration) {
   // on repeated interface registration.
   DialectRegistry secondRegistry;
   secondRegistry.insert<TestDialect>();
-  secondRegistry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+  secondRegistry.addExtension(
+      "TEST_DIALECT", +[](MLIRContext *ctx, TestDialect *dialect) {
     dialect->addInterfaces<TestDialectInterface>();
   });
   context.appendDialectRegistry(secondRegistry);
@@ -143,12 +146,18 @@ struct DummyExtension : DialectExtension<DummyExtension, TestDialect> {
   DummyExtension(int *counter, int numRecursive)
       : DialectExtension(), counter(counter), numRecursive(numRecursive) {}
 
+  inline static std::vector<std::string> extensionIDs;
+
   void apply(MLIRContext *ctx, TestDialect *dialect) const final {
     ++(*counter);
     DialectRegistry nestedRegistry;
-    for (int i = 0; i < numRecursive; ++i)
+    extensionIDs.reserve(extensionIDs.size() + numRecursive);
+    for (int i = 0; i < numRecursive; ++i) {
+      extensionIDs.push_back("DUMMY_" + std::to_string(i));
       nestedRegistry.addExtension(
+          extensionIDs.back(),
           std::make_unique<DummyExtension>(counter, /*numRecursive=*/0));
+    }
     // Adding additional extensions may trigger a reallocation of the
     // `extensions` vector in the dialect registry.
     ctx->appendDialectRegistry(nestedRegistry);
@@ -166,10 +175,12 @@ TEST(Dialect, NestedDialectExtension) {
 
   // Add an extension that adds 100 more extensions.
   int counter1 = 0;
-  registry.addExtension(std::make_unique<DummyExtension>(&counter1, 100));
+  registry.addExtension("DUMMY",
+                        std::make_unique<DummyExtension>(&counter1, 100));
   // Add one more extension. This should not crash.
   int counter2 = 0;
-  registry.addExtension(std::make_unique<DummyExtension>(&counter2, 0));
+  registry.addExtension("DUMMY2",
+                        std::make_unique<DummyExtension>(&counter2, 0));
 
   // Load dialect and apply extensions.
   MLIRContext context(registry);
diff --git a/mlir/unittests/IR/InterfaceAttachmentTest.cpp b/mlir/unittests/IR/InterfaceAttachmentTest.cpp
index b6066dd5685dc..a5b835153b636 100644
--- a/mlir/unittests/IR/InterfaceAttachmentTest.cpp
+++ b/mlir/unittests/IR/InterfaceAttachmentTest.cpp
@@ -103,9 +103,11 @@ TEST(InterfaceAttachment, TypeDelayedContextConstruct) {
   // Put the interface in the registry.
   DialectRegistry registry;
   registry.insert<test::TestDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
-    test::TestType::attachInterface<TestTypeModel>(*ctx);
-  });
+  registry.addExtension(
+      "TYPE_DELAYED_CONSTRUCT",
+      +[](MLIRContext *ctx, test::TestDialect *dialect) {
+        test::TestType::attachInterface<TestTypeModel>(*ctx);
+      });
 
   // Check that when a context is constructed with the given registry, the type
   // interface gets registered.
@@ -122,9 +124,10 @@ TEST(InterfaceAttachment, TypeDelayedContextAppend) {
   // Put the interface in the registry.
   DialectRegistry registry;
   registry.insert<test::TestDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
-    test::TestType::attachInterface<TestTypeModel>(*ctx);
-  });
+  registry.addExtension(
+      "TYPE_DELAYED_APPEND", +[](MLIRContext *ctx, test::TestDialect *dialect) {
+        test::TestType::attachInterface<TestTypeModel>(*ctx);
+      });
 
   // Check that when the registry gets appended to the context, the interface
   // becomes available for objects in loaded dialects.
@@ -138,9 +141,10 @@ TEST(InterfaceAttachment, TypeDelayedContextAppend) {
 
 TEST(InterfaceAttachment, RepeatedRegistration) {
   DialectRegistry registry;
-  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
-    IntegerType::attachInterface<Model>(*ctx);
-  });
+  registry.addExtension(
+      "REPEATED", +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+        IntegerType::attachInterface<Model>(*ctx);
+      });
   MLIRContext context(registry);
 
   // Should't fail on repeated registration through the dialect registry.
@@ -151,9 +155,10 @@ TEST(InterfaceAttachment, TypeBuiltinDelayed) {
   // Builtin dialect needs to registration or loading, but delayed interface
   // registration must still work.
   DialectRegistry registry;
-  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
-    IntegerType::attachInterface<Model>(*ctx);
-  });
+  registry.addExtension(
+      "BUILTIN_DELAYED", +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+        IntegerType::attachInterface<Model>(*ctx);
+      });
 
   MLIRContext context(registry);
   IntegerType i16 = IntegerType::get(&context, 16);
@@ -246,9 +251,10 @@ TEST(InterfaceAttachmentTest, AttributeDelayed) {
   // that the delayed registration work for attributes.
   DialectRegistry registry;
   registry.insert<test::TestDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
-    test::SimpleAAttr::attachInterface<TestExternalSimpleAAttrModel>(*ctx);
-  });
+  registry.addExtension(
+      "ATTRIBUTE_DELAYED", +[](MLIRContext *ctx, test::TestDialect *dialect) {
+        test::SimpleAAttr::attachInterface<TestExternalSimpleAAttrModel>(*ctx);
+      });
 
   MLIRContext context(registry);
   context.loadDialect<test::TestDialect>();
@@ -352,13 +358,17 @@ struct TestExternalTestOpModel
 TEST(InterfaceAttachment, OperationDelayedContextConstruct) {
   DialectRegistry registry;
   registry.insert<test::TestDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
-    ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
-  });
-  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
-    test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
-    test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
-  });
+  registry.addExtension(
+      "OPERATION_DELAYED_BUILTIN",
+      +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+        ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
+      });
+  registry.addExtension(
+      "OPERATION_DELAYED_TEST",
+      +[](MLIRContext *ctx, test::TestDialect *dialect) {
+        test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
+        test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
+      });
 
   // Construct the context directly from a registry. The interfaces are
   // expected to be readily available on operations.
@@ -383,13 +393,17 @@ TEST(InterfaceAttachment, OperationDelayedContextConstruct) {
 TEST(InterfaceAttachment, OperationDelayedContextAppend) {
   DialectRegistry registry;
   registry.insert<test::TestDialect>();
-  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
-    ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
-  });
-  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
-    test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
-    test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
-  });
+  registry.addExtension(
+      "OPERATION_DELAYED_BUILTIN",
+      +[](MLIRContext *ctx, BuiltinDialect *dialect) {
+        ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
+      });
+  registry.addExtension(
+      "OPERATION_DELAYED_TEST",
+      +[](MLIRContext *ctx, test::TestDialect *dialect) {
+        test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
+        test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
+      });
 
   // Construct the context, create ops, and only then append the registry. The
   // interfaces are expected to be available after appending the registry.

>From 0cf68e846b20027c7903551c6a0d2321f4c06f88 Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Wed, 31 Jul 2024 19:36:45 -0700
Subject: [PATCH 3/3] add dialect registry extension comparison by key

---
 mlir/lib/IR/Dialect.cpp           | 14 ++++++++++--
 mlir/unittests/IR/DialectTest.cpp | 37 ++++++++++++++++++++++++++-----
 2 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/IR/Dialect.cpp b/mlir/lib/IR/Dialect.cpp
index 0bdc34bbc5ae8..c8b55bf751091 100644
--- a/mlir/lib/IR/Dialect.cpp
+++ b/mlir/lib/IR/Dialect.cpp
@@ -291,9 +291,19 @@ void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
 }
 
 bool DialectRegistry::isSubsetOf(const DialectRegistry &rhs) const {
-  // Treat any extensions conservatively.
-  if (!extensions.empty())
+  // Check that all extension keys are present in 'rhs'.
+  llvm::DenseSet<llvm::StringRef> rhsExtensionKeys;
+  {
+    auto rhsKeys = llvm::map_range(rhs.extensions,
+                                   [](const auto &item) { return item.first; });
+    rhsExtensionKeys.insert(rhsKeys.begin(), rhsKeys.end());
+  }
+
+  if (!llvm::all_of(extensions, [&rhsExtensionKeys](const auto &extension) {
+        return rhsExtensionKeys.contains(extension.first);
+      }))
     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); });
diff --git a/mlir/unittests/IR/DialectTest.cpp b/mlir/unittests/IR/DialectTest.cpp
index b79d7941a12a0..2623978639674 100644
--- a/mlir/unittests/IR/DialectTest.cpp
+++ b/mlir/unittests/IR/DialectTest.cpp
@@ -77,8 +77,8 @@ TEST(Dialect, DelayedInterfaceRegistration) {
   // Delayed registration of an interface for TestDialect.
   registry.addExtension(
       "TEST_DIALECT_DELAYED", +[](MLIRContext *ctx, TestDialect *dialect) {
-    dialect->addInterfaces<TestDialectInterface>();
-  });
+        dialect->addInterfaces<TestDialectInterface>();
+      });
 
   MLIRContext context(registry);
 
@@ -116,8 +116,8 @@ TEST(Dialect, RepeatedDelayedRegistration) {
   registry.insert<TestDialect>();
   registry.addExtension(
       "TEST_DIALECT", +[](MLIRContext *ctx, TestDialect *dialect) {
-    dialect->addInterfaces<TestDialectInterface>();
-  });
+        dialect->addInterfaces<TestDialectInterface>();
+      });
   MLIRContext context(registry);
 
   // Load the TestDialect and check that the interface got registered for it.
@@ -132,8 +132,8 @@ TEST(Dialect, RepeatedDelayedRegistration) {
   secondRegistry.insert<TestDialect>();
   secondRegistry.addExtension(
       "TEST_DIALECT", +[](MLIRContext *ctx, TestDialect *dialect) {
-    dialect->addInterfaces<TestDialectInterface>();
-  });
+        dialect->addInterfaces<TestDialectInterface>();
+      });
   context.appendDialectRegistry(secondRegistry);
   testDialectInterface = dyn_cast<TestDialectInterfaceBase>(testDialect);
   EXPECT_TRUE(testDialectInterface != nullptr);
@@ -193,4 +193,29 @@ TEST(Dialect, NestedDialectExtension) {
   EXPECT_GE(counter2, 1);
 }
 
+TEST(Dialect, SubsetWithExtensions) {
+  DialectRegistry registry1, registry2;
+  registry1.insert<TestDialect>();
+  registry2.insert<TestDialect>();
+
+  // Validate that the registries are equivalent.
+  ASSERT_TRUE(registry1.isSubsetOf(registry2));
+  ASSERT_TRUE(registry2.isSubsetOf(registry1));
+
+  // Add extensions to registry2.
+  int counter;
+  registry2.addExtension("EXT", std::make_unique<DummyExtension>(&counter, 0));
+
+  // Expect that (1) is a subset of (2) but not the other way around.
+  ASSERT_TRUE(registry1.isSubsetOf(registry2));
+  ASSERT_FALSE(registry2.isSubsetOf(registry1));
+
+  // Add extensions to registry1.
+  registry1.addExtension("EXT", std::make_unique<DummyExtension>(&counter, 0));
+
+  // Expect that (1) and (2) are equivalent.
+  ASSERT_TRUE(registry1.isSubsetOf(registry2));
+  ASSERT_TRUE(registry2.isSubsetOf(registry1));
+}
+
 } // namespace



More information about the Mlir-commits mailing list