[mlir] Fix crash when adding nested dialect extensions

[springerm]

Summary

A dialect extension can add additional dialect extensions in its apply function. This used to crash when the vector of extensions was internally reallocated while it is being iterated over.

[nicolasvasilache]

thanks!

mlir/lib/IR/Dialect.cpp

 MLIRContext *DialectInterface::getContext() const {
   return dialect->getContext();
 }
 
 DialectInterfaceCollectionBase::DialectInterfaceCollectionBase(
     MLIRContext *ctx, TypeID interfaceKind, StringRef interfaceName) {
   for (auto *dialect : ctx->getLoadedDialects()) {
 #ifndef NDEBUG
-  dialect->handleUseOfUndefinedPromisedInterface(interfaceKind, interfaceName);
+    dialect->handleUseOfUndefinedPromisedInterface(interfaceKind,
+                                                   interfaceName);
 #endif
     if (auto *interface = dialect->getRegisteredInterface(interfaceKind)) {
       interfaces.insert(interface);
       orderedInterfaces.push_back(interface);
     }
   }
 }
 
for (auto it = dialectNames.begin(), e = dialectNames.end(); it != e;
       Dialect *loadedDialect = ctx->getLoadedDialect(*it);
       if (!loadedDialect)
         return;
       requiredDialects.push_back(loadedDialect);
     }
     extension.apply(ctx, requiredDialects);
   };
 
-  for (const auto &extension : extensions)
-    applyExtension(*extension);
+  // Note: Additional extensions may be added while applying an extension.
+  for (int i = 0; i < static_cast<int>(extensions.size()); ++i)
+    applyExtension(*extensions[i]);
 }
 
 void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
   // Functor used to try to apply the given extension.
   auto applyExtension = [&](const DialectExtensionBase &extension) {
     ArrayRef<StringRef> dialectNames = extension.getRequiredDialects();
     if (dialectNames.empty()) {
       auto loadedDialects = ctx->getLoadedDialects();
       extension.apply(ctx, loadedDialects);
       return;
     }
 
     // Check to see if all of the dialects for this extension are loaded.
     SmallVector<Dialect *> requiredDialects;
     requiredDialects.reserve(dialectNames.size());
     for (StringRef dialectName : dialectNames) {
       Dialect *loadedDialect = ctx->getLoadedDialect(dialectName);
       if (!loadedDialect)
         return;
       requiredDialects.push_back(loadedDialect);
     }
     extension.apply(ctx, requiredDialects);
   };
 
-  for (const auto &extension : extensions)
-    applyExtension(*extension);
+  // Note: Additional extensions may be added while applying an extension.
+  for (int i = 0; i < static_cast<int>(extensions.size()); ++i)
+    applyExtension(*extensions[i]);
 }
 
 bool DialectRegistry::isSubsetOf(const DialectRegistry &rhs) const {
   // Treat any extensions conservatively.
   if (!extensions.empty())
     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

TEST(Dialect, RepeatedDelayedRegistration) {
   secondRegistry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
     dialect->addInterfaces<TestDialectInterface>();
   });
   context.appendDialectRegistry(secondRegistry);
   testDialectInterface = dyn_cast<TestDialectInterfaceBase>(testDialect);
   EXPECT_TRUE(testDialectInterface != nullptr);
 }
 
+namespace {
+/// A dummy extension that increases a counter when being applied and
+/// recursively adds additional extensions.
+struct DummyExtension : DialectExtension<DummyExtension, TestDialect> {
+  DummyExtension(int *counter, int numRecursive)
+      : DialectExtension(), counter(counter), numRecursive(numRecursive) {}
+
+  void apply(MLIRContext *ctx, TestDialect *dialect) const final {
+    ++(*counter);
+    DialectRegistry nestedRegistry;
+    for (int i = 0; i < numRecursive; ++i)
+      nestedRegistry.addExtension(
+          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);
+  }
+
+private:
+  int *counter;
+  int numRecursive;
+};
+} // namespace
+
+TEST(Dialect, NestedDialectExtension) {
+  DialectRegistry registry;
+  registry.insert<TestDialect>();
+
+  // Add an extension that adds 100 more extensions.
+  int counter1 = 0;
+  registry.addExtension(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));
+
+  // Load dialect and apply extensions.
+  MLIRContext context(registry);
+  Dialect *testDialect = context.getOrLoadDialect<TestDialect>();
+  ASSERT_TRUE(testDialect != nullptr);
+
+  // Extensions may be applied multiple times. Make sure that each expected
+  // extension was applied at least once.
+  EXPECT_GE(counter1, 101);
+  EXPECT_GE(counter2, 1);
+}
+
 } // namespace