[mlir] Fix circular dialect initialization

[springerm]

Summary

This change fixes a bug where a dialect is initialized multiple times. This triggers an assertion when the ops of the dialect are registered (error: operation named ... is already registered).

This bug can be triggered as follows:

1. Dialect A depends on dialect B (as per ADialect.td).

2. Somewhere there is an extension of dialect B that depends on dialect A (e.g., it defines external models create ops from dialect A). E.g.:

registry.addExtension(+[](MLIRContext *ctx, BDialect *dialect) {
  BDialectOp::attachInterface ...
  ctx->loadDialect<ADialect>();
});

3. When dialect A is loaded, its initialize function is called twice:

ADialect::ADialect()
   |     |
   |     v
   |   ADialect::initialize()
   v
getOrLoadDialect<BDialect>()
   |
   v
(load extension of BDialect)
   |
   v
ctx->loadDialect<ADialect>()  // user wrote this in the extension
   |
   v
getOrLoadDialect<ADialect>()  // the dialect is not "fully" loaded yet
   |
   v
ADialect::ADialect()
   |
   v
ADialect::initialize()

An example of a dialect extension that depends on other dialects is Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp. That particular dialect extension does not trigger this bug. (It would trigger this bug if the SCF dialect would depend on the Tensor dialect.)

This change introduces a new dialect state: dialects that are currently being loaded. Same as dialects that were already fully loaded (and initialized), dialects that are in the process of being loaded are not loaded a second time.

[springerm]

Note: I am not an expert on dialect loading/initialization. I can't tell whether this is the right way to fix the bug or not.

[mehdi_amini]

mlir/tools/mlir-tblgen/DialectGen.cpp

262

I think this require much more explanations: call out the recursive initialization through get-dependent.

But more importantly I don't actually understand how initialize() gets called twice, where is the second call coming from?
You mention in the description that there is a call to initialize() from ctx->loadDialect<ADialect>(), but how does this call initialize()? It may call another constructor, but we'd have the constructor called twice? Does not seem right, so I'm missing something :)

[springerm]

mlir/tools/mlir-tblgen/DialectGen.cpp

262

Yes, the same constructor is getting called twice.

When ADialect is loaded for the first time, an instance of ADialect is created, so the constructor is called. But the ADialect instance is not yet added to the list. This happens after the constructor is finished executing:

std::unique_ptr<Dialect> &dialect =
    impl.loadedDialects.insert({dialectNamespace, ctor()}).first->second;
assert(dialect && "dialect ctor failed");

The constructor (ADialect::ADialect) then indirectly calls ctx->loadDialect<ADialect>(), which calls getOrLoadDialect again. But we are still in the process of constructing the other ADialect instance, so it has not been added to loadedDialects yet. And therefore we create a second instance.

Note: impl.loadedDialects.insert is called twice. The second insert fails; i.e., it returns the already inserted ADialect instance, so one of the two ADialect objects is discarded.

[springerm]

mlir/tools/mlir-tblgen/DialectGen.cpp

262

so one of the two ADialect objects is discarded

This makes me wonder, are dialect objects allowed to have state? In that case, this change is wrong (as it may discard state).

However, it's still not ideal that ADialect is created twice....

[springerm]

a better fix: do not load/create the same dialect multiple times

[springerm]

I uploaded a better fix that avoid creating two versions of ADialect in the first place.

[ftynse]

On a technical level, the issue is with the fact that there's no entry in the dialect map when the second call is issued. Maybe we can proactively add a null pointer to that map before calling the constructor to work around this? It will result in the second getOrLoad call returning a null pointer. Does it happen only in extensions? If so, it may be okay as long as we document that extensions must not attempt to load dialects and use them, instead they should actually "extend" the dialect they want to use (since extensions may depend on multiple dialects).

On the conceptual level, loading dialects in extensions let us create cycles in the dialect dependency graph that are not visible in the dialect itself. I wonder if we can trigger a similar cycle without extensions, by having two mutually dependent dialects. If so, I'd go with whatever solution that addresses both problems.

mlir/lib/IR/MLIRContext.cpp

178

[mehdi_amini]

mlir/lib/IR/MLIRContext.cpp

180

I find this a bit annoying to have this separate data structure, what about Alex's suggestion of inserting a nullptr in the loadedDialects map?

mlir/tools/mlir-tblgen/DialectGen.cpp

262

Yes, the same constructor is getting called twice.

Right let's avoid that :)

This makes me wonder, are dialect objects allowed to have state?

Yes.

[springerm]

address comments

[springerm]

Update: Using nullptr to indicate "dialect loading" and also keeping isDialectLoading for error handling (report_fatal_error).

[ftynse]

mlir/lib/IR/MLIRContext.cpp

462

Nit: I'd expand this auto.

465

I'd rather do == nullptr to better match with the comment.

[springerm]

address comments

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

101–102

This doesn't look like something that should be publicly exposed.

mlir/lib/IR/MLIRContext.cpp

436–438

Can you avoid the double lookup here? It'd be nicer to just assign std::unique_ptr<Dialect> &dialect = impl.loadedDialects[dialectNamespace]; first, and then call the constructor. The lookup should already default initialize it to null.

mlir/include/mlir/IR/MLIRContext.h

public:
   T *getOrLoadDialect() {
     return static_cast<T *>(
         getOrLoadDialect(T::getDialectNamespace(), TypeID::get<T>(), [this]() {
           std::unique_ptr<T> dialect(new T(this));
           return dialect;
         }));
   }
 
+  /// Return whether the given dialect is loaded or not.
+  bool isDialectLoaded(StringRef dialectNamespace);
+
+  /// Return whether the given dialect is loaded or not.
+  template <typename T> bool isDialectLoaded() {
+    return isDialectLoaded(T::getDialectNamespace());
+  }
+
   /// Load a dialect in the context.
   template <typename Dialect>
   void loadDialect() {
     getOrLoadDialect<Dialect>();
   }
 
   /// Load a list dialects in the context.
   template <typename Dialect, typename OtherDialect, typename... MoreDialects>

mlir/lib/IR/MLIRContext.cpp

#endif
   /// destruction with the context.
   std::unique_ptr<llvm::ThreadPool> ownedThreadPool;
 
   /// This is a list of dialects that are created referring to this context.
   /// The MLIRContext owns the objects.
   DenseMap<StringRef, std::unique_ptr<Dialect>> loadedDialects;
   DialectRegistry dialectsRegistry;
 
   /// An allocator used for AbstractAttribute and AbstractType objects.
   llvm::BumpPtrAllocator abstractDialectSymbolAllocator;
 
   /// This is a mapping from operation name to the operation info describing it.
   llvm::StringMap<OperationName::Impl> operations;
 
   /// A vector of operation info specifically for registered operations.
   llvm::StringMap<RegisteredOperationName> registeredOperations;
 
   /// This is a sorted container of registered operations for a deterministic
   /// and efficient `getRegisteredOperations` implementation.
if (impl.multiThreadedExecutionContext != 0)
           "Loading a dialect (" + dialectNamespace +
           ") while in a multi-threaded execution context (maybe "
           "the PassManager): this can indicate a "
           "missing `dependentDialects` in a pass for example.");
 #endif
     std::unique_ptr<Dialect> &dialect =
         impl.loadedDialects.insert({dialectNamespace, ctor()}).first->second;
     assert(dialect && "dialect ctor failed");
 
     // Refresh all the identifiers dialect field, this catches cases where a
     // dialect may be loaded after identifier prefixed with this dialect name
     // were already created.
     auto stringAttrsIt = impl.dialectReferencingStrAttrs.find(dialectNamespace);
     if (stringAttrsIt != impl.dialectReferencingStrAttrs.end()) {
       for (StringAttrStorage *storage : stringAttrsIt->second)
         storage->referencedDialect = dialect.get();
       impl.dialectReferencingStrAttrs.erase(stringAttrsIt);
     }
 
     // Apply any extensions to this newly loaded dialect.
     impl.dialectsRegistry.applyExtensions(dialect.get());
     return dialect.get();
   }
 
   // Abort if dialect with namespace has already been registered.
   std::unique_ptr<Dialect> &dialect = dialectIt->second;
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");
 
   return dialect.get();
 }
 
+bool MLIRContext::isDialectLoaded(StringRef dialectNamespace) {
+  auto &impl = getImpl();
+  return impl.loadedDialects.find(dialectNamespace) !=
+         impl.loadedDialects.end();
+}
+
 DynamicDialect *MLIRContext::getOrLoadDynamicDialect(
     StringRef dialectNamespace, function_ref<void(DynamicDialect *)> ctor) {
   auto &impl = getImpl();
   // Get the correct insertion position sorted by namespace.
   auto dialectIt = impl.loadedDialects.find(dialectNamespace);
 
   if (dialectIt != impl.loadedDialects.end()) {
     if (auto dynDialect = dyn_cast<DynamicDialect>(dialectIt->second.get()))

mlir/tools/mlir-tblgen/DialectGen.cpp

 /// {0}: The name of the dialect class.
 /// {1}: initialization code that is emitted in the ctor body before calling
 ///      initialize().
 /// {2}: The dialect parent class.
 static const char *const dialectConstructorStr = R"(
 {0}::{0}(::mlir::MLIRContext *context) 
     : ::mlir::{2}(getDialectNamespace(), context, ::mlir::TypeID::get<{0}>()) {{
   {1}
+
+  // Do not initialize if this dialect was just loaded.
+  if (!context->isDialectLoaded(getDialectNamespace()))
     initialize();
 }
 )";
 
 /// The code block to generate a default desturctor definition.
 ///
 /// {0}: The name of the dialect class.
 static const char *const dialectDestructorStr = R"(
 {0}::~{0}() = default;