[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

473

Nit: I'd expand this auto.

476

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

446–448

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 true if the given dialect is currently loading.
+  bool isDialectLoading(StringRef dialectNamespace);
+
   /// Load a dialect in the context.
   template <typename Dialect>
   void loadDialect() {
+    // Do not load the dialect if it is currently loading. This can happen if a
+    // dialect initializer triggers loading the same dialect recursively.
+    if (!isDialectLoading(Dialect::getDialectNamespace()))
       getOrLoadDialect<Dialect>();
   }
 
   /// Load a list dialects in the context.
   template <typename Dialect, typename OtherDialect, typename... MoreDialects>
   void loadDialect() {
-    getOrLoadDialect<Dialect>();
+    loadDialect<Dialect>();
     loadDialect<OtherDialect, MoreDialects...>();
   }
 
   /// Get (or create) a dynamic dialect for the given name.
   DynamicDialect *
   getOrLoadDynamicDialect(StringRef dialectNamespace,
                           function_ref<void(DynamicDialect *)> ctor);
 

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;
 
+  /// This is a list of dialect that are currently in the process of loading.
+  /// I.e., their constructor/initializer is still executing.
+  DenseSet<StringRef> loadingDialects;
+
   /// 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;
LLVM_DEBUG(llvm::dbgs()
                << "Load new dialect in Context " << dialectNamespace << "\n");
 #ifndef NDEBUG
     if (impl.multiThreadedExecutionContext != 0)
       llvm::report_fatal_error(
           "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.");
+    if (impl.loadingDialects.contains(dialectNamespace))
+      llvm::report_fatal_error(
+          "Loading (and getting) a dialect (" + dialectNamespace +
+          ") while the same dialect is still loading: use loadDialect instead "
+          "of getOrLoadDialect.");
 #endif
+    auto it = impl.loadingDialects.insert(dialectNamespace);
     std::unique_ptr<Dialect> &dialect =
         impl.loadedDialects.insert({dialectNamespace, ctor()}).first->second;
     assert(dialect && "dialect ctor failed");
+    impl.loadingDialects.erase(it.first);
 
     // 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::isDialectLoading(StringRef dialectNamespace) {
+  return getImpl().loadingDialects.contains(dialectNamespace);
+}
+
 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()))
       return dynDialect;

mlir/tools/mlir-tblgen/DialectGen.cpp

public:
   static constexpr ::llvm::StringLiteral getDialectNamespace() {
     return ::llvm::StringLiteral("{1}");
   }
 )";
 
 /// Registration for a single dependent dialect: to be inserted in the ctor
 /// above for each dependent dialect.
 const char *const dialectRegistrationTemplate = R"(
-    getContext()->getOrLoadDialect<{0}>();
+    getContext()->loadDialect<{0}>();
 )";
 
 /// The code block for the attribute parser/printer hooks.
 static const char *const attrParserDecl = R"(
   /// Parse an attribute registered to this dialect.
   ::mlir::Attribute parseAttribute(::mlir::DialectAsmParser &parser,
                                    ::mlir::Type type) const override;
 

 /// {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}
   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;