add dialect registry extension comparison by key

Author: nikalra

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); });

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