[RFC] Dialect::registerDependentDialects

Passes have registerDependentDialects to ensure that the necessary dialects are registered before a pass is run. Dialects can indicate dependent dialects for “in-dialect” stuff like op folding, constant materialize, and attribute/type handling, by sticking a call to getOrLoadDialect in the constructor.

Not only can this lead to confusing behaviour:

DialectA::DialectA(MLIRContext *ctx) : Dialect("a", ctx) {}
DialectB::DialectB(MLIRContext *ctx) : Dialect("b", ctx) {
  ctx->getOrLoadDialect<DialectA>();
}

// ...
MLIRContext ctx;
DialectRegistry registry.
registry.insert<DialectB>();
ctx.appendDialectRegistry(registry);

// ...
mlir::parseSourceString(R"(
  "a.op"() {} // error, unregistered dialect "a"
  "b.op"() {}
)", &ctx);

// ...
mlir::parseSourceString(R"(
  "b.op"() {} // dialect B is loaded which loads dialect A
  "a.op"() {} // works fine, no error
)", &ctx);

But it also means that dialects that aren’t self-contained and need other dialects to describe valid IR must also register those dependent dialects everywhere they are registered. Why not have a Dialect::registerDependentDialects so this isn’t needed?

But it also means that dialects that aren’t self-contained and need other dialects to describe valid IR

Do you have examples of such dialects?

must also register those dependent dialects everywhere they are registered

So far the idea has been that the registration is done by the API exposing the ability to parse an input. Such an API, by doing explicit registration, controls kind of dialect that are allowed to be present in the input.
The fact that a dialect can load another one as dependent kind of creates a “loophole”, even though not a good one as you noticed. I would say it is more like “hiding some bugs” than anything else.

1 Like

I guess no dialect is not “self-contained”, because inputs can always be provided through block arguments. I’m thinking, for example, scf.for which typically uses arith.constant to indicate loop bounds, or some other non-SCF op.

Or maybe when ops get folded and produce constants, which are then hoisted to the top of the function. If ArithmeticDialect isn’t registered, passing the resultant IR through the tool again will result in an unknown op error.

I guess this is true. I’m fine if the intended behaviour is “register what you use” all the way down the dependency graph.

We were discussing the other day having modes to make this easier to check rather than getting mysterious errors some of the time. Jacques suggested a mode where the passmanager unloads all the dialects after every pass, which would ensure all the dialects you use are registered in the main function (so the same tool can always process IR after each pass and reproducers work) and that each pass correctly declares its dependent dialects. I think it actually wouldn’t work for the particular case you’re describing though, since the loading happens while parsing. Would a method on Dialect help here? My understanding is that the intent behind loading other dialects in the constructor was to enable cases where canonicalization patterns depended on other dialects. If it was in a separate method that loading could be delayed until a pass was about to be run (which I think is the only time that something should be creating new operations of a dependent dialect). This might involve an annoying amount of bookkeeping, but I think it would help avoid these hidden bugs that appear at inconvenient moments.

I like this idea!

2 Likes

That’s a good point, but we shouldn’t mix up “debugging tools” with APIs for loading standalone inputs: adding an API here for mlir-opt purpose would end up impacting every IR loader. I’m not sure I see a good layering for that.

I think something like the above specifically for opt-style debugging tools (and probably behind a flag and/or build setting) would be a nice improvement. I wonder if there’s a way to achieve the same sort of thing for the case here with something like the delayed loading. I think that would require a separate method and mlir-opt could delay the loading of dialects from the populated registry (although now that I think of it, loading after completing all of parsing doesn’t seem bad for the general case either). IMO having an extra registerDependentDialects is actually a clearer contract than loading them in the constructor. I frequently see people asking what they’re supposed to do if a canonicalization pattern depends on another dialect.

You may have missed it, but that’s already implemented (example)?

Oh :smiley: I thought that’s what Jeff was asking for. I guess this inserts the loads in the constructor?

1 Like

Yes that’s what it does. I’m thinking of a mechanism that registers dependent dialects when the dialect is registered, instead of loading dependent dialects when the dialect is loaded.

E.g. registry.insert<DialectA>() would also call DialectA::registerDependentDialects(registry).

Are you referring to the parsing scenario or to the in-memory IR processing scenario? Why should a dialect register other dialects? It’s a pass that needs to load zero or more other dialects for it to access those dialects. I didn’t understand what it means for one dialect to require other dialects to be registered or loaded.

I can’t find registerDependentDialects anywhere in the codebase! (current official tree). There’s only getDependentDialects and it only loads dialects and does not register. (Some of the doc and code comments may be stale and referring to register instead of load.)

Oh yeah, I meant getDependentDialects. My bad.

And this method registers dialects, not load. A dialect is registered when its name and constructor are added to the MLIRContext. It is loaded when the constructor is called and the dialect is instantiated, registering all of its ops/types/attributes. If a dialect is loaded it must be registered.

Dialects already load (and register) other dialects when they are loaded, so that the ops in those dialects are accessible by op folders, materialize constants, etc. The problem is that if the depender dialect is registered but not loaded, its dependent dialects are neither registered nor loaded (resulting in the confusing IR error I showed above).

What I mean is that by adding calls to getOrLoadDialect on dependent dialects inside the constructor, dependent dialects are loaded when the depender dialect is loaded, but there is no way to register dependent dialects when the depender dialect is registered.

Essentially, if DialectB depends on DialectA, I want to be able to write registry.insert<DialectB>() instead of registry.insert<DialectA, DialectB>().

I think I was missing the dependentDialects on the Dialect definition itself. I’m still missing where this information is used to add the specified dependent dialects to the registry. I only see a Pass's dependentDialects getting looked at in the codebase.

In tablegen, there is a let dependentDialects = ["MyDialect"] under the Dialect tablegen class. This generates calls to ctx->getOrLoadDialect<MyDialect> inside the constructor.

1 Like

But this would lead to repeated unloading of even dialects that a user would have explicitly loaded via context.loadDialect<....> before a pass pipeline is run. This would both go against a user’s (compiler developer’s) choice and lead to repeated construction and destruction of dialect instances used throughout the pipeline for example.

Looks like I was confused between “register” and “load” here - thanks.

I’m missing this still. You can always do a context.loadDialect<...> without adding anything in the dialect registry. The registry is only used by parsing tools. The following method from MLIRContext.cpp doesn’t care if a dialect is registered or not:

Dialect *
  MLIRContext::getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
                                function_ref<std::unique_ptr<Dialect>()> ctor) {`

One can emit MLIR and transform it programmatically without registering any dialects – the dialects only need to be loaded.
https://mlir.llvm.org/getting_started/Faq/#registered-loaded-dependent-whats-up-with-dialects-management

Are you referring to the “mode where the passmanager unloads all the dialects after every pass”?
If so it is intended to be a debugging tool, like a verifier ++. The point isn’t to enable this in production, but to catch more issues ahead of time.

I’m talking about registered in a context. The “registry” is just a convenient collection of dialect constructors. What I mean is context.appendDialectRegistry(registry) registers the dialects in the context, but does not load them. They are instantiated as needed.

Yes - I was referring to that proposal.

But that looks quite confusing – to have that kind of different behavior in debug mode vs non-debug mode – a bit subjective though.

The terminology I’m using is consistent with the FAQ linked below. You don’t need a dialect to be registered with the context in order to load it or when it is loaded. Could you see the last paragraph on that FAQ section? Is it inconsistent? You stated "If a dialect is loaded it must be registered." However, a getOrLoadDialect never appends anything to the registry. Registering dialects with the context is only needed for parsing tools in order to control what can be loaded on demand as you mention further below.

https://mlir.llvm.org/getting_started/Faq/#registered-loaded-dependent-whats-up-with-dialects-management

Yes, this part is already pretty clear.

It depends: do you see a case where the different behavior in debug mode wouldn’t be a bug?
(similar to the fatal error messages we have when someone creates an operation with a builder without loading the dialect first).
Basically these are cases that “appear to work” without assert but are actually hiding latent bugs.

To be clear, I don’t think such a debugging mode should be added as part of NDEBUG (in general I think LLVM actually already overuses NDEBUG, which makes it impossible to build with asserts :frowning_face:), but rather hidden behind a different macro like MLIR_EXPENSIVE_CHECKS